Proposal for reworking host memory - #1928
Conversation
and resolve merge conflicts
and resolve merge conflicts
adayton1
changed the title
| axom::setDefaultAllocator(axom::MemorySpace::Unified); | ||
|
|
||
| // set Axom global allocator to an explicitly chosen Umpire allocator | ||
| int allocId = |
There was a problem hiding this comment.
axom::getAllocatorIDFromMemorySpace(axom::MemorySpace::Pinned) seems like a better example.
There was a problem hiding this comment.
Yeah, I didn't think too hard about the example. As I noted in the PR summary, many of the changes involve overloads to preserve the pre-existing interfaces, but show how the new explicit host allocator API works.
If we agree that this is an improvement and the way to move forward, then we would deprecate the "implicit" host allocator API and remove all references to global host allocator state, including in the docs.
There was a problem hiding this comment.
@BradWhitlock I added your suggestion as an alternative example.
| : | ||
| #else | ||
|
|
||
| const axom::HostAllocator hostAllocator {axom::execution_space<axom::SEQ_EXEC>::allocatorID()}; |
There was a problem hiding this comment.
This leaves me wanting syntactic sugar: const auto hostAllocator = axom::execution_space<axom::SEQ_EXEC>::hostAllocator();
There was a problem hiding this comment.
I agree that we can smooth this further.
| static int allocatorID() noexcept | ||
| { | ||
| return axom::getUmpireResourceAllocatorID(umpire::resource::Device); | ||
| return axom::getAllocatorIDFromMemorySpace(memory_space); |
| void check_device(axom::Array<T, DIM, SPACE>& v) | ||
| { | ||
| const axom::IndexType size = v.size(); | ||
| const int explicit_host_alloc = |
There was a problem hiding this comment.
Why this and not axom::getAllocatorIDFromMemorySpace(axom::MemorySpace::Host)?
There was a problem hiding this comment.
This is intended to be an explicit check that the default allocator and the Umpire Host allocator are consistent. In particular, axom::MemorySpace::Host could refer to either the Umpire Host allocator or Axom MALLOC, depending on what the default is set to.
|
|
||
| // Then check the contents by assigning to an explicitly Host Array | ||
| axom::Array<T, 1, axom::MemorySpace::Host> check_raw_array_host = v; | ||
| axom::Array<T, 1, axom::MemorySpace::Host> check_raw_array_host(v, explicit_host_alloc); |
There was a problem hiding this comment.
In this statement and some others in this test it switched from assigment operator to copy constructor. Is that intentional?
There was a problem hiding this comment.
This is testing the mechanics of the ctor overloads, that takes a host allocator id explicitly.
| ::check_alloc(v_int_device, axom::getUmpireResourceAllocatorID(umpire::resource::Device)); | ||
| axom::Array<double, 1, axom::MemorySpace::Device> v_double_device(capacity, capacity); | ||
| ::check_alloc(v_double_device, axom::getUmpireResourceAllocatorID(umpire::resource::Device)); | ||
| if(runtimeMemorySpaceAvailable(axom::MemorySpace::Device)) |
There was a problem hiding this comment.
I've seen other efforts in this PR to minimize use of the Umpire enums and related functions and instead request allocators for the axom::MemorySpace::Device (whichever enum value). Here they are being mixed. Could axom::getAllocatorIDFromMemorySpace(axom::MemorySpace::Device) be used?
There was a problem hiding this comment.
These tests are checking that the axom::MemorySpace enum values match up with the umpire::resource:: enum values.
| // in check_alloc_realloc_free when reallocating to 3 * ARRAY_SIZE. | ||
| constexpr int ARRAY_SIZE = 5345; | ||
|
|
||
| struct ScopedDefaultAllocatorState |
There was a problem hiding this comment.
This RAII pattern is repeated from a different test. Would it be useful to have in axom::core?
| buffer_size *= 3; | ||
| buffer = axom::reallocate<int>(buffer, buffer_size); | ||
| #ifdef AXOM_USE_UMPIRE | ||
| #if defined(AXOM_USE_UMPIRE) |
There was a problem hiding this comment.
Is the AXOM_USE_UMPIRE check still needed since you switched to not using the umpire::ResourceManager?
There was a problem hiding this comment.
It may be redundant here since I think (hope) that if AXOM_USE_UMPIRE is not defined, then UMPIRE_ENABLE_* will not be defined. However, I found the use of AXOM_USE_UMPIRE and UMPIRE_ENABLE_* macro constants inconsistent and confusing. It was there in the original, so I left it.
| } | ||
| }; | ||
|
|
||
| inline HostAllocator hostAllocatorForPrimaryAllocator(int allocator_id) |
There was a problem hiding this comment.
It seems like these functions might be useful in other places too. Do they belong in memory_management.hpp?
| , m_host_allocator(other.m_host_allocator) | ||
| , m_arrayOps(m_allocator_id, m_executeOnGPU, m_host_allocator) | ||
| { | ||
| #if defined(AXOM_DEVICE_CODE) |
There was a problem hiding this comment.
Making this AXOM_HOST_DEVICE and adding the trap/abort behavior suggests to me that we'll now get a runtime error if we accidentally capture axom::Array in device code. Nice. It might be good to test this in CI if it isn't already (codex did not see a test).
There was a problem hiding this comment.
There is a test for it, but it's disabled on AMD; however, AMD's host/device checks are stricter at compile-time so it's less of a concern.
|
|
||
| AXOM_HOST_DEVICE inline T* data() { return m_data; } | ||
| AXOM_HOST_DEVICE inline const T* data() const { return m_data; } | ||
| inline int getHostAllocatorID() const { return m_host_allocator.getID(); } |
There was a problem hiding this comment.
It looks like some other properties can be queried from the host allocator. Would be make sense to instead provide a const HostAllocator& getHostAllocator() const method? It would simplify some of the FlatMap initializers too.
|
I have questions / comments for Monday. |
| template <IndexType SFINAE_DIM = DIM, | ||
| MemorySpace SFINAE_SPACE = SPACE, | ||
| typename std::enable_if<SFINAE_DIM == 1>::type* = nullptr> | ||
| Array(IndexType num_elements, IndexType capacity, int allocator_id, HostAllocator host_allocator); |
There was a problem hiding this comment.
Big picture: do we expect users to be frequently setting different "host allocators" between different Array instances? I'm a little worried about the expansion in the array API surface; wondering if a global interface for a "default" value for all Arrays, plus a setter interface for specific instances, might tame things a little.
There was a problem hiding this comment.
I share your concern. My PR attempt before this tried to do that, but it seemed brittle and subtle w.r.t. maintaining global host allocator state (default or user-selected) and what to do when the host allocator is implicit.
What I think we want is a simple way to define a host allocator default and allow users to change it if/where they want. Unfortunately, trying to make this robust requires a lot of subtle checking for what had been done in the past (like host allocation, then changing the host allocator default).
I'm open to other ideas. This is the sort of input/feedback I'm looking for by posting this PR.
There was a problem hiding this comment.
As it stands now in develop branch, the host allocator is determined based on how an Axom build is configured. If Umpire is enabled, Umpire Host allocator is used/assumed. If Umpire is NOT enabled, Axom MALLOC is used/assumed.
In this PR, Axom MALLOC is the default in either case, which is preferable IMO.
3 participants
Summary
This PR touches a lot of files, but the changes are all closely related and limited in scope. I tried to stack this as a sequence of PRs but struggled to reconcile subsets of cherry picked commits with the develop branch due to all the changes there.
The bulk of the changes are in Axom core. Other components are modified as the changes in core percolate through to higher level components.
The main goal of this PR is to remove implicit dependence on Axom's global default host allocator for runtime memory allocations. The most substantial changes in this PR are:
getDefaultHostAllocatorID()or resolveMemorySpace::Hostthrough global state unless it is in a legacy compatibility path.The pre-existing global host allocator API still exists, especially in the core container classes (Array, etc.) because they are user-facing. Thus, most API changes are done via method overloads so that the pre-existing APIs still function as before. In particular, the
getDefaultHostAllocatorID()andsetDefaultHostAllocator(...)methods are retained as user-facing compatibility APIs, but they are no longer used in Axom internals.Rather than thread raw
intvalues (allocator ids) everywhere, aaxom::HostAllocatorwrapper was introduced to:allocator_idparameters that refer to primary storage rather than host scratch memory.Many APIs already accepted a primary
allocator_idfor owned storage. However, that is not enough for code paths that may also allocate temporary host ("staging") buffers. This PR takes the following approach:Tests have been adjusted/added and documentation has been adjusted.
I will update the release notes if/when this PR is approved for merging.
If this PR is approved and merged. The global allocator dependent code and APIs would be deprecated and eventually removed in a future PR.