Remove all RMM singletons except for the default resource, fix rmmInitialize/rmmFinalize

Author: harrism

include/rmm/mr/default_memory_resource.hpp

@@ -54,69 +54,5 @@ device_memory_resource* get_default_resource();
 device_memory_resource* set_default_resource(
     device_memory_resource* new_resource);
 
-namespace detail{
-
-
-/**---------------------------------------------------------------------------*
- * @brief gets the default memory_resource when none is set
- *
- * A static function which will return a singleton cuda_memory_resource
- *
- *
- * @return device_memory_resource* a pointer to the static
- * cuda_memory_resource
- *---------------------------------------------------------------------------**/
-device_memory_resource* initial_resource();
-
-/**---------------------------------------------------------------------------*
- * @brief gets a cuda_memory_resource
- *
- * A static function which will return a singleton cuda_memory_resource
- *
- *
- *
- * @return device_memory_resource* a pointer to the static
- * cuda_memory_resource
- *---------------------------------------------------------------------------**/
-device_memory_resource* cuda_resource();
-
-/**---------------------------------------------------------------------------*
- * @brief gets a cnmem_memory_resource
- *
- * A static function which will return a singleton cnmem_memory_resource which
- * manages a pool
- *
- *
- * @param pool_size The initial size of the pool
- * @return device_memory_resource* a pointer to the static
- * cnmem_memory_resource
- *---------------------------------------------------------------------------**/
-device_memory_resource* pool_resource(std::size_t pool_size = 0);
-
-/**---------------------------------------------------------------------------*
- * @brief gets a cnmem_managed_memory_resource
- *
- * A static function which will return a singleton cnmem_memory_resource which
- * manages a pool of UVM memory
- *
- * @param pool_size The initial size of the pool
- * @return device_memory_resource* a pointer to the static
- * cnmem_managed_memory_resource
- *---------------------------------------------------------------------------**/
-device_memory_resource* managed_pool_resource(std::size_t pool_size = 0);
-
-
-/**---------------------------------------------------------------------------*
- * @brief gets a managed_memory_resource
- *
- * A static function that returns a singleton managed_memory_resource.
- *
- *
- *
- * @return device_memory_resource* a pointer to the static
- * managed_memory_resource
- *---------------------------------------------------------------------------**/
-device_memory_resource* managed_resource();
-}  // namespace detail
 }  // namespace mr
 }  // namespace rmm

src/mr/default_memory_resource.cpp

@@ -25,49 +25,20 @@
 namespace rmm {
 namespace mr {
 
-namespace detail{
-device_memory_resource* initial_resource() {
-  return cuda_resource();
-}
-
-device_memory_resource* cuda_resource() {
-  static cuda_memory_resource resource{};
-  return &resource;
-}
-
-device_memory_resource* pool_resource(std::size_t pool_size){
-   static cnmem_memory_resource mr{pool_size};
-   return &mr;
-}
-
-device_memory_resource* managed_pool_resource(std::size_t pool_size){
-   static cnmem_managed_memory_resource mr{pool_size};
-   return &mr;
-}
-
-device_memory_resource* managed_resource(){
-   static managed_memory_resource mr{};
-   return &mr;
-}
-} //namespace detail
 namespace {
-
-
-
-
-
 // Use an atomic to guarantee thread safety
 std::atomic<device_memory_resource*>& get_default() {
-  static std::atomic<device_memory_resource*> res{detail::initial_resource()};
+  static std::atomic<device_memory_resource*> res{new cuda_memory_resource{}};
   return res;
 }
-}  // namespace
+}  // namespace anonymous
 
 device_memory_resource* get_default_resource() { return get_default().load(); }
 
 device_memory_resource* set_default_resource(
-    device_memory_resource* new_resource) {
-  new_resource = (new_resource == nullptr) ? detail::initial_resource() : new_resource;
+  device_memory_resource* new_resource) {
+  new_resource = (new_resource == nullptr) ?
+    new cuda_memory_resource() : new_resource;
   return get_default().exchange(new_resource);
 }
 }  // namespace mr

src/rmm.cpp

