From 66682d156936fab3a19b915d8004a0ccba5f131e Mon Sep 17 00:00:00 2001
From: Michal Pandy <michalpandy@fb.com>
Date: Mon, 2 Sep 2019 07:06:09 -0700
Subject: [PATCH] Tbb scalable allocator for containers

---
 setup.py                              | 9 ++++++---
 sparseconvnet/SCN/Metadata/Metadata.h | 9 +++++----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/setup.py b/setup.py
index 7b189ab..f1cf3bb 100644
--- 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
index 2266aa9..ecbcfa0 100644
--- 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,9 +23,9 @@
 #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;
@@ -32,7 +33,7 @@ template <Int dimension> class SparseGrid {
   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,