Tbb scalable allocator for containers

facebookresearch · Sep 2, 2019 · 66682d1 · 66682d1
1 parent 1171aae
commit 66682d1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/setup.py b/setup.py
@@ -14,8 +14,10 @@
 this_dir = os.path.dirname(os.path.realpath(__file__))
 torch_dir = os.path.dirname(torch.__file__)
 conda_include_dir = '/'.join(torch_dir.split('/')[:-4]) + '/include'
+conda_lib_dir = os.environ["CONDA_PREFIX"] +  '/lib/'
 
 extra = {'cxx': ['-std=c++11', '-fopenmp'], 'nvcc': ['-std=c++11', '-Xcompiler', '-fopenmp']}
+extra_link = ['-ltbb', '-ltbbmalloc']
 
 setup(
     name='sparseconvnet',
@@ -29,13 +31,14 @@
       CUDAExtension('sparseconvnet.SCN',
         [
          'sparseconvnet/SCN/cuda.cu', 'sparseconvnet/SCN/sparseconvnet_cuda.cpp', 'sparseconvnet/SCN/pybind.cpp'],
-        include_dirs=[conda_include_dir, this_dir+'/sparseconvnet/SCN/'],
+        include_dirs=[conda_include_dir, this_dir+'/sparseconvnet/SCN/', conda_lib_dir],
         extra_compile_args=extra)
-      if torch.cuda.is_available()  else
+      if False else
       CppExtension('sparseconvnet.SCN',
         ['sparseconvnet/SCN/pybind.cpp', 'sparseconvnet/SCN/sparseconvnet_cpu.cpp'],
         include_dirs=[conda_include_dir, this_dir+'/sparseconvnet/SCN/'],
-        extra_compile_args=extra['cxx'])],
+        extra_compile_args=extra['cxx'], 
+        extra_link_args=extra_link)],
     cmdclass={'build_ext': BuildExtension},
     zip_safe=False,
 )
diff --git a/sparseconvnet/SCN/Metadata/Metadata.h b/sparseconvnet/SCN/Metadata/Metadata.h
@@ -7,6 +7,7 @@
 #ifndef Metadata_H
 #define Metadata_H
 #include "32bits.h"
+#include "tbb/scalable_allocator.h"
 #include <algorithm>
 #include <array>
 #include <cassert>
@@ -22,17 +23,17 @@
 #include <vector>
 
 template <Int dimension>
-using SparseGridMap =
-    google::dense_hash_map<Point<dimension>, Int, IntArrayHash<dimension>,
-                           std::equal_to<Point<dimension>>>;
+using SparseGridMap = google::dense_hash_map<Point<dimension>, Int, IntArrayHash<dimension>,
+    std::equal_to<Point<dimension>>, tbb::scalable_allocator<std::pair<Point<dimension>, Int>>>;
+
 template <Int dimension> class SparseGrid {
 public:
   Int ctr;
   SparseGridMap<dimension> mp;
   SparseGrid();
 };
 template <Int dimension> using SparseGrids = std::vector<SparseGrid<dimension>>;
-using RuleBook = std::vector<std::vector<Int>>;
+using RuleBook = std::vector<std::vector<Int, tbb::scalable_allocator<Int>>>;
 
 template <Int dimension>
 void addPointToSparseGridMapAndFeatures(SparseGridMap<dimension> &mp,