Allocation of MPI buffers

[pmarguinaud]

Is your feature request related to a problem? Please describe.

Memory used in MPI communications has to be pinned. On CPU based architectures, this pinning has a cost (not too expensive), but on accelerated architectures, vendors (eg NVIDIA) insist that we reserve memory and reuse it for communications.
This led us to allocate persistent buffers in ectrans and in the semi-lagrangian; at some point, other parts of the code (IO, semi-implicit) requiring communications would be affected.

Describe the solution you'd like

I would like to have a mechanism to manage a memory buffer dedicated to MPI communications. This buffer could be allocated in the setup after querying the different parts of the code requiring MPI communications. We could use a heap based algorithm to manage this buffer.

This would avoid creating persistent buffers in different parts of the code and reduce the memory footprint of the application, at least for GPU architectures (where memory is expensive).

This, of course, has to be prototyped and benchmarked (at least on CPU architectures).

Describe alternatives you've considered

No response

Additional context

No response

Organisation

Météo-France

[wdeconinck]

Hi Philippe,

Memory used in MPI communications has to be pinned

This statement only applies for CPU-memory with CPU <-> GPU communication right? I don't think this applies for CPU <-> CPU communication, and I don't think GPU memory buffers need persistent memory?

It sounds like a good idea, and a "must" for performance.

We could use a heap based algorithm to manage this buffer.

What do you mean exactly with this?

[marsdeno]

Actually I believe the MPI library does typically pin buffers even for CPU <-> CPU, so that the NIC can access the buffer safely and efficiently?

[pmarguinaud]

Communications between CPUs require pinned memory. For GPUs, I do not remember what is exactly happening (NVIDIA explained me once), but the issue is similar.

heap based algorithm : we allocate a buffer in the setup, and then when we need memory for MPI, we allocate memory with functions like malloc/free which will return a pointer in the big buffer.

[pmarguinaud]

I would like something like this :

https://github.com/pmarguinaud/fiat/blob/14b2f07354186ae2a41034f2710592d91b6ab7b9/src/fiat/mpl/mpl_allocate_mod.fypp
https://github.com/pmarguinaud/fiat/blob/14b2f07354186ae2a41034f2710592d91b6ab7b9/src/fiat/heap/fiat_heap.c

I would call MPL_ALLOCATE_INIT in the setup to reserve memory for communications, and then allocate/deallocate MPI buffers when I need it (eg in the semi-Lagrangian).

[pmarguinaud]

mpi-memory.pdf

A short document describing a possible implementation, its usage in the code and the impact on performance.

[amckinstry]

@wdeconinck I am starting work on ACCORD (on behalf of Met Éireann).
I propose to try this implementation, testing on Mexluxina and Lumi. If I prepare and test a MR, would it be considered by ECMWF?

Thanks
Alastair McKinstry

[wdeconinck]

Dear @pmarguinaud, thank you for bringing up the issue of memory allocations, which is a very interesting one and also something close to my heart. In the proposal's example code using the buddy_alloc, I am not seeing where pinning or page-locking is applied and was this not the point? It is not immediately obvious to me or perhaps it is not yet implemented?

In the context of (de)allocation strategies it would be good to adopt a more general approach which is not limited to fiat/MPL. For the past years we have been developing a completely standalone project called pluto , which is forming the basis now for Atlas' field allocations and deallocations. It is embedded within the atlas repository currently here, and can be installed simply:

git clone https://github.com/ecmwf/atlas
cmake -S atlas/pluto -B build --install-prefix $PWD/install
cmake --build build
cmake --install build

The principle is an abstract concept of a pluto::allocator and pluto::memory_resource (typedef of c++ std::pmr::memory_resource). So it is completely compatible with C++ polymorphic memory resources and its documentation is worth a read.

Currently these are the pluto-built-in implementations of memory_resource available:

• host : new/delete
• device : cudaMalloc
• pinned : new/delete + cudaHostRegister (or hipHostRegister or no-op)
• managed : cudaMallocManaged (or hipMallocManaged or new/delete)
• host_pool : memory pool of host
• device_pool : memory pool of device
• pinned_pool : memory pool of pinned
• managed_pool : memory pool of managed

Custom resources can be easily created to extend the existing ones. However for all the pluto-built-in resources, statistics on allocations may be retrieved, and debugging output is also available upon request.

I have last week improved the pluto further, currently in a branch still, and will be merged in the next release soon.
New in this branch are also two new memory resources which use MPI_Alloc_mem and MPI_Free_mem

• mpi
• mpi_pool

Very likely this mpi_pool memory_resource is what we will want to be using for MPL as its actual implementation of MPI_Alloc_mem should be optimised by the HPC vendor / HPC stack to provide the best possible performance (pinned or not..?). However, the proposed design based on pluto it flexible enough to change at runtime even which allocation mechanism to use.

