Skip to content

Commit

Permalink
Add differentiable volumetric rendering ops for SPC (#485)
Browse files Browse the repository at this point in the history 
* Add differentiable volumetric rendering ops for SPC

Signed-off-by: Towaki Takikawa <ttakikawa@nvidia.com>

* Flake 8 issues

Signed-off-by: Towaki Takikawa <ttakikawa@nvidia.com>

* Fix documentation issues

Signed-off-by: Towaki Takikawa <ttakikawa@nvidia.com>

Co-authored-by: Towaki Takikawa <ttakikawa@nvidia.com>
  • Loading branch information
@tovacinni
tovacinni and Towaki Takikawa committed Dec 8, 2021
1 parent d387158 commit 3ade542
Show file tree
Hide file tree
Showing 6 changed files with 943 additions and 31 deletions.
5 changes: 5 additions & 0 deletions kaolin/csrc/bindings.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
render_spc.def("generate_primary_rays_cuda", &generate_primary_rays_cuda); // Deprecate soon
render_spc.def("mark_pack_boundary_cuda", &mark_pack_boundary_cuda);
render_spc.def("generate_shadow_rays_cuda", &generate_shadow_rays_cuda); // Deprecate soon
render_spc.def("inclusive_sum_cuda", &inclusive_sum_cuda);
render_spc.def("diff_cuda", &diff_cuda);
render_spc.def("sum_reduce_cuda", &sum_reduce_cuda);
render_spc.def("cumsum_cuda", &cumsum_cuda);
render_spc.def("cumprod_cuda", &cumprod_cuda);
}

} // namespace kaolin
158 changes: 158 additions & 0 deletions kaolin/csrc/render/spc/raytrace.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,44 @@ uint generate_shadow_rays_cuda_impl(
uint* info,
uint* prefix_sum);

void diff_cuda_impl(
int64_t num_packs,
int64_t num_feats,
int64_t feat_dim,
at::Tensor feats_in,
at::Tensor feats_out,
at::Tensor pack_indices);

void inclusive_sum_cuda_impl(
int64_t num,
at::Tensor info,
at::Tensor inclusive_sum);

uint sum_reduce_cuda_impl(
int64_t num_feats,
int64_t feat_dim,
at::Tensor feats_in,
at::Tensor feats_out,
at::Tensor inclusive_sum);

void cumsum_cuda_impl(
int64_t num_feats,
int64_t feat_dim,
at::Tensor feats_in,
at::Tensor feats_out,
at::Tensor pack_indices,
bool exclusive,
bool reverse);

void cumprod_cuda_impl(
int64_t num_feats,
int64_t feat_dim,
at::Tensor feats_in,
at::Tensor feats_out,
at::Tensor pack_indices,
bool exclusive,
bool reverse);

#endif

