Merge pull request #1894 from llnl/bugfix/whitlock/axom_copy_umpire

Fix axom::copy so it can be called from OpenMP loop

Author: BradWhitlock

RELEASE-NOTES.md

@@ -65,6 +65,7 @@ The Axom project release numbers follow [Semantic Versioning](http://semver.org/
 - Core: A moved-from Array is valid, e.g. it can be pushed to
 - Core: Improved `axom::FlatMap` insertion performance by fusing duplicate-key lookup with empty-slot probing.
 - Core: Updated DeviceHash to use 64-bit hash results and improved coverage for integer and floating-point hashing.
+- Core: Avoids a first-use race in `axom::copy()` when multiple OpenMP threads concurrently trigger Umpire fallback host-copy initialization.
 - Python: Improves lifetime handling for python wrapped sidre entities, including support for external views into numpy arrays.
 
 ## [Version 0.14.0] - Release date 2026-03-31

src/axom/core/memory_management.hpp

@@ -27,11 +27,49 @@
 #include <cstddef>
 #include <cstdlib>
 #include <iostream>
+#include <mutex>
 #include <string>
 #include <type_traits>
 
 namespace axom
 {
+#ifdef AXOM_USE_UMPIRE
+namespace detail
+{
+/*!
+ * \brief Cache for Umpire data used in axom::copy.
+ */
+struct UmpireCopyContext
+{
+  umpire::strategy::AllocationStrategy* hostStrategy {nullptr};
+  umpire::op::MemoryOperationRegistry* operationRegistry {nullptr};
+};
+
+/*!
+ * \brief Gets a reference to an UmpireCopyContext object, initializing it on demand.
+ *        The static UmpireCopyContext is initialized via std::call_once so multiple
+ *        threads can call this function and only initialize the object once.
+ *
+ * \return A reference to the cached UmpireCopyContext.
+ */
+inline const UmpireCopyContext& getUmpireCopyContext() noexcept
+{
+  static std::once_flag once;
+  static UmpireCopyContext context {};
+
+  // Resolve Umpire's fallback HOST path once so the first threaded axom::copy()
+  // cannot race through lazy resource creation.
+  std::call_once(once, []() {
+    auto& rm = umpire::ResourceManager::getInstance();
+    context.hostStrategy = rm.getAllocator("HOST").getAllocationStrategy();
+    context.operationRegistry = &umpire::op::MemoryOperationRegistry::getInstance();
+  });
+
+  return context;
+}
+}  // namespace detail
+#endif
+
 // To co-exist with Umpire allocator ids, use negative values here.
 constexpr int INVALID_ALLOCATOR_ID = -1;  //!< Place holder for no/unknown allocator
 constexpr int MALLOC_ALLOCATOR_ID = -3;   //!< Refers to MemorySpace::Malloc
@@ -462,11 +500,11 @@ inline T* reallocate(T* pointer, std::size_t n, int allocID) noexcept
 inline void copy(void* dst, const void* src, std::size_t numbytes) noexcept
 {
 #ifdef AXOM_USE_UMPIRE
+  const auto& copyContext = detail::getUmpireCopyContext();
   umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
-  umpire::op::MemoryOperationRegistry& op_registry =
-    umpire::op::MemoryOperationRegistry::getInstance();
+  umpire::op::MemoryOperationRegistry& op_registry = *copyContext.operationRegistry;
 
-  auto dstStrategy = rm.getAllocator("HOST").getAllocationStrategy();
+  auto dstStrategy = copyContext.hostStrategy;
   auto srcStrategy = dstStrategy;
 
   using AllocationRecord = umpire::util::AllocationRecord;

src/axom/core/tests/CMakeLists.txt

@@ -153,6 +153,18 @@ if (AXOM_ENABLE_OPENMP AND RAJA_FOUND)
                    COMMAND         core_openmp_tests --gtest_filter=${test_suite}*
                    NUM_OMP_THREADS ${AXOM_TEST_NUM_OMP_THREADS} )
   endforeach()
+
+  if (UMPIRE_FOUND)
+    axom_add_executable(NAME       core_copy_openmp_race_test
+                        SOURCES    core_copy_openmp_race.cpp
+                        OUTPUT_DIR ${TEST_OUTPUT_DIRECTORY}
+                        DEPENDS_ON ${core_test_depends}
+                        FOLDER     axom/core/tests )
+
+    axom_add_test( NAME            core_copy_openmp_race
+                   COMMAND         core_copy_openmp_race_test
+                   NUM_OMP_THREADS ${AXOM_TEST_NUM_OMP_THREADS} )
+  endif()
 endif()
 
 #------------------------------------------------------------------------------

src/axom/core/tests/core_copy_openmp_race.cpp

@@ -0,0 +1,55 @@
+// Copyright (c) Lawrence Livermore National Security, LLC and other
+// Axom Project Contributors. See top-level LICENSE and COPYRIGHT
+// files for dates and other details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+
+#include "gtest/gtest.h"
+
+#include "axom/core/Array.hpp"
+#include "axom/core/execution/for_all.hpp"
+#include "axom/core/memory_management.hpp"
+
+namespace
+{
+template <typename ExecSpace>
+struct CopySlices
+{
+  static void run(int* dst, const int* src, axom::IndexType nslices, axom::IndexType sliceSize)
+  {
+    axom::for_all<ExecSpace>(nslices, [=](axom::IndexType slice) {
+      axom::copy(dst + slice * sliceSize, src + slice * sliceSize, sliceSize * sizeof(int));
+    });
+  }
+};
+}  // namespace
+
+TEST(core_copy_openmp_race, first_use_inside_for_all)
+{
+  constexpr axom::IndexType nslices = 128;
+  constexpr axom::IndexType sliceSize = 64;
+  constexpr axom::IndexType size = nslices * sliceSize;
+
+  axom::Array<int> src(size);
+  axom::Array<int> dst(size);
+
+  for(axom::IndexType i = 0; i < size; ++i)
+  {
+    src[i] = static_cast<int>(i);
+    dst[i] = -1;
+  }
+
+  CopySlices<axom::OMP_EXEC>::run(dst.data(), src.data(), nslices, sliceSize);
+
+  for(axom::IndexType i = 0; i < size; ++i)
+  {
+    EXPECT_EQ(dst[i], src[i]);
+  }
+}
+
+int main(int argc, char** argv)
+{
+  ::testing::InitGoogleTest(&argc, argv);
+  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
+  return RUN_ALL_TESTS();
+}