convert.cpp 647 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
#include <torch/script.h>
rusty1s's avatar
rusty1s committed
2

rusty1s's avatar
rusty1s committed
3
4
#define CHECK_CUDA(x)                                                          \
  AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
rusty1s's avatar
rusty1s committed
5

rusty1s's avatar
rusty1s committed
6
7
torch::Tensor ind2ptr_cuda(torch::Tensor ind, int64_t M);
torch::Tensor ptr2ind_cuda(torch::Tensor ptr, int64_t E);
rusty1s's avatar
rusty1s committed
8

rusty1s's avatar
rusty1s committed
9
torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
rusty1s's avatar
rusty1s committed
10
11
12
13
  CHECK_CUDA(ind);
  return ind2ptr_cuda(ind, M);
}

rusty1s's avatar
rusty1s committed
14
torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) {
rusty1s's avatar
rusty1s committed
15
16
17
18
  CHECK_CUDA(ptr);
  return ptr2ind_cuda(ptr, E);
}

rusty1s's avatar
rusty1s committed
19
20
21
static auto registry =
    torch::RegisterOperators("torch_sparse_cuda::ind2ptr", &ind2ptr)
        .op("torch_sparse_cuda::ptr2ind", &ptr2ind);