From 244bee4754f50baa3f888cbb99dedc6ad2a68142 Mon Sep 17 00:00:00 2001
From: Kristi Belcher <belcher6@llnl.gov>
Date: Fri, 11 Oct 2024 14:45:23 -0700
Subject: [PATCH] updates to RAP

---
 examples/rap_example.cpp                        | 10 +++++-----
 src/umpire/strategy/ResourceAwarePool.cpp       |  7 +------
 src/umpire/strategy/ResourceAwarePool.hpp       |  3 +++
 tests/integration/resource_aware_pool_tests.cpp |  2 +-
 4 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/examples/rap_example.cpp b/examples/rap_example.cpp
index fa42ecd4d..be7446670 100644
--- a/examples/rap_example.cpp
+++ b/examples/rap_example.cpp
@@ -1,5 +1,4 @@
 #include <stdio.h>
-
 #include <iostream>
 
 #include "camp/camp.hpp"
@@ -68,8 +67,9 @@ int main(int, char**)
   auto pool = rm.makeAllocator<umpire::strategy::ResourceAwarePool>("rap-pool", rm.getAllocator("HOST"));
 #endif
 
-  resource_type d1, d2, d3;
-  Resource r1{d1}, r2{d2}, r3{d3};
+  // Create camp resources for RAP
+  resource_type d1, d2;
+  Resource r1{d1}, r2{d2};
 
   // allocate memory in the pool with r1
   double* a = static_cast<double*>(pool.allocate(r1, NUM_THREADS * sizeof(double)));
@@ -89,14 +89,14 @@ int main(int, char**)
 
   // deallocate memory with r1 and reallocate using a different stream r2
   pool.deallocate(r1, a);
-
   a = static_cast<double*>(pool.allocate(r2, NUM_THREADS * sizeof(double)));
   double* ptr2 = a;
 
+  // Make sure resource was correctly tracked
   UMPIRE_ASSERT(getResource(pool, a) == r2);
 
   // Use Camp resource to synchronize devices
-  r1.get_event().wait();
+  r2.get_event().wait();
 
 #if defined(UMPIRE_ENABLE_CUDA) || defined(UMPIRE_ENABLE_HIP)
   UMPIRE_ASSERT(ptr1 != ptr2);
diff --git a/src/umpire/strategy/ResourceAwarePool.cpp b/src/umpire/strategy/ResourceAwarePool.cpp
index 715ee9994..e5ae92d0b 100644
--- a/src/umpire/strategy/ResourceAwarePool.cpp
+++ b/src/umpire/strategy/ResourceAwarePool.cpp
@@ -65,8 +65,6 @@ void* ResourceAwarePool::allocate_resource(camp::resources::Resource r, std::siz
 
   Chunk* chunk{nullptr};
 
-  // auto pending_chunks_exist = m_pending_map.find(r);
-
   if (!m_pending_map.empty()) {
     for (auto pending_chunk : m_pending_map) {
       if (pending_chunk->free == false && pending_chunk->m_event.check()) // no longer pending
@@ -270,12 +268,9 @@ void ResourceAwarePool::deallocate_resource(camp::resources::Resource r, void* p
 
   auto my_r = getResource(ptr);
   if (my_r != r) {
-    UMPIRE_LOG(
-        Warning,
+    UMPIRE_ERROR(runtime_error,
         fmt::format("Called deallocate with different resource than what is returned by getResource. Called with {},",
                     "but getResource returned: {}", camp::resources::to_string(r), camp::resources::to_string(my_r)));
-    UMPIRE_LOG(Debug, fmt::format("getResource doesn't match resource passed to deallocate. Resource used: {} .",
-                                  camp::resources::to_string(r)));
   }
 
   // Chunk is now pending, add to list
diff --git a/src/umpire/strategy/ResourceAwarePool.hpp b/src/umpire/strategy/ResourceAwarePool.hpp
index f8739b196..957dbc851 100644
--- a/src/umpire/strategy/ResourceAwarePool.hpp
+++ b/src/umpire/strategy/ResourceAwarePool.hpp
@@ -96,6 +96,9 @@ class ResourceAwarePool : public AllocationStrategy, private mixins::AlignedAllo
   std::size_t getPendingSize() const noexcept;
 
   Platform getPlatform() noexcept override;
+  /*!
+   * \brief get the (generic) camp resource associated with a ptr
+   */
   Resource getResource(void* ptr) const;
 
   MemoryResourceTraits getTraits() const noexcept override;
diff --git a/tests/integration/resource_aware_pool_tests.cpp b/tests/integration/resource_aware_pool_tests.cpp
index 44ed06645..7dea1db21 100644
--- a/tests/integration/resource_aware_pool_tests.cpp
+++ b/tests/integration/resource_aware_pool_tests.cpp
@@ -141,7 +141,7 @@ TEST_P(ResourceAwarePoolTest, Check_States)
   hipLaunchKernelGGL(do_sleep, 1, 32, 0, d1.get_stream(), ptr);
 #endif
 
-  m_pool.deallocate(ptr);
+  m_pool.deallocate(r1, ptr);
   EXPECT_EQ(getPendingSize(m_pool), 1);
 
   double* ptr2 = static_cast<double*>(m_pool.allocate(r2, 2048));