Merge branch 'master' of https://github.com/rusty1s/pytorch_spline_conv

aten/THCC/THCCWeighting.c

-#include <THC/THC.h>
-
-#include "THC.h"
-
-#define THCCTensor_(NAME) TH_CONCAT_4(THCC,Real,Tensor_,NAME)
-
-extern THCState *state;
-
-#include "generic/THCCWeighting.c"
-#include "THCGenerateFloatTypes.h"
-

aten/THCC/THCCWeighting.h

-void  THCCFloatTensor_weightingForward(      THCudaTensor *self,       THCudaTensor *src,       THCudaTensor *weight,       THCudaTensor *basis, THCudaLongTensor *weightIndex);
-void THCCDoubleTensor_weightingForward(THCudaDoubleTensor *self, THCudaDoubleTensor *src, THCudaDoubleTensor *weight, THCudaDoubleTensor *basis, THCudaLongTensor *weightIndex);
-
-void  THCCFloatTensor_weightingBackwardSrc(      THCudaTensor *self,       THCudaTensor *gradOutput,       THCudaTensor *weight,       THCudaTensor *basis, THCudaLongTensor *weightIndex);
-void THCCDoubleTensor_weightingBackwardSrc(THCudaDoubleTensor *self, THCudaDoubleTensor *gradOutput, THCudaDoubleTensor *weight, THCudaDoubleTensor *basis, THCudaLongTensor *weightIndex);
-
-void  THCCFloatTensor_weightingBackwardWeight(      THCudaTensor *self,      THCudaTensor *gradOutput,        THCudaTensor *src,       THCudaTensor *basis, THCudaLongTensor *weightIndex);
-void THCCDoubleTensor_weightingBackwardWeight(THCudaDoubleTensor *self, THCudaDoubleTensor *gradOutput, THCudaDoubleTensor *src, THCudaDoubleTensor *basis, THCudaLongTensor *weightIndex);
-
-void  THCCFloatTensor_weightingBackwardBasis(      THCudaTensor *self,       THCudaTensor *gradOutput,       THCudaTensor *src,       THCudaTensor *weight, THCudaLongTensor *weightIndex);
-void THCCDoubleTensor_weightingBackwardBasis(THCudaDoubleTensor *self, THCudaDoubleTensor *gradOutput, THCudaDoubleTensor *src, THCudaDoubleTensor *weight, THCudaLongTensor *weightIndex);

aten/THCC/generic/THCCBasis.c

-#ifndef THC_GENERIC_FILE
-#define THC_GENERIC_FILE "generic/THCCBasis.c"
-#else
-
-void THCCTensor_(linearBasisForward)(THCTensor *basis, THCudaLongTensor *weightIndex,
-                                     THCTensor *pseudo, THCudaLongTensor *kernelSize,
-                                     THCudaByteTensor *isOpenSpline) {
-  THCTensor_(linearBasisForward)(state, basis, weightIndex, pseudo, kernelSize, isOpenSpline);
-}
-
-void THCCTensor_(quadraticBasisForward)(THCTensor *basis, THCudaLongTensor *weightIndex,
-                                        THCTensor *pseudo, THCudaLongTensor *kernelSize,
-                                        THCudaByteTensor *isOpenSpline) {
-  THCTensor_(quadraticBasisForward)(state, basis, weightIndex, pseudo, kernelSize, isOpenSpline);
-}
-
-void THCCTensor_(cubicBasisForward)(THCTensor *basis, THCudaLongTensor *weightIndex,
-                                    THCTensor *pseudo, THCudaLongTensor *kernelSize,
-                                    THCudaByteTensor *isOpenSpline) {
-  THCTensor_(cubicBasisForward)(state, basis, weightIndex, pseudo, kernelSize, isOpenSpline);
-}
-
-void THCCTensor_(linearBasisBackward)(THCTensor *self, THCTensor *gradBasis, THCTensor *pseudo,
-                                      THCudaLongTensor *kernelSize,
-                                      THCudaByteTensor *isOpenSpline) {
-  THCTensor_(linearBasisBackward)(state, self, gradBasis, pseudo, kernelSize, isOpenSpline);
-}
-
-void THCCTensor_(quadraticBasisBackward)(THCTensor *self, THCTensor *gradBasis, THCTensor *pseudo,
-                                         THCudaLongTensor *kernelSize,
-                                         THCudaByteTensor *isOpenSpline) {
-  THCTensor_(quadraticBasisBackward)(state, self, gradBasis, pseudo, kernelSize, isOpenSpline);
-}
-
-void THCCTensor_(cubicBasisBackward)(THCTensor *self, THCTensor *gradBasis, THCTensor *pseudo,
-                                     THCudaLongTensor *kernelSize,
-                                     THCudaByteTensor *isOpenSpline) {
-  THCTensor_(cubicBasisBackward)(state, self, gradBasis, pseudo, kernelSize, isOpenSpline);
-}
-
-#endif  // THC_GENERIC_FILE

