fix CPU only build with memman

NVIDIA · May 20, 2020 · be5c3c1 · be5c3c1
1 parent 1ddf529
commit be5c3c1
Show file tree

Hide file tree

Showing 7 changed files with 24 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,6 @@
 # Change Log
 
-## [master] - 2020-05-07
+## [master] - 2020-05-19
 
 ### Changed
 

diff --git a/README.md b/README.md
@@ -49,7 +49,7 @@ We visualized a sparse tensor network operation on a sparse tensor, convolution,
 
 - Ubuntu 14.04 or higher
 - CUDA 10.1 or higher
-- pytorch 1.3 or higher
+- pytorch 1.5 or higher
 - python 3.6 or higher
 - GCC 7 or higher
 

diff --git a/setup.py b/setup.py
@@ -131,7 +131,9 @@ def _argparse(pattern, argv, is_flag=True):
 
 # extra_compile_args+=['-g']  # Uncomment for debugging
 if CPU_ONLY and not FORCE_CUDA:
-    print("\nCPU_ONLY build set")
+    print("--------------------------------")
+    print("| WARNING: CPU_ONLY build set  |")
+    print("--------------------------------")
     compile_args += ["CPU_ONLY=1"]
     extra_compile_args += ["-DCPU_ONLY"]
     Extension = CppExtension

diff --git a/src/coords_manager.hpp b/src/coords_manager.hpp
@@ -122,18 +122,18 @@ template <typename MapType = CoordsToIndexMap> class CoordsManager {
   unordered_map<InOutMapKey, InOutMaps<int>, InOutMapKeyHash> in_maps;
   unordered_map<InOutMapKey, InOutMaps<int>, InOutMapKeyHash> out_maps;
 
-  CoordsManager(){
-    gpu_memory_manager = std::make_shared<GPUMemoryManager>();
-  };
-  CoordsManager(int num_threads) {
-    omp_set_dynamic(0);
-    omp_set_num_threads(num_threads);
-  }
   CoordsManager(int num_threads, MemoryManagerBackend backend) {
-    omp_set_dynamic(0);
-    omp_set_num_threads(num_threads);
+    if (num_threads > 0) {
+      omp_set_dynamic(0);
+      omp_set_num_threads(num_threads);
+    }
+#ifndef CPU_ONLY
     gpu_memory_manager = std::make_shared<GPUMemoryManager>(backend);
+#endif
   }
+  CoordsManager(int num_threads): CoordsManager(num_threads, PYTORCH) {}
+  CoordsManager(): CoordsManager(-1, PYTORCH) {}
+
   ~CoordsManager() { clear(); }
 
   void printDiagnostics(py::object py_coords_key) const;

diff --git a/src/gpu_memory_manager.hpp b/src/gpu_memory_manager.hpp
@@ -41,8 +41,6 @@ namespace minkowski {
 
 using std::vector;
 
-enum MemoryManagerBackend { CUDA = 0, PYTORCH = 1 };
-
 class GPUMemoryManager {
 private:
   int initial_size = 256;

diff --git a/src/types.hpp b/src/types.hpp
@@ -97,6 +97,9 @@ using InOutMapsRefPair = pair<InOutMaps<Itype> &, InOutMaps<Itype> &>;
 template <typename Itype>
 using pInOutMapsRefPair = pair<pInOutMaps<Itype> &, pInOutMaps<Itype> &>;
 
+// GPU memory manager backend. No effect with CPU_ONLY build
+enum MemoryManagerBackend { CUDA = 0, PYTORCH = 1 };
+
 // FNV64-1a
 // uint64_t for unsigned long, must use CXX -m64
 template <typename T> uint64_t hash_vec(T p) {

diff --git a/tests/coords.py b/tests/coords.py
@@ -169,11 +169,14 @@ def test_batch_size_initialize(self):
         self.assertTrue(cm.get_batch_size() == 2)
 
     def test_memory_manager_backend(self):
-        CoordsManager(memory_manager_backend=MemoryManagerBackend.CUDA, D=2)
-        CoordsManager(memory_manager_backend=MemoryManagerBackend.PYTORCH, D=2)
-
+        # Set the global GPU memory manager backend. By default PYTORCH.
         ME.set_memory_manager_backend(MemoryManagerBackend.PYTORCH)
-        CoordsManager(D=2)
+        ME.set_memory_manager_backend(MemoryManagerBackend.CUDA)
+
+        # Create a coords man with the specified GPU memory manager backend.
+        # No effect with CPU_ONLY build
+        cm = CoordsManager(memory_manager_backend=MemoryManagerBackend.CUDA, D=2)
+        cm = CoordsManager(memory_manager_backend=MemoryManagerBackend.PYTORCH, D=2)
 
 
 if __name__ == '__main__':