Skip to content

Rework host memory allocation mechanics - #1964

Open
rhornung67 wants to merge 24 commits into
developfrom
task/rhornung67/host-memory-api
Open

Rework host memory allocation mechanics#1964
rhornung67 wants to merge 24 commits into
developfrom
task/rhornung67/host-memory-api

Conversation

@rhornung67

@rhornung67 rhornung67 commented Aug 25, 2026

Copy link
Copy Markdown
Member

Summary

  • This PR makes axom::MALLOC_ALLOCATOR_ID the default host allocator when Axom is configured with Umpire enabled and when it is not.
  • The CMake variable AXOM_DEFAULT_HOST_ALLOCATOR was added as a compile-time option for users who need to make a different default host execution-space policy. Its valid values are MALLOC (default) and UMPIRE_HOST (available when Axom is configured with Umpire), which was previous default.
    • We could make this an environment variable to allow users to change the default without recompilation. However, it could make testing more complicated, especially for core library allocator semantics.
    • One GitLab CI job was added to dane to test the UMPIRE_HOST case.
  • Explicit allocator-ID arguments can still be used for runtime per-use control in Axom APIs that supported it previously.
  • Also, the preexisting role of setDefaultAllocator() as the interface for Umpire default allocator state has been preserved.

No substantial changes to Axom public APIs were made and no static state was introduced.

This PR addresses #1816

@rhornung67 rhornung67 changed the title Rework host memory allocator Rework host memory allocation mechanics Aug 25, 2026

@rhornung67 rhornung67 left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@lc-hubcast approved

@rhornung67

rhornung67 commented Aug 26, 2026

Copy link
Copy Markdown
Member Author

@kennyweiss @publixsubfan @BradWhitlock @bmhan12 this is a much smaller, more focused PR that addresses much of what I have been talking about regarding the default host memory allocation scheme in Axom. If you think we should include it in the Axom release, please review. Thank you.

@kennyweiss kennyweiss left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for working though this @rhornung67 (and reworking it, and reworking it).

Overall, this seems like a nice compromise given the constraints and improves our memory management, but might need to be integrated with more of our infrastructure, e.g. axom::Array, sidre::Group, ...

There are some things that we should investigate before merging
(caveat: I haven't built the code, so this is based on staring at the code w/ claude):

  • we might need to also update memory_management's reallocate(MALLOC_ALLOCATOR_ID) to move the malloc handling outside the #else.
    The handling of allocate/deallocate and reallocate was already inconsistent in axom@develop and not introduced by your changes, but the consequences are perhaps more significant now.
  • With AXOM_DEFAULT_HOST_ALLOCATOR=MALLOC, Array<double> a(n) still uses the Host allocator rather than the malloc allocator since Array's default is Dynamic, (defined as Umpire's current default allocator). I'm pretty sure we'd want the default Array<double> to follow AXOM_DEFAULT_HOST_ALLOCATOR
  • With AXOM_DEFAULT_HOST_ALLOCATOR=MALLOC, sidre still defaults to umpire Host (getDefaultAllocatorID() rather than detail::getDefaultHostAllocatorID())

