File size: 9,059 Bytes
3a2242c c0e0c9d 3a2242c c0e0c9d 3a2242c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | #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");
}
} // namespace
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)
|