| #include <torch/library.h> |
|
|
| #include <ATen/cuda/CUDAContext.h> |
| #include <c10/cuda/CUDAGuard.h> |
|
|
| #include "../vox_cuda/vox_kernels.h" |
| #include "registration.h" |
| #include "torch_binding.h" |
|
|
| namespace { |
|
|
| void check_dims(int64_t dx, int64_t dy, int64_t dz) { |
| TORCH_CHECK(dx >= 1 && dy >= 1 && dz >= 1 && dx <= 8192 && dy <= 8192 && |
| dz <= 8192, |
| "grid dims must be in [1, 8192]"); |
| } |
|
|
| void check_grid(const at::Tensor& grid, int64_t dx, int64_t dy, int64_t dz) { |
| const int64_t words = ((dz + 31) / 32) * dy * dx; |
| TORCH_CHECK(grid.is_cuda() && grid.scalar_type() == at::kInt && |
| grid.is_contiguous() && grid.numel() == words, |
| "grid must be a contiguous CUDA int32 tensor with " |
| "ceil(dz/32)*dy*dx elements"); |
| } |
|
|
| void check_mesh(const at::Tensor& verts, const at::Tensor& faces, |
| const at::Tensor& tris) { |
| TORCH_CHECK(verts.is_cuda() && verts.scalar_type() == at::kFloat && |
| verts.dim() == 2 && verts.size(1) == 3 && |
| verts.is_contiguous(), |
| "verts must be a contiguous CUDA float32 [V, 3] tensor"); |
| TORCH_CHECK(faces.is_cuda() && faces.scalar_type() == at::kInt && |
| faces.dim() == 2 && faces.size(1) == 3 && |
| faces.is_contiguous(), |
| "faces must be a contiguous CUDA int32 [F, 3] tensor"); |
| TORCH_CHECK(tris.is_cuda() && tris.scalar_type() == at::kInt && |
| tris.dim() == 1 && tris.is_contiguous(), |
| "tris must be a contiguous CUDA int32 1-D tensor"); |
| } |
|
|
| void check_points(const at::Tensor& points) { |
| TORCH_CHECK(points.is_cuda() && points.scalar_type() == at::kFloat && |
| points.dim() == 2 && points.size(1) == 3 && |
| points.is_contiguous(), |
| "points must be a contiguous CUDA float32 [N, 3] tensor"); |
| } |
|
|
| } |
|
|
| void vox_surface(at::Tensor verts, at::Tensor faces, at::Tensor tris, |
| at::Tensor grid, int64_t dx, int64_t dy, int64_t dz, |
| double ox, double oy, double oz, double inv_vs, |
| int64_t group) { |
| check_dims(dx, dy, dz); |
| check_mesh(verts, faces, tris); |
| check_grid(grid, dx, dy, dz); |
| TORCH_CHECK(group == 1 || group == 32 || group == 256, |
| "group must be 1, 32, or 256"); |
| const c10::cuda::CUDAGuard guard(grid.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(); |
| const cudaError_t err = vox_surface_launch( |
| verts.data_ptr<float>(), faces.data_ptr<int>(), tris.data_ptr<int>(), |
| tris.numel(), reinterpret_cast<unsigned int*>(grid.data_ptr<int>()), |
| (int)dx, (int)dy, (int)dz, (float)ox, (float)oy, (float)oz, |
| (float)inv_vs, (int)group, stream.stream()); |
| TORCH_CHECK(err == cudaSuccess, "vox_surface launch failed: ", |
| cudaGetErrorString(err)); |
| } |
|
|
| void vox_solid_cross(at::Tensor verts, at::Tensor faces, at::Tensor tris, |
| at::Tensor grid, int64_t dx, int64_t dy, int64_t dz, |
| double ox, double oy, double oz, double inv_vs, |
| int64_t group) { |
| check_dims(dx, dy, dz); |
| check_mesh(verts, faces, tris); |
| check_grid(grid, dx, dy, dz); |
| TORCH_CHECK(group == 1 || group == 32 || group == 256, |
| "group must be 1, 32, or 256"); |
| const c10::cuda::CUDAGuard guard(grid.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(); |
| const cudaError_t err = vox_solid_cross_launch( |
| verts.data_ptr<float>(), faces.data_ptr<int>(), tris.data_ptr<int>(), |
| tris.numel(), reinterpret_cast<unsigned int*>(grid.data_ptr<int>()), |
| (int)dx, (int)dy, (int)dz, (float)ox, (float)oy, (float)oz, |
| (float)inv_vs, (int)group, stream.stream()); |
| TORCH_CHECK(err == cudaSuccess, "vox_solid_cross launch failed: ", |
| cudaGetErrorString(err)); |
| } |
|
|
| void vox_parity(at::Tensor grid, int64_t dx, int64_t dy, int64_t dz) { |
| check_dims(dx, dy, dz); |
| check_grid(grid, dx, dy, dz); |
| const c10::cuda::CUDAGuard guard(grid.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(); |
| const cudaError_t err = vox_parity_launch( |
| reinterpret_cast<unsigned int*>(grid.data_ptr<int>()), (int)dx, (int)dy, |
| (int)dz, stream.stream()); |
| TORCH_CHECK(err == cudaSuccess, "vox_parity launch failed: ", |
| cudaGetErrorString(err)); |
| } |
|
|
| void vox_point_keys(at::Tensor points, at::Tensor keys, int64_t dx, |
| int64_t dy, int64_t dz, double ox, double oy, double oz, |
| double inv_vs) { |
| check_dims(dx, dy, dz); |
| check_points(points); |
| TORCH_CHECK(keys.is_cuda() && keys.scalar_type() == at::kLong && |
| keys.is_contiguous() && keys.numel() == points.size(0), |
| "keys must be a contiguous CUDA int64 tensor with one element " |
| "per point"); |
| const c10::cuda::CUDAGuard guard(points.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(); |
| const cudaError_t err = vox_point_keys_launch( |
| points.data_ptr<float>(), points.size(0), |
| reinterpret_cast<long long*>(keys.data_ptr<int64_t>()), (int)dx, |
| (int)dy, (int)dz, (float)ox, (float)oy, (float)oz, (float)inv_vs, |
| stream.stream()); |
| TORCH_CHECK(err == cudaSuccess, "vox_point_keys launch failed: ", |
| cudaGetErrorString(err)); |
| } |
|
|
| void vox_point_occ(at::Tensor points, at::Tensor grid, int64_t dx, int64_t dy, |
| int64_t dz, double ox, double oy, double oz, |
| double inv_vs) { |
| check_dims(dx, dy, dz); |
| check_points(points); |
| check_grid(grid, dx, dy, dz); |
| const c10::cuda::CUDAGuard guard(grid.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(); |
| const cudaError_t err = vox_point_occ_launch( |
| points.data_ptr<float>(), points.size(0), |
| reinterpret_cast<unsigned int*>(grid.data_ptr<int>()), (int)dx, (int)dy, |
| (int)dz, (float)ox, (float)oy, (float)oz, (float)inv_vs, |
| stream.stream()); |
| TORCH_CHECK(err == cudaSuccess, "vox_point_occ launch failed: ", |
| cudaGetErrorString(err)); |
| } |
|
|
| void vox_seg_reduce(at::Tensor feats, at::Tensor order, at::Tensor segoff, |
| at::Tensor out, int64_t mode) { |
| TORCH_CHECK(feats.is_cuda() && feats.scalar_type() == at::kFloat && |
| feats.dim() == 2 && feats.is_contiguous(), |
| "feats must be a contiguous CUDA float32 [N, C] tensor"); |
| TORCH_CHECK(order.is_cuda() && order.scalar_type() == at::kLong && |
| order.dim() == 1 && order.is_contiguous(), |
| "order must be a contiguous CUDA int64 1-D tensor"); |
| TORCH_CHECK(segoff.is_cuda() && segoff.scalar_type() == at::kLong && |
| segoff.dim() == 1 && segoff.is_contiguous() && |
| segoff.numel() >= 1, |
| "segoff must be a contiguous CUDA int64 [S+1] tensor"); |
| const int64_t nseg = segoff.numel() - 1; |
| TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::kFloat && |
| out.dim() == 2 && out.is_contiguous() && |
| out.size(0) == nseg && out.size(1) == feats.size(1), |
| "out must be a contiguous CUDA float32 [S, C] tensor"); |
| TORCH_CHECK(mode >= 0 && mode <= 2, "mode must be 0 (mean), 1 (max), or " |
| "2 (sum)"); |
| const c10::cuda::CUDAGuard guard(feats.device()); |
| auto stream = at::cuda::getCurrentCUDAStream(); |
| const cudaError_t err = vox_seg_reduce_launch( |
| feats.data_ptr<float>(), |
| reinterpret_cast<const long long*>(order.data_ptr<int64_t>()), |
| reinterpret_cast<const long long*>(segoff.data_ptr<int64_t>()), nseg, |
| feats.size(1), out.data_ptr<float>(), (int)mode, stream.stream()); |
| TORCH_CHECK(err == cudaSuccess, "vox_seg_reduce launch failed: ", |
| cudaGetErrorString(err)); |
| } |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def("vox_surface(Tensor verts, Tensor faces, Tensor tris, Tensor grid, " |
| "int dx, int dy, int dz, float ox, float oy, float oz, " |
| "float inv_vs, int group) -> ()"); |
| ops.def("vox_solid_cross(Tensor verts, Tensor faces, Tensor tris, " |
| "Tensor grid, int dx, int dy, int dz, float ox, float oy, " |
| "float oz, float inv_vs, int group) -> ()"); |
| ops.def("vox_parity(Tensor grid, int dx, int dy, int dz) -> ()"); |
| ops.def("vox_point_keys(Tensor points, Tensor keys, int dx, int dy, " |
| "int dz, float ox, float oy, float oz, float inv_vs) -> ()"); |
| ops.def("vox_point_occ(Tensor points, Tensor grid, int dx, int dy, int dz, " |
| "float ox, float oy, float oz, float inv_vs) -> ()"); |
| ops.def("vox_seg_reduce(Tensor feats, Tensor order, Tensor segoff, " |
| "Tensor out, int mode) -> ()"); |
|
|
| #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL) |
| ops.impl("vox_surface", torch::kCUDA, &vox_surface); |
| ops.impl("vox_solid_cross", torch::kCUDA, &vox_solid_cross); |
| ops.impl("vox_parity", torch::kCUDA, &vox_parity); |
| ops.impl("vox_point_keys", torch::kCUDA, &vox_point_keys); |
| ops.impl("vox_point_occ", torch::kCUDA, &vox_point_occ); |
| ops.impl("vox_seg_reduce", torch::kCUDA, &vox_seg_reduce); |
| #endif |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
