conquer3d 0.2.4__tar.gz → 0.2.6__tar.gz

{conquer3d-0.2.4/conquer3d.egg-info → conquer3d-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: conquer3d
-Version: 0.2.4
+Version: 0.2.6
 Summary: Geometric Cuda Tool Box
 Author-email: Do Hoang Khoi <khoido8899@gmail.com>
 License-Expression: MIT
@@ -68,17 +68,54 @@ twine upload dist/* --verbose
 ## Research Paper
 ```bibtex
 @inproceedings{2383795.2383801,
-author = {Karras, Tero},
-title = {Maximizing parallelism in the construction of BVHs, octrees, and k-d trees},
-year = {2012},
-isbn = {9783905674415},
-publisher = {Eurographics Association},
-address = {Goslar, DEU},
-booktitle = {Proceedings of the Fourth ACM SIGGRAPH / Eurographics Conference on High-Performance Graphics},
-pages = {33–37},
-numpages = {5},
-location = {Paris, France},
-series = {EGGH-HPG'12}
+    author = {Karras, Tero},
+    title = {Maximizing parallelism in the construction of BVHs, octrees, and k-d trees},
+    year = {2012},
+    isbn = {9783905674415},
+    publisher = {Eurographics Association},
+    address = {Goslar, DEU},
+    booktitle = {Proceedings of the Fourth ACM SIGGRAPH / Eurographics Conference on High-Performance Graphics},
+    pages = {33–37},
+    numpages = {5},
+    location = {Paris, France},
+    series = {EGGH-HPG'12}
+}
+```
+
+```bibtex
+@article{10.1080/10867651.1997.10487472,
+    author = {M\"{o}ller, Tomas},
+    title = {A fast triangle-triangle intersection test},
+    year = {1997},
+    issue_date = {1997},
+    publisher = {A. K. Peters, Ltd.},
+    address = {USA},
+    volume = {2},
+    number = {2},
+    issn = {1086-7651},
+    url = {https://doi.org/10.1080/10867651.1997.10487472},
+    doi = {10.1080/10867651.1997.10487472},
+    journal = {J. Graph. Tools},
+    month = nov,
+    pages = {25–30},
+    numpages = {6}
+}
+```
+
+```bibtex
+@inproceedings{10.1145/1198555.1198746,
+    author = {M\"{o}ller, Tomas and Trumbore, Ben},
+    title = {Fast, minimum storage ray/triangle intersection},
+    year = {2005},
+    isbn = {9781450378338},
+    publisher = {Association for Computing Machinery},
+    address = {New York, NY, USA},
+    url = {https://doi.org/10.1145/1198555.1198746},
+    doi = {10.1145/1198555.1198746},
+    pages = {7–es},
+    keywords = {base transformation, intersection, ray tracing, ray/triangle-intersection},
+    location = {Los Angeles, California},
+    series = {SIGGRAPH '05}
 }
 ```
 

{conquer3d-0.2.4 → conquer3d-0.2.6}/README.md

@@ -44,17 +44,54 @@ twine upload dist/* --verbose
 ## Research Paper
 ```bibtex
 @inproceedings{2383795.2383801,
-author = {Karras, Tero},
-title = {Maximizing parallelism in the construction of BVHs, octrees, and k-d trees},
-year = {2012},
-isbn = {9783905674415},
-publisher = {Eurographics Association},
-address = {Goslar, DEU},
-booktitle = {Proceedings of the Fourth ACM SIGGRAPH / Eurographics Conference on High-Performance Graphics},
-pages = {33–37},
-numpages = {5},
-location = {Paris, France},
-series = {EGGH-HPG'12}
+    author = {Karras, Tero},
+    title = {Maximizing parallelism in the construction of BVHs, octrees, and k-d trees},
+    year = {2012},
+    isbn = {9783905674415},
+    publisher = {Eurographics Association},
+    address = {Goslar, DEU},
+    booktitle = {Proceedings of the Fourth ACM SIGGRAPH / Eurographics Conference on High-Performance Graphics},
+    pages = {33–37},
+    numpages = {5},
+    location = {Paris, France},
+    series = {EGGH-HPG'12}
+}
+```
+
+```bibtex
+@article{10.1080/10867651.1997.10487472,
+    author = {M\"{o}ller, Tomas},
+    title = {A fast triangle-triangle intersection test},
+    year = {1997},
+    issue_date = {1997},
+    publisher = {A. K. Peters, Ltd.},
+    address = {USA},
+    volume = {2},
+    number = {2},
+    issn = {1086-7651},
+    url = {https://doi.org/10.1080/10867651.1997.10487472},
+    doi = {10.1080/10867651.1997.10487472},
+    journal = {J. Graph. Tools},
+    month = nov,
+    pages = {25–30},
+    numpages = {6}
+}
+```
+
+```bibtex
+@inproceedings{10.1145/1198555.1198746,
+    author = {M\"{o}ller, Tomas and Trumbore, Ben},
+    title = {Fast, minimum storage ray/triangle intersection},
+    year = {2005},
+    isbn = {9781450378338},
+    publisher = {Association for Computing Machinery},
+    address = {New York, NY, USA},
+    url = {https://doi.org/10.1145/1198555.1198746},
+    doi = {10.1145/1198555.1198746},
+    pages = {7–es},
+    keywords = {base transformation, intersection, ray tracing, ray/triangle-intersection},
+    location = {Los Angeles, California},
+    series = {SIGGRAPH '05}
 }
 ```
 

{conquer3d-0.2.4 → conquer3d-0.2.6}/conquer3d/_C.pyi

@@ -14,6 +14,14 @@ class BVH:
         """
         Query the BVH with bounding boxes or segments. Returns (query_ids, object_ids).
         """
+    def query_point(self, query_points: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Find the closest leaf AABB to each query point. Returns (query_ids, object_ids, distances)
+        """
+    def query_ray(self, ray_origins: torch.Tensor, ray_dirs: torch.Tensor, max_capacity: int = 10000000) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Find all ray-AABB intersections. Returns (ray_ids, object_ids)
+        """
     def query_self(self) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Query the BVH against itself. Returns unique overlapping (query_ids, object_ids).
@@ -49,6 +57,10 @@ class MeshBVH(BVH):
         """
         Construct MeshBVH from Triangle AABBs.
         """
+    def get_ray_intersection(self, ray_origins: torch.Tensor, ray_dirs: torch.Tensor, vertices: torch.Tensor, triangles: torch.Tensor, return_distance: bool = False) -> typing.Any:
+        """
+        Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])
+        """
     def get_self_intersection(self, vertices: torch.Tensor, triangles: torch.Tensor) -> torch.Tensor:
         """
         Find all self-intersecting triangle pairs
@@ -57,6 +69,10 @@ class MeshBVH(BVH):
         """
         Check if there are any self-intersecting triangle pairs
         """
+    def query_point(self, query_points: torch.Tensor, vertices: torch.Tensor, triangles: torch.Tensor, return_sdf: bool = False, return_prj_pts: bool = True, sign_mode: int = 0) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Find the closest triangle to each query point. Returns (query_ids, triangle_ids, projected_points, distances)
+        """
 class PGSBVH(BVH):
     def __init__(self, in_aabb_mins: torch.Tensor, in_aabb_maxs: torch.Tensor) -> None:
         """
@@ -87,8 +103,14 @@ class TriangleMesh:
         ...
     def get_non_manifold_vertices(self) -> torch.Tensor:
         ...
+    def get_ray_intersection(self, ray_origins: torch.Tensor, ray_dirs: torch.Tensor, return_distance: bool = False) -> typing.Any:
+        """
+        Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])
+        """
     def get_self_intersection(self) -> torch.Tensor:
-        ...
+        """
+        Find all self-intersecting triangle pairs
+        """
     def get_vertices_to_triangle_map(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         ...
     def is_edge_manifold(self, allow_boundary_edge: bool = True) -> bool:
@@ -96,9 +118,15 @@ class TriangleMesh:
     def is_manifold(self, allow_boundary_edge: bool = True) -> bool:
         ...
     def is_self_intersection(self) -> bool:
-        ...
+        """
+        Check if there are any self-intersecting triangle pairs
+        """
     def is_vertex_manifold(self) -> bool:
         ...
+    def query_points(self, query_pts: torch.Tensor, return_sdf: bool = False, return_prj_pts: bool = True, sign_mode: int = 0, distance_mode: int = 0) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Query points against the mesh. Returns (query_ids, triangle_ids, projected_points, distances)
+        """
     def remove_triangles_by_mask(self, keep_mask: torch.Tensor) -> None:
         ...
     @property

{conquer3d-0.2.4 → conquer3d-0.2.6}/conquer3d/csrc/binds/data_structure/bvh.cpp

@@ -153,6 +153,85 @@ std::tuple<torch::Tensor, torch::Tensor> BVH::query_self()
     );
 }
 
+std::tuple<torch::Tensor, torch::Tensor> BVH::query_ray(
+    const torch::Tensor &ray_origins,
+    const torch::Tensor &ray_dirs,
+    int64_t max_capacity)
+{
+    CHECK_INPUT(ray_origins);
+    CHECK_INPUT(ray_dirs);
+    TORCH_CHECK(ray_origins.size(0) == ray_dirs.size(0), "ray_origins and ray_dirs must have the same number of rays");
+
+    int num_queries = ray_origins.size(0);
+
+    auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(ray_origins.device());
+    torch::Tensor out_query_ids = torch::empty({max_capacity}, options_i64);
+    torch::Tensor out_object_ids = torch::empty({max_capacity}, options_i64);
+    torch::Tensor hit_counter = torch::zeros({1}, options_i64);
+
+    bvh::query_ray(
+        num_queries,
+        this->num_objects,
+        (const float3 *)ray_origins.data_ptr<float>(),
+        (const float3 *)ray_dirs.data_ptr<float>(),
+        (const float3 *)this->aabb_mins.data_ptr<float>(),
+        (const float3 *)this->aabb_maxs.data_ptr<float>(),
+        (const int2 *)this->bvh_children.data_ptr<int>(),
+        this->object_ids.data_ptr<int>(),
+        out_query_ids.data_ptr<int64_t>(),
+        out_object_ids.data_ptr<int64_t>(),
+        hit_counter.data_ptr<int64_t>(),
+        max_capacity);
+
+    int64_t h_hit_counter = hit_counter.item<int64_t>();
+
+    if (h_hit_counter >= max_capacity)
+    {
+        TORCH_WARN("BVH ray query capacity exceeded. Some hits were dropped.");
+        h_hit_counter = max_capacity;
+    }
+
+    if (h_hit_counter == 0) {
+        return std::make_tuple(
+            torch::empty({0}, options_i64),
+            torch::empty({0}, options_i64)
+        );
+    }
+
+    return std::make_tuple(
+        out_query_ids.slice(0, 0, h_hit_counter),
+        out_object_ids.slice(0, 0, h_hit_counter));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> BVH::query_point(
+    const torch::Tensor &query_points)
+{
+    CHECK_INPUT(query_points);
+
+    int num_queries = query_points.size(0);
+    
+    auto options_i64 = torch::TensorOptions().dtype(torch::kInt64).device(query_points.device());
+    auto options_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(query_points.device());
+    
+    torch::Tensor out_query_ids = torch::empty({num_queries}, options_i64);
+    torch::Tensor out_object_ids = torch::empty({num_queries}, options_i64);
+    torch::Tensor out_distances = torch::empty({num_queries}, options_f32);
+
+    bvh::query_point(
+        num_queries,
+        this->num_objects,
+        (const float3 *)query_points.data_ptr<float>(),
+        (const float3 *)this->aabb_mins.data_ptr<float>(),
+        (const float3 *)this->aabb_maxs.data_ptr<float>(),
+        (const int2 *)this->bvh_children.data_ptr<int>(),
+        this->object_ids.data_ptr<int>(),
+        out_query_ids.data_ptr<int64_t>(),
+        out_object_ids.data_ptr<int64_t>(),
+        out_distances.data_ptr<float>());
+
+    return std::make_tuple(out_query_ids, out_object_ids, out_distances);
+}
+
 void bind_ds_bvh(py::module_& m)
 {
     py::class_<BVH>(m, "BVH")
@@ -167,5 +246,14 @@ void bind_ds_bvh(py::module_& m)
              "Query the BVH with bounding boxes or segments. Returns (query_ids, object_ids).")
              
         .def("query_self", &BVH::query_self,
-             "Query the BVH against itself. Returns unique overlapping (query_ids, object_ids).");
+             "Query the BVH against itself. Returns unique overlapping (query_ids, object_ids).")
+
+        .def("query_ray", &BVH::query_ray,
+             py::arg("ray_origins"),
+             py::arg("ray_dirs"),
+             py::arg("max_capacity") = BVH_MAX_CAPACITY,
+             "Find all ray-AABB intersections. Returns (ray_ids, object_ids)")
+        .def("query_point", &BVH::query_point,
+             py::arg("query_points"),
+             "Find the closest leaf AABB to each query point. Returns (query_ids, object_ids, distances)");
 }

conquer3d-0.2.6/conquer3d/csrc/binds/data_structure/mesh_bvh.cpp

@@ -0,0 +1,44 @@
+#include <torch/extension.h>
+#include "../../data_structure/mesh_bvh.h"
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+void bind_ds_mesh_bvh(py::module_& m)
+{
+    py::class_<MeshBVH, BVH>(m, "MeshBVH")
+        .def(py::init<const torch::Tensor &, const torch::Tensor &>(),
+             py::arg("in_aabb_mins"), 
+             py::arg("in_aabb_maxs"),
+             "Construct MeshBVH from Triangle AABBs.")
+        .def("get_self_intersection", &MeshBVH::get_self_intersection,
+             py::arg("vertices"), 
+             py::arg("triangles"),
+             "Find all self-intersecting triangle pairs")
+        .def("is_self_intersection", &MeshBVH::is_self_intersection,
+             py::arg("vertices"), 
+             py::arg("triangles"),
+             "Check if there are any self-intersecting triangle pairs")
+        .def("get_ray_intersection", [](MeshBVH& self, const torch::Tensor& ray_origins, const torch::Tensor& ray_dirs, const torch::Tensor& vertices, const torch::Tensor& triangles, bool return_distance) -> py::object {
+                 auto result = self.get_ray_intersection(ray_origins, ray_dirs, vertices, triangles, return_distance);
+                 if (return_distance) {
+                     return py::cast(result);
+                 } else {
+                     return py::cast(std::make_tuple(std::get<0>(result), std::get<1>(result), std::get<2>(result)));
+                 }
+             },
+             py::arg("ray_origins"),
+             py::arg("ray_dirs"),
+             py::arg("vertices"),
+             py::arg("triangles"),
+             py::arg("return_distance") = false,
+             "Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])")
+        .def("query_point", &MeshBVH::query_point,
+             py::arg("query_points"),
+             py::arg("vertices"),
+             py::arg("triangles"),
+             py::arg("return_sdf") = false,
+             py::arg("return_prj_pts") = true,
+             py::arg("sign_mode") = 0,
+             "Find the closest triangle to each query point. Returns (query_ids, triangle_ids, projected_points, distances)");
+}

{conquer3d-0.2.4 → conquer3d-0.2.6}/conquer3d/csrc/binds/data_structure/triangle_mesh.cpp

@@ -97,6 +97,30 @@ bool TriangleMesh::is_self_intersection()
     return this->bvh.value().is_self_intersection(this->vertices, this->triangles);
 }
 
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> TriangleMesh::query_points(
+    const torch::Tensor &query_pts,
+    bool return_sdf,
+    bool return_prj_pts,
+    int sign_mode,
+    int distance_mode)
+{
+    if (distance_mode == 0) {
+        this->build_bvh();
+        return this->bvh.value().query_point(query_pts, this->vertices, this->triangles, return_sdf, return_prj_pts, sign_mode);
+    } else {
+        throw std::runtime_error("distance_mode != 0 not implemented yet");
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> TriangleMesh::get_ray_intersection(
+    const torch::Tensor &ray_origins,
+    const torch::Tensor &ray_dirs,
+    bool return_distance)
+{
+    this->build_bvh();
+    return this->bvh.value().get_ray_intersection(ray_origins, ray_dirs, this->vertices, this->triangles, return_distance);
+}
+
 torch::Tensor TriangleMesh::get_triangle_areas()
 {
     if (!this->triangle_areas.defined())
@@ -313,8 +337,29 @@ void bind_ds_triangle_mesh(py::module_ &m)
         .def_property_readonly("surface_area", &TriangleMesh::get_surface_area)
         .def_property_readonly("bvh", &TriangleMesh::build_bvh)
         .def("build_bvh", &TriangleMesh::build_bvh)
-        .def("get_self_intersection", &TriangleMesh::get_self_intersection)
-        .def("is_self_intersection", &TriangleMesh::is_self_intersection)
+        .def("get_self_intersection", &TriangleMesh::get_self_intersection,
+             "Find all self-intersecting triangle pairs")
+        .def("is_self_intersection", &TriangleMesh::is_self_intersection,
+             "Check if there are any self-intersecting triangle pairs")
+        .def("get_ray_intersection", [](TriangleMesh& self, const torch::Tensor& ray_origins, const torch::Tensor& ray_dirs, bool return_distance) -> py::object {
+                 auto result = self.get_ray_intersection(ray_origins, ray_dirs, return_distance);
+                 if (return_distance) {
+                     return py::cast(result);
+                 } else {
+                     return py::cast(std::make_tuple(std::get<0>(result), std::get<1>(result), std::get<2>(result)));
+                 }
+             },
+             py::arg("ray_origins"),
+             py::arg("ray_dirs"),
+             py::arg("return_distance") = false,
+             "Find all ray-triangle intersections. Returns (ray_ids, triangle_ids, intersect_points, [distances])")
+        .def("query_points", &TriangleMesh::query_points,
+             py::arg("query_pts"),
+             py::arg("return_sdf") = false,
+             py::arg("return_prj_pts") = true,
+             py::arg("sign_mode") = 0,
+             py::arg("distance_mode") = 0,
+             "Query points against the mesh. Returns (query_ids, triangle_ids, projected_points, distances)")
         .def_property_readonly("edges", &TriangleMesh::get_edges)
         .def_property_readonly("edge_to_triangle_offsets", &TriangleMesh::get_edge_to_triangle_offsets)
         .def_property_readonly("edge_to_triangle_counts", &TriangleMesh::get_edge_to_triangle_counts)

{conquer3d-0.2.4 → conquer3d-0.2.6}/conquer3d/csrc/data_structure/bvh.cu

@@ -1,5 +1,6 @@
 #include "bvh.h"
 #include "zcurve.h"
+#include "../primitive/ray.h"
 
 #include <thrust/device_vector.h>
 #include <thrust/device_ptr.h>
@@ -453,4 +454,185 @@ namespace bvh
             max_capacity
         );
     }
+
+    __global__ void query_ray_bvh_kernel(
+        const uint32_t num_queries,
+        const uint32_t num_objects,
+        const float3* __restrict__ ray_origins,
+        const float3* __restrict__ ray_dirs,
+        const float3* __restrict__ bvh_aabb_mins,
+        const float3* __restrict__ bvh_aabb_maxs,
+        const int2* __restrict__ bvh_children,
+        const int* __restrict__ object_ids,
+        int64_t* __restrict__ out_query_ids,
+        int64_t* __restrict__ out_object_ids,
+        int64_t* __restrict__ hit_counter,
+        const int64_t max_capacity)
+    {
+        uint32_t q_idx = blockIdx.x * blockDim.x + threadIdx.x;
+        if (q_idx >= num_queries) return;
+
+        Ray ray(ray_origins[q_idx], ray_dirs[q_idx]);
+
+        int stack[BVH_STACK_SIZE];
+        int stack_ptr = 0;
+        stack[0] = 0; 
+
+        while (stack_ptr >= 0) 
+        {
+            int node_idx = stack[stack_ptr--]; 
+            float t_hit;
+
+            if (ray.is_intersect_aabb(bvh_aabb_mins[node_idx], bvh_aabb_maxs[node_idx], t_hit)) 
+            {
+                if (node_idx >= num_objects - 1) 
+                {
+                    int leaf_idx = node_idx - (num_objects - 1);
+                    int original_obj_id = object_ids[leaf_idx];
+
+                    uint64_t write_idx = (uint64_t)atomicAdd((unsigned long long int*)hit_counter, 1ULL);
+
+                    if (write_idx < max_capacity) {
+                        out_query_ids[write_idx] = q_idx;
+                        out_object_ids[write_idx] = original_obj_id;
+                    }
+                } 
+                else 
+                {
+                    if (stack_ptr + 2 < BVH_STACK_SIZE)
+                    {
+                        int2 children = bvh_children[node_idx];
+                        stack[++stack_ptr] = children.x;
+                        stack[++stack_ptr] = children.y;
+                    }
+                }
+            }
+        }
+    }
+
+    void query_ray(
+        const uint32_t num_queries,
+        const uint32_t num_objects,
+        const float3* __restrict__ ray_origins,
+        const float3* __restrict__ ray_dirs,
+        const float3* __restrict__ bvh_aabb_mins,
+        const float3* __restrict__ bvh_aabb_maxs,
+        const int2* __restrict__ bvh_children,
+        const int* __restrict__ object_ids,
+        int64_t* __restrict__ out_query_ids,
+        int64_t* __restrict__ out_object_ids,
+        int64_t* __restrict__ hit_counter,
+        const int64_t max_capacity)
+    {
+        uint32_t threads = NTHREADS;
+        uint32_t blocks = (num_queries + threads - 1) / threads;
+
+        query_ray_bvh_kernel<<<blocks, threads>>>(
+            num_queries,
+            num_objects,
+            ray_origins,
+            ray_dirs,
+            bvh_aabb_mins,
+            bvh_aabb_maxs,
+            bvh_children,
+            object_ids,
+            out_query_ids,
+            out_object_ids,
+            hit_counter,
+            max_capacity
+        );
+    }
+
+    __global__ void query_point_bvh_kernel(
+        const uint32_t num_queries,
+        const uint32_t num_objects,
+        const float3* __restrict__ query_points,
+        const float3* __restrict__ bvh_aabb_mins,
+        const float3* __restrict__ bvh_aabb_maxs,
+        const int2* __restrict__ bvh_children,
+        const int* __restrict__ object_ids,
+        int64_t* __restrict__ out_query_ids,
+        int64_t* __restrict__ out_object_ids,
+        float* __restrict__ out_distances)
+    {
+        uint32_t q_idx = blockIdx.x * blockDim.x + threadIdx.x;
+        if (q_idx >= num_queries) return;
+
+        float3 p = query_points[q_idx];
+
+        int stack[BVH_STACK_SIZE];
+        int stack_ptr = 0;
+        stack[0] = 0; 
+        
+        float best_dist_sq = FLT_MAX;
+        int best_leaf_original_id = -1;
+
+        while (stack_ptr >= 0) 
+        {
+            int node_idx = stack[stack_ptr--]; 
+
+            float dist_sq = aabb::compute_squared_distance(p, bvh_aabb_mins[node_idx], bvh_aabb_maxs[node_idx]);
+            
+            if (dist_sq > best_dist_sq) continue; // Prune branch
+
+            if (node_idx >= num_objects - 1) 
+            {
+                if (dist_sq < best_dist_sq) {
+                    best_dist_sq = dist_sq;
+                    best_leaf_original_id = object_ids[node_idx - (num_objects - 1)];
+                }
+            } 
+            else 
+            {
+                if (stack_ptr + 2 < BVH_STACK_SIZE)
+                {
+                    int2 children = bvh_children[node_idx];
+                    
+                    float dist_l = aabb::compute_squared_distance(p, bvh_aabb_mins[children.x], bvh_aabb_maxs[children.x]);
+                    float dist_r = aabb::compute_squared_distance(p, bvh_aabb_mins[children.y], bvh_aabb_maxs[children.y]);
+                    
+                    if (dist_l > dist_r) {
+                        stack[++stack_ptr] = children.x; 
+                        stack[++stack_ptr] = children.y;
+                    } else {
+                        stack[++stack_ptr] = children.y;
+                        stack[++stack_ptr] = children.x;
+                    }
+                }
+            }
+        }
+        
+        out_query_ids[q_idx] = q_idx;
+        out_object_ids[q_idx] = best_leaf_original_id;
+        out_distances[q_idx] = sqrtf(best_dist_sq);
+    }
+
+    void query_point(
+        const uint32_t num_queries,
+        const uint32_t num_objects,
+        const float3* __restrict__ query_points,
+        const float3* __restrict__ bvh_aabb_mins,
+        const float3* __restrict__ bvh_aabb_maxs,
+        const int2* __restrict__ bvh_children,
+        const int* __restrict__ object_ids,
+        int64_t* __restrict__ out_query_ids,
+        int64_t* __restrict__ out_object_ids,
+        float* __restrict__ out_distances)
+    {
+        uint32_t threads = NTHREADS;
+        uint32_t blocks = (num_queries + threads - 1) / threads;
+
+        query_point_bvh_kernel<<<blocks, threads>>>(
+            num_queries,
+            num_objects,
+            query_points,
+            bvh_aabb_mins,
+            bvh_aabb_maxs,
+            bvh_children,
+            object_ids,
+            out_query_ids,
+            out_object_ids,
+            out_distances
+        );
+    }
 }

{conquer3d-0.2.4 → conquer3d-0.2.6}/conquer3d/csrc/data_structure/bvh.h

@@ -4,6 +4,7 @@
 #include "../maths/maths.h"
 #include "../constants.h"
 #include "../primitive/aabb.h"
+#include "../primitive/ray.h"
 
 #include <torch/extension.h>
 #include <cuda_runtime.h>
@@ -29,6 +30,14 @@ public:
         const torch::Tensor &query_aabb_maxs);
 
     std::tuple<torch::Tensor, torch::Tensor> query_self();
+
+    std::tuple<torch::Tensor, torch::Tensor> query_ray(
+        const torch::Tensor &ray_origins,
+        const torch::Tensor &ray_dirs,
+        int64_t max_capacity = BVH_MAX_CAPACITY);
+
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> query_point(
+        const torch::Tensor &query_points);
 };
 
 namespace bvh
@@ -69,6 +78,32 @@ namespace bvh
         int64_t *__restrict__ out_object_ids,
         int64_t *__restrict__ hit_counter,
         const int64_t max_capacity);
+
+    __host__ void query_ray(
+        const uint32_t num_queries,
+        const uint32_t num_objects,
+        const float3 *__restrict__ ray_origins,
+        const float3 *__restrict__ ray_dirs,
+        const float3 *__restrict__ bvh_aabb_mins,
+        const float3 *__restrict__ bvh_aabb_maxs,
+        const int2 *__restrict__ bvh_children,
+        const int *__restrict__ object_ids,
+        int64_t *__restrict__ out_query_ids,
+        int64_t *__restrict__ out_object_ids,
+        int64_t *__restrict__ hit_counter,
+        const int64_t max_capacity);
+
+    __host__ void query_point(
+        const uint32_t num_queries,
+        const uint32_t num_objects,
+        const float3 *__restrict__ query_points,
+        const float3 *__restrict__ bvh_aabb_mins,
+        const float3 *__restrict__ bvh_aabb_maxs,
+        const int2 *__restrict__ bvh_children,
+        const int *__restrict__ object_ids,
+        int64_t *__restrict__ out_query_ids,
+        int64_t *__restrict__ out_object_ids,
+        float *__restrict__ out_distances);
 }
 
 #endif // BVH_H