Comment thread RELEASE-NOTES.md Outdated
Comment thread src/cmake/AxomConfig.cmake
Comment thread src/axom/core/memory_management.hpp Outdated
{
#ifdef AXOM_USE_UMPIRE
umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
if(allocId == MALLOC_ALLOCATOR_ID)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't setDefaultAllocator also be influenced by AXOM_DEFAULT_HOST_ALLOCATOR_USES_UMPIRE_HOST ?

It seems that if we're in an Umpire build, we can never get back to the malloc allocator as the default.
E.g., consider the following in a build with Umpire:

axom::setDefaultAllocator(Host);       // starts at Malloc, ends at host
axom::setDefaultAllocator(Malloc);   // still ends at Host

@rhornung67 rhornung67 Aug 27, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This method is only setting the default allocator in the Umpire ResourceManager. As long as all Axom usage honors the Axom default (Malloc or Host) when appropriate this should be OK.

Do we need to handle the case where a user calls setDefaultAllocator() with something that is neither Malloc or Host, but is still valid on the host, such as Pinned when Umpire is enabled? As I understand it, that is the main purpose of this method.

I think only the method documentation needs to be clarified, which I did.

Does that make sense now?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This method is only setting the default allocator in the Umpire ResourceManager. As long as all Axom usage honors the Axom default (Malloc or Host) when appropriate this should be OK.

If that's the case, perhaps it should be renamed setUmpireDefaultAllocator ?
(and be a no-op in non-umpire configs)

This is the current doxygen brief for the setDefaultAllocator

\brief Sets the default memory allocator to use.

Until now, we've been using it as Axom's default allocator
(which points to Umpire when Axom is configured against Umpire).

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I modified the doxygen brief comment to indicate that it sets the default allocator in the Umpire ResourceManager, which it did before my changes IIRC. It is a no-op in non-umpire configs.

// Device memory: fill on host, then copy to device
const auto num_bytes = n * sizeof(T);
T* src = allocate<T>(num_bytes, rm.getDefaultAllocator().getId());
T* src = allocate<T>(n, axom::detail::getDefaultHostAllocatorID());

@BradWhitlock BradWhitlock Aug 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch.

It's not here but in allocate() if Umpire is enabled then we check whether the passed allocator Id is a valid Umpire allocator before the malloc check. Since we're defaulting the host allocator id to MALLOC, even for Umpire-enabled builds (unless they select UMPIRE_HOST), should we swap the order of the umpire/malloc allocation checks so we can skip checking whether the allocator Id is valid for Umpire? Or, will allocate() more likely be called with allocators that are associated with Umpire?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good thought. I swapped the allocation check ordering.

@BradWhitlock BradWhitlock left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work @rhornung67

@rhornung67

rhornung67 commented Sep 2, 2026

Copy link
Copy Markdown
Member Author

@kennyweiss I think I addressed all the concerns in your review comments. I have done the following:

  • Made the cases and usage of all Axom core host allocation internally consistent and explicit w.r.t. to the configuration value of AXOM_DEFAULT_HOST_ALLOCATOR.
  • Did the same for Sidre Group, View, Buffer default allocations.

However, I maintained the distinction between MemorySpace::Dynamic, which follows axom::getDefaultAllocatorID(), and the host defaults which follow axom::detail::getDefaultHostAllocatorID(). Specifically,

  • axom::Array<T> uses MemorySpace::Dynamic by default (I did not change this), so in Umpire builds it uses Umpire’s current default allocator, not AXOM_DEFAULT_HOST_ALLOCATOR.
  • axom::Array<T, DIM, MemorySpace::Malloc> always uses MALLOC_ALLOCATOR_ID.
  • axom::Array<T, DIM, MemorySpace::Host> uses Umpire HOST when configured with Umpire.
  • Arrays constructed with axom::policyToDefaultAllocatorID(seq/omp) use the configured host default:
    - AXOM_DEFAULT_HOST_ALLOCATOR=MALLOC -> MALLOC_ALLOCATOR_ID
    - AXOM_DEFAULT_HOST_ALLOCATOR=UMPIRE_HOST -> Umpire HOST

Therefore, AXOM_DEFAULT_HOST_ALLOCATOR controls host execution-space defaults, not the generic dynamic default allocator. The staging paths in ArrayBase also use getDefaultHostAllocatorID(), so they line up with the host-default policy.

The goal was to preserve existing semantics and not change any user facing interfaces. The only potential confusion here is that Array<T> with no memory-space or allocator template argument still means MemorySpace::Dynamic, not “configured default host allocator.”

Please take a look when you have time and let me know if this satisfies your concerns. Thank you.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants