spmm.cpp 1.29 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
#include <torch/extension.h>

#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")

rusty1s's avatar
rusty1s committed
5
6
7
std::tuple<at::Tensor, at::optional<at::Tensor>>
spmm_cuda(at::Tensor rowptr, at::Tensor col, at::optional<at::Tensor> value_opt,
          at::Tensor mat, std::string reduce);
rusty1s's avatar
rusty1s committed
8

rusty1s's avatar
rusty1s committed
9
10
11
at::Tensor spmm_val_bw_cuda(at::Tensor index, at::Tensor rowptr, at::Tensor mat,
                            at::Tensor grad, std::string reduce);

rusty1s's avatar
rusty1s committed
12
13
14
std::tuple<at::Tensor, at::optional<at::Tensor>>
spmm(at::Tensor rowptr, at::Tensor col, at::optional<at::Tensor> value_opt,
     at::Tensor mat, std::string reduce) {
rusty1s's avatar
rusty1s committed
15
16
  CHECK_CUDA(rowptr);
  CHECK_CUDA(col);
rusty1s's avatar
rusty1s committed
17
18
  if (value_opt.has_value())
    CHECK_CUDA(value_opt.value());
rusty1s's avatar
rusty1s committed
19
  CHECK_CUDA(mat);
rusty1s's avatar
rusty1s committed
20
  return spmm_cuda(rowptr, col, value_opt, mat, reduce);
rusty1s's avatar
rusty1s committed
21
22
}

rusty1s's avatar
rusty1s committed
23
24
25
at::Tensor spmm_val_bw(at::Tensor row, at::Tensor rowptr, at::Tensor col,
                       at::Tensor mat, at::Tensor grad, std::string reduce) {
  CHECK_CUDA(row);
rusty1s's avatar
rusty1s committed
26
  CHECK_CUDA(rowptr);
rusty1s's avatar
rusty1s committed
27
  CHECK_CUDA(col);
rusty1s's avatar
rusty1s committed
28
29
  CHECK_CUDA(mat);
  CHECK_CUDA(grad);
rusty1s's avatar
rusty1s committed
30
  return spmm_val_bw_cuda(row, rowptr, col, mat, grad, reduce);
rusty1s's avatar
rusty1s committed
31
32
}

rusty1s's avatar
rusty1s committed
33
34
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("spmm", &spmm, "Sparse Matrix Multiplication (CUDA)");
rusty1s's avatar
rusty1s committed
35
36
  m.def("spmm_val_bw", &spmm_val_bw,
        "Sparse-Dense Matrix Multiplication Value Backward (CPU)");
rusty1s's avatar
rusty1s committed
37
}