File size: 2,147 Bytes
6b92ff7 | 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 | #pragma once
#include <Eigen/Dense>
#include <ATen/cuda/CUDAContext.h>
#include <torch/torch.h>
#include <cuda_runtime.h>
#include <memory>
#include <tuple>
using namespace Eigen;
using Verts = Matrix<float, Dynamic, 3, RowMajor>;
using Trigs = Matrix<uint32_t, Dynamic, 3, RowMajor>;
namespace cubvh {
// abstract class of raytracer
class cuBVH {
public:
cuBVH() {}
virtual ~cuBVH() {}
virtual void ray_trace(at::Tensor rays_o, at::Tensor rays_d, at::Tensor positions, at::Tensor face_id, at::Tensor depth) = 0;
virtual void unsigned_distance(at::Tensor positions, at::Tensor distances, at::Tensor face_id, at::optional<at::Tensor> uvw) = 0;
virtual void signed_distance(at::Tensor positions, at::Tensor distances, at::Tensor face_id, at::optional<at::Tensor> uvw, uint32_t mode) = 0;
virtual std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> export_state() const = 0;
};
// function to create an implementation of cuBVH
cuBVH* create_cuBVH(Ref<const Verts> vertices, Ref<const Trigs> triangles);
cuBVH* create_cuBVH_from_state(at::Tensor triangle_vertices, at::Tensor triangle_ids, at::Tensor node_mins, at::Tensor node_maxs, at::Tensor node_children);
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> build_cuBVH_state(Ref<const Verts> vertices, Ref<const Trigs> triangles);
// floodfill
at::Tensor floodfill(at::Tensor grid);
// sparse marching cubes
std::tuple<at::Tensor, at::Tensor> sparse_marching_cubes(at::Tensor coords, at::Tensor corners, double iso_d, bool ensure_consistency = false);
// GPU ND integer hash table (default D=3) - abstract interface
class cuHashTable {
public:
cuHashTable() {}
virtual ~cuHashTable() {}
virtual void set_num_dims(int d) = 0;
virtual int get_num_dims() const = 0;
virtual void resize(int capacity) = 0;
virtual void prepare() = 0;
virtual void insert(at::Tensor coords) = 0;
virtual void build(at::Tensor coords) = 0;
virtual at::Tensor search(at::Tensor queries) const = 0;
};
// Factory to create an implementation of cuHashTable
cuHashTable* create_cuHashTable();
} // namespace cubvh |