From 5137c5a005f1ab1665d98226a82b375e6c9544fb Mon Sep 17 00:00:00 2001
From: Daniel Arndt <arndtd@ornl.gov>
Date: Wed, 11 Sep 2024 15:02:21 -0400
Subject: [PATCH] Update memory space concept

---
 docs/source/API/core/KokkosConcepts.rst |  7 +++++-
 docs/source/API/core/memory_spaces.rst  | 33 ++++++++++++++++++-------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/docs/source/API/core/KokkosConcepts.rst b/docs/source/API/core/KokkosConcepts.rst
index 7351b1018..92029b196 100644
--- a/docs/source/API/core/KokkosConcepts.rst
+++ b/docs/source/API/core/KokkosConcepts.rst
@@ -258,12 +258,17 @@ Looking at the common functionality in the current implementations of ``CudaSpac
             std::is_same_v<Mem, typename Mem::memory_space>;
             Kokkos::is_execution_space<typename Mem::execution_space>::value;
             typename Mem::device_type;
+            typename Mem::size_type;
         }
         // Required methods:
-        requires(Mem m, size_t size, void* ptr) {
+        requires(Mem m, size_t size, void* ptr, const char* label, typename Mem::execution_space exec) {
             { m.name() } -> const char*;
+            { m.allocate(exec, size) } -> void*;
+            { m.allocate(exec, label, size) } -> void*;
             { m.allocate(size) } -> void*;
+            { m.allocate(label, size) } -> void*;
             { m.deallocate(ptr, size) };
+            { m.deallocate(label, ptr, size) };
         };
 
 Implementation Requirements
diff --git a/docs/source/API/core/memory_spaces.rst b/docs/source/API/core/memory_spaces.rst
index 22488da0f..6639e59ca 100644
--- a/docs/source/API/core/memory_spaces.rst
+++ b/docs/source/API/core/memory_spaces.rst
@@ -116,12 +116,28 @@ Synopsis
         typedef MemorySpaceConcept memory_space;
         typedef ... execution_space;
         typedef Device<execution_space, memory_space> device_type;
+        typedef ... size_type;
 
         MemorySpaceConcept();
-        MemorySpaceConcept(const MemorySpaceConcept& src);
-        const char* name() const;
-        void * allocate(ptrdiff_t size) const;
-        void deallocate(void* ptr, ptrdiff_t size) const;
+
+        template <typename AccessibleExecutionSpace>
+        void* allocate(const AccessibleExecutionSpace& exec_space,
+                             const std::size_t arg_alloc_size) const;
+        template <typename AccessibleExecutionSpace>
+        void* allocate(const AccessibleExecutionSpace& exec_space, const char* arg_label,
+                 const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const;
+        void* allocate(const std::size_t arg_alloc_size) const;
+        void* allocate(const char* arg_label, const size_t arg_alloc_size,
+                 const size_t arg_logical_size = 0) const;
+
+        void deallocate(void* const arg_alloc_ptr,
+                  const std::size_t arg_alloc_size) const;
+        void deallocate(const char* arg_label, void* const arg_alloc_ptr,
+                  const size_t arg_alloc_size,
+                  const size_t arg_logical_size = 0) const;
+
+        static constexpr const char* name();
     };
 
     template<class MS>
@@ -132,7 +148,7 @@ Synopsis
     template<>
     struct is_memory_space<MemorySpaceConcept> {
     enum { value = true };
-    };  
+    };
 
 Typedefs
 ~~~~~~~~
@@ -157,17 +173,16 @@ Constructors
 ~~~~~~~~~~~~
 
 * ``MemorySpaceConcept()``: Default constructor.
-* ``MemorySpaceConcept(const MemorySpaceConcept& src)``: Copy constructor.
 
 Functions
 ~~~~~~~~~
 
 * ``const char* name() const;``: Returns the label of the memory space instance.
-* ``void * allocate(ptrdiff_t size) const;``: Allocates a buffer of at least ``size`` bytes using the memory resource that ``MemorySpaceConcept`` represents.
-* ``void deallocate(void* ptr, ptrdiff_t size) const;``: Frees the buffer starting at ``ptr`` (of type ``void*``) previously allocated with exactly ``allocate(size)``.
+* ``void * allocate(ptrdiff_t size) const;``: Allocates a buffer of at least ``size`` bytes using the memory resource that ``MemorySpaceConcept`` represents. An optional execution space argument allows enqueuing the allocation operation on a particular execution space instance that can access the memory space. An optional label argument allows customizing the event reported to Kokkos Tools.
+* ``void deallocate(void* ptr, ptrdiff_t size) const;``: Frees the buffer starting at ``ptr`` (of type ``void*``) previously allocated with exactly ``allocate(size)``. An optional label argument allows customizing the event reported to Kokkos Tools.
 
 Non Member Facilities
 ~~~~~~~~~~~~~~~~~~~~~
 
 * ``template<class MS> struct is_memory_space;``: typetrait to check whether a class is a memory space.
-* ``template<class S1, class S2> struct SpaceAccessibility;``: typetraits to check whether two spaces are compatible (assignable, deep_copy-able, accessible). 
+* ``template<class S1, class S2> struct SpaceAccessibility;``: typetraits to check whether two spaces are compatible (assignable, deep_copy-able, accessible).