In this branch I also propose an implementation and example program on how an MPL allocator could be created using the mpi_pool memory resource with a lot of comments serving as a tutorial.
https://github.com/ecmwf/atlas/blob/pluto/pluto/examples/mpl_allocator.F90
The underlying mpi_pool memory resource could be used not only by MPL, but also for other programs making use of Atlas to create temporary MPI buffers. This should be a very nice to have feature.

Not shown in this example nor present yet in the Fortran API are asynchronous allocations using cudaStreams (or hipStreams).
Pluto also offers host-device data transfer, prefetching and data introspection mechanisms (currently only in the C++ API), using the low-level CUDA or HIP api directly.
The underlying C++ API already supports all of this and it is only a matter of making the Fortran bindings for the async_memory_resource. It is however very good to already keep this in mind when thinking about next steps for a generic GPU-allocation mechanism.

I realise it is more than just the few files present in @pmarguinaud 's proposal, but pluto should be providing us with many more features that will be useful in interoperability with the entire software stack.

Finally, I don't think we really need an API for GROW. Though it is convenient, the given case on FOUBUF is lacking the zero-initialisation that was present, and which is not needed in a general case.
This GROW pattern should not be that common, and a few more lines of code here are not so bad.
In fact better would be to use a pluto memory_pool resource, which would also remove the need or benefit of LALLOPERM implicitly.

@amckinstry I would refrain from doing too many developments right now that would be incompatible with the pluto alternative, although I believe the changes are minimal.

I would be very grateful for any review of this of course. Currently pluto is mainly (or only perhaps?) used within Atlas as it was still maturing, and especially the Fortran API is still very flexible to change. We can turn pluto into a standalone repository in due course as well.

[pmarguinaud]

Hello Willem,

Thank you for these interesting explanations. My problem is that I have been waiting for an MPL ready for GPU for two years now, and it is not there yet. And I need to have this working this year.

The solution I suggest (MPL_ALLOCATE/MPL_DEALLOCATE) is easy to implement, and could be backed by something else at some point, keeping the very same interface.

Would it be possible to have that while the pluto backend is being implemented ?

We would immediately benefit from it, as it brings about 10% on the forecast model running on AMD CPUs, and we could start implement it in ARPEGE/IFS and ECTRANS.

To answer one of your question, you do not need to do the pinning yourself; the MPI library will do it for you. What is important is to reuse already pinned memory. It is the same on GPU; at least this is what I have been told by NVIDIA.

[wdeconinck]

Hello Philippe,

I don't think the MPL-GPU readiness is related? Two years waiting for that is definitely excessive 😔. I was not aware of the requirement for you. The prerequisite move to MPI_F08 was very tedious, and now there's the final remaining issue #90 still in progress, which I believe is almost finished by @marsdeno. Thank you for your patience!

I like your suggestion to have an API for MPL_ALLOCATE/MPL_DEALLOCATE which could be flexible in adapting to pluto later, or even optionally.
For this purpose I have a fiat branch here, with an example API that could hopefully satisfy all?
It includes buddy_alloc and fiat_heap from your proposal, but I hope you're fine with the omission of the GROW api. For simplicity in prototyping, I only implemented the REAL64_R2 cases, avoiding fypp preprocessing during prototyping. Key is the "RESOURCE=STRING" argument to MPL_ALLOCATOR_INIT to be able to make runtime choices, and the addition of a label argument to allocate/deallocate for tracing. I have tested the compatibility with pluto backend as well, and a standard "FORTRAN" backend.

To avoid any future incompatibilities I think the fiat_heap should not be public at the moment, what do you think?

I have adapted the buddy_alloc.h code very slightly to add a configurable C++ namespace. I should try to feed this back upstream. This allows to also use buddy_alloc in other codes packages like field_api without having duplicate symbols.

Thanks for the insight that it is the MPI library that possibly does the pinning, it is clear to me now. I guess it will still be better to use something like MPI_Alloc_mem for a large memory-pool chunks immediately rather than have MPI pin select chunks.
With this MPL_ALLOCATE pluto backend it will be very easy to test and benchmark various strategies.

I'm happy to declare a fiat release without this for now, or do you want to wait for this to go in as well?

[pmarguinaud]

Hello Willem, yes it suits me.

Just one thing, I would really prefer not modifying buddy_alloc.h as it comes from some other repository, and we might have to update it in the future.

It is already possible to make it private and avoid symbol collision, as you can see here :

https://github.com/pmarguinaud/fiat/blob/mpl_allocate/src/fiat/heap/fiat_heap.cc

[wdeconinck]

That's a good point. I thought it would be problematic to #include "buddy_alloc.h" within the anonymous namespace as it also #includes further C standard headers. Some of my testing seems that it indeed is no issue, and all linkage is internal.

[wdeconinck]

Hello @pmarguinaud I have reworked my branch and made a draft PR #109 . You could already check it out and leave comments and suggestions.