Use a more object-oriented approach to the allocator API

devreal · Mar 15, 2023 · 2a195f2 · 2a195f2
1 parent 8204248
commit 2a195f2
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 103 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -87,7 +87,7 @@ include(FindOrFetchBoost)
 find_package(CXXStdCoroutine MODULE REQUIRED COMPONENTS Final Experimental)
 
 ##########################
-#### CUDA
+#### CUDA -- CUDA language must be added before PaRSEC detects CUDA
 ##########################
 option(TTG_DISABLE_CUDA "True iff support for CUDA files is disabled, even if it is possible" OFF)
 if( NOT TTG_DISABLE_CUDA )

diff --git a/ttg/ttg/device.h b/ttg/ttg/device.h
@@ -2,20 +2,12 @@
 #define TTG_DEVICE_H
 
 #include "ttg/fwd.h"
+#include "ttg/execution.h"
 
 namespace ttg {
     namespace device {
-        class DeviceAllocator {
-            public:
-            virtual DeviceAllocator(int did) = 0;
-            virtual ~DeviceAllocator() = 0;
-            virtual void *allocate(std::size_t size) = 0;
-            virtual void  free(void *ptr) = 0;
-            virtual ttg::ExecutionSpace executionSpace() = 0;
-        };
-
+        using DeviceAllocator = TTG_IMPL_NS::device::DeviceAllocator;
         std::size_t nb_devices() { return TTG_IMPL_NS::device::nb_devices(); }
-        const DeviceAllocator &allocator(int did) { return TTG_IMPL_NS::device::get_device_allocator(did); }
     }
 }
 

diff --git a/ttg/ttg/madness/fwd.h b/ttg/ttg/madness/fwd.h
@@ -46,17 +46,8 @@ namespace ttg_madness {
   inline void ttg_broadcast(ttg::World world, T &data, int source_rank);
 
   namespace device {
-    class DeviceAllocator : ttg::device::DeviceAllocator {
-      public:
-        DeviceAllocator(int did);
-        ~DeviceAllocator();
-        void *allocate(std::size_t size);
-        void  free(void *ptr);
-        ttg::ExecutionSpace executionSpace();
-    };
-
+    class DeviceAllocator;
     std::size_t nb_devices();
-    const DeviceAllocator &allocator(int did);
   }
 
 }  // namespace ttg_madness