aten/THCC/generic/THCCWeighting.c

-#ifndef THC_GENERIC_FILE
-#define THC_GENERIC_FILE "generic/THCCWeighting.c"
-#else
-
-void THCCTensor_(weightingForward)(THCTensor *self, THCTensor *src, THCTensor *weight,
-                                   THCTensor *basis, THCudaLongTensor *weightIndex) {
-  THCTensor_(weightingForward)(state, self, src, weight, basis, weightIndex);
-}
-
-void THCCTensor_(weightingBackwardSrc)(THCTensor *self, THCTensor *gradOutput, THCTensor *weight,
-                                       THCTensor *basis, THCudaLongTensor *weightIndex) {
-  THCTensor_(weightingBackwardSrc)(state, self, gradOutput, weight, basis, weightIndex);
-}
-
-void THCCTensor_(weightingBackwardWeight)(THCTensor *self, THCTensor *gradOutput, THCTensor *src,
-                                          THCTensor *basis, THCudaLongTensor *weightIndex) {
-  THCTensor_(weightingBackwardWeight)(state, self, gradOutput, src, basis, weightIndex);
-}
-
-void THCCTensor_(weightingBackwardBasis)(THCTensor *self, THCTensor *gradOutput, THCTensor *src,
-                                         THCTensor *weight, THCudaLongTensor *weightIndex) {
-  THCTensor_(weightingBackwardBasis)(state, self, gradOutput, src, weight, weightIndex);
-}
-
-#endif  // THC_GENERIC_FILE

cpu/basis.cpp

@@ -142,12 +142,11 @@ inline scalar_t grad_cubic(scalar_t v, int64_t k_mod) {
             tmp = v;                                                           \
                                                                                \
             for (ptrdiff_t d_it = 1; d_it < D; d_it++) {                       \
-              auto d_other = d_it - (d >= d_it);                               \
-              k_mod = (s / (int64_t)(pow(M + 1, d_other) + 0.5)) % (M + 1);    \
+              auto d_new = d_it - (d >= d_it);                                 \
+              k_mod = (s / (int64_t)(pow(M + 1, d_new) + 0.5)) % (M + 1);      \
               v = pseudo_data[e * pseudo.stride(0) +                           \
-                              d_other * pseudo.stride(1)];                     \
-              v *= kernel_size_data[d_other] -                                 \
-                   M * is_open_spline_data[d_other];                           \
+                              d_new * pseudo.stride(1)];                       \
+              v *= kernel_size_data[d_new] - M * is_open_spline_data[d_new];   \
               v -= floor(v);                                                   \
               v = FUNC<scalar_t>(v, k_mod);                                    \
               tmp *= v;                                                        \

cpu/weighting.cpp

@@ -142,6 +142,6 @@ at::Tensor weighting_bw_b(at::Tensor grad_out, at::Tensor x, at::Tensor weight,
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("weighting_fw", &weighting_fw, "Weighting Forward (CPU)");
   m.def("weighting_bw_x", &weighting_bw_x, "Weighting Backward X (CPU)");
-  m.def("weighting_bw_w", &weighting_bw_w, "Weighting Backward W (CPU)");
-  m.def("weighting_bw_b", &weighting_bw_b, "Weighting Backward B (CPU)");
+  m.def("weighting_bw_w", &weighting_bw_w, "Weighting Backward Weight (CPU)");
+  m.def("weighting_bw_b", &weighting_bw_b, "Weighting Backward Basis (CPU)");
 }

cuda/atomics.cuh

+#pragma once
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || CUDA_VERSION < 8000)
+static inline __device__ void atomicAdd(double *address, double val) {
+  unsigned long long int *address_as_ull = (unsigned long long int *)address;
+  unsigned long long int old = *address_as_ull;
+  unsigned long long int assumed;
+
+  do {
+    assumed = old;
+    old = atomicCAS(address_as_ull, assumed,
+                    __double_as_longlong(val + __longlong_as_double(assumed)));
+  } while (assumed != old);
+}
+#endif

