From 8b1235838e6256dbf1a209069f56640679c50bc6 Mon Sep 17 00:00:00 2001
From: Arturo Vargas <vargas45@llnl.gov>
Date: Sun, 13 Apr 2025 23:42:38 -0700
Subject: [PATCH 1/2] memory arena for launch

---
 include/RAJA/pattern/launch/launch_core.hpp | 37 +++++++++++++++++++--
 1 file changed, 34 insertions(+), 3 deletions(-)

diff --git a/include/RAJA/pattern/launch/launch_core.hpp b/include/RAJA/pattern/launch/launch_core.hpp
index bfce94057c..1ef1f02eda 100644
--- a/include/RAJA/pattern/launch/launch_core.hpp
+++ b/include/RAJA/pattern/launch/launch_core.hpp
@@ -182,9 +182,11 @@ class LaunchContext
   // Bump style allocator used to
   // get memory from the pool
   size_t shared_mem_offset;
-
   void* shared_mem_ptr;
 
+  size_t mem_arena_offset;
+  void* mem_arena_ptr;
+
 #if defined(RAJA_ENABLE_SYCL)
   mutable ::sycl::nd_item<3>* itm;
 #endif
@@ -208,6 +210,26 @@ class LaunchContext
     return static_cast<T*>(mem_ptr);
   }
 
+  template<typename T>
+  RAJA_HOST_DEVICE createMemoryArena(T *ptr, size_t bytes)
+  {
+    void * mem_arena_ptr = static_cast<char*>(&ptr[0]);
+  }
+
+  // TODO handle alignment
+  template<typename T>
+  RAJA_HOST_DEVICE T* getArenaMemory(size_t bytes)
+  {
+
+    // Calculate offset in bytes with a char pointer
+    void* mem_ptr = static_cast<char*>(arena_mem_ptr) + shared_arena_offset;
+
+    arena_mem_offset += bytes * sizeof(T);
+
+    // convert to desired type
+    return static_cast<T*>(mem_ptr);
+  }
+
   /*
   //Odd dependecy with atomics is breaking CI builds
   template<typename T, size_t DIM, typename IDX_T=RAJA::Index_type, ptrdiff_t
@@ -222,12 +244,21 @@ class LaunchContext
   }
   */
 
-  RAJA_HOST_DEVICE void releaseSharedMemory()
+  //default is to release all
+  RAJA_HOST_DEVICE void releaseSharedMemory(size_t offset=-1)
   {
     // On the cpu/gpu we want to restart the count
-    shared_mem_offset = 0;
+    shared_mem_offset = offset < 0 ? 0 : shared_mem_offset - offset;
   }
 
+  //default is to release all
+  RAJA_HOST_DEVICE void releaseArenaMemory(size_t offset=-1)
+  {
+    // On the cpu/gpu we want to restart the count
+    arena_mem_offset = offset < 0 ? 0 : arena_mem_offset - offset;
+  }
+
+
   RAJA_HOST_DEVICE
   void teamSync()
   {

From 5bfd2c84b730b536678a09b21c6595627add3685 Mon Sep 17 00:00:00 2001
From: Arturo Vargas <vargas45@llnl.gov>
Date: Tue, 15 Apr 2025 10:06:49 -0700
Subject: [PATCH 2/2] Update include/RAJA/pattern/launch/launch_core.hpp

Co-authored-by: Yohann <yohann.dudouit@gmail.com>
---
 include/RAJA/pattern/launch/launch_core.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/include/RAJA/pattern/launch/launch_core.hpp b/include/RAJA/pattern/launch/launch_core.hpp
index 1ef1f02eda..ce3530004e 100644
--- a/include/RAJA/pattern/launch/launch_core.hpp
+++ b/include/RAJA/pattern/launch/launch_core.hpp
@@ -222,7 +222,7 @@ class LaunchContext
   {
 
     // Calculate offset in bytes with a char pointer
-    void* mem_ptr = static_cast<char*>(arena_mem_ptr) + shared_arena_offset;
+    void* mem_ptr = static_cast<char*>(mem_arena_ptr) + mem_arena_offset;
 
     arena_mem_offset += bytes * sizeof(T);