Skip to content

Commit 1c4fdfe

Browse files
rusty1srusty1s
committed
pytorch 1.3 support
1 parent 573ad11 commit 1c4fdfe

File tree

8 files changed

+65
-37
lines changed

8 files changed

+65
-37
lines changed

‎.travis.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ before_install:
1717
- export CXX="g++-4.9"
1818
install:
1919
- pip install numpy
20-
- pip install -q torch -f https://download.pytorch.org/whl/nightly/cpu/torch.html
20+
- pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
2121
- pip install pycodestyle
2222
- pip install flake8
2323
- pip install codecov

‎cpu/compat.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
#ifdef VERSION_GE_1_3
2+
#define DATA_PTR data_ptr
3+
#else
4+
#define DATA_PTR data
5+
#endif

‎cpu/spspmm.cpp‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include <torch/extension.h>
22

3+
#include "compat.h"
4+
35
at::Tensor degree(at::Tensor row, int64_t num_nodes) {
46
auto zero = at::zeros(num_nodes, row.options());
57
auto one = at::ones(row.size(0), row.options());
@@ -18,23 +20,23 @@ at::Tensor spspmm_bw(at::Tensor index, at::Tensor indexA, at::Tensor valueA,
1820
at::Tensor indexB, at::Tensor valueB, size_t rowA_max,
1921
size_t rowB_max) {
2022

21-
int64_t *index_data = index.data<int64_t>();
23+
int64_t *index_data = index.DATA_PTR<int64_t>();
2224
auto value = at::zeros(index.size(1), valueA.options());
2325

2426
at::Tensor rowA, colA;
2527
std::tie(rowA, colA) = to_csr(indexA[0], indexA[1], rowA_max);
26-
int64_t *rowA_data = rowA.data<int64_t>();
27-
int64_t *colA_data = colA.data<int64_t>();
28+
int64_t *rowA_data = rowA.DATA_PTR<int64_t>();
29+
int64_t *colA_data = colA.DATA_PTR<int64_t>();
2830

2931
at::Tensor rowB, colB;
3032
std::tie(rowB, colB) = to_csr(indexB[0], indexB[1], rowB_max);
31-
int64_t *rowB_data = rowB.data<int64_t>();
32-
int64_t *colB_data = colB.data<int64_t>();
33+
int64_t *rowB_data = rowB.DATA_PTR<int64_t>();
34+
int64_t *colB_data = colB.DATA_PTR<int64_t>();
3335

3436
AT_DISPATCH_FLOATING_TYPES(valueA.scalar_type(), "spspmm_bw", [&] {
35-
scalar_t *value_data = value.data<scalar_t>();
36-
scalar_t *valueA_data = valueA.data<scalar_t>();
37-
scalar_t *valueB_data = valueB.data<scalar_t>();
37+
scalar_t *value_data = value.DATA_PTR<scalar_t>();
38+
scalar_t *valueA_data = valueA.DATA_PTR<scalar_t>();
39+
scalar_t *valueB_data = valueB.DATA_PTR<scalar_t>();
3840

3941
for (int64_t e = 0; e < value.size(0); e++) {
4042
int64_t i = index_data[e], j = index_data[value.size(0) + e];

‎cuda/compat.cuh‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
#ifdef VERSION_GE_1_3
2+
#define DATA_PTR data_ptr
3+
#else
4+
#define DATA_PTR data
5+
#endif

‎cuda/spspmm_kernel.cu‎

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
#include <ATen/ATen.h>
2-
32
#include <cusparse.h>
43

4+
#include "compat.cuh"
5+
56
#define THREADS 1024
67
#define BLOCKS(N) (N + THREADS - 1) / THREADS
78

@@ -51,18 +52,18 @@ spspmm_cuda(at::Tensor indexA, at::Tensor valueA, at::Tensor indexB,
5152

5253
// Convert A to CSR format.
5354
auto row_ptrA = at::empty(m + 1, indexA.options());
54-
cusparseXcoo2csr(cusparse_handle, indexA[0].data<int>(), nnzA, k,
55-
row_ptrA.data<int>(), CUSPARSE_INDEX_BASE_ZERO);
55+
cusparseXcoo2csr(cusparse_handle, indexA[0].DATA_PTR<int>(), nnzA, k,
56+
row_ptrA.DATA_PTR<int>(), CUSPARSE_INDEX_BASE_ZERO);
5657
auto colA = indexA[1];
57-
cudaMemcpy(row_ptrA.data<int>() + m, &nnzA, sizeof(int),
58+
cudaMemcpy(row_ptrA.DATA_PTR<int>() + m, &nnzA, sizeof(int),
5859
cudaMemcpyHostToDevice);
5960

6061
// Convert B to CSR format.
6162
auto row_ptrB = at::empty(k + 1, indexB.options());
62-
cusparseXcoo2csr(cusparse_handle, indexB[0].data<int>(), nnzB, k,
63-
row_ptrB.data<int>(), CUSPARSE_INDEX_BASE_ZERO);
63+
cusparseXcoo2csr(cusparse_handle, indexB[0].DATA_PTR<int>(), nnzB, k,
64+
row_ptrB.DATA_PTR<int>(), CUSPARSE_INDEX_BASE_ZERO);
6465
auto colB = indexB[1];
65-
cudaMemcpy(row_ptrB.data<int>() + k, &nnzB, sizeof(int),
66+
cudaMemcpy(row_ptrB.DATA_PTR<int>() + k, &nnzB, sizeof(int),
6667
cudaMemcpyHostToDevice);
6768

6869
cusparseMatDescr_t descr = 0;
@@ -74,22 +75,23 @@ spspmm_cuda(at::Tensor indexA, at::Tensor valueA, at::Tensor indexB,
7475
auto row_ptrC = at::empty(m + 1, indexB.options());
7576
cusparseXcsrgemmNnz(cusparse_handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
7677
CUSPARSE_OPERATION_NON_TRANSPOSE, m, n, k, descr, nnzA,
77-
row_ptrA.data<int>(), colA.data<int>(), descr, nnzB,
78-
row_ptrB.data<int>(), colB.data<int>(), descr,
79-
row_ptrC.data<int>(), &nnzC);
78+
row_ptrA.DATA_PTR<int>(), colA.DATA_PTR<int>(), descr,
79+
nnzB, row_ptrB.DATA_PTR<int>(), colB.DATA_PTR<int>(),
80+
descr, row_ptrC.DATA_PTR<int>(), &nnzC);
8081
auto colC = at::empty(nnzC, indexA.options());
8182
auto valueC = at::empty(nnzC, valueA.options());
8283

8384
CSRGEMM(valueC.scalar_type(), cusparse_handle,
8485
CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE, m,
85-
n, k, descr, nnzA, valueA.data<scalar_t>(), row_ptrA.data<int>(),
86-
colA.data<int>(), descr, nnzB, valueB.data<scalar_t>(),
87-
row_ptrB.data<int>(), colB.data<int>(), descr,
88-
valueC.data<scalar_t>(), row_ptrC.data<int>(), colC.data<int>());
86+
n, k, descr, nnzA, valueA.DATA_PTR<scalar_t>(),
87+
row_ptrA.DATA_PTR<int>(), colA.DATA_PTR<int>(), descr, nnzB,
88+
valueB.DATA_PTR<scalar_t>(), row_ptrB.DATA_PTR<int>(),
89+
colB.DATA_PTR<int>(), descr, valueC.DATA_PTR<scalar_t>(),
90+
row_ptrC.DATA_PTR<int>(), colC.DATA_PTR<int>());
8991

9092
auto rowC = at::empty(nnzC, indexA.options());
91-
cusparseXcsr2coo(cusparse_handle, row_ptrC.data<int>(), nnzC, m,
92-
rowC.data<int>(), CUSPARSE_INDEX_BASE_ZERO);
93+
cusparseXcsr2coo(cusparse_handle, row_ptrC.DATA_PTR<int>(), nnzC, m,
94+
rowC.DATA_PTR<int>(), CUSPARSE_INDEX_BASE_ZERO);
9395

9496
auto indexC = at::stack({rowC, colC}, 0).toType(at::kLong);
9597

@@ -154,9 +156,10 @@ at::Tensor spspmm_bw_cuda(at::Tensor index, at::Tensor indexA,
154156

155157
AT_DISPATCH_FLOATING_TYPES(valueA.scalar_type(), "spspmm_bw", [&] {
156158
spspmm_bw_kernel<scalar_t><<<BLOCKS(value.numel()), THREADS>>>(
157-
index.data<int64_t>(), value.data<scalar_t>(), rowA.data<int64_t>(),
158-
colA.data<int64_t>(), valueA.data<scalar_t>(), rowB.data<int64_t>(),
159-
colB.data<int64_t>(), valueB.data<scalar_t>(), value.numel());
159+
index.DATA_PTR<int64_t>(), value.DATA_PTR<scalar_t>(),
160+
rowA.DATA_PTR<int64_t>(), colA.DATA_PTR<int64_t>(),
161+
valueA.DATA_PTR<scalar_t>(), rowB.DATA_PTR<int64_t>(),
162+
colB.DATA_PTR<int64_t>(), valueB.DATA_PTR<scalar_t>(), value.numel());
160163
});
161164

162165
return value;

‎cuda/unique_kernel.cu‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include <ATen/ATen.h>
22

3+
#include "compat.cuh"
4+
35
#define THREADS 1024
46
#define BLOCKS(N) (N + THREADS - 1) / THREADS
57

@@ -23,7 +25,7 @@ std::tuple<at::Tensor, at::Tensor> unique_cuda(at::Tensor src) {
2325
auto mask = at::zeros(src.numel(), src.options().dtype(at::kByte));
2426
AT_DISPATCH_ALL_TYPES(src.scalar_type(), "grid_cuda_kernel", [&] {
2527
unique_cuda_kernel<scalar_t><<<BLOCKS(src.numel()), THREADS>>>(
26-
src.data<scalar_t>(), mask.data<uint8_t>(), src.numel());
28+
src.DATA_PTR<scalar_t>(), mask.DATA_PTR<uint8_t>(), src.numel());
2729
});
2830

2931
src = src.masked_select(mask);

‎setup.py‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,17 @@
33
import torch
44
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
55

6-
ext_modules = [CppExtension('torch_sparse.spspmm_cpu', ['cpu/spspmm.cpp'])]
6+
TORCH_MAJOR = int(torch.__version__.split('.')[0])
7+
TORCH_MINOR = int(torch.__version__.split('.')[1])
8+
9+
extra_compile_args = []
10+
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
11+
extra_compile_args += ['-DVERSION_GE_1_3']
12+
13+
ext_modules = [
14+
CppExtension('torch_sparse.spspmm_cpu', ['cpu/spspmm.cpp'],
15+
extra_compile_args=extra_compile_args)
16+
]
717
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
818

919
if CUDA_HOME is not None:
@@ -13,15 +23,16 @@
1323
extra_link_args = ['-lcusparse', '-l', 'cusparse']
1424

1525
ext_modules += [
16-
CUDAExtension(
17-
'torch_sparse.spspmm_cuda',
18-
['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'],
19-
extra_link_args=extra_link_args),
26+
CUDAExtension('torch_sparse.spspmm_cuda',
27+
['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'],
28+
extra_link_args=extra_link_args,
29+
extra_compile_args=extra_compile_args),
2030
CUDAExtension('torch_sparse.unique_cuda',
21-
['cuda/unique.cpp', 'cuda/unique_kernel.cu']),
31+
['cuda/unique.cpp', 'cuda/unique_kernel.cu'],
32+
extra_compile_args=extra_compile_args),
2233
]
2334

24-
__version__ = '0.4.0'
35+
__version__ = '0.4.1'
2536
url = 'https://github.com/rusty1s/pytorch_sparse'
2637

2738
install_requires = ['scipy']

‎torch_sparse/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from .spmm import spmm
66
from .spspmm import spspmm
77

8-
__version__ = '0.4.0'
8+
__version__ = '0.4.1'
99

1010
__all__ = [
1111
'__version__',

0 commit comments

Comments
 (0)