cuda/basis_kernel.cu

@@ -185,6 +185,32 @@ template <typename scalar_t> struct BasisBackward {
     const size_t index = blockIdx.x * blockDim.x + threadIdx.x;                \
     const size_t stride = blockDim.x * gridDim.x;                              \
     for (ptrdiff_t i = index; i < NUMEL; i += stride) {                        \
+      int64_t e = i / GRAD_PSEUDO.sizes[1], d = i % GRAD_PSEUDO.sizes[1];      \
+      scalar_t g = 0, tmp;                                                     \
+                                                                               \
+      for (ptrdiff_t s = 0; s < GRAD_BASIS.sizes[1]; s++) {                    \
+        auto k_mod = (s / (int64_t)(pow(M + 1, d) + 0.5)) % (M + 1);           \
+        auto v = PSEUDO.data[e * PSEUDO.strides[0] + d * PSEUDO.strides[1]];   \
+        v *= KERNEL_SIZE[d] - M * IS_OPEN_SPLINE[d];                           \
+        v -= floor(v);                                                         \
+        v = GRAD_CODE;                                                         \
+        tmp = v;                                                               \
+                                                                               \
+        for (ptrdiff_t d_it = 1; d_it < GRAD_PSEUDO.sizes[1]; d_it++) {        \
+          auto d_new = d_it - (d >= d_it);                                     \
+          k_mod = (s / (int64_t)(pow(M + 1, d_new) + 0.5)) % (M + 1);          \
+          v = PSEUDO.data[e * pseudo.strides[0] + d_new * PSEUDO.strides[1]];  \
+          v *= KERNEL_SIZE[d_new] - M * IS_OPEN_SPLINE[d_new];                 \
+          v -= floor(v);                                                       \
+          v = CODE;                                                            \
+          tmp *= v;                                                            \
+        }                                                                      \
+        g += tmp *                                                             \
+             GRAD_BASIS                                                        \
+                 .data[e * GRAD_BASIS.strides[0] + s * GRAD_BASIS.strides[1]]; \
+      }                                                                        \
+      g *= KERNEL_SIZE[d] - M * IS_OPEN_SPLINE[d];                             \
+      GRAD_PSEUDO.data[e * GRAD_PSEUDO.sizes[1] + d] = g;                      \
     }                                                                          \
   }()
 

cuda/weighting.cpp

