Proposal for reworking host memory by rhornung67 · Pull Request #1928 · llnl/axom

rhornung67 · 2026-07-22T22:12:04Z

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:

Routines that allocate host memory now take an explicit host allocator argument.
Internal helper code does not call getDefaultHostAllocatorID() or resolve MemorySpace::Host through global state unless it is in a legacy compatibility path.
Startup behavior no longer depends on whether some earlier static constructor happened to allocate host memory.

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() and setDefaultHostAllocator(...) methods are retained as user-facing compatibility APIs, but they are no longer used in Axom internals.

Rather than thread raw int values (allocator ids) everywhere, a axom::HostAllocator wrapper was introduced to:

Better communicate intent throughout the code.
Prevent accidental passage of a device allocator where host staging memory is needed.
Avoid overloading ambiguity with existing allocator_id parameters that refer to primary storage rather than host scratch memory.

Many APIs already accepted a primary allocator_id for owned storage. However, that is not enough for code paths that may also allocate temporary host ("staging") buffers. This PR takes the following approach:

If an object allocates only one owned buffer, keep the existing explicit allocator parameter.
If an object may allocate both primary storage and host staging/scratch, add a second explicit host allocator parameter.

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.

… usage

and resolve merge conflicts

BradWhitlock · 2026-07-31T17:52:09Z

+Changing the default host and global allocators (legacy)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+THe contents of this section refer to Axom *legacy* convenience routines,


D'oh! Thanks. Fixed.

BradWhitlock · 2026-07-31T18:21:51Z

+  axom::setDefaultAllocator(axom::MemorySpace::Unified);
+
+  // set Axom global allocator to an explicitly chosen Umpire allocator
+  int allocId =


axom::getAllocatorIDFromMemorySpace(axom::MemorySpace::Pinned) seems like a better example.

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.

@BradWhitlock I added your suggestion as an alternative example.

BradWhitlock · 2026-07-31T18:25:28Z

-    :
+#else
+
+  const axom::HostAllocator hostAllocator {axom::execution_space<axom::SEQ_EXEC>::allocatorID()};


This leaves me wanting syntactic sugar: const auto hostAllocator = axom::execution_space<axom::SEQ_EXEC>::hostAllocator();

I agree that we can smooth this further.

BradWhitlock · 2026-07-31T18:34:08Z

  static int allocatorID() noexcept
  {
-    return axom::getUmpireResourceAllocatorID(umpire::resource::Device);
+    return axom::getAllocatorIDFromMemorySpace(memory_space);


Good change.

BradWhitlock · 2026-07-31T18:38:14Z

 namespace
 {
+#if defined(AXOM_USE_GPU) && defined(AXOM_USE_UMPIRE)
+bool runtimeMemorySpaceAvailable(axom::MemorySpace space)


This looks like it belongs in the new MemoryTesting.hpp file.

Probably. I'll look into it.

BradWhitlock · 2026-07-31T18:44:00Z

 void check_device(axom::Array<T, DIM, SPACE>& v)
 {
  const axom::IndexType size = v.size();
+  const int explicit_host_alloc =


Why this and not axom::getAllocatorIDFromMemorySpace(axom::MemorySpace::Host)?

BradWhitlock · 2026-07-31T18:44:40Z


  // 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);


In this statement and some others in this test it switched from assigment operator to copy constructor. Is that intentional?

BradWhitlock · 2026-07-31T18:49:33Z

-      ::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))


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?

BradWhitlock · 2026-07-31T19:02:09Z

 {
+namespace
+{
+bool runtimeMemorySpaceAvailable(axom::MemorySpace space)


I think this is repeated in another test.

BradWhitlock · 2026-07-31T19:05:48Z

 // in check_alloc_realloc_free when reallocating to 3 * ARRAY_SIZE.
 constexpr int ARRAY_SIZE = 5345;

+struct ScopedDefaultAllocatorState


This RAII pattern is repeated from a different test. Would it be useful to have in axom::core?

BradWhitlock · 2026-07-31T19:08:46Z

    buffer_size *= 3;
-    buffer = axom::reallocate<int>(buffer, buffer_size);
-#ifdef AXOM_USE_UMPIRE
+#if defined(AXOM_USE_UMPIRE)


Is the AXOM_USE_UMPIRE check still needed since you switched to not using the umpire::ResourceManager?

BradWhitlock · 2026-07-31T19:11:49Z

+};
+
+#if defined(AXOM_USE_UMPIRE)
+bool coreMemoryManagementRuntimeMemorySpaceAvailable(axom::MemorySpace space)


This may be repeated, this time with a different name.

BradWhitlock · 2026-07-31T19:25:03Z

  }
 };

