Fix bug when using partial space list

LLNL · Jul 14, 2021 · df3e8e0 · df3e8e0
1 parent 84c08e9
commit df3e8e0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/chai/ManagedArray.inl b/src/chai/ManagedArray.inl
@@ -39,6 +39,10 @@ CHAI_HOST_DEVICE ManagedArray<T>::ManagedArray(
 #if !defined(CHAI_DEVICE_COMPILE)
   m_pointer_record = new PointerRecord();
   int i = 0;
+  for (int s = CPU; s < NUM_EXECUTION_SPACES; ++s) {
+    m_pointer_record->m_allocators[s] = m_resource_manager->getAllocatorId(ExecutionSpace(s));
+  }
+
   for (const auto& space : spaces) {
     m_pointer_record->m_allocators[space] = allocators.begin()[i++].getId();
   }

diff --git a/tests/integration/managed_array_tests.cpp b/tests/integration/managed_array_tests.cpp
@@ -29,6 +29,8 @@
 
 #include "chai/ManagedArray.hpp"
 
+#include "umpire/ResourceManager.hpp"
+
 
 struct my_point {
   double x;
@@ -1973,6 +1975,23 @@ GPU_TEST(ManagedArray, NoAllocationGPU)
   array.free();
 }
 
+GPU_TEST(ManagedArray, NoAllocationGPUList)
+{
+  auto& rm = umpire::ResourceManager::getInstance();
+  chai::ManagedArray<double> array(10,
+      std::initializer_list<chai::ExecutionSpace>{chai::CPU},
+      std::initializer_list<umpire::Allocator>{rm.getAllocator("HOST")}
+  );
+
+  forall(gpu(), 0, 10, [=] __device__ (int i) {
+    array[i] = i;
+  });
+
+  forall(sequential(), 0, 10, [=](int i) { ASSERT_EQ(array[i], i); });
+
+  array.free();
+}
+
 GPU_TEST(ManagedArray, NoAllocationNullGPU)
 {
   chai::ManagedArray<double> array;