diff --git a/ttg/ttg/madness/ttg.h b/ttg/ttg/madness/ttg.h
@@ -1277,27 +1277,26 @@ namespace ttg_madness {
 
 #include "ttg/make_tt.h"
 
-  namespace device {
-    std::size_t get_nb() { return 1; }
-    void *memory_allocate(int did, std::size_t size) { 
-      if(did != 0) {
-        throw std::out_of_range("TTG MADNESS Backend does not support other devices than CPU at this time");
-      }
-      return ::malloc(size);
-    }
-    void  memory_free(int did, void *ptr) {
-      if(did != 0) {
-        throw std::out_of_range("TTG MADNESS Backend does not support other devices than CPU at this time");
-      }
-      ::free(ptr);
-    }
-    ttg::ExecutionSpace execution_space(int did) {
-      if(did != 0) {
-        throw std::out_of_range("TTG MADNESS Backend does not support other devices than CPU at this time");
-      }
-      return ttg::ExecutionSpace::Host;
-    }
+   namespace device {
+    class DeviceAllocator {
+      public:
+        DeviceAllocator(int did) { 
+          if(did != 0) {
+            throw std::out_of_range("TTG MADNESS Backend: current implementation only supports CPU devices")
+          }
+        }
+        void *allocate(std::size_t size) {
+          return ::malloc(size);
+        };
+        void  free(void *ptr) {
+          ::free(ptr);
+        }
+        ::ttg::ExecutionSpace executionSpace() {
+          return ::ttg::ExecutionSpace::Host;
+        }
+    };
 
+    std::size_t nb_devices() { return 1; }
   }
 
 }  // namespace ttg_madness

diff --git a/ttg/ttg/parsec/fwd.h b/ttg/ttg/parsec/fwd.h
@@ -70,17 +70,8 @@ namespace ttg_parsec {
   static void ttg_broadcast(ttg::World world, T &data, int source_rank);
 
   namespace device {
-    class DeviceAllocator : ttg::device::DeviceAllocator {
-      public:
-        virtual DeviceAllocator(int did);
-        virtual ~DeviceAllocator();
-        virtual void *allocate(std::size_t size);
-        virtual void  free(void *ptr);
-        virtual ttg::ExecutionSpace executionSpace();
-    };
-
+    class DeviceAllocator;
     std::size_t nb_devices();
-    const DeviceAllocator &allocator(int did);
   }
 
 #if 0

diff --git a/ttg/ttg/parsec/ttg.h b/ttg/ttg/parsec/ttg.h
@@ -3484,78 +3484,68 @@ namespace ttg_parsec {
 #include "ttg/make_tt.h"
 
   namespace device {
-    namespace detail {
-      static std::size_t nb_ttg_devices = 0;
-      static std::vector<const DeviceAllocator &>device_allocators;
-    }
 
-    class DeviceAllocator : ttg::device::DeviceAllocator {
+    class DeviceAllocator {
       private:
           int ttg_did, parsec_did;
-          struct zone_malloc_s *zone;
-          ttg::ExecutionSpace exec_space;
-
+          struct ::zone_malloc_s *zone;
+          ::ttg::ExecutionSpace exec_space;
       public:
-      virtual DeviceAllocator(int did) : ttg_did(-1), parsec_did(-1), zone(nullptr), exec_space(ttg::ExecutionSpace::Invalid) {
-        parsec_did = -1;
-        for(int i = 0; i < parsec_nb_devices; i++) {
-          parsec_device_module_t *m = parsec_mca_device_get(i);
-          if(m->type == PARSEC_DEV_CPU || m->type == PARSEC_DEV_CUDA) { 
-            if(did == 0) {
-              parsec_did = i;
-              ttg_did = did;
-              if(m->ype == PARSEC_DEV_CUDA) {
-                parsec_device_gpu_module_t *gm = reinterpret_cast<parsec_device_gpu_module_t*>(m);
-                zone = gm->memory;
-                exec_space = ttg::ExecutionSpace::CUDA;
-              } else {
-                exec_space = ttg::ExecutionSpace::Host;
-              }
-              return;
+        DeviceAllocator(int did);
+        void *allocate(std::size_t size);
+        void  free(void *ptr);
+        ::ttg::ExecutionSpace executionSpace();
+    };
+
+    DeviceAllocator::DeviceAllocator(int did) : ttg_did(-1), parsec_did(-1), zone(nullptr), exec_space(::ttg::ExecutionSpace::Invalid) {
+      for(int i = 0; i < parsec_nb_devices; i++) {
+        parsec_device_module_t *m = parsec_mca_device_get(i);
+        if(m->type == PARSEC_DEV_CPU || m->type == PARSEC_DEV_CUDA) { 
+          if(did == 0) {
+            parsec_did = i;
+            ttg_did = did;
+            if(m->type == PARSEC_DEV_CUDA) {
+              parsec_device_gpu_module_t *gm = reinterpret_cast<parsec_device_gpu_module_t*>(m);
+              zone = gm->memory;
+              exec_space = ::ttg::ExecutionSpace::CUDA;
+            } else {
+              exec_space = ::ttg::ExecutionSpace::Host;
             }
-            did--;
+            return;
           }
+          did--;
         }
-        throw std::out_of_range("Device identifier is out of range");
       }
+      throw std::out_of_range("Device identifier is out of range");
+    }
 
-      virtual ~DeviceAllocator() = default;
+    void *DeviceAllocator::allocate(std::size_t size) {
+      if(nullptr == zone) return malloc(size);
+      return zone_malloc(zone, size);
+    }
 
-      virtual void *allocate(std::size_t size) {
-        if(nullptr == zone) return malloc(size);
-        return zone_malloc(zone, size);
-      }
-      virtual void  free(void *ptr) {
-        if(nullptr == zone) {
-          free(ptr);
-          return;
-        }
-        zone_free(zone, ptr);
-      }
-      virtual ttg::ExecutionSpace executionSpace() {
-        return exec_space;
+    void DeviceAllocator::free(void *ptr) {
+      if(nullptr == zone) {
+        free(ptr);
+        return;
       }
+      zone_free(zone, ptr);
+    }
+
+    ::ttg::ExecutionSpace DeviceAllocator::executionSpace() {
+      return exec_space;
     }
 
     std::size_t nb_devices() { 
-      if( detail::nb_ttg_devices > 0 ) return detail::nb_ttg_devices;
+      std::size_t nb = 0;
       for(int i = 0; i < parsec_nb_devices; i++) {
         parsec_device_module_t *m = parsec_mca_device_get(i);
         if(m->type == PARSEC_DEV_CPU || m->type == PARSEC_DEV_CUDA) {
-          device_allocators.append( std::move(DeviceAllocator(detail::nb_ttg_devices)) );
-          detail::nb_ttg_devices++;
+          nb++;
         }
       }
-      return detail::nb_ttg_devices;
+      return nb;
     }
-
-    const DeviceAllocator &allocator(int did) { 
-      if( did >= nb_devices() ) {
-        throw std::out_of_range("Device identifier is out of range");
-      }
-      return detail::device_allocators.at(did);
-    }
-
   } // namespace ttg_parsec::device
 
 }  // namespace ttg_parsec