Add the implementation of the Gridding layer (arXiv 2006.03761).

torch-points3d · Dec 22, 2020 · 8488978 · 8488978
1 parent 48e3a16
commit 8488978
Show file tree

Hide file tree

Showing 6 changed files with 390 additions and 14 deletions.
diff --git a/cuda/include/cuda_utils.h b/cuda/include/cuda_utils.h
@@ -32,23 +32,22 @@ inline dim3 opt_block_config(int x, int y)
 // from https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
 #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 600
 #else
-    __device__ double 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, assumed;
+__device__ double 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, assumed;
 
-        do {
-            assumed = old;
-            old = atomicCAS(address_as_ull, assumed,
-                            __double_as_longlong(val +
-                                __longlong_as_double(assumed)));
+    do
+    {
+        assumed = old;
+        old = atomicCAS(address_as_ull, assumed,
+                        __double_as_longlong(val + __longlong_as_double(assumed)));
 
         // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
-        } while (assumed != old);
+    } while (assumed != old);
 
-        return __longlong_as_double(old);
-    }
+    return __longlong_as_double(old);
+}
 #endif
 
 #define CUDA_CHECK_ERRORS()                                                                        \

diff --git a/cuda/include/gridding.h b/cuda/include/gridding.h
@@ -0,0 +1,18 @@
+#include <vector>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> gridding_kernel_warpper(float min_x, float max_x, float min_y,
+                                                   float max_y, float min_z, float max_z,
+                                                   torch::Tensor ptcloud, cudaStream_t stream);
+
+torch::Tensor gridding_grad_kernel_warpper(torch::Tensor grid_pt_weights,
+                                           torch::Tensor grid_pt_indexes, torch::Tensor grad_grid,
+                                           cudaStream_t stream);
+
+std::vector<torch::Tensor> gridding(float min_x, float max_x, float min_y, float max_y, float min_z,
+                                    float max_z, torch::Tensor ptcloud);
+
+torch::Tensor gridding_grad(torch::Tensor grid_pt_weights, torch::Tensor grid_pt_indexes,
+                            torch::Tensor grad_grid);
diff --git a/cuda/src/bindings.cpp b/cuda/src/bindings.cpp
@@ -1,6 +1,7 @@
 #include "ball_query.h"
 #include "chamfer_dist.h"
 #include "cubic_feature_sampling.h"
+#include "gridding.h"
 #include "interpolate.h"
 #include "metrics.h"
 #include "sampling.h"
@@ -23,4 +24,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 
     m.def("cubic_feature_sampling", &cubic_feature_sampling);
     m.def("cubic_feature_sampling_grad", &cubic_feature_sampling_grad);
+
+    m.def("gridding", &gridding);
+    m.def("gridding_grad", &gridding_grad);
 }
diff --git a/cuda/src/chamfer_dist_gpu.cu b/cuda/src/chamfer_dist_gpu.cu
@@ -2,8 +2,8 @@
 #include <cuda_runtime.h>
 #include <torch/extension.h>
 
-#include <vector>
 #include "cuda_utils.h"
+#include <vector>
 
 template <typename scalar_t>
 __global__ void chamfer_dist_kernel(int batch_size, int n, const scalar_t* __restrict__ xyz1, int m,

diff --git a/cuda/src/gridding.cpp b/cuda/src/gridding.cpp
@@ -0,0 +1,26 @@
+#include "gridding.h"
+#include "utils.h"
+
+std::vector<torch::Tensor> gridding(float min_x, float max_x, float min_y, float max_y, float min_z,
+                                    float max_z, torch::Tensor ptcloud)
+{
+    CHECK_CUDA(ptcloud);
+    CHECK_CONTIGUOUS(ptcloud);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    return gridding_kernel_warpper(min_x, max_x, min_y, max_y, min_z, max_z, ptcloud, stream);
+}
+
+torch::Tensor gridding_grad(torch::Tensor grid_pt_weights, torch::Tensor grid_pt_indexes,
+                            torch::Tensor grad_grid)
+{
+    CHECK_CUDA(grid_pt_weights);
+    CHECK_CONTIGUOUS(grid_pt_weights);
+    CHECK_CUDA(grid_pt_indexes);
+    CHECK_CONTIGUOUS(grid_pt_indexes);
+    CHECK_CUDA(grad_grid);
+    CHECK_CONTIGUOUS(grad_grid);
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    return gridding_grad_kernel_warpper(grid_pt_weights, grid_pt_indexes, grad_grid, stream);
+}