Skip to content
This repository was archived by the owner on Mar 1, 2025. It is now read-only.

Commit 66682d1

Browse files
committed
Tbb scalable allocator for containers
1 parent 1171aae commit 66682d1

File tree

2 files changed

+11
-7
lines changed

2 files changed

+11
-7
lines changed

setup.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
this_dir = os.path.dirname(os.path.realpath(__file__))
1515
torch_dir = os.path.dirname(torch.__file__)
1616
conda_include_dir = '/'.join(torch_dir.split('/')[:-4]) + '/include'
17+
conda_lib_dir = os.environ["CONDA_PREFIX"] + '/lib/'
1718

1819
extra = {'cxx': ['-std=c++11', '-fopenmp'], 'nvcc': ['-std=c++11', '-Xcompiler', '-fopenmp']}
20+
extra_link = ['-ltbb', '-ltbbmalloc']
1921

2022
setup(
2123
name='sparseconvnet',
@@ -29,13 +31,14 @@
2931
CUDAExtension('sparseconvnet.SCN',
3032
[
3133
'sparseconvnet/SCN/cuda.cu', 'sparseconvnet/SCN/sparseconvnet_cuda.cpp', 'sparseconvnet/SCN/pybind.cpp'],
32-
include_dirs=[conda_include_dir, this_dir+'/sparseconvnet/SCN/'],
34+
include_dirs=[conda_include_dir, this_dir+'/sparseconvnet/SCN/', conda_lib_dir],
3335
extra_compile_args=extra)
34-
if torch.cuda.is_available() else
36+
if False else
3537
CppExtension('sparseconvnet.SCN',
3638
['sparseconvnet/SCN/pybind.cpp', 'sparseconvnet/SCN/sparseconvnet_cpu.cpp'],
3739
include_dirs=[conda_include_dir, this_dir+'/sparseconvnet/SCN/'],
38-
extra_compile_args=extra['cxx'])],
40+
extra_compile_args=extra['cxx'],
41+
extra_link_args=extra_link)],
3942
cmdclass={'build_ext': BuildExtension},
4043
zip_safe=False,
4144
)

sparseconvnet/SCN/Metadata/Metadata.h

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#ifndef Metadata_H
88
#define Metadata_H
99
#include "32bits.h"
10+
#include "tbb/scalable_allocator.h"
1011
#include <algorithm>
1112
#include <array>
1213
#include <cassert>
@@ -22,17 +23,17 @@
2223
#include <vector>
2324

2425
template <Int dimension>
25-
using SparseGridMap =
26-
google::dense_hash_map<Point<dimension>, Int, IntArrayHash<dimension>,
27-
std::equal_to<Point<dimension>>>;
26+
using SparseGridMap = google::dense_hash_map<Point<dimension>, Int, IntArrayHash<dimension>,
27+
std::equal_to<Point<dimension>>, tbb::scalable_allocator<std::pair<Point<dimension>, Int>>>;
28+
2829
template <Int dimension> class SparseGrid {
2930
public:
3031
Int ctr;
3132
SparseGridMap<dimension> mp;
3233
SparseGrid();
3334
};
3435
template <Int dimension> using SparseGrids = std::vector<SparseGrid<dimension>>;
35-
using RuleBook = std::vector<std::vector<Int>>;
36+
using RuleBook = std::vector<std::vector<Int, tbb::scalable_allocator<Int>>>;
3637

3738
template <Int dimension>
3839
void addPointToSparseGridMapAndFeatures(SparseGridMap<dimension> &mp,

0 commit comments

Comments
 (0)