+#include <torch/torch.h>
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
+
+at::Tensor weighting_fw_cuda(at::Tensor x, at::Tensor weight, at::Tensor basis,
+                             at::Tensor weight_index);
+
+at::Tensor weighting_bw_x_cuda(at::Tensor grad_out, at::Tensor weight,
+                               at::Tensor basis, at::Tensor weight_index);
+
+at::Tensor weighting_bw_w_cuda(at::Tensor grad_out, at::Tensor x,
+                               at::Tensor basis, at::Tensor weight_index,
+                               int64_t K);
+
+at::Tensor weighting_bw_b_cuda(at::Tensor grad_out, at::Tensor x,
+                               at::Tensor weight, at::Tensor weight_index);
+
+at::Tensor weighting_fw(at::Tensor x, at::Tensor weight, at::Tensor basis,
+                        at::Tensor weight_index) {
+  CHECK_CUDA(x);
+  CHECK_CUDA(weight);
+  CHECK_CUDA(basis);
+  CHECK_CUDA(weight_index);
+  return weighting_fw_cuda(x, weight, basis, weight_index);
+}
+
+at::Tensor weighting_bw_x(at::Tensor grad_out, at::Tensor weight,
+                          at::Tensor basis, at::Tensor weight_index) {
+  CHECK_CUDA(grad_out);
+  CHECK_CUDA(weight);
+  CHECK_CUDA(basis);
+  CHECK_CUDA(weight_index);
+  return weighting_bw_x_cuda(grad_out, weight, basis, weight_index);
+}
+
+at::Tensor weighting_bw_w(at::Tensor grad_out, at::Tensor x, at::Tensor basis,
+                          at::Tensor weight_index, int64_t K) {
+  CHECK_CUDA(grad_out);
+  CHECK_CUDA(x);
+  CHECK_CUDA(basis);
+  CHECK_CUDA(weight_index);
+  return weighting_bw_w_cuda(grad_out, x, basis, weight_index, K);
+}
+
+at::Tensor weighting_bw_b(at::Tensor grad_out, at::Tensor x, at::Tensor weight,
+                          at::Tensor weight_index) {
+  CHECK_CUDA(grad_out);
+  CHECK_CUDA(x);
+  CHECK_CUDA(weight);
+  CHECK_CUDA(weight_index);
+  return weighting_bw_b_cuda(grad_out, x, weight, weight_index);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("weighting_fw", &weighting_fw, "Weighting Forward (CUDA)");
+  m.def("weighting_bw_x", &weighting_bw_x, "Weighting Backward X (CUDA)");
+  m.def("weighting_bw_w", &weighting_bw_w, "Weighting Backward Weight (CUDA)");
+  m.def("weighting_bw_b", &weighting_bw_b, "Weighting Backward Basis (CUDA)");
+}
+#define BLOCKS(N) (N + THREADS - 1) / THREADS

cuda/weighting_kernel.cu

+#include <ATen/ATen.h>
+#include <ATen/cuda/detail/IndexUtils.cuh>
+#include <ATen/cuda/detail/TensorInfo.cuh>
+
+#include "atomics.cuh"
+
+#define THREADS 1024
+#define BLOCKS(N) (N + THREADS - 1) / THREADS
+
+template <typename scalar_t>
+__global__ void
+weighting_fw_kernel(at::cuda::detail::TensorInfo<scalar_t, int64_t> out,
+                    at::cuda::detail::TensorInfo<scalar_t, int64_t> x,
+                    at::cuda::detail::TensorInfo<scalar_t, int64_t> weight,
+                    at::cuda::detail::TensorInfo<scalar_t, int64_t> basis,
+                    at::cuda::detail::TensorInfo<int64_t, int64_t> weight_index,
+                    size_t numel) {
+  const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+  for (ptrdiff_t i = index; i < numel; i += stride) {
+    int64_t e = i / out.sizes[1], m_out = i % out.sizes[1];
+    auto S = basis.sizes[1];
+    scalar_t v = 0;
+
+    for (ptrdiff_t s = 0; s < S; s++) {
+      auto b = basis.data[e * S + s];
+      auto wi = weight_index.data[e * S + s];
+      for (ptrdiff_t m_in = 0; m_in < x.sizes[1]; m_in++) {
+        auto tmp =
+            weight.data[wi * weight.strides[0] + m_in * weight.strides[1] +
+                        m_out * weight.strides[2]];
+        tmp *= b * x.data[e * x.strides[0] + m_in * x.strides[1]];
+        v += tmp;
+      }
+    }
+    out.data[e * out.sizes[1] + m_out] = v;
+  }
+}
+
+at::Tensor weighting_fw_cuda(at::Tensor x, at::Tensor weight, at::Tensor basis,
+                             at::Tensor weight_index) {
+  auto E = x.size(0), M_out = weight.size(2);
+  auto out = at::empty({E, M_out}, x.type());
+  AT_DISPATCH_FLOATING_TYPES(out.type(), "weighting_fw", [&] {
+    weighting_fw_kernel<scalar_t><<<BLOCKS(out.numel()), THREADS>>>(
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(out),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(x),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(weight),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(basis),
+        at::cuda::detail::getTensorInfo<int64_t, int64_t>(weight_index),
+        out.numel());
+  });
+  return out;
+}
+
+template <typename scalar_t>
+__global__ void weighting_bw_x_kernel(
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> grad_x,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> grad_out,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> weight,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> basis,
+    at::cuda::detail::TensorInfo<int64_t, int64_t> weight_index, size_t numel) {
+  const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+  for (ptrdiff_t i = index; i < numel; i += stride) {
+    int64_t e = i / grad_x.sizes[1], m_in = i % grad_x.sizes[1];
+    auto S = basis.sizes[1];
+    scalar_t v = 0;
+
+    for (ptrdiff_t s = 0; s < S; s++) {
+      auto b = basis.data[e * S + s];
+      auto wi = weight_index.data[e * S + s];
+      for (ptrdiff_t m_out = 0; m_out < grad_out.sizes[1]; m_out++) {
+        auto tmp =
+            weight.data[wi * weight.strides[0] + m_out * weight.strides[1] +
+                        m_in * weight.strides[2]];
+        tmp *= b *
+               grad_out
+                   .data[e * grad_out.strides[0] + m_out * grad_out.strides[1]];
+        v += tmp;
+      }
+    }
+    grad_x.data[e * grad_x.sizes[1] + m_in] = v;
+  }
+}
+
+at::Tensor weighting_bw_x_cuda(at::Tensor grad_out, at::Tensor weight,
+                               at::Tensor basis, at::Tensor weight_index) {
+  auto E = grad_out.size(0), M_in = weight.size(1);
+  auto grad_x = at::empty({E, M_in}, grad_out.type());
+  weight = weight.transpose(1, 2).contiguous();
+  AT_DISPATCH_FLOATING_TYPES(grad_x.type(), "weighting_bw_x", [&] {
+    weighting_bw_x_kernel<scalar_t><<<BLOCKS(grad_x.numel()), THREADS>>>(
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_x),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_out),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(weight),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(basis),
+        at::cuda::detail::getTensorInfo<int64_t, int64_t>(weight_index),
+        grad_x.numel());
+  });
+  return grad_x;
+}
+
+template <typename scalar_t>
+__global__ void weighting_bw_w_kernel(
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> grad_weight,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> grad_out,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> x,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> basis,
+    at::cuda::detail::TensorInfo<int64_t, int64_t> weight_index, size_t numel) {
+  const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+  for (ptrdiff_t i = index; i < numel; i += stride) {
+    int64_t e = i / grad_out.sizes[1], m_out = i % grad_out.sizes[1];
+    int64_t S = basis.sizes[1], M_in = x.sizes[1], M_out = grad_out.sizes[1];
+
+    auto g =
+        grad_out.data[e * grad_out.strides[0] + m_out * grad_out.strides[1]];
+    for (ptrdiff_t s = 0; s < S; s++) {
+      auto b = basis.data[e * S + s];
+      auto wi = weight_index.data[e * S + s];
+      for (ptrdiff_t m_in = 0; m_in < M_in; m_in++) {
+        auto v = g * b * x.data[e * x.strides[0] + m_in * x.strides[1]];
+        atomicAdd(&grad_weight.data[wi * M_in * M_out + m_in * M_out + m_out],
+                  v);
+      }
+    }
+  }
+}
+
+at::Tensor weighting_bw_w_cuda(at::Tensor grad_out, at::Tensor x,
+                               at::Tensor basis, at::Tensor weight_index,
+                               int64_t K) {
+  auto M_in = x.size(1), M_out = grad_out.size(1);
+  auto grad_weight = at::zeros({K, M_in, M_out}, grad_out.type());
+  AT_DISPATCH_FLOATING_TYPES(grad_out.type(), "weighting_bw_w", [&] {
+    weighting_bw_w_kernel<scalar_t><<<BLOCKS(grad_out.numel()), THREADS>>>(
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_weight),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_out),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(x),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(basis),
+        at::cuda::detail::getTensorInfo<int64_t, int64_t>(weight_index),
+        grad_out.numel());
+  });
+  return grad_weight;
+}
+
+template <typename scalar_t>
+__global__ void weighting_bw_b_kernel(
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> grad_basis,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> grad_out,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> x,
+    at::cuda::detail::TensorInfo<scalar_t, int64_t> weight,
+    at::cuda::detail::TensorInfo<int64_t, int64_t> weight_index, size_t numel) {
+  const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+  for (ptrdiff_t i = index; i < numel; i += stride) {
+    int64_t e = i / grad_out.sizes[1], m_out = i % grad_out.sizes[1];
+    auto S = grad_basis.sizes[1];
+
+    auto g =
+        grad_out.data[e * grad_out.strides[0] + m_out * grad_out.strides[1]];
+    for (ptrdiff_t s = 0; s < S; s++) {
+      scalar_t v = 0;
+      auto wi = weight_index.data[e * S + s];
+      for (ptrdiff_t m_in = 0; m_in < x.sizes[1]; m_in++) {
+        auto w = weight.data[wi * weight.strides[0] + m_in * weight.strides[1] +
+                             m_out * weight.strides[2]];
+        v += g * w * x.data[e * x.strides[0] + m_in * x.strides[1]];
+      }
+      atomicAdd(&grad_basis.data[e * S + s], v);
+    }
+  }
+}
+
+at::Tensor weighting_bw_b_cuda(at::Tensor grad_out, at::Tensor x,
+                               at::Tensor weight, at::Tensor weight_index) {
+  auto E = x.size(0), S = weight_index.size(1);
+  auto grad_basis = at::zeros({E, S}, grad_out.type());
+  AT_DISPATCH_FLOATING_TYPES(grad_out.type(), "weighting_bw_b", [&] {
+    weighting_bw_b_kernel<scalar_t><<<BLOCKS(grad_out.numel()), THREADS>>>(
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_basis),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_out),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(x),
+        at::cuda::detail::getTensorInfo<scalar_t, int64_t>(weight),
+        at::cuda::detail::getTensorInfo<int64_t, int64_t>(weight_index),
+        grad_out.numel());
+  });
+  return grad_basis;
+}

setup.py

@@ -10,7 +10,10 @@ cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
 
 if torch.cuda.is_available():
     ext_modules += [
-        CUDAExtension('basis_cuda', ['cuda/basis.cpp', 'cuda/basis_kernel.cu'])
+        CUDAExtension('basis_cuda',
+                      ['cuda/basis.cpp', 'cuda/basis_kernel.cu']),
+        CUDAExtension('weighting_cuda',
+                      ['cuda/weighting.cpp', 'cuda/weighting_kernel.cu']),
     ]
 
 __version__ = '1.0.4'

torch_spline_conv/weighting.py

 import torch
 import weighting_cpu
 
+if torch.cuda.is_available():
+    import weighting_cuda
+
 
 def get_func(name, tensor):
-    # module = weighting_cuda if tensor.is_cuda else weighting_cpu
-    module = weighting_cpu
+    module = weighting_cuda if tensor.is_cuda else weighting_cpu
     return getattr(module, name)