@@ -59,39 +59,48 @@ const char * rmmGetErrorString(rmmError_t errcode) {
 // Initialize memory manager state and storage.
 rmmError_t rmmInitialize(rmmOptions_t *options)
 {
+  rmm::Manager::getInstance().initialize(options);
+
+  rmm::mr::device_memory_resource * memory_resource = nullptr;
+
+  if (rmm::Manager::usePoolAllocator()) {
+    if (rmm::Manager::useManagedMemory()) {
+      memory_resource = new rmm::mr::cnmem_managed_memory_resource{
+        rmm::Manager::getOptions().initial_pool_size};
+    } else {
+      memory_resource = new rmm::mr::cnmem_memory_resource{
+        rmm::Manager::getOptions().initial_pool_size};
+    } 
+  } else if (rmm::Manager::useManagedMemory()) {
+    memory_resource = new rmm::mr::managed_memory_resource();
+  }
 
-    rmm::Manager::getInstance().initialize(options);
-    rmm::mr::device_memory_resource * memory_resource = rmm::mr::detail::initial_resource();
-
-    if (rmm::Manager::usePoolAllocator())
-    {
-        memory_resource = rmm::mr::detail::pool_resource(rmm::Manager::getOptions().initial_pool_size);
-
-    }else if(rmm::Manager::useManagedMemory()){
-      memory_resource = rmm::mr::detail::managed_resource();
-
-    }else if(rmm::Manager::useManagedMemory() && rmm::Manager::usePoolAllocator()){
-        memory_resource =  rmm::mr::detail::managed_pool_resource(rmm::Manager::getOptions().initial_pool_size);
-    }
-    rmm::mr::set_default_resource(memory_resource);
+  rmm::mr::set_default_resource(memory_resource);
 
-    return RMM_SUCCESS;
+  return RMM_SUCCESS;
 }
 
 // Shutdown memory manager.
 rmmError_t rmmFinalize()
 {
-    rmm::Manager::getInstance().finalize();
-    return RMM_SUCCESS;
+  // delete the current default resource to reset it, and replace with the 
+  // initial resource, so that subsequent calls to initialize will act 
+  // just like the first call to initialize
+  rmm::mr::device_memory_resource *mr = rmm::mr::get_default_resource();
+  rmm::mr::set_default_resource(nullptr);
+  delete mr;
+
+  rmm::Manager::getInstance().finalize();
+  return RMM_SUCCESS;
 }
 
 // Query the initialization state of RMM.
 bool rmmIsInitialized(rmmOptions_t *options)
 {
-    if (nullptr != options) {
-        *options = rmm::Manager::getOptions();
-    }
-    return rmm::Manager::getInstance().isInitialized();
+  if (nullptr != options) {
+    *options = rmm::Manager::getOptions();
+  }
+  return rmm::Manager::getInstance().isInitialized();
 }
 
 // Allocate memory and return a pointer to device memory.
@@ -109,72 +118,70 @@ rmmError_t rmmAlloc(void **ptr, size_t size, cudaStream_t stream, const char* fi
 // Release device memory and recycle the associated memory.
 rmmError_t rmmFree(void *ptr, cudaStream_t stream, const char* file, unsigned int line)
 {
-    return rmm::free(ptr, stream, file, line);
+  return rmm::free(ptr, stream, file, line);
 }
 
 // Get the offset of ptr from its base allocation
 rmmError_t rmmGetAllocationOffset(ptrdiff_t *offset,
                                   void *ptr,
                                   cudaStream_t stream)
 {
-    void *base = (void*)0xffffffff;
-    CUresult res = cuMemGetAddressRange((CUdeviceptr*)&base, nullptr,
-                                        (CUdeviceptr)ptr);
-    if (res != CUDA_SUCCESS)
-        return RMM_ERROR_INVALID_ARGUMENT;
-    *offset = reinterpret_cast<ptrdiff_t>(ptr) -
-              reinterpret_cast<ptrdiff_t>(base);
-    return RMM_SUCCESS;
+  void *base = (void*)0xffffffff;
+  CUresult res = cuMemGetAddressRange((CUdeviceptr*)&base, nullptr,
+                                      (CUdeviceptr)ptr);
+  if (res != CUDA_SUCCESS)
+    return RMM_ERROR_INVALID_ARGUMENT;
+  *offset = reinterpret_cast<ptrdiff_t>(ptr) -
+            reinterpret_cast<ptrdiff_t>(base);
+  return RMM_SUCCESS;
 }
 
 // Get amounts of free and total memory managed by a manager associated
 // with the stream.
 rmmError_t rmmGetInfo(size_t *freeSize, size_t *totalSize, cudaStream_t stream)
 {
-    try{
-      std::pair<size_t,size_t> memInfo = rmm::mr::get_default_resource()->get_mem_info( stream);
-      *freeSize = memInfo.first;
-      *totalSize = memInfo.second;
-    }catch(std::runtime_error){
-      return RMM_ERROR_CUDA_ERROR;
-    }
-	return RMM_SUCCESS;
+  try {
+    std::pair<size_t,size_t> memInfo = rmm::mr::get_default_resource()->get_mem_info( stream);
+    *freeSize = memInfo.first;
+    *totalSize = memInfo.second;
+  } catch(std::runtime_error){
+    return RMM_ERROR_CUDA_ERROR;
+  }
+  return RMM_SUCCESS;
 }
 
 // Write the memory event stats log to specified path/filename
 rmmError_t rmmWriteLog(const char* filename)
 {
-    try
-    {
-        std::ofstream csv;
-        csv.open(filename);
-        rmm::Manager::getLogger().to_csv(csv);
-    }
-    catch (const std::ofstream::failure& e) {
-        return RMM_ERROR_IO;
-    }
-    return RMM_SUCCESS;
+  try {
+    std::ofstream csv;
+    csv.open(filename);
+    rmm::Manager::getLogger().to_csv(csv);
+  }
+  catch (const std::ofstream::failure& e) {
+    return RMM_ERROR_IO;
+  }
+  return RMM_SUCCESS;
 }
 
 // Get the size opf the CSV log
 size_t rmmLogSize()
 {
-    std::ostringstream csv;
-    rmm::Manager::getLogger().to_csv(csv);
-    return csv.str().size();
+  std::ostringstream csv;
+  rmm::Manager::getLogger().to_csv(csv);
+  return csv.str().size();
 }
 
 // Get the CSV log as a string
 rmmError_t rmmGetLog(char *buffer, size_t buffer_size)
 {
-    try
-    {
-        std::ostringstream csv;
-        rmm::Manager::getLogger().to_csv(csv);
-        csv.str().copy(buffer, std::min(buffer_size, csv.str().size()));
-    }
-    catch (const std::ofstream::failure& e) {
-        return RMM_ERROR_IO;
-    }
-    return RMM_SUCCESS;
+  try {
+    std::ostringstream csv;
+    rmm::Manager::getLogger().to_csv(csv);
+    csv.str().copy(buffer, std::min(buffer_size, csv.str().size()));
+  }
+  catch (const std::ofstream::failure& e) {
+    return RMM_ERROR_IO;
+  }
+  return RMM_SUCCESS;
 }