+inline HostAllocator hostAllocatorForPrimaryAllocator(int allocator_id)


It seems like these functions might be useful in other places too. Do they belong in memory_management.hpp?

BradWhitlock · 2026-07-31T19:56:12Z

+  , m_host_allocator(other.m_host_allocator)
+  , m_arrayOps(m_allocator_id, m_executeOnGPU, m_host_allocator)
 {
+#if defined(AXOM_DEVICE_CODE)


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).

BradWhitlock · 2026-07-31T20:16:31Z


  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(); }


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.

BradWhitlock · 2026-07-31T23:13:13Z

I have questions / comments for Monday.

rhornung67 added 30 commits July 8, 2026 13:45

Rework plumbing of host allocator in core so it is more explicit

7d5d307

Merge branch 'develop' into task/rhornung67/rework-host_memory

e4c1ea1

Use HostAllocator in examples/tests, mark default-host APIs as legacy

2662d51

More conversion to explicit host allocator

af2b6be

Merge branch 'develop' into task/rhornung67/rework-host_memory

fdf74aa

Thread host allocator aware exection paths through quest.

a598f75

Merge branch 'develop' into task/rhornung67/rework-host_memory

02fe33b

Thread most explicit host allocator in quest.

79a4de5

Merge branch 'develop' into task/rhornung67/rework-host_memory

36f63ea

Merge branch 'develop' into task/rhornung67/rework-host_memory

5798798

Clean up pass to make things consistent with preferred host allocator…

873728e

… usage

run clang-format

85c0fd9

Merge branch 'develop' into task/rhornung67/rework-host_memory

86ef9f3

Remove remaining internal reliance on implicit host allocator

0eb62b2

Run clang-format

b68bde6

Updated host allocator fallback logic and Umpire-forcused test

7e5d6e4

Make quest host-allocator usage more rigorous and robust

ae1344a

Thread host allocator mechanics through quest clipper operations.

aa081bd

More polishin host allocation mechanics in quest clipping stuff

cdbd2bf

remove more implicit host default allocator use

5f9e117

Add explcit host allocator overleads for BVH

bb0ae9e

Work explicit host allocator through intersection shaper

be0b34b

Thread most explicit host allocator mechanics through quest

5e4d953

Add most explicit host allocator state and usage

3436a98

Add more explicit host allocator overloads

7a1cf9f

Add explicit host allocator overload for "reduce tree"

4ffa685

Thread more host allocator through quest point finding stuff

ab3726a

Merge branch 'develop' into task/rhornung67/rework-host_memory

de39b8e

Add explicit host allocator to mint helpers

6cd80d0

Added default ctor with host allocator

b775f9e

rhornung67 requested review from Arlie-Capps, BradWhitlock, bmhan12, kennyweiss, nselliott and publixsubfan July 27, 2026 15:55

rhornung67 added 4 commits July 28, 2026 08:44

Merge branch 'develop' into task/rhornung67/rework-host_memory

a122c6e

Merge branch 'develop' into task/rhornung67/rework-host_memory

e82fcbe

Merge branch 'develop' into task/rhornung67/rework-host_memory

e0d6da8

and resolve merge conflicts

Merge branch 'develop' into task/rhornung67/rework-host_memory

d017041

and resolve merge conflicts

adayton1 changed the title ~~Proposal for reworking host memory memory~~ Proposal for reworking host memory Jul 31, 2026

BradWhitlock reviewed Jul 31, 2026

View reviewed changes

rhornung67 added 2 commits July 31, 2026 11:01

Fix typo

b8e9364

clang-format

4833931

BradWhitlock reviewed Jul 31, 2026

View reviewed changes

Based on PR comment, add alternative example.

67c47a4

BradWhitlock reviewed Jul 31, 2026

View reviewed changes

Uh oh!

Conversation

rhornung67 commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BradWhitlock Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BradWhitlock Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BradWhitlock Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BradWhitlock Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

BradWhitlock commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

rhornung67 commented Jul 22, 2026 •

edited

Loading

BradWhitlock Jul 31, 2026 •

edited

Loading

BradWhitlock Jul 31, 2026 •

edited

Loading

BradWhitlock Jul 31, 2026 •

edited

Loading

BradWhitlock Jul 31, 2026 •

edited

Loading