mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 11:24:00 +08:00
Ethan Weber
676beebb14
* some initial cuda code as I learn. using git to save progress. * first cuda kernel for ray sampling that mostly works * cleaning up code. working on lr. structs for scene bounds aabb * added scene bounds collider code * updated setup * setup * setup * black * including some files * exporting env variable * new file * git workflows * github actions * github actions * github actions * removing env var
30 lines
732 B
C++
30 lines
732 B
C++
/*
|
|
This file uses pybind to make CUDA calls.
|
|
*/
|
|
|
|
#include <torch/extension.h>
|
|
|
|
#include <vector>
|
|
|
|
// CUDA forward declarations
|
|
|
|
std::vector<torch::Tensor> sample_uniformly_along_ray_bundle(
|
|
torch::Tensor origins,
|
|
torch::Tensor directions,
|
|
torch::Tensor nears,
|
|
torch::Tensor fars,
|
|
torch::Tensor offsets,
|
|
int max_num_samples);
|
|
|
|
// C++ interface
|
|
|
|
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
|
|
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
|
#define CHECK_INPUT(x) \
|
|
CHECK_CUDA(x); \
|
|
CHECK_CONTIGUOUS(x)
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
m.def("sample_uniformly_along_ray_bundle", &sample_uniformly_along_ray_bundle);
|
|
} |