std::vector<at::Tensor> generate_primary_rays_cuda(
Expand Down Expand Up @@ -263,4 +301,124 @@ std::vector<at::Tensor> generate_shadow_rays_cuda(
#endif // WITH_CUDA
}

at::Tensor diff_cuda(
at::Tensor feats,
at::Tensor pack_indices) {
#ifdef WITH_CUDA
at::TensorArg feats_arg{feats, "feats", 1};
at::TensorArg pack_indices_arg{pack_indices, "pack_indices", 2};
at::checkDim(__func__, feats_arg, 2);
at::checkDim(__func__, pack_indices_arg, 1);
at::checkAllSameGPU(__func__, {feats_arg, pack_indices_arg});
at::checkAllContiguous(__func__, {feats_arg, pack_indices_arg});
at::checkScalarTypes(__func__, feats_arg, {at::kHalf, at::kFloat, at::kDouble});
at::checkScalarType(__func__, pack_indices_arg, at::kLong);

int64_t num_feats = feats.size(0);
int64_t feat_dim = feats.size(1);
at::Tensor feats_out = at::zeros({num_feats, feat_dim}, feats.options());
int64_t num_packs = pack_indices.size(0);

diff_cuda_impl(num_packs, num_feats, feat_dim, feats, feats_out, pack_indices);

return feats_out;
#else
KAOLIN_NO_CUDA_ERROR(__func__);
#endif // WITH_CUDA
}

at::Tensor inclusive_sum_cuda(at::Tensor info) {
#ifdef WITH_CUDA
at::TensorArg info_arg{info, "info", 1};
at::checkDim(__func__, info_arg, 1);
at::checkAllSameGPU(__func__, {info_arg});
at::checkAllContiguous(__func__, {info_arg});
at::checkScalarType(__func__, info_arg, at::kInt);

int num = info.size(0);
at::Tensor inclusive_sum = at::zeros({num}, info.options().dtype(at::kInt));
inclusive_sum_cuda_impl(num, info, inclusive_sum);
return inclusive_sum;
#else
KAOLIN_NO_CUDA_ERROR(__func__);
#endif // WITH_CUDA
}

at::Tensor sum_reduce_cuda(
at::Tensor feats,
at::Tensor inclusive_sum) {
#ifdef WITH_CUDA
at::TensorArg feats_arg{feats, "feats", 1};
at::TensorArg inclusive_sum_arg{inclusive_sum, "inclusive_sum", 2};
at::checkDim(__func__, feats_arg, 2);
at::checkDim(__func__, inclusive_sum_arg, 1);
at::checkAllSameGPU(__func__, {feats_arg, inclusive_sum_arg});
at::checkAllContiguous(__func__, {feats_arg, inclusive_sum_arg});
at::checkScalarTypes(__func__, feats_arg, {at::kHalf, at::kFloat, at::kDouble});
at::checkScalarType(__func__, inclusive_sum_arg, at::kInt);

int num_feats = feats.size(0);
int feat_dim = feats.size(1);

at::Tensor feats_out = at::zeros({num_feats, feat_dim}, feats.options());

int cnt = sum_reduce_cuda_impl(num_feats, feat_dim, feats, feats_out, inclusive_sum);

return feats_out.index({Slice(None, cnt)}).contiguous();
#else
KAOLIN_NO_CUDA_ERROR(__func__);
#endif // WITH_CUDA
}

at::Tensor cumsum_cuda(
at::Tensor feats,
at::Tensor pack_indices,
bool exclusive,
bool reverse) {
#ifdef WITH_CUDA
at::TensorArg feats_arg{feats, "feats", 1};
at::TensorArg pack_indices_arg{pack_indices, "pack_indices", 2};
at::checkDim(__func__, feats_arg, 2);
at::checkDim(__func__, pack_indices_arg, 1);
at::checkAllSameGPU(__func__, {feats_arg, pack_indices_arg});
at::checkAllContiguous(__func__, {feats_arg, pack_indices_arg});
at::checkScalarTypes(__func__, feats_arg, {at::kHalf, at::kFloat, at::kDouble});
at::checkScalarType(__func__, pack_indices_arg, at::kInt);

//int num_feats = feats.size(0);
int num_feats = feats.size(0);
int feat_dim = feats.size(1);
at::Tensor feats_out = at::zeros({num_feats, feat_dim}, feats.options());
cumsum_cuda_impl(num_feats, feat_dim, feats, feats_out, pack_indices, exclusive, reverse);
return feats_out.contiguous();
#else
KAOLIN_NO_CUDA_ERROR(__func__);
#endif // WITH_CUDA
}

at::Tensor cumprod_cuda(
at::Tensor feats,
at::Tensor pack_indices,
bool exclusive,
bool reverse) {
#ifdef WITH_CUDA
at::TensorArg feats_arg{feats, "feats", 1};
at::TensorArg pack_indices_arg{pack_indices, "pack_indices", 2};
at::checkDim(__func__, feats_arg, 2);
at::checkDim(__func__, pack_indices_arg, 1);
at::checkAllSameGPU(__func__, {feats_arg, pack_indices_arg});
at::checkAllContiguous(__func__, {feats_arg, pack_indices_arg});
at::checkScalarTypes(__func__, feats_arg, {at::kHalf, at::kFloat, at::kDouble});
at::checkScalarType(__func__, pack_indices_arg, at::kInt);

int num_feats = feats.size(0);
int feat_dim = feats.size(1);
at::Tensor feats_out = at::ones({num_feats, feat_dim}, feats.options());
cumprod_cuda_impl(num_feats, feat_dim, feats, feats_out, pack_indices, exclusive, reverse);
return feats_out.contiguous();
#else
KAOLIN_NO_CUDA_ERROR(__func__);
#endif // WITH_CUDA
}

} // namespace kaolin
23 changes: 23 additions & 0 deletions kaolin/csrc/render/spc/raytrace.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,29 @@ std::vector<at::Tensor> generate_shadow_rays_cuda(
at::Tensor light,
at::Tensor plane);

at::Tensor diff_cuda(
at::Tensor feats,
at::Tensor pack_indices);

at::Tensor inclusive_sum_cuda(
at::Tensor info);

at::Tensor sum_reduce_cuda(
at::Tensor feats,
at::Tensor inclusive_sum);

at::Tensor cumsum_cuda(
at::Tensor feats,
at::Tensor pack_indices,
bool exclusive,
bool reverse);

at::Tensor cumprod_cuda(
at::Tensor feats,
at::Tensor pack_indices,
bool exclusive,
bool reverse);

} // namespace kaolin

#endif // KAOLIN_OPS_RENDER_SPC_RAYTRACE_H_

0 comments on commit 3ade542

Please sign in to comment.