Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "a9a47402dfc7b6ff6286ac13343a0335a35121fd"
Commit 9824c5f1 authored by rusty1s's avatar rusty1s
Browse files

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

parents 46040140 bdfdb070
Show whitespace changes
Inline Side-by-side
Showing with 305 additions and 98 deletions
aten/THCC/THCCWeighting.c deleted 100644 → 0
View file @ 46040140
#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 deleted 100644 → 0
View file @ 46040140
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 deleted 100644 → 0
View file @ 46040140
#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 deleted 100644 → 0
View file @ 46040140
#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
View file @ 9824c5f1
......@@ -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
View file @ 9824c5f1
......@@ -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 0 → 100644
View file @ 9824c5f1
#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
View file @ 9824c5f1
......@@ -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 0 → 100644
View file @ 9824c5f1
#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 0 → 100644
View file @ 9824c5f1
#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
View file @ 9824c5f1
......@@ -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
View file @ 9824c5f1
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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment