diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a127cb3..f38a0938 100644 --- 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 index a6fe19f1..6809e34e 100644 --- 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 index 7f0c5e4d..deb20ea1 100644 --- 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 index 3fe4fa89..e2ca7ba7 100644 --- a/src/coords_manager.hpp +++ b/src/coords_manager.hpp @@ -122,18 +122,18 @@ template class CoordsManager { unordered_map, InOutMapKeyHash> in_maps; unordered_map, InOutMapKeyHash> out_maps; - CoordsManager(){ - gpu_memory_manager = std::make_shared(); - }; - 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(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 index 94c6e508..255c32da 100644 --- 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 index dff59317..7c8e3318 100644 --- a/src/types.hpp +++ b/src/types.hpp @@ -97,6 +97,9 @@ using InOutMapsRefPair = pair &, InOutMaps &>; template using pInOutMapsRefPair = pair &, pInOutMaps &>; +// 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 uint64_t hash_vec(T p) { diff --git a/tests/coords.py b/tests/coords.py index 2e24ba64..df33daa9 100644 --- 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__':