init

LICENSE

+MIT License
+
+Copyright (c) 2022 kiui
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LICENSE_NVIDIA

+Copyright (c) 2022, NVIDIA Corporation & affiliates. All rights reserved.
+
+
+NVIDIA Source Code License for instant neural graphics primitives
+
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. Notwithstanding
+    the foregoing, NVIDIA and its affiliates may use the Work and any
+    derivative works commercially. As used herein, "non-commercially"
+    means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor’s or its affiliates’ names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

cubvh/__init__.py

+from .api import cuBVH
\ No newline at end of file

cubvh/api.py

+
+import numpy as np
+import torch
+
+# CUDA extension
+import _cubvh as _backend
+
+class cuBVH():
+    def __init__(self, vertices, triangles):
+        # vertices: np.ndarray, [N, 3]
+        # triangles: np.ndarray, [M, 3]
+
+        if torch.is_tensor(vertices): vertices = vertices.detach().cpu().numpy()
+        if torch.is_tensor(triangles): triangles = triangles.detach().cpu().numpy()
+
+        assert triangles.shape[0] > 8, "BVH needs at least 8 triangles."
+        
+        # implementation
+        self.impl = _backend.create_cuBVH(vertices, triangles)
+
+    def ray_trace(self, rays_o, rays_d):
+        # rays_o: torch.Tensor, float, [N, 3]
+        # rays_d: torch.Tensor, float, [N, 3]
+
+        rays_o = rays_o.float().contiguous()
+        rays_d = rays_d.float().contiguous()
+
+        if not rays_o.is_cuda: rays_o = rays_o.cuda()
+        if not rays_d.is_cuda: rays_d = rays_d.cuda()
+
+        prefix = rays_o.shape[:-1]
+        rays_o = rays_o.view(-1, 3)
+        rays_d = rays_d.view(-1, 3)
+
+        N = rays_o.shape[0]
+
+        # init output buffers
+        positions = torch.empty(N, 3, dtype=torch.float32, device=rays_o.device)
+        face_id = torch.empty(N, dtype=torch.int64, device=rays_o.device)
+        depth = torch.empty(N, dtype=torch.float32, device=rays_o.device)
+        
+        self.impl.ray_trace(rays_o, rays_d, positions, face_id, depth) # [N, 3]
+
+        positions = positions.view(*prefix, 3)
+        face_id = face_id.view(*prefix)
+        depth = depth.view(*prefix)
+
+        return positions, face_id, depth
+
+    def unsigned_distance(self, positions, return_uvw=False):
+        # positions: torch.Tensor, float, [N, 3]
+
+        positions = positions.float().contiguous()
+
+        if not positions.is_cuda: positions = positions.cuda()
+
+        prefix = positions.shape[:-1]
+        positions = positions.view(-1, 3)
+
+        N = positions.shape[0]
+
+        # init output buffers
+        distances = torch.empty(N, dtype=torch.float32, device=positions.device)
+        face_id = torch.empty(N, dtype=torch.int64, device=positions.device)
+
+        if return_uvw:
+            uvw = torch.empty(N, 3, dtype=torch.float32, device=positions.device)
+        else:
+            uvw = None
+        
+        self.impl.unsigned_distance(positions, distances, face_id, uvw) # [N, 3]
+
+        distances = distances.view(*prefix)
+        face_id = face_id.view(*prefix)
+        if uvw is not None:
+            uvw = uvw.view(*prefix, 3)
+
+        return distances, face_id, uvw
+
+    
+    # def signed_distance():

include/cubvh/api.h

+#pragma once
+
+#include <Eigen/Dense>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/torch.h>
+#include <cuda_runtime.h>
+
+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;
+};
+
+// function to create an implementation of cuBVH
+cuBVH* create_cuBVH(Ref<const Verts> vertices, Ref<const Trigs> triangles);
+    
+} // namespace cubvh
\ No newline at end of file

include/cubvh/bounding_box.cuh

+#pragma once
+
+#include <cubvh/common.h>
+#include <cubvh/triangle.cuh>
+
+namespace cubvh {
+
+template <int N_POINTS>
+__host__ __device__ inline void project(Eigen::Vector3f points[N_POINTS], const Eigen::Vector3f& axis, float& min, float& max) {
+    min = std::numeric_limits<float>::infinity();
+    max = -std::numeric_limits<float>::infinity();
+
+    #pragma unroll
+    for (uint32_t i = 0; i < N_POINTS; ++i) {
+        float val = axis.dot(points[i]);
+
+        if (val < min) {
+            min = val;
+        }
+
+        if (val > max) {
+            max = val;
+        }
+    }
+}
+
+struct BoundingBox {
+
+    Eigen::Vector3f min = Eigen::Vector3f::Constant(std::numeric_limits<float>::infinity());
+    Eigen::Vector3f max = Eigen::Vector3f::Constant(-std::numeric_limits<float>::infinity());
+
+    __host__ __device__ BoundingBox() {}
+
+    __host__ __device__ BoundingBox(const Eigen::Vector3f& a, const Eigen::Vector3f& b) : min{a}, max{b} {}
+
+    __host__ __device__ explicit BoundingBox(const Triangle& tri) {
+        min = max = tri.a;
+        enlarge(tri.b);
+        enlarge(tri.c);
+    }
+
+    BoundingBox(std::vector<Triangle>::iterator begin, std::vector<Triangle>::iterator end) {
+        min = max = begin->a;
+        for (auto it = begin; it != end; ++it) {
+            enlarge(*it);
+        }
+    }
+
+    __host__ __device__ void enlarge(const BoundingBox& other) {
+        min = min.cwiseMin(other.min);
+        max = max.cwiseMax(other.max);
+    }
+
+    __host__ __device__ void enlarge(const Triangle& tri) {
+        enlarge(tri.a);
+        enlarge(tri.b);
+        enlarge(tri.c);
+    }
+
+    __host__ __device__ void enlarge(const Eigen::Vector3f& point) {
+        min = min.cwiseMin(point);
+        max = max.cwiseMax(point);
+    }
+
+    __host__ __device__ void inflate(float amount) {
+        min -= Eigen::Vector3f::Constant(amount);
+        max += Eigen::Vector3f::Constant(amount);
+    }
+
+    __host__ __device__ Eigen::Vector3f diag() const {
+        return max - min;
+    }
+
+    __host__ __device__ Eigen::Vector3f relative_pos(const Eigen::Vector3f& pos) const {
+        return (pos - min).cwiseQuotient(diag());
+    }
+
+    __host__ __device__ Eigen::Vector3f center() const {
+        return 0.5f * (max + min);
+    }
+
+    __host__ __device__ BoundingBox intersection(const BoundingBox& other) const {
+        BoundingBox result = *this;
+        result.min = result.min.cwiseMax(other.min);
+        result.max = result.max.cwiseMin(other.max);
+        return result;
+    }
+
+    __host__ __device__ bool intersects(const BoundingBox& other) const {
+        return !intersection(other).is_empty();
+    }
+
+    // Based on the separating axis theorem
+    // (https://fileadmin.cs.lth.se/cs/Personal/Tomas_Akenine-Moller/code/tribox_tam.pdf)
+    // Code adapted from a C# implementation at stack overflow
+    // https://stackoverflow.com/a/17503268
+    __host__ __device__ bool intersects(const Triangle& triangle) const {
+        float triangle_min, triangle_max;
+        float box_min, box_max;
+
+        // Test the box normals (x-, y- and z-axes)
+        Eigen::Vector3f box_normals[3] = {
+            Eigen::Vector3f{1.0f, 0.0f, 0.0f},
+            Eigen::Vector3f{0.0f, 1.0f, 0.0f},
+            Eigen::Vector3f{0.0f, 0.0f, 1.0f},
+        };
+
+        Eigen::Vector3f triangle_normal = triangle.normal();
+        Eigen::Vector3f triangle_verts[3];
+        triangle.get_vertices(triangle_verts);
+
+        for (int i = 0; i < 3; i++) {
+            project<3>(triangle_verts, box_normals[i], triangle_min, triangle_max);
+            if (triangle_max < min[i] || triangle_min > max[i]) {
+                return false; // No intersection possible.
+            }
+        }
+
+        Eigen::Vector3f verts[8];
+        get_vertices(verts);
+
+        // Test the triangle normal
+        float triangle_offset = triangle_normal.dot(triangle.a);
+        project<8>(verts, triangle_normal, box_min, box_max);
+        if (box_max < triangle_offset || box_min > triangle_offset) {
+            return false; // No intersection possible.
+        }
+
+        // Test the nine edge cross-products
+        Eigen::Vector3f edges[3] = {
+            triangle.a - triangle.b,
+            triangle.a - triangle.c,
+            triangle.b - triangle.c,
+        };
+
+        for (int i = 0; i < 3; i++) {
+            for (int j = 0; j < 3; j++) {
+                // The box normals are the same as it's edge tangents
+                Eigen::Vector3f axis = edges[i].cross(box_normals[j]);
+                project<8>(verts, axis, box_min, box_max);
+                project<3>(triangle_verts, axis, triangle_min, triangle_max);
+                if (box_max < triangle_min || box_min > triangle_max)
+                    return false; // No intersection possible
+            }
+        }
+
+        // No separating axis found.
+        return true;
+    }
+
+    __host__ __device__ Eigen::Vector2f ray_intersect(Eigen::Ref<const Eigen::Vector3f> pos, Eigen::Ref<const Eigen::Vector3f> dir) const {
+        float tmin = (min.x() - pos.x()) / dir.x();
+        float tmax = (max.x() - pos.x()) / dir.x();
+
+        if (tmin > tmax) {
+            host_device_swap(tmin, tmax);
+        }
+
+        float tymin = (min.y() - pos.y()) / dir.y();
+        float tymax = (max.y() - pos.y()) / dir.y();
+
+        if (tymin > tymax) {
+            host_device_swap(tymin, tymax);
+        }
+
+        if (tmin > tymax || tymin > tmax) {
+            return { std::numeric_limits<float>::max(), std::numeric_limits<float>::max() };
+        }
+
+        if (tymin > tmin) {
+            tmin = tymin;
+        }
+
+        if (tymax < tmax) {
+            tmax = tymax;
+        }
+
+        float tzmin = (min.z() - pos.z()) / dir.z();
+        float tzmax = (max.z() - pos.z()) / dir.z();
+
+        if (tzmin > tzmax) {
+            host_device_swap(tzmin, tzmax);
+        }
+
+        if (tmin > tzmax || tzmin > tmax) {
+            return { std::numeric_limits<float>::max(), std::numeric_limits<float>::max() };
+        }
+
+        if (tzmin > tmin) {
+            tmin = tzmin;
+        }
+
+        if (tzmax < tmax) {
+            tmax = tzmax;
+        }
+
+        return { tmin, tmax };
+    }
+
+    __host__ __device__ bool is_empty() const {
+        return (max.array() < min.array()).any();
+    }
+
+    __host__ __device__ bool contains(const Eigen::Vector3f& p) const {
+        return
+            p.x() >= min.x() && p.x() <= max.x() &&
+            p.y() >= min.y() && p.y() <= max.y() &&
+            p.z() >= min.z() && p.z() <= max.z();
+    }
+
+    /// Calculate the squared point-AABB distance
+    __host__ __device__ float distance(const Eigen::Vector3f& p) const {
+        return sqrt(distance_sq(p));
+    }
+
+    __host__ __device__ float distance_sq(const Eigen::Vector3f& p) const {
+        return (min - p).cwiseMax(p - max).cwiseMax(0.0f).squaredNorm();
+    }
+
+    __host__ __device__ float signed_distance(const Eigen::Vector3f& p) const {
+        Eigen::Vector3f q = (p - min).cwiseAbs() - diag();
+        return q.cwiseMax(0.0f).norm() + std::min(std::max(q.x(), std::max(q.y(), q.z())), 0.0f);
+    }
+
+    __host__ __device__ void get_vertices(Eigen::Vector3f v[8]) const {
+        v[0] = {min.x(), min.y(), min.z()};
+        v[1] = {min.x(), min.y(), max.z()};
+        v[2] = {min.x(), max.y(), min.z()};
+        v[3] = {min.x(), max.y(), max.z()};
+        v[4] = {max.x(), min.y(), min.z()};
+        v[5] = {max.x(), min.y(), max.z()};
+        v[6] = {max.x(), max.y(), min.z()};
+        v[7] = {max.x(), max.y(), max.z()};
+    }
+
+};
+
+inline std::ostream& operator<< (std::ostream& os, const cubvh::BoundingBox& bb) {
+    os << "[";
+    os << "min=[" << bb.min.x() << "," << bb.min.y() << "," << bb.min.z() << "], ";
+    os << "max=[" << bb.max.x() << "," << bb.max.y() << "," << bb.max.z() << "]";
+    os << "]";
+    return os;
+}
+
+}
\ No newline at end of file

include/cubvh/bvh.cuh

+#pragma once
+
+#include <cubvh/common.h>
+#include <cubvh/triangle.cuh>
+#include <cubvh/bounding_box.cuh>
+#include <cubvh/gpu_memory.h>
+
+#include <memory>
+
+namespace cubvh {
+
+struct TriangleBvhNode {
+    BoundingBox bb;
+    int left_idx; // negative values indicate leaves
+    int right_idx;
+};
+
+
+template <typename T, int MAX_SIZE=32>
+class FixedStack {
+public:
+    __host__ __device__ void push(T val) {
+        if (m_count >= MAX_SIZE-1) {
+            printf("WARNING TOO BIG\n");
+        }
+        m_elems[m_count++] = val;
+    }
+
+    __host__ __device__ T pop() {
+        return m_elems[--m_count];
+    }
+
+    __host__ __device__ bool empty() const {
+        return m_count <= 0;
+    }
+
+private:
+    T m_elems[MAX_SIZE];
+    int m_count = 0;
+};
+
+using FixedIntStack = FixedStack<int>;
+
+
+class TriangleBvh {
+
+protected:
+    std::vector<TriangleBvhNode> m_nodes;
+    GPUMemory<TriangleBvhNode> m_nodes_gpu;
+    TriangleBvh() {};
+
+public:
+    virtual void build(std::vector<Triangle>& triangles, uint32_t n_primitives_per_leaf) = 0;
+
+    virtual void unsigned_distance_gpu(uint32_t n_elements, const float* positions, float* distances, int64_t* face_id, float* uvw, const Triangle* gpu_triangles, cudaStream_t stream) = 0;
+    virtual void ray_trace_gpu(uint32_t n_elements, const float* rays_o, const float* rays_d, float* positions, int64_t* face_id, float* depth, const Triangle* gpu_triangles, cudaStream_t stream) = 0;
+
+    // KIUI: not supported now.
+    // virtual void signed_distance_gpu(uint32_t n_elements, EMeshSdfMode mode, const Eigen::Vector3f* gpu_positions, float* gpu_distances, const Triangle* gpu_triangles, bool use_existing_distances_as_upper_bounds, cudaStream_t stream) = 0;
+    // virtual bool touches_triangle(const BoundingBox& bb, const Triangle* __restrict__ triangles) const = 0;
+    // virtual void build_optix(const GPUMemory<Triangle>& triangles, cudaStream_t stream) = 0;
+
+    static std::unique_ptr<TriangleBvh> make();
+
+    TriangleBvhNode* nodes_gpu() const {
+        return m_nodes_gpu.data();
+    }
+
+};
+
+}
\ No newline at end of file

include/cubvh/common.h

+#pragma once
+
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include <cstdint>
+#include <cmath>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+#include <Eigen/Dense>
+
+namespace cubvh {
+
+static constexpr float PI = 3.14159265358979323846f;
+static constexpr float SQRT2 = 1.41421356237309504880f;
+
+
+// enum class EMeshSdfMode : int {
+// 	Watertight,
+// 	Raystab,
+// 	PathEscape,
+// };
+// static constexpr const char* MeshSdfModeStr = "Watertight\0Raystab\0PathEscape\0\0";
+
+template <typename T>
+__host__ __device__ T div_round_up(T val, T divisor) {
+    return (val + divisor - 1) / divisor;
+}
+
+constexpr uint32_t n_threads_linear = 128;
+
+template <typename T>
+constexpr uint32_t n_blocks_linear(T n_elements) {
+    return (uint32_t)div_round_up(n_elements, (T)n_threads_linear);
+}
+
+template <typename K, typename T, typename ... Types>
+inline void linear_kernel(K kernel, uint32_t shmem_size, cudaStream_t stream, T n_elements, Types ... args) {
+    if (n_elements <= 0) {
+        return;
+    }
+    kernel<<<n_blocks_linear(n_elements), n_threads_linear, shmem_size, stream>>>((uint32_t)n_elements, args...);
+}
+
+inline __host__ __device__ float sign(float x) {
+    return copysignf(1.0, x);
+}
+
+template <typename T>
+__host__ __device__ T clamp(T val, T lower, T upper) {
+    return val < lower ? lower : (upper < val ? upper : val);
+}
+
+template <typename T>
+__host__ __device__ void host_device_swap(T& a, T& b) {
+    T c(a); a=b; b=c;
+}
+
+}
\ No newline at end of file

include/cubvh/gpu_memory.h

+/*
+ * Copyright (c) 2020-2022, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *//*
+ */
+
+/** @file   gpu_memory.h
+ *  @author Nikolaus Binder and Thomas Müller, NVIDIA
+ *  @brief  Managed memory on the GPU. Like a std::vector, memory is alocated  either explicitly (resize/enlarge)
+ *          or implicitly (resize_and_copy_from_host etc). Memory is always and automatically released in the destructor.
+ */
+
+#pragma once
+
+#include <cubvh/common.h>
+#include <atomic>
+#include <stdexcept>
+#include <stdint.h>
+#include <string>
+#include <vector>
+
+/// Checks the result of a cudaXXXXXX call and throws an error on failure
+#define CUDA_CHECK_THROW(x)                                                                                               \
+    do {                                                                                                                  \
+        cudaError_t result = x;                                                                                           \
+        if (result != cudaSuccess)                                                                                        \
+            throw std::runtime_error(std::string("CUDA Error: " #x " failed with error ") + cudaGetErrorString(result));  \
+    } while(0)
+
+
+namespace cubvh {
+
+#define DEBUG_GUARD_SIZE 0
+
+inline std::atomic<size_t>& total_n_bytes_allocated() {
+    static std::atomic<size_t> s_total_n_bytes_allocated{0};
+    return s_total_n_bytes_allocated;
+}
+
+/// Managed memory on the Device
+template<class T>
+class GPUMemory {
+private:
+    T* m_data = nullptr;
+    size_t m_size = 0; // Number of elements
+    bool m_owned = true;
+
+public:
+    GPUMemory() {}
+
+    GPUMemory<T>& operator=(GPUMemory<T>&& other) {
+        std::swap(m_data, other.m_data);
+        std::swap(m_size, other.m_size);
+        return *this;
+    }
+
+    GPUMemory(GPUMemory<T>&& other) {
+        *this = std::move(other);
+    }
+
+    __host__ __device__ GPUMemory(const GPUMemory<T> &other) : m_data{other.m_data}, m_size{other.m_size}, m_owned{false} {}
+
+    void check_guards() const {
+#if DEBUG_GUARD_SIZE > 0
+        if (!m_data)
+            return;
+        uint8_t buf[DEBUG_GUARD_SIZE];
+        const uint8_t *rawptr=(const uint8_t *)m_data;
+        cudaMemcpy(buf, rawptr-DEBUG_GUARD_SIZE, DEBUG_GUARD_SIZE, cudaMemcpyDeviceToHost);
+        for (int i=0;i<DEBUG_GUARD_SIZE;++i) if (buf[i] != 0xff) {
+            printf("TRASH BEFORE BLOCK offset %d data %p, read 0x%02x expected 0xff!\n", i, m_data, buf[i] );
+            break;
+        }
+        cudaMemcpy(buf, rawptr+m_size*sizeof(T), DEBUG_GUARD_SIZE, cudaMemcpyDeviceToHost);
+        for (int i=0;i<DEBUG_GUARD_SIZE;++i) if (buf[i] != 0xfe) {
+            printf("TRASH AFTER BLOCK offset %d data %p, read 0x%02x expected 0xfe!\n", i, m_data, buf[i] );
+            break;
+        }
+#endif
+    }
+
+    void allocate_memory(size_t n_bytes) {
+        if (n_bytes == 0) {
+            return;
+        }
+
+#ifdef TCNN_VERBOSE_MEMORY_ALLOCS
+        std::cout << "GPUMemory: Allocating " << bytes_to_string(n_bytes) << "." << std::endl;
+#endif
+
+        uint8_t *rawptr = nullptr;
+        CUDA_CHECK_THROW(cudaMalloc(&rawptr, n_bytes+DEBUG_GUARD_SIZE*2));
+#if DEBUG_GUARD_SIZE > 0
+        CUDA_CHECK_THROW(cudaMemset(rawptr , 0xff, DEBUG_GUARD_SIZE));
+        CUDA_CHECK_THROW(cudaMemset(rawptr+n_bytes+DEBUG_GUARD_SIZE , 0xfe, DEBUG_GUARD_SIZE));
+#endif
+        if (rawptr) rawptr+=DEBUG_GUARD_SIZE;
+        m_data=(T*)(rawptr);
+        total_n_bytes_allocated() += n_bytes;
+    }
+
+    void free_memory() {
+        if (!m_data) {
+            return;
+        }
+
+        uint8_t *rawptr = (uint8_t*)m_data;
+        if (rawptr) rawptr-=DEBUG_GUARD_SIZE;
+        CUDA_CHECK_THROW(cudaFree(rawptr));
+
+        total_n_bytes_allocated() -= get_bytes();
+
+        m_data = nullptr;
+    }
+
+    /// Allocates memory for size items of type T
+    GPUMemory(const size_t size) {
+        resize(size);
+    }
+
+    /// Frees memory again
+    __host__ __device__ ~GPUMemory() {
+#ifndef __CUDA_ARCH__
+        if (!m_owned) {
+            return;
+        }
+
+        try {
+            if (m_data) {
+                free_memory();
+                m_size = 0;
+            }
+        } catch (std::runtime_error error) {
+            // Don't need to report on memory-free problems when the driver is shutting down.
+            if (std::string{error.what()}.find("driver shutting down") == std::string::npos) {
+                fprintf(stderr, "Could not free memory: %s\n", error.what());
+            }
+        }
+#endif
+    }
+
+    /** @name Resizing/enlargement
+     *  @{
+     */
+    /// Resizes the array to the exact new size, even if it is already larger
+    void resize(const size_t size) {
+        if (!m_owned) {
+            throw std::runtime_error("Cannot resize non-owned memory.");
+        }
+
+        if (m_size != size) {
+            if (m_size) {
+                try {
+                    free_memory();
+                } catch (std::runtime_error error) {
+                    throw std::runtime_error(std::string("Could not free memory: ") + error.what());
+                }
+            }
+
+            if (size > 0) {
+                try {
+                    allocate_memory(size * sizeof(T));
+                } catch (std::runtime_error error) {
+                    throw std::runtime_error(std::string("Could not allocate memory: ") + error.what());
+                }
+            }
+
+            m_size = size;
+        }
+    }
+
+    /// Enlarges the array if its size is smaller
+    void enlarge(const size_t size) {
+        if (size > m_size) {
+            resize(size);
+        }
+    }
+    /** @} */
+
+    /** @name Memset
+     *  @{
+     */
+    /// Sets the memory of the first num_elements to value
+    void memset(const int value, const size_t num_elements, const size_t offset = 0) {
+        if (num_elements + offset > m_size) {
+            throw std::runtime_error("Could not set memory: Number of elements larger than allocated memory");
+        }
+
+        try {
+            CUDA_CHECK_THROW(cudaMemset(m_data + offset, value, num_elements * sizeof(T)));
+        } catch (std::runtime_error error) {
+            throw std::runtime_error(std::string("Could not set memory: ") + error.what());
+        }
+    }
+
+    /// Sets the memory of the all elements to value
+    void memset(const int value) {
+        memset(value, m_size);
+    }
+    /** @} */
+
+    /** @name Copy operations
+     *  @{
+     */
+    /// Copy data of num_elements from the raw pointer on the host
+    void copy_from_host(const T* host_data, const size_t num_elements) {
+        try {
+            CUDA_CHECK_THROW(cudaMemcpy(data(), host_data, num_elements * sizeof(T), cudaMemcpyHostToDevice));
+        } catch (std::runtime_error error) {
+            throw std::runtime_error(std::string("Could not copy from host: ") + error.what());
+        }
+    }
+
+    /// Copy num_elements from the host vector
+    void copy_from_host(const std::vector<T>& data, const size_t num_elements) {
+        if (data.size() < num_elements) {
+            throw std::runtime_error(std::string("Trying to copy ") + std::to_string(num_elements) + std::string(" elements, but vector size is only ") + std::to_string(data.size()));
+        }
+        copy_from_host(data.data(), num_elements);
+    }
+
+    /// Copies data from the raw host pointer to fill the entire array
+    void copy_from_host(const T* data) {
+        copy_from_host(data, m_size);
+    }
+
+    /// Copies num_elements of data from the raw host pointer after enlarging the array so that everything fits in
+    void enlarge_and_copy_from_host(const T* data, const size_t num_elements) {
+        enlarge(num_elements);
+        copy_from_host(data, num_elements);
+    }
+
+    /// Copies num_elements from the host vector after enlarging the array so that everything fits in
+    void enlarge_and_copy_from_host(const std::vector<T>& data, const size_t num_elements) {
+        enlarge_and_copy_from_host(data.data(), num_elements);
+    }
+
+    /// Copies the entire host vector after enlarging the array so that everything fits in
+    void enlarge_and_copy_from_host(const std::vector<T>& data) {
+        enlarge_and_copy_from_host(data.data(), data.size());
+    }
+
+    /// Copies num_elements of data from the raw host pointer after resizing the array
+    void resize_and_copy_from_host(const T* data, const size_t num_elements) {
+        resize(num_elements);
+        copy_from_host(data, num_elements);
+    }
+
+    /// Copies num_elements from the host vector after resizing the array
+    void resize_and_copy_from_host(const std::vector<T>& data, const size_t num_elements) {
+        resize_and_copy_from_host(data.data(), num_elements);
+    }
+
+    /// Copies the entire host vector after resizing the array
+    void resize_and_copy_from_host(const std::vector<T>& data) {
+        resize_and_copy_from_host(data.data(), data.size());
+    }
+
+    /// Copies the entire host vector to the device. Fails if there is not enough space available.
+    void copy_from_host(const std::vector<T>& data) {
+        if (data.size() < m_size) {
+            throw std::runtime_error(std::string("Trying to copy ") + std::to_string(m_size) + std::string(" elements, but vector size is only ") + std::to_string(data.size()));
+        }
+        copy_from_host(data.data(), m_size);
+    }
+
+    /// Copies num_elements of data from the raw host pointer to the device. Fails if there is not enough space available.
+    void copy_to_host(T* host_data, const size_t num_elements) const {
+        if (num_elements > m_size) {
+            throw std::runtime_error(std::string("Trying to copy ") + std::to_string(num_elements) + std::string(" elements, but vector size is only ") + std::to_string(m_size));
+        }
+        try {
+            CUDA_CHECK_THROW(cudaMemcpy(host_data, data(), num_elements * sizeof(T), cudaMemcpyDeviceToHost));
+        } catch (std::runtime_error error) {
+            throw std::runtime_error(std::string("Could not copy to host: ") + error.what());
+        }
+    }
+
+    /// Copies num_elements from the device to a vector on the host
+    void copy_to_host(std::vector<T>& data, const size_t num_elements) const {
+        if (data.size() < num_elements) {
+            throw std::runtime_error(std::string("Trying to copy ") + std::to_string(num_elements) + std::string(" elements, but vector size is only ") + std::to_string(data.size()));
+        }
+        copy_to_host(data.data(), num_elements);
+    }
+
+    /// Copies num_elements from the device to a raw pointer on the host
+    void copy_to_host(T* data) const {
+        copy_to_host(data, m_size);
+    }
+
+    /// Copies all elements from the device to a vector on the host
+    void copy_to_host(std::vector<T>& data) const {
+        if (data.size() < m_size) {
+            throw std::runtime_error(std::string("Trying to copy ") + std::to_string(m_size) + std::string(" elements, but vector size is only ") + std::to_string(data.size()));
+        }
+        copy_to_host(data.data(), m_size);
+    }
+
+    /// Copies data from another device array to this one, automatically resizing it
+    void copy_from_device(const GPUMemory<T> &other) {
+        if (m_size != other.m_size) {
+            resize(other.m_size);
+        }
+
+        try {
+            CUDA_CHECK_THROW(cudaMemcpy(m_data, other.m_data, m_size * sizeof(T), cudaMemcpyDeviceToDevice));
+        } catch (std::runtime_error error) {
+            throw std::runtime_error(std::string("Could not copy from device: ") + error.what());
+        }
+    }
+
+    /// Copies size elements from another device array to this one, automatically resizing it
+    void copy_from_device(const GPUMemory<T> &other, const size_t size) {
+        if (m_size < size) {
+            resize(size);
+        }
+
+        try {
+            CUDA_CHECK_THROW(cudaMemcpy(m_data, other.m_data, size * sizeof(T), cudaMemcpyDeviceToDevice));
+        } catch (std::runtime_error error) {
+            throw std::runtime_error(std::string("Could not copy from device: ") + error.what());
+        }
+    }
+
+    // Created an (owned) copy of the data
+    GPUMemory<T> copy() const {
+        GPUMemory<T> result{m_size};
+        result.copy_from_device(*this);
+        return result;
+    }
+
+    T* data() const {
+        check_guards();
+        return m_data;
+    }
+
+    __host__ __device__ T& operator[](size_t idx) const {
+#ifdef DEBUG_BUFFER_OVERRUN
+        if (idx > m_size) {
+            printf("WARNING: buffer overrun of %p at idx %zu\n", idx);
+        }
+#endif
+        return m_data[idx];
+    }
+
+    __host__ __device__ T& operator[](uint32_t idx) const {
+#ifdef DEBUG_BUFFER_OVERRUN
+        if (idx > m_size) {
+            printf("WARNING: buffer overrun of %p at idx %u\n", idx);
+        }
+#endif
+        return m_data[idx];
+    }
+
+    size_t get_num_elements() const {
+        return m_size;
+    }
+
+    size_t size() const {
+        return get_num_elements();
+    }
+
+    size_t get_bytes() const {
+        return m_size * sizeof(T);
+    }
+
+    size_t bytes() const {
+        return get_bytes();
+    }
+};
+
+}
\ No newline at end of file

include/cubvh/triangle.cuh

+#pragma once
+
+#include <cubvh/common.h>
+
+namespace cubvh {
+
+// Triangle data structure
+struct Triangle {
+
+    __host__ __device__ Eigen::Vector3f sample_uniform_position(const Eigen::Vector2f& sample) const {
+        float sqrt_x = std::sqrt(sample.x());
+        float factor0 = 1.0f - sqrt_x;
+        float factor1 = sqrt_x * (1.0f - sample.y());
+        float factor2 = sqrt_x * sample.y();
+
+        return factor0 * a + factor1 * b + factor2 * c;
+    }
+
+    __host__ __device__ float surface_area() const {
+        return 0.5f * Eigen::Vector3f((b - a).cross(c - a)).norm();
+    }
+
+    __host__ __device__ Eigen::Vector3f normal() const {
+        return (b - a).cross(c - a).normalized();
+    }
+
+    __host__ __device__ float ray_intersect(const Eigen::Vector3f &ro, const Eigen::Vector3f &rd, Eigen::Vector3f& n) const { // based on https://www.iquilezles.org/www/articles/intersectors/intersectors.htm
+        Eigen::Vector3f v1v0 = b - a;
+        Eigen::Vector3f v2v0 = c - a;
+        Eigen::Vector3f rov0 = ro - a;
+        n = v1v0.cross( v2v0 );
+        Eigen::Vector3f q = rov0.cross( rd );
+        float d = 1.0f/rd.dot( n );
+        float u = d*-q.dot( v2v0 );
+        float v = d* q.dot( v1v0 );
+        float t = d*-n.dot( rov0 );
+        if( u<0.0f || u>1.0f || v<0.0f || (u+v)>1.0f || t<0.0f) t = 1e6f;
+        return t; // Eigen::Vector3f( t, u, v );
+    }
+
+    __host__ __device__ float ray_intersect(const Eigen::Vector3f &ro, const Eigen::Vector3f &rd) const {
+        Eigen::Vector3f n;
+        return ray_intersect(ro, rd, n);
+    }
+
+    __host__ __device__ float distance_sq(const Eigen::Vector3f& pos) const {
+        // prepare data
+        Eigen::Vector3f v21 = b - a; Eigen::Vector3f p1 = pos - a;
+        Eigen::Vector3f v32 = c - b; Eigen::Vector3f p2 = pos - b;
+        Eigen::Vector3f v13 = a - c; Eigen::Vector3f p3 = pos - c;
+        Eigen::Vector3f nor = v21.cross(v13);
+
+        return
+            // inside/outside test
+            (sign(v21.cross(nor).dot(p1)) + sign(v32.cross(nor).dot(p2)) + sign(v13.cross(nor).dot(p3)) < 2.0f)
+            ?
+            // 3 edges
+            std::min(
+                std::min(
+                    (v21 * clamp(v21.dot(p1) / v21.squaredNorm(), 0.0f, 1.0f)-p1).squaredNorm(),
+                    (v32 * clamp(v32.dot(p2) / v32.squaredNorm(), 0.0f, 1.0f)-p2).squaredNorm()
+                ),
+                (v13 * clamp(v13.dot(p3) / v13.squaredNorm(), 0.0f, 1.0f)-p3).squaredNorm()
+            )
+            :
+            // 1 face
+            nor.dot(p1)*nor.dot(p1)/nor.squaredNorm();
+    }
+
+    __host__ __device__ float distance(const Eigen::Vector3f& pos) const {
+        return std::sqrt(distance_sq(pos));
+    }
+
+    __host__ __device__ bool point_in_triangle(const Eigen::Vector3f& p) const {
+        // Move the triangle so that the point becomes the
+        // triangles origin
+        Eigen::Vector3f local_a = a - p;
+        Eigen::Vector3f local_b = b - p;
+        Eigen::Vector3f local_c = c - p;
+
+        // The point should be moved too, so they are both
+        // relative, but because we don't use p in the
+        // equation anymore, we don't need it!
+        // p -= p;
+
+        // Compute the normal vectors for triangles:
+        // u = normal of PBC
+        // v = normal of PCA
+        // w = normal of PAB
+
+        Eigen::Vector3f u = local_b.cross(local_c);
+        Eigen::Vector3f v = local_c.cross(local_a);
+        Eigen::Vector3f w = local_a.cross(local_b);
+
+        // Test to see if the normals are facing
+        // the same direction, return false if not
+        if (u.dot(v) < 0.0f) {
+            return false;
+        }
+        if (u.dot(w) < 0.0f) {
+            return false;
+        }
+
+        // All normals facing the same way, return true
+        return true;
+    }
+
+    __host__ __device__ Eigen::Vector3f closest_point_to_line(const Eigen::Vector3f& a, const Eigen::Vector3f& b, const Eigen::Vector3f& c) const {
+        float t = (c - a).dot(b - a) / (b - a).dot(b - a);
+        t = std::max(std::min(t, 1.0f), 0.0f);
+        return a + t * (b - a);
+    }
+
+    __host__ __device__ Eigen::Vector3f closest_point(Eigen::Vector3f point) const {
+        point -= normal().dot(point - a) * normal();
+
+        if (point_in_triangle(point)) {
+            return point;
+        }
+
+        Eigen::Vector3f c1 = closest_point_to_line(a, b, point);
+        Eigen::Vector3f c2 = closest_point_to_line(b, c, point);
+        Eigen::Vector3f c3 = closest_point_to_line(c, a, point);
+
+        float mag1 = (point - c1).squaredNorm();
+        float mag2 = (point - c2).squaredNorm();
+        float mag3 = (point - c3).squaredNorm();
+
+        float min = std::min(mag1, mag2);
+        min = std::min(min, mag3);
+
+        if (min == mag1) {
+            return c1;
+        }
+        else if (min == mag2) {
+            return c2;
+        }
+        return c3;
+    }
+
+    __host__ __device__ Eigen::Vector3f barycentric(const Eigen::Vector3f& p) const {
+        Eigen::Vector3f v0 = b - a;
+        Eigen::Vector3f v1 = c - a;
+        Eigen::Vector3f v2 = p - a;
+
+        float d00 = v0.dot(v0);
+        float d01 = v0.dot(v1);
+        float d11 = v1.dot(v1);
+        float d20 = v2.dot(v0);
+        float d21 = v2.dot(v1);
+
+        float denom = d00 * d11 - d01 * d01;
+        float v = (d11 * d20 - d01 * d21) / denom;
+        float w = (d00 * d21 - d01 * d20) / denom;
+        float u = 1.0 - v - w;
+        
+        return Eigen::Vector3f(u, v, w);
+    }
+
+    __host__ __device__ Eigen::Vector3f centroid() const {
+        return (a + b + c) / 3.0f;
+    }
+
+    __host__ __device__ float centroid(int axis) const {
+        return (a[axis] + b[axis] + c[axis]) / 3;
+    }
+
+    __host__ __device__ void get_vertices(Eigen::Vector3f v[3]) const {
+        v[0] = a;
+        v[1] = b;
+        v[2] = c;
+    }
+
+    Eigen::Vector3f a, b, c;
+    int64_t id;
+};
+
+
+inline std::ostream& operator<<(std::ostream& os, const Triangle& triangle) {
+    os << "[";
+    os << "a=[" << triangle.a.x() << "," << triangle.a.y() << "," << triangle.a.z() << "], ";
+    os << "b=[" << triangle.b.x() << "," << triangle.b.y() << "," << triangle.b.z() << "], ";
+    os << "c=[" << triangle.c.x() << "," << triangle.c.y() << "," << triangle.c.z() << "]";
+    os << "]";
+    return os;
+}
+
+
+}
\ No newline at end of file

readme.md

+# cuBVH
+
+A CUDA Mesh BVH acceleration toolkit.
+
+### Install
+
+```python
+git clone https://github.com/ashawkey/cubvh
+cd cubvh
+pip install .
+```
+
+### Usage
+
+Example for a mesh normal renderer:
+
+```bash
+python test/renderer.py # default, show a dodecahedron
+python test/renderer.py --mesh example.ply # show any mesh file
+```
+
+https://user-images.githubusercontent.com/25863658/183238748-7ac82808-6cd3-4bb6-867a-9c22f8e3f7dd.mp4
+
+Example code:
+
+```python
+import numpy as np
+import trimesh
+
+import torch
+import cubvh
+
+# build BVH from mesh
+mesh = trimesh.load('example.ply')
+BVH = cubvh.cuBVH(mesh.vertices, mesh.faces) # build with numpy.ndarray/torch.Tensor
+
+
+# query ray-mesh intersection
+rays_o, rays_d = get_ray(pose, intrinsics, H, W) # [N, 3], [N, 3], query with torch.Tensor (cuda)
+intersections, face_normals, depth = BVH.ray_trace(rays_o, rays_d) # [N, 3], [N, 3], [N,]
+
+# query unsigned distance
+points # [N, 3]
+distances, face_id, uvw = BVH.unsigned_distance(points, return_uvw=True) # [N], [N], [N, 3]
+```
+
+
+### Acknowledgement
+
+* Credits to [Thomas Müller](https://tom94.net/)'s amazing [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) and [instant-ngp](https://github.com/NVlabs/instant-ngp)!

setup.py

+import os
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+_src_path = os.path.dirname(os.path.abspath(__file__))
+
+# ref: https://github.com/sxyu/sdf/blob/master/setup.py
+def find_eigen(min_ver=(3, 3, 0)):
+    """Helper to find or download the Eigen C++ library"""
+    import re, os
+    try_paths = [
+        '/usr/include/eigen3',
+        '/usr/local/include/eigen3',
+        os.path.expanduser('~/.local/include/eigen3'),
+        'C:/Program Files/eigen3',
+        'C:/Program Files (x86)/eigen3',
+    ]
+    WORLD_VER_STR = "#define EIGEN_WORLD_VERSION"
+    MAJOR_VER_STR = "#define EIGEN_MAJOR_VERSION"
+    MINOR_VER_STR = "#define EIGEN_MINOR_VERSION"
+    EIGEN_WEB_URL = 'https://gitlab.com/libeigen/eigen/-/archive/3.3.7/eigen-3.3.7.tar.bz2'
+    TMP_EIGEN_FILE = 'tmp_eigen.tar.bz2'
+    TMP_EIGEN_DIR = 'eigen-3.3.7'
+    min_ver_str = '.'.join(map(str, min_ver))
+
+    eigen_path = None
+    for path in try_paths:
+        macros_path = os.path.join(path, 'Eigen/src/Core/util/Macros.h')
+        if os.path.exists(macros_path):
+            macros = open(macros_path, 'r').read().split('\n')
+            world_ver, major_ver, minor_ver = None, None, None
+            for line in macros:
+                if line.startswith(WORLD_VER_STR):
+                    world_ver = int(line[len(WORLD_VER_STR):])
+                elif line.startswith(MAJOR_VER_STR):
+                    major_ver = int(line[len(MAJOR_VER_STR):])
+                elif line.startswith(MINOR_VER_STR):
+                    minor_ver = int(line[len(MAJOR_VER_STR):])
+            if not world_ver or not major_ver or not minor_ver:
+                print('Failed to parse macros file', macros_path)
+            else:
+                ver = (world_ver, major_ver, minor_ver)
+                ver_str = '.'.join(map(str, ver))
+                if ver < min_ver:
+                    print('Found unsuitable Eigen version', ver_str, 'at',
+                          path, '(need >= ' + min_ver_str + ')')
+                else:
+                    print('Found Eigen version', ver_str, 'at', path)
+                    eigen_path = path
+                    break
+
+    if eigen_path is None:
+        try:
+            import urllib.request
+            print("Couldn't find Eigen locally, downloading...")
+            req = urllib.request.Request(
+                EIGEN_WEB_URL,
+                data=None,
+                headers={
+                    'User-Agent':
+                    'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.135 Safari/537.36'
+                })
+
+            with urllib.request.urlopen(req) as resp,\
+                  open(TMP_EIGEN_FILE, 'wb') as file:
+                data = resp.read()
+                file.write(data)
+            import tarfile
+            tar = tarfile.open(TMP_EIGEN_FILE)
+            tar.extractall()
+            tar.close()
+
+            eigen_path = TMP_EIGEN_DIR
+            os.remove(TMP_EIGEN_FILE)
+        except:
+            print('Download failed, failed to find Eigen')
+
+    if eigen_path is not None:
+        print('Found eigen at', eigen_path)
+
+    return eigen_path
+
+nvcc_flags = [
+    '-O3', '-std=c++14',
+    "--expt-extended-lambda",
+	"--expt-relaxed-constexpr",
+    '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__',
+]
+
+if os.name == "posix":
+    c_flags = ['-O3', '-std=c++14']
+elif os.name == "nt":
+    c_flags = ['/O2', '/std:c++17']
+
+    # find cl.exe
+    def find_cl_path():
+        import glob
+        for edition in ["Enterprise", "Professional", "BuildTools", "Community"]:
+            paths = sorted(glob.glob(r"C:\\Program Files (x86)\\Microsoft Visual Studio\\*\\%s\\VC\\Tools\\MSVC\\*\\bin\\Hostx64\\x64" % edition), reverse=True)
+            if paths:
+                return paths[0]
+
+    # If cl.exe is not on path, try to find it.
+    if os.system("where cl.exe >nul 2>nul") != 0:
+        cl_path = find_cl_path()
+        if cl_path is None:
+            raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation")
+        os.environ["PATH"] += ";" + cl_path
+
+'''
+Usage:
+python setup.py build_ext --inplace # build extensions locally, do not install (only can be used from the parent directory)
+python setup.py install # build extensions and install (copy) to PATH.
+pip install . # ditto but better (e.g., dependency & metadata handling)
+python setup.py develop # build extensions and install (symbolic) to PATH.
+pip install -e . # ditto but better (e.g., dependency & metadata handling)
+'''
+setup(
+    name='cubvh', # package name, import this to use python API
+    version='0.1.0',
+    description='CUDA BVH implementation',
+    url='https://github.com/ashawkey/cubvh',
+    author='kiui',
+    author_email='ashawkey1999@gmail.com',
+    ext_modules=[
+        CUDAExtension(
+            name='_cubvh', # extension name, import this to use CUDA API
+            sources=[os.path.join(_src_path, 'src', f) for f in [
+                'bvh.cu',
+                'api.cu',
+                'bindings.cpp',
+            ]],
+            include_dirs=[
+                os.path.join(_src_path, 'include'),
+                find_eigen(),
+            ],
+            extra_compile_args={
+                'cxx': c_flags,
+                'nvcc': nvcc_flags,
+            }
+        ),
+    ],
+    cmdclass={
+        'build_ext': BuildExtension,
+    },
+    install_requires=[
+        'ninja',
+        'trimesh',
+        'opencv-python',
+        'torch',
+        'numpy ',
+        'tqdm',
+        'matplotlib',
+        'dearpygui',
+    ],
+)
\ No newline at end of file

src/api.cu

+#pragma once
+
+#include <cubvh/api.h>
+
+#include <cubvh/common.h>
+#include <cubvh/bvh.cuh>
+
+#include <Eigen/Dense>
+
+using namespace Eigen;
+
+using Verts = Matrix<float, Dynamic, 3, RowMajor>;
+using Trigs = Matrix<uint32_t, Dynamic, 3, RowMajor>;
+
+namespace cubvh {
+
+class cuBVHImpl : public cuBVH {
+public:
+
+    // accept numpy array (cpu) to init 
+    cuBVHImpl(Ref<const Verts> vertices, Ref<const Trigs> triangles) : cuBVH() {
+
+        const size_t n_vertices = vertices.rows();
+        const size_t n_triangles = triangles.rows();
+
+        triangles_cpu.resize(n_triangles);
+
+        for (size_t i = 0; i < n_triangles; i++) {
+            triangles_cpu[i] = {vertices.row(triangles(i, 0)), vertices.row(triangles(i, 1)), vertices.row(triangles(i, 2)), (int64_t)i};
+        }
+
+        if (!triangle_bvh) {
+            triangle_bvh = TriangleBvh::make();
+        }
+
+        triangle_bvh->build(triangles_cpu, 8);
+
+        triangles_gpu.resize_and_copy_from_host(triangles_cpu);
+
+        // TODO: need OPTIX
+        // triangle_bvh->build_optix(triangles_gpu, m_inference_stream);
+
+    }
+
+    // accept torch tensor (gpu) to init
+    void ray_trace(at::Tensor rays_o, at::Tensor rays_d, at::Tensor positions, at::Tensor face_id, at::Tensor depth) {
+
+        const uint32_t n_elements = rays_o.size(0);
+        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+        triangle_bvh->ray_trace_gpu(n_elements, rays_o.data_ptr<float>(), rays_d.data_ptr<float>(), positions.data_ptr<float>(), face_id.data_ptr<int64_t>(), depth.data_ptr<float>(), triangles_gpu.data(), stream);
+    }
+
+    // optionally return uvw (barycentric) and face_id
+    void unsigned_distance(at::Tensor positions, at::Tensor distances, at::Tensor face_id, at::optional<at::Tensor> uvw) {
+
+        const uint32_t n_elements = positions.size(0);
+        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+        triangle_bvh->unsigned_distance_gpu(n_elements, positions.data_ptr<float>(), distances.data_ptr<float>(), face_id.data_ptr<int64_t>(), uvw.has_value() ? uvw.value().data_ptr<float>() : nullptr, triangles_gpu.data(), stream);
+
+    }
+
+    std::vector<Triangle> triangles_cpu;
+    GPUMemory<Triangle> triangles_gpu;
+    std::shared_ptr<TriangleBvh> triangle_bvh;
+};
+    
+cuBVH* create_cuBVH(Ref<const Verts> vertices, Ref<const Trigs> triangles) {
+    return new cuBVHImpl{vertices, triangles};
+}
+
+} // namespace cubvh
\ No newline at end of file

src/bindings.cpp

+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/eigen.h>
+
+#include <torch/extension.h>
+
+#include <cubvh/api.h>
+
+
+namespace py = pybind11;
+using namespace cubvh;
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+
+py::class_<cuBVH>(m, "cuBVH")
+    .def("ray_trace", &cuBVH::ray_trace)
+    .def("unsigned_distance", &cuBVH::unsigned_distance);
+
+m.def("create_cuBVH", &create_cuBVH);
+
+}
\ No newline at end of file

src/bvh.cu

+// #include <Eigen/Dense>
+#include <cubvh/common.h>
+#include <cubvh/triangle.cuh>
+#include <cubvh/bvh.cuh>
+
+// #ifdef NGP_OPTIX
+// #  include <optix.h>
+// #  include <optix_stubs.h>
+// #  include <optix_function_table_definition.h>
+// #  include <optix_stack_size.h>
+
+// // Custom optix toolchain stuff
+// #  include "optix/pathescape.h"
+// #  include "optix/raystab.h"
+// #  include "optix/raytrace.h"
+
+// #  include "optix/program.h"
+
+// // Compiled optix program PTX generated by cmake and wrapped in a C
+// // header by bin2c.
+// namespace optix_ptx {
+// 	#include <optix_ptx.h>
+// }
+// #endif //NGP_OPTIX
+
+#include <stack>
+#include <iostream>
+#include <cstdio>
+
+using namespace Eigen;
+using namespace cubvh;
+
+
+namespace cubvh {
+
+constexpr float MAX_DIST = 10.0f;
+constexpr float MAX_DIST_SQ = MAX_DIST*MAX_DIST;
+
+
+// #ifdef NGP_OPTIX
+// OptixDeviceContext g_optix;
+
+// namespace optix {
+// 	bool initialize() {
+// 		static bool ran_before = false;
+// 		static bool is_optix_initialized = false;
+// 		if (ran_before) {
+// 			return is_optix_initialized;
+// 		}
+
+// 		ran_before = true;
+
+// 		// Initialize CUDA with a no-op call to the the CUDA runtime API
+// 		CUDA_CHECK_THROW(cudaFree(nullptr));
+
+// 		try {
+// 			// Initialize the OptiX API, loading all API entry points
+// 			OPTIX_CHECK_THROW(optixInit());
+
+// 			// Specify options for this context. We will use the default options.
+// 			OptixDeviceContextOptions options = {};
+
+// 			// Associate a CUDA context (and therefore a specific GPU) with this
+// 			// device context
+// 			CUcontext cuCtx = 0; // NULL means take the current active context
+
+// 			OPTIX_CHECK_THROW(optixDeviceContextCreate(cuCtx, &options, &g_optix));
+// 		} catch (std::exception& e) {
+// 			tlog::warning() << "OptiX failed to initialize: " << e.what();
+// 			return false;
+// 		}
+
+// 		is_optix_initialized = true;
+// 		return true;
+// 	}
+
+// 	class Gas {
+// 	public:
+// 		Gas(const GPUMemory<Triangle>& triangles, OptixDeviceContext optix, cudaStream_t stream) {
+// 			// Specify options for the build. We use default options for simplicity.
+// 			OptixAccelBuildOptions accel_options = {};
+// 			accel_options.buildFlags = OPTIX_BUILD_FLAG_NONE;
+// 			accel_options.operation = OPTIX_BUILD_OPERATION_BUILD;
+
+// 			// Populate the build input struct with our triangle data as well as
+// 			// information about the sizes and types of our data
+// 			const uint32_t triangle_input_flags[1] = { OPTIX_GEOMETRY_FLAG_NONE };
+// 			OptixBuildInput triangle_input = {};
+
+// 			CUdeviceptr d_triangles = (CUdeviceptr)(uintptr_t)triangles.data();
+
+// 			triangle_input.type = OPTIX_BUILD_INPUT_TYPE_TRIANGLES;
+// 			triangle_input.triangleArray.vertexFormat = OPTIX_VERTEX_FORMAT_FLOAT3;
+// 			triangle_input.triangleArray.numVertices = (uint32_t)triangles.size()*3;
+// 			triangle_input.triangleArray.vertexBuffers = &d_triangles;
+// 			triangle_input.triangleArray.flags = triangle_input_flags;
+// 			triangle_input.triangleArray.numSbtRecords = 1;
+
+// 			// Query OptiX for the memory requirements for our GAS
+// 			OptixAccelBufferSizes gas_buffer_sizes;
+// 			OPTIX_CHECK_THROW(optixAccelComputeMemoryUsage(optix, &accel_options, &triangle_input, 1, &gas_buffer_sizes));
+
+// 			// Allocate device memory for the scratch space buffer as well
+// 			// as the GAS itself
+// 			GPUMemory<char> gas_tmp_buffer{gas_buffer_sizes.tempSizeInBytes};
+// 			m_gas_gpu_buffer.resize(gas_buffer_sizes.outputSizeInBytes);
+
+// 			OPTIX_CHECK_THROW(optixAccelBuild(
+// 				optix,
+// 				stream,
+// 				&accel_options,
+// 				&triangle_input,
+// 				1,           // num build inputs
+// 				(CUdeviceptr)(uintptr_t)gas_tmp_buffer.data(),
+// 				gas_buffer_sizes.tempSizeInBytes,
+// 				(CUdeviceptr)(uintptr_t)m_gas_gpu_buffer.data(),
+// 				gas_buffer_sizes.outputSizeInBytes,
+// 				&m_gas_handle, // Output handle to the struct
+// 				nullptr,       // emitted property list
+// 				0              // num emitted properties
+// 			));
+// 		}
+
+// 		OptixTraversableHandle handle() const {
+// 			return m_gas_handle;
+// 		}
+
+// 	private:
+// 		OptixTraversableHandle m_gas_handle;
+// 		GPUMemory<char> m_gas_gpu_buffer;
+// 	};
+// }
+// #endif //NGP_OPTIX
+
+// __global__ void signed_distance_watertight_kernel(uint32_t n_elements, const Vector3f* __restrict__ positions, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, float* __restrict__ distances, bool use_existing_distances_as_upper_bounds = false);
+// __global__ void signed_distance_raystab_kernel(uint32_t n_elements, const Vector3f* __restrict__ positions, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, float* __restrict__ distances, bool use_existing_distances_as_upper_bounds = false);
+__global__ void unsigned_distance_kernel(uint32_t n_elements, const Vector3f* __restrict__ positions, float* __restrict__ distances, int64_t* __restrict__ face_id, Vector3f* __restrict__ uvw, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, bool use_existing_distances_as_upper_bounds);
+__global__ void raytrace_kernel(uint32_t n_elements, const Vector3f* __restrict__ rays_o, const Vector3f* __restrict__ rays_d, Vector3f* __restrict__ positions, int64_t* __restrict__ face_id, float* __restrict__ depth, const TriangleBvhNode* __restrict__ nodes, const Triangle* __restrict__ triangles);
+
+struct DistAndIdx {
+    float dist;
+    uint32_t idx;
+
+    // Sort in descending order!
+    __host__ __device__ bool operator<(const DistAndIdx& other) {
+        return dist < other.dist;
+    }
+};
+
+template <typename T>
+__host__ __device__ void inline compare_and_swap(T& t1, T& t2) {
+    if (t1 < t2) {
+        T tmp{t1}; t1 = t2; t2 = tmp;
+    }
+}
+
+// Sorting networks from http://users.telenet.be/bertdobbelaere/SorterHunter/sorting_networks.html#N4L5D3
+template <uint32_t N, typename T>
+__host__ __device__ void sorting_network(T values[N]) {
+    static_assert(N <= 8, "Sorting networks are only implemented up to N==8");
+    if (N <= 1) {
+        return;
+    } else if (N == 2) {
+        compare_and_swap(values[0], values[1]);
+    } else if (N == 3) {
+        compare_and_swap(values[0], values[2]);
+        compare_and_swap(values[0], values[1]);
+        compare_and_swap(values[1], values[2]);
+    } else if (N == 4) {
+        compare_and_swap(values[0], values[2]);
+        compare_and_swap(values[1], values[3]);
+        compare_and_swap(values[0], values[1]);
+        compare_and_swap(values[2], values[3]);
+        compare_and_swap(values[1], values[2]);
+    } else if (N == 5) {
+        compare_and_swap(values[0], values[3]);
+        compare_and_swap(values[1], values[4]);
+
+        compare_and_swap(values[0], values[2]);
+        compare_and_swap(values[1], values[3]);
+
+        compare_and_swap(values[0], values[1]);
+        compare_and_swap(values[2], values[4]);
+
+        compare_and_swap(values[1], values[2]);
+        compare_and_swap(values[3], values[4]);
+
+        compare_and_swap(values[2], values[3]);
+    } else if (N == 6) {
+        compare_and_swap(values[0], values[5]);
+        compare_and_swap(values[1], values[3]);
+        compare_and_swap(values[2], values[4]);
+
+        compare_and_swap(values[1], values[2]);
+        compare_and_swap(values[3], values[4]);
+
+        compare_and_swap(values[0], values[3]);
+        compare_and_swap(values[2], values[5]);
+
+        compare_and_swap(values[0], values[1]);
+        compare_and_swap(values[2], values[3]);
+        compare_and_swap(values[4], values[5]);
+
+        compare_and_swap(values[1], values[2]);
+        compare_and_swap(values[3], values[4]);
+    } else if (N == 7) {
+        compare_and_swap(values[0], values[6]);
+        compare_and_swap(values[2], values[3]);
+        compare_and_swap(values[4], values[5]);
+
+        compare_and_swap(values[0], values[2]);
+        compare_and_swap(values[1], values[4]);
+        compare_and_swap(values[3], values[6]);
+
+        compare_and_swap(values[0], values[1]);
+        compare_and_swap(values[2], values[5]);
+        compare_and_swap(values[3], values[4]);
+
+        compare_and_swap(values[1], values[2]);
+        compare_and_swap(values[4], values[6]);
+
+        compare_and_swap(values[2], values[3]);
+        compare_and_swap(values[4], values[5]);
+
+        compare_and_swap(values[1], values[2]);
+        compare_and_swap(values[3], values[4]);
+        compare_and_swap(values[5], values[6]);
+    } else if (N == 8) {
+        compare_and_swap(values[0], values[2]);
+        compare_and_swap(values[1], values[3]);
+        compare_and_swap(values[4], values[6]);
+        compare_and_swap(values[5], values[7]);
+
+        compare_and_swap(values[0], values[4]);
+        compare_and_swap(values[1], values[5]);
+        compare_and_swap(values[2], values[6]);
+        compare_and_swap(values[3], values[7]);
+
+        compare_and_swap(values[0], values[1]);
+        compare_and_swap(values[2], values[3]);
+        compare_and_swap(values[4], values[5]);
+        compare_and_swap(values[6], values[7]);
+
+        compare_and_swap(values[2], values[4]);
+        compare_and_swap(values[3], values[5]);
+
+        compare_and_swap(values[1], values[4]);
+        compare_and_swap(values[3], values[6]);
+
+        compare_and_swap(values[1], values[2]);
+        compare_and_swap(values[3], values[4]);
+        compare_and_swap(values[5], values[6]);
+    }
+}
+
+template <uint32_t BRANCHING_FACTOR>
+class TriangleBvhWithBranchingFactor : public TriangleBvh {
+public:
+    __host__ __device__ static std::pair<int, float> ray_intersect(Ref<const Vector3f> ro, Ref<const Vector3f> rd, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles) {
+        FixedIntStack query_stack;
+        query_stack.push(0);
+
+        float mint = MAX_DIST;
+        int shortest_idx = -1;
+
+        while (!query_stack.empty()) {
+            int idx = query_stack.pop();
+
+            const TriangleBvhNode& node = bvhnodes[idx];
+
+            if (node.left_idx < 0) {
+                int end = -node.right_idx-1;
+                for (int i = -node.left_idx-1; i < end; ++i) {
+                    float t = triangles[i].ray_intersect(ro, rd);
+                    if (t < mint) {
+                        mint = t;
+                        shortest_idx = i;
+                    }
+                }
+            } else {
+                DistAndIdx children[BRANCHING_FACTOR];
+
+                uint32_t first_child = node.left_idx;
+
+                #pragma unroll
+                for (uint32_t i = 0; i < BRANCHING_FACTOR; ++i) {
+                    children[i] = {bvhnodes[i+first_child].bb.ray_intersect(ro, rd).x(), i+first_child};
+                }
+
+                sorting_network<BRANCHING_FACTOR>(children);
+
+                #pragma unroll
+                for (uint32_t i = 0; i < BRANCHING_FACTOR; ++i) {
+                    if (children[i].dist < mint) {
+                        query_stack.push(children[i].idx);
+                    }
+                }
+            }
+        }
+
+        return {shortest_idx, mint};
+    }
+
+    __host__ __device__ static std::pair<int, float> closest_triangle(const Vector3f& point, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, float max_distance_sq = MAX_DIST_SQ) {
+        FixedIntStack query_stack;
+        query_stack.push(0);
+
+        float shortest_distance_sq = max_distance_sq;
+        int shortest_idx = -1;
+
+        while (!query_stack.empty()) {
+            int idx = query_stack.pop();
+
+            const TriangleBvhNode& node = bvhnodes[idx];
+
+            if (node.left_idx < 0) {
+                int end = -node.right_idx-1;
+                for (int i = -node.left_idx-1; i < end; ++i) {
+                    float dist_sq = triangles[i].distance_sq(point);
+                    if (dist_sq <= shortest_distance_sq) {
+                        shortest_distance_sq = dist_sq;
+                        shortest_idx = i;
+                    }
+                }
+            } else {
+                DistAndIdx children[BRANCHING_FACTOR];
+
+                uint32_t first_child = node.left_idx;
+
+                #pragma unroll
+                for (uint32_t i = 0; i < BRANCHING_FACTOR; ++i) {
+                    children[i] = {bvhnodes[i+first_child].bb.distance_sq(point), i+first_child};
+                }
+
+                sorting_network<BRANCHING_FACTOR>(children);
+
+                #pragma unroll
+                for (uint32_t i = 0; i < BRANCHING_FACTOR; ++i) {
+                    if (children[i].dist <= shortest_distance_sq) {
+                        query_stack.push(children[i].idx);
+                    }
+                }
+            }
+        }
+
+        if (shortest_idx == -1) {
+            // printf("No closest triangle found. This must be a bug! %d\n", BRANCHING_FACTOR);
+            shortest_idx = 0;
+            shortest_distance_sq = 0.0f;
+        }
+
+        return {shortest_idx, std::sqrt(shortest_distance_sq)};
+    }
+
+    // Assumes that "point" is a location on a triangle
+    __host__ __device__ static Vector3f avg_normal_around_point(const Vector3f& point, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles) {
+        FixedIntStack query_stack;
+        query_stack.push(0);
+
+        static constexpr float EPSILON = 1e-6f;
+
+        float total_weight = 0;
+        Vector3f result = Vector3f::Zero();
+
+        while (!query_stack.empty()) {
+            int idx = query_stack.pop();
+
+            const TriangleBvhNode& node = bvhnodes[idx];
+
+            if (node.left_idx < 0) {
+                int end = -node.right_idx-1;
+                for (int i = -node.left_idx-1; i < end; ++i) {
+                    if (triangles[i].distance_sq(point) < EPSILON) {
+                        float weight = 1; // TODO: cot weight
+                        result += triangles[i].normal();
+                        total_weight += weight;
+                    }
+                }
+            } else {
+                uint32_t first_child = node.left_idx;
+
+                #pragma unroll
+                for (uint32_t i = 0; i < BRANCHING_FACTOR; ++i) {
+                    if (bvhnodes[i+first_child].bb.distance_sq(point) < EPSILON) {
+                        query_stack.push(i+first_child);
+                    }
+                }
+            }
+        }
+
+        return result / total_weight;
+    }
+
+    // __host__ __device__ static float signed_distance_watertight(const Vector3f& point, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, float max_distance_sq = MAX_DIST_SQ) {
+    //     auto p = closest_triangle(point, bvhnodes, triangles, max_distance_sq);
+
+    //     const Triangle& tri = triangles[p.first];
+    //     Vector3f closest_point = tri.closest_point(point);
+    //     Vector3f avg_normal = avg_normal_around_point(closest_point, bvhnodes, triangles);
+
+    //     return std::copysignf(p.second, avg_normal.dot(point - closest_point));
+    // }
+
+    // __host__ __device__ static float signed_distance_raystab(const Vector3f& point, const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, float max_distance_sq = MAX_DIST_SQ, default_rng_t rng={}) {
+    //     float distance = unsigned_distance(point, bvhnodes, triangles, max_distance_sq);
+
+    //     Vector2f offset = random_val_2d(rng);
+
+    //     static constexpr uint32_t N_STAB_RAYS = 32;
+    //     for (uint32_t i = 0; i < N_STAB_RAYS; ++i) {
+    //         // Use a Fibonacci lattice (with random offset) to regularly
+    //         // distribute the stab rays over the sphere.
+    //         Vector3f d = fibonacci_dir<N_STAB_RAYS>(i, offset);
+
+    //         // If any of the stab rays goes outside the mesh, the SDF is positive.
+    //         if (ray_intersect(point, -d, bvhnodes, triangles).first < 0 || ray_intersect(point, d, bvhnodes, triangles).first < 0) {
+    //             return distance;
+    //         }
+    //     }
+
+    //     return -distance;
+    // }
+
+    // // // Assumes that "point" is a location on a triangle
+    // // Vector3f avg_normal_around_point(const Vector3f& point, const Triangle* __restrict__ triangles) const {
+    // //     return avg_normal_around_point(point, m_nodes.data(), triangles);
+    // // }
+
+    // // float signed_distance(EMeshSdfMode mode, const Vector3f& point, const std::vector<Triangle>& triangles) const {
+    // //     if (mode == EMeshSdfMode::Watertight) {
+    // //         return signed_distance_watertight(point, m_nodes.data(), triangles.data());
+    // //     } else {
+    // //         return signed_distance_raystab(point, m_nodes.data(), triangles.data());
+    // //     }
+    // // }
+
+    // void signed_distance_gpu(uint32_t n_elements, EMeshSdfMode mode, const Vector3f* gpu_positions, float* gpu_distances, const Triangle* gpu_triangles, bool use_existing_distances_as_upper_bounds, cudaStream_t stream) override {
+    //     if (mode == EMeshSdfMode::Watertight) {
+    //         linear_kernel(signed_distance_watertight_kernel, 0, stream,
+    //             n_elements,
+    //             gpu_positions,
+    //             m_nodes_gpu.data(),
+    //             gpu_triangles,
+    //             gpu_distances,
+    //             use_existing_distances_as_upper_bounds
+    //         );
+    //     } else {
+    //         {
+    //             if (mode == EMeshSdfMode::Raystab) {
+    //                 linear_kernel(signed_distance_raystab_kernel, 0, stream,
+    //                     n_elements,
+    //                     gpu_positions,
+    //                     m_nodes_gpu.data(),
+    //                     gpu_triangles,
+    //                     gpu_distances,
+    //                     use_existing_distances_as_upper_bounds
+    //                 );
+    //             } else if (mode == EMeshSdfMode::PathEscape) {
+    //                 throw std::runtime_error{"TriangleBvh: EMeshSdfMode::PathEscape is only supported with OptiX enabled."};
+    //             }
+    //         }
+    //     }
+    // }
+
+    void unsigned_distance_gpu(uint32_t n_elements, const float* positions, float* distances, int64_t* face_id, float* uvw, const Triangle* gpu_triangles, cudaStream_t stream) override {
+
+        const Vector3f* positions_vec = (const Vector3f*)positions;
+        Vector3f* uvw_vec = (Vector3f*)uvw;
+
+        linear_kernel(unsigned_distance_kernel, 0u, stream,
+            n_elements,
+            positions_vec,
+            distances,
+            face_id,
+            uvw_vec,
+            m_nodes_gpu.data(),
+            gpu_triangles,
+            false
+        );
+
+    }
+
+    void ray_trace_gpu(uint32_t n_elements, const float* rays_o, const float* rays_d, float* positions, int64_t* face_id, float* depth, const Triangle* gpu_triangles, cudaStream_t stream) override {
+
+        // cast float* to Vector3f*
+        const Vector3f* rays_o_vec = (const Vector3f*)rays_o;
+        const Vector3f* rays_d_vec = (const Vector3f*)rays_d;
+        Vector3f* positions_vec = (Vector3f*)positions;
+        
+        linear_kernel(raytrace_kernel, 0u, stream,
+            n_elements,
+            rays_o_vec,
+            rays_d_vec,
+            positions_vec,
+            face_id,
+            depth,
+            m_nodes_gpu.data(),
+            gpu_triangles
+        );
+
+    }
+
+    void build(std::vector<Triangle>& triangles, uint32_t n_primitives_per_leaf) override {
+        m_nodes.clear();
+
+        // Root
+        m_nodes.emplace_back();
+        m_nodes.front().bb = BoundingBox(std::begin(triangles), std::end(triangles));
+
+        struct BuildNode {
+            int node_idx;
+            std::vector<Triangle>::iterator begin;
+            std::vector<Triangle>::iterator end;
+        };
+
+        std::stack<BuildNode> build_stack;
+        build_stack.push({0, std::begin(triangles), std::end(triangles)});
+
+        while (!build_stack.empty()) {
+            const BuildNode& curr = build_stack.top();
+            size_t node_idx = curr.node_idx;
+
+            std::array<BuildNode, BRANCHING_FACTOR> children;
+            children[0].begin = curr.begin;
+            children[0].end = curr.end;
+
+            build_stack.pop();
+
+            // Partition the triangles into the children
+            int n_children = 1;
+            while (n_children < BRANCHING_FACTOR) {
+                for (int i = n_children - 1; i >= 0; --i) {
+                    auto& child = children[i];
+
+                    // Choose axis with maximum standard deviation
+                    Vector3f mean = Vector3f::Zero();
+                    for (auto it = child.begin; it != child.end; ++it) {
+                        mean += it->centroid();
+                    }
+                    mean /= (float)std::distance(child.begin, child.end);
+
+                    Vector3f var = Vector3f::Zero();
+                    for (auto it = child.begin; it != child.end; ++it) {
+                        Vector3f diff = it->centroid() - mean;
+                        var += diff.cwiseProduct(diff);
+                    }
+                    var /= (float)std::distance(child.begin, child.end);
+
+                    Vector3f::Index axis;
+                    var.maxCoeff(&axis);
+
+                    auto m = child.begin + std::distance(child.begin, child.end)/2;
+                    std::nth_element(child.begin, m, child.end, [&](const Triangle& tri1, const Triangle& tri2) { return tri1.centroid(axis) < tri2.centroid(axis); });
+
+                    children[i*2].begin = children[i].begin;
+                    children[i*2+1].end = children[i].end;
+                    children[i*2].end = children[i*2+1].begin = m;
+                }
+
+                n_children *= 2;
+            }
+
+            // Create next build nodes
+            m_nodes[node_idx].left_idx = (int)m_nodes.size();
+            for (uint32_t i = 0; i < BRANCHING_FACTOR; ++i) {
+                auto& child = children[i];
+                assert(child.begin != child.end);
+                child.node_idx = (int)m_nodes.size();
+
+                m_nodes.emplace_back();
+                m_nodes.back().bb = BoundingBox(child.begin, child.end);
+
+                if (std::distance(child.begin, child.end) <= n_primitives_per_leaf) {
+                    m_nodes.back().left_idx = -(int)std::distance(std::begin(triangles), child.begin)-1;
+                    m_nodes.back().right_idx = -(int)std::distance(std::begin(triangles), child.end)-1;
+                } else {
+                    build_stack.push(child);
+                }
+            }
+            m_nodes[node_idx].right_idx = (int)m_nodes.size();
+        }
+
+        m_nodes_gpu.resize_and_copy_from_host(m_nodes);
+
+        // std::cout << "[INFO] Built TriangleBvh: nodes=" << m_nodes.size() << std::endl;
+    }
+
+//     void build_optix(const GPUMemory<Triangle>& triangles, cudaStream_t stream) override {
+// #ifdef NGP_OPTIX
+//         m_optix.available = optix::initialize();
+//         if (m_optix.available) {
+//             m_optix.gas = std::make_unique<optix::Gas>(triangles, g_optix, stream);
+//             m_optix.raystab = std::make_unique<optix::Program<Raystab>>((const char*)optix_ptx::raystab_ptx, sizeof(optix_ptx::raystab_ptx), g_optix);
+//             m_optix.raytrace = std::make_unique<optix::Program<Raytrace>>((const char*)optix_ptx::raytrace_ptx, sizeof(optix_ptx::raytrace_ptx), g_optix);
+//             m_optix.pathescape = std::make_unique<optix::Program<PathEscape>>((const char*)optix_ptx::pathescape_ptx, sizeof(optix_ptx::pathescape_ptx), g_optix);
+//             tlog::success() << "Built OptiX GAS and shaders";
+//         } else {
+//             tlog::warning() << "Falling back to slower TriangleBVH::ray_intersect.";
+//         }
+// #else //NGP_OPTIX
+//         tlog::warning() << "OptiX was not built. Falling back to slower TriangleBVH::ray_intersect.";
+// #endif //NGP_OPTIX
+//     }
+
+    TriangleBvhWithBranchingFactor() {}
+
+// private:
+// #ifdef NGP_OPTIX
+//     struct {
+//         std::unique_ptr<optix::Gas> gas;
+//         std::unique_ptr<optix::Program<Raystab>> raystab;
+//         std::unique_ptr<optix::Program<Raytrace>> raytrace;
+//         std::unique_ptr<optix::Program<PathEscape>> pathescape;
+//         bool available = false;
+//     } m_optix;
+// #endif //NGP_OPTIX
+};
+
+using TriangleBvh4 = TriangleBvhWithBranchingFactor<4>;
+
+std::unique_ptr<TriangleBvh> TriangleBvh::make() {
+    return std::unique_ptr<TriangleBvh>(new TriangleBvh4());
+}
+
+// __global__ void signed_distance_watertight_kernel(uint32_t n_elements,
+//     const Vector3f* __restrict__ positions,
+//     const TriangleBvhNode* __restrict__ bvhnodes,
+//     const Triangle* __restrict__ triangles,
+//     float* __restrict__ distances,
+//     bool use_existing_distances_as_upper_bounds
+// ) {
+//     uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+//     if (i >= n_elements) return;
+
+//     float max_distance = use_existing_distances_as_upper_bounds ? distances[i] : MAX_DIST;
+//     distances[i] = TriangleBvh4::signed_distance_watertight(positions[i], bvhnodes, triangles, max_distance*max_distance);
+// }
+
+// __global__ void signed_distance_raystab_kernel(
+//     uint32_t n_elements,
+//     const Vector3f* __restrict__ positions,
+//     const TriangleBvhNode* __restrict__ bvhnodes,
+//     const Triangle* __restrict__ triangles,
+//     float* __restrict__ distances,
+//     bool use_existing_distances_as_upper_bounds
+// ) {
+//     uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+//     if (i >= n_elements) return;
+
+//     float max_distance = use_existing_distances_as_upper_bounds ? distances[i] : MAX_DIST;
+//     default_rng_t rng;
+//     rng.advance(i * 2);
+
+//     distances[i] = TriangleBvh4::signed_distance_raystab(positions[i], bvhnodes, triangles, max_distance*max_distance, rng);
+// }
+
+__global__ void unsigned_distance_kernel(
+    uint32_t n_elements, const Vector3f* __restrict__ positions,
+    float* __restrict__ distances, int64_t* __restrict__ face_id, Vector3f* __restrict__ uvw,
+    const TriangleBvhNode* __restrict__ bvhnodes, const Triangle* __restrict__ triangles, bool use_existing_distances_as_upper_bounds
+) {
+    uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n_elements) return;
+
+    float max_distance = use_existing_distances_as_upper_bounds ? distances[i] : MAX_DIST;
+
+    Vector3f point = positions[i];
+
+    auto res = TriangleBvh4::closest_triangle(point, bvhnodes, triangles, max_distance*max_distance);
+
+    // write 
+    distances[i] = res.second;
+    face_id[i] = triangles[res.first].id;
+
+    // optional output
+    if (uvw) {
+        // get closest point
+        Vector3f cpoint = triangles[res.first].closest_point(point);
+        // query uvw
+        uvw[i] = triangles[res.first].barycentric(cpoint);
+    }
+}
+
+__global__ void raytrace_kernel(
+    uint32_t n_elements, const Vector3f* __restrict__ rays_o, const Vector3f* __restrict__ rays_d, 
+    Vector3f* __restrict__ positions, int64_t* __restrict__ face_id, float* __restrict__ depth, 
+    const TriangleBvhNode* __restrict__ nodes, const Triangle* __restrict__ triangles
+) {
+    uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= n_elements) return;
+
+    Vector3f ro = rays_o[i];
+    Vector3f rd = rays_d[i];
+
+    auto res = TriangleBvh4::ray_intersect(ro, rd, nodes, triangles);
+
+    // write depth
+    depth[i] = res.second;
+ 
+    // intersection point is written back to positions.
+    // non-intersect point reaches at most 10 depth
+    positions[i] = ro + res.second * rd;
+
+    // write face_id
+    if (res.first >= 0) {
+        face_id[i] = triangles[res.first].id;
+    } else {
+        face_id[i] = -1;
+    }
+}
+    
+}
\ No newline at end of file

test/renderer.py

+import argparse
+import numpy as np
+import trimesh
+
+import torch
+import cubvh
+
+import dearpygui.dearpygui as dpg
+from scipy.spatial.transform import Rotation as R
+
+def create_dodecahedron(radius=1, center=np.array([0, 0, 0])):
+
+    vertices = np.array([
+        -0.57735,  -0.57735,  0.57735,
+        0.934172,  0.356822,  0,
+        0.934172,  -0.356822,  0,
+        -0.934172,  0.356822,  0,
+        -0.934172,  -0.356822,  0,
+        0,  0.934172,  0.356822,
+        0,  0.934172,  -0.356822,
+        0.356822,  0,  -0.934172,
+        -0.356822,  0,  -0.934172,
+        0,  -0.934172,  -0.356822,
+        0,  -0.934172,  0.356822,
+        0.356822,  0,  0.934172,
+        -0.356822,  0,  0.934172,
+        0.57735,  0.57735,  -0.57735,
+        0.57735,  0.57735,  0.57735,
+        -0.57735,  0.57735,  -0.57735,
+        -0.57735,  0.57735,  0.57735,
+        0.57735,  -0.57735,  -0.57735,
+        0.57735,  -0.57735,  0.57735,
+        -0.57735,  -0.57735,  -0.57735,
+        ]).reshape((-1,3), order="C")
+
+    faces = np.array([
+        19, 3, 2,
+        12, 19, 2,
+        15, 12, 2,
+        8, 14, 2,
+        18, 8, 2,
+        3, 18, 2,
+        20, 5, 4,
+        9, 20, 4,
+        16, 9, 4,
+        13, 17, 4,
+        1, 13, 4,
+        5, 1, 4,
+        7, 16, 4,
+        6, 7, 4,
+        17, 6, 4,
+        6, 15, 2,
+        7, 6, 2,
+        14, 7, 2,
+        10, 18, 3,
+        11, 10, 3,
+        19, 11, 3,
+        11, 1, 5,
+        10, 11, 5,
+        20, 10, 5,
+        20, 9, 8,
+        10, 20, 8,
+        18, 10, 8,
+        9, 16, 7,
+        8, 9, 7,
+        14, 8, 7,
+        12, 15, 6,
+        13, 12, 6,
+        17, 13, 6,
+        13, 1, 11,
+        12, 13, 11,
+        19, 12, 11,
+        ]).reshape((-1, 3), order="C")-1
+
+    length = np.linalg.norm(vertices, axis=1).reshape((-1, 1))
+    vertices = vertices / length * radius + center
+
+    return trimesh.Trimesh(vertices=vertices, faces=faces)
+
+
+class OrbitCamera:
+    def __init__(self, W, H, r=2, fovy=60):
+        self.W = W
+        self.H = H
+        self.radius = r # camera distance from center
+        self.fovy = fovy # in degree
+        self.center = np.array([0, 0, 0], dtype=np.float32) # look at this point
+        self.rot = R.from_quat([1, 0, 0, 0]) # init camera matrix: [[1, 0, 0], [0, -1, 0], [0, 0, 1]] (to suit ngp convention)
+        self.up = np.array([0, 1, 0], dtype=np.float32) # need to be normalized!
+
+    # pose
+    @property
+    def pose(self):
+        # first move camera to radius
+        res = np.eye(4, dtype=np.float32)
+        res[2, 3] -= self.radius
+        # rotate
+        rot = np.eye(4, dtype=np.float32)
+        rot[:3, :3] = self.rot.as_matrix()
+        res = rot @ res
+        # translate
+        res[:3, 3] -= self.center
+        return res
+    
+    # intrinsics
+    @property
+    def intrinsics(self):
+        focal = self.H / (2 * np.tan(np.radians(self.fovy) / 2))
+        return np.array([focal, focal, self.W // 2, self.H // 2])
+    
+    def orbit(self, dx, dy):
+        # rotate along camera up/side axis!
+        side = self.rot.as_matrix()[:3, 0] # why this is side --> ? # already normalized.
+        rotvec_x = self.up * np.radians(-0.05 * dx)
+        rotvec_y = side * np.radians(-0.05 * dy)
+        self.rot = R.from_rotvec(rotvec_x) * R.from_rotvec(rotvec_y) * self.rot
+
+    def scale(self, delta):
+        self.radius *= 1.1 ** (-delta)
+
+    def pan(self, dx, dy, dz=0):
+        # pan in camera coordinate system (careful on the sensitivity!)
+        self.center += 0.0005 * self.rot.as_matrix()[:3, :3] @ np.array([dx, dy, dz])
+
+
+@torch.cuda.amp.autocast(enabled=False)
+def get_rays(poses, intrinsics, H, W, N=-1, error_map=None):
+    ''' get rays
+    Args:
+        poses: [B, 4, 4], cam2world
+        intrinsics: [4]
+        H, W, N: int
+        error_map: [B, 128 * 128], sample probability based on training error
+    Returns:
+        rays_o, rays_d: [B, N, 3]
+        inds: [B, N]
+    '''
+
+    device = poses.device
+    B = poses.shape[0]
+    fx, fy, cx, cy = intrinsics
+
+    i, j = torch.meshgrid(torch.linspace(0, W-1, W, device=device), torch.linspace(0, H-1, H, device=device))
+    i = i.t().reshape([1, H*W]).expand([B, H*W]) + 0.5
+    j = j.t().reshape([1, H*W]).expand([B, H*W]) + 0.5
+
+    results = {}
+
+    if N > 0:
+        N = min(N, H*W)
+
+        if error_map is None:
+            inds = torch.randint(0, H*W, size=[N], device=device) # may duplicate
+            inds = inds.expand([B, N])
+        else:
+
+            # weighted sample on a low-reso grid
+            inds_coarse = torch.multinomial(error_map.to(device), N, replacement=False) # [B, N], but in [0, 128*128)
+
+            # map to the original resolution with random perturb.
+            inds_x, inds_y = inds_coarse // 128, inds_coarse % 128 # `//` will throw a warning in torch 1.10... anyway.
+            sx, sy = H / 128, W / 128
+            inds_x = (inds_x * sx + torch.rand(B, N, device=device) * sx).long().clamp(max=H - 1)
+            inds_y = (inds_y * sy + torch.rand(B, N, device=device) * sy).long().clamp(max=W - 1)
+            inds = inds_x * W + inds_y
+
+            results['inds_coarse'] = inds_coarse # need this when updating error_map
+
+        i = torch.gather(i, -1, inds)
+        j = torch.gather(j, -1, inds)
+
+        results['inds'] = inds
+
+    else:
+        inds = torch.arange(H*W, device=device).expand([B, H*W])
+
+    zs = torch.ones_like(i)
+    xs = (i - cx) / fx * zs
+    ys = (j - cy) / fy * zs
+    directions = torch.stack((xs, ys, zs), dim=-1)
+    directions = directions / torch.norm(directions, dim=-1, keepdim=True)
+    rays_d = directions @ poses[:, :3, :3].transpose(-1, -2) # (B, N, 3)
+
+    rays_o = poses[..., :3, 3] # [B, 3]
+    rays_o = rays_o[..., None, :].expand_as(rays_d) # [B, N, 3]
+
+    results['rays_o'] = rays_o
+    results['rays_d'] = rays_d
+
+    return results
+    
+
+class GUI:
+    def __init__(self, opt, debug=True):
+        self.opt = opt # shared with the trainer's opt to support in-place modification of rendering parameters.
+        self.W = opt.W
+        self.H = opt.H
+        self.debug = debug
+        self.cam = OrbitCamera(opt.W, opt.H, r=opt.radius, fovy=opt.fovy)
+        self.bg_color = torch.ones(3, dtype=torch.float32) # default white bg
+
+        self.render_buffer = np.zeros((self.W, self.H, 3), dtype=np.float32)
+        self.need_update = True # camera moved, should reset accumulation
+        
+        self.mode = 'face_id' # choose from ['position', 'depth', 'face_id']?
+
+        # Set3
+        self.cmap = np.array([
+            (0.5529411764705883, 0.8274509803921568, 0.7803921568627451),
+            (1.0, 1.0, 0.7019607843137254),
+            (0.7450980392156863, 0.7294117647058823, 0.8549019607843137),
+            (0.984313725490196, 0.5019607843137255, 0.4470588235294118),
+            (0.5019607843137255, 0.6941176470588235, 0.8274509803921568),
+            (0.9921568627450981, 0.7058823529411765, 0.3843137254901961),
+            (0.7019607843137254, 0.8705882352941177, 0.4117647058823529),
+            (0.9882352941176471, 0.803921568627451, 0.8980392156862745),
+            (0.8509803921568627, 0.8509803921568627, 0.8509803921568627),
+            (0.7372549019607844, 0.5019607843137255, 0.7411764705882353),
+            (0.8, 0.9215686274509803, 0.7725490196078432),
+            (1.0, 0.9294117647058824, 0.43529411764705883)], dtype=np.float32)
+
+        # load mesh
+        if opt.mesh == '':
+            self.mesh = create_dodecahedron()
+        else:
+            self.mesh = trimesh.load(opt.mesh, force='mesh', skip_material=True)
+
+        # normalize
+        center = self.mesh.vertices.mean(axis=0)
+        length = (self.mesh.vertices.max(axis=0) - self.mesh.vertices.min(axis=0)).max()
+        self.mesh.vertices = (self.mesh.vertices - center) / (length + 1e-5)
+        print(f'[INFO] load mesh {self.mesh.vertices.shape}, {self.mesh.faces.shape}')
+
+        # prepare raytracer
+        self.RT = cubvh.cuBVH(self.mesh.vertices, self.mesh.faces)
+
+        dpg.create_context()
+        self.register_dpg()
+        self.step()
+        
+
+    def __del__(self):
+        dpg.destroy_context()
+
+
+    def prepare_buffer(self, outputs):
+        positions, face_id, depth = outputs
+
+        if self.mode == 'position':
+            # outputs is the actual 3D point, how to visualize them ???
+            # naive normalize...
+            positions = positions.detach().cpu().numpy().reshape(self.H, self.W, 3)
+            positions = (positions - positions.min(axis=0, keepdims=True)) / (positions.max(axis=0, keepdims=True) - positions.min(axis=0, keepdims=True) + 1e-8)
+            return positions
+        elif self.mode == 'face_id':
+            # already normalized to [-1, 1]
+            face_id = face_id.detach().cpu().numpy().reshape(self.H, self.W)
+            mask = face_id < 0 # the bg
+            face_id = self.cmap[face_id % self.cmap.shape[0]]
+            face_id[mask] = 0
+            return face_id
+        elif self.mode == 'depth':
+            depth = depth.detach().cpu().numpy().reshape(self.H, self.W, 1)
+            mask = depth >= 10
+            mn = depth[~mask].min()
+            mx = depth[~mask].max()
+            depth = (depth - mn) / (mx - mn + 1e-5)
+            depth[mask] = 0
+            depth = depth.repeat(3, -1)
+            return depth
+        else:
+            raise NotImplementedError()
+
+    
+    def step(self):
+
+        if self.need_update:
+        
+            starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+            starter.record()
+
+            # outputs = self.trainer.test_gui(self.cam.pose, self.cam.intrinsics, self.W, self.H, self.bg_color, self.spp, self.downscale)
+
+            pose = torch.from_numpy(self.cam.pose).unsqueeze(0).cuda()
+            rays = get_rays(pose, self.cam.intrinsics, self.H, self.W, -1)
+            rays_o = rays['rays_o'].contiguous().view(-1, 3)
+            rays_d = rays['rays_d'].contiguous().view(-1, 3)
+            outputs = self.RT.ray_trace(rays_o, rays_d)
+            
+            ender.record()
+            torch.cuda.synchronize()
+            t = starter.elapsed_time(ender)
+
+            if self.need_update:
+                self.render_buffer = self.prepare_buffer(outputs)
+                self.need_update = False
+            else:
+                self.render_buffer = (self.render_buffer * self.spp + self.prepare_buffer(outputs)) / (self.spp + 1)
+
+            dpg.set_value("_log_infer_time", f'{t:.4f}ms ({int(1000/t)} FPS)')
+            dpg.set_value("_texture", self.render_buffer)
+
+        
+    def register_dpg(self):
+
+        ### register texture 
+
+        with dpg.texture_registry(show=False):
+            dpg.add_raw_texture(self.W, self.H, self.render_buffer, format=dpg.mvFormat_Float_rgb, tag="_texture")
+
+        ### register window
+
+        # the rendered image, as the primary window
+        with dpg.window(tag="_primary_window", width=self.W, height=self.H):
+
+            # add the texture
+            dpg.add_image("_texture")
+
+        dpg.set_primary_window("_primary_window", True)
+
+        # control window
+        with dpg.window(label="Control", tag="_control_window", width=300, height=200):
+
+            # button theme
+            with dpg.theme() as theme_button:
+                with dpg.theme_component(dpg.mvButton):
+                    dpg.add_theme_color(dpg.mvThemeCol_Button, (23, 3, 18))
+                    dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (51, 3, 47))
+                    dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (83, 18, 83))
+                    dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 5)
+                    dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 3, 3)              
+
+            with dpg.group(horizontal=True):
+                dpg.add_text("Infer time: ")
+                dpg.add_text("no data", tag="_log_infer_time")
+            
+            # rendering options
+            with dpg.collapsing_header(label="Options", default_open=True):
+
+                # mode combo
+                def callback_change_mode(sender, app_data):
+                    self.mode = app_data
+                    self.need_update = True
+                
+                dpg.add_combo(('position', 'face_id', 'depth'), label='mode', default_value=self.mode, callback=callback_change_mode)
+
+                # # bg_color picker
+                # def callback_change_bg(sender, app_data):
+                #     self.bg_color = torch.tensor(app_data[:3], dtype=torch.float32) # only need RGB in [0, 1]
+                #     self.need_update = True
+
+                # dpg.add_color_edit((255, 255, 255), label="Background Color", width=200, tag="_color_editor", no_alpha=True, callback=callback_change_bg)
+
+                # fov slider
+                def callback_set_fovy(sender, app_data):
+                    self.cam.fovy = app_data
+                    self.need_update = True
+
+                dpg.add_slider_int(label="FoV (vertical)", min_value=1, max_value=120, format="%d deg", default_value=self.cam.fovy, callback=callback_set_fovy)
+
+            # debug info
+            if self.debug:
+                with dpg.collapsing_header(label="Debug"):
+                    # pose
+                    dpg.add_separator()
+                    dpg.add_text("Camera Pose:")
+                    dpg.add_text(str(self.cam.pose), tag="_log_pose")
+
+
+        ### register camera handler
+
+        def callback_camera_drag_rotate(sender, app_data):
+
+            if not dpg.is_item_focused("_primary_window"):
+                return
+
+            dx = app_data[1]
+            dy = app_data[2]
+
+            self.cam.orbit(dx, dy)
+            self.need_update = True
+
+            if self.debug:
+                dpg.set_value("_log_pose", str(self.cam.pose))
+
+
+        def callback_camera_wheel_scale(sender, app_data):
+
+            if not dpg.is_item_focused("_primary_window"):
+                return
+
+            delta = app_data
+
+            self.cam.scale(delta)
+            self.need_update = True
+
+            if self.debug:
+                dpg.set_value("_log_pose", str(self.cam.pose))
+
+
+        def callback_camera_drag_pan(sender, app_data):
+
+            if not dpg.is_item_focused("_primary_window"):
+                return
+
+            dx = app_data[1]
+            dy = app_data[2]
+
+            self.cam.pan(dx, dy)
+            self.need_update = True
+
+            if self.debug:
+                dpg.set_value("_log_pose", str(self.cam.pose))
+
+
+        with dpg.handler_registry():
+            dpg.add_mouse_drag_handler(button=dpg.mvMouseButton_Left, callback=callback_camera_drag_rotate)
+            dpg.add_mouse_wheel_handler(callback=callback_camera_wheel_scale)
+            dpg.add_mouse_drag_handler(button=dpg.mvMouseButton_Middle, callback=callback_camera_drag_pan)
+
+        
+        dpg.create_viewport(title='mesh viewer', width=self.W, height=self.H, resizable=False)
+
+        ### global theme
+        with dpg.theme() as theme_no_padding:
+            with dpg.theme_component(dpg.mvAll):
+                # set all padding to 0 to avoid scroll bar
+                dpg.add_theme_style(dpg.mvStyleVar_WindowPadding, 0, 0, category=dpg.mvThemeCat_Core)
+                dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 0, 0, category=dpg.mvThemeCat_Core)
+                dpg.add_theme_style(dpg.mvStyleVar_CellPadding, 0, 0, category=dpg.mvThemeCat_Core)
+        
+        dpg.bind_item_theme("_primary_window", theme_no_padding)
+
+        dpg.setup_dearpygui()
+
+        #dpg.show_metrics()
+
+        dpg.show_viewport()
+
+
+    def render(self):
+        while dpg.is_dearpygui_running():
+            self.step()
+            dpg.render_dearpygui_frame()
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--mesh', default='', type=str)
+    parser.add_argument('--W', type=int, default=1920, help="GUI width")
+    parser.add_argument('--H', type=int, default=1080, help="GUI height")
+    parser.add_argument('--radius', type=float, default=5, help="default GUI camera radius from center")
+    parser.add_argument('--fovy', type=float, default=50, help="default GUI camera fovy")
+
+    opt = parser.parse_args()
+
+    gui = GUI(opt)
+    gui.render()

test/unsigned_distance.py

+import numpy as np
+import trimesh
+import argparse
+import torch
+import cubvh
+import time
+
+
+def create_dodecahedron(radius=1, center=np.array([0, 0, 0])):
+
+    vertices = np.array([
+        -0.57735,  -0.57735,  0.57735,
+        0.934172,  0.356822,  0,
+        0.934172,  -0.356822,  0,
+        -0.934172,  0.356822,  0,
+        -0.934172,  -0.356822,  0,
+        0,  0.934172,  0.356822,
+        0,  0.934172,  -0.356822,
+        0.356822,  0,  -0.934172,
+        -0.356822,  0,  -0.934172,
+        0,  -0.934172,  -0.356822,
+        0,  -0.934172,  0.356822,
+        0.356822,  0,  0.934172,
+        -0.356822,  0,  0.934172,
+        0.57735,  0.57735,  -0.57735,
+        0.57735,  0.57735,  0.57735,
+        -0.57735,  0.57735,  -0.57735,
+        -0.57735,  0.57735,  0.57735,
+        0.57735,  -0.57735,  -0.57735,
+        0.57735,  -0.57735,  0.57735,
+        -0.57735,  -0.57735,  -0.57735,
+        ]).reshape((-1,3), order="C")
+
+    faces = np.array([
+        19, 3, 2,
+        12, 19, 2,
+        15, 12, 2,
+        8, 14, 2,
+        18, 8, 2,
+        3, 18, 2,
+        20, 5, 4,
+        9, 20, 4,
+        16, 9, 4,
+        13, 17, 4,
+        1, 13, 4,
+        5, 1, 4,
+        7, 16, 4,
+        6, 7, 4,
+        17, 6, 4,
+        6, 15, 2,
+        7, 6, 2,
+        14, 7, 2,
+        10, 18, 3,
+        11, 10, 3,
+        19, 11, 3,
+        11, 1, 5,
+        10, 11, 5,
+        20, 10, 5,
+        20, 9, 8,
+        10, 20, 8,
+        18, 10, 8,
+        9, 16, 7,
+        8, 9, 7,
+        14, 8, 7,
+        12, 15, 6,
+        13, 12, 6,
+        17, 13, 6,
+        13, 1, 11,
+        12, 13, 11,
+        19, 12, 11,
+        ]).reshape((-1, 3), order="C")-1
+
+    length = np.linalg.norm(vertices, axis=1).reshape((-1, 1))
+    vertices = vertices / length * radius + center
+
+    return trimesh.Trimesh(vertices=vertices, faces=faces)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--N', default=1000, type=int)
+    parser.add_argument('--mesh', default='', type=str)
+    
+    opt = parser.parse_args()
+
+    if opt.mesh == '':
+        mesh = create_dodecahedron()
+    else:
+        mesh = trimesh.load(opt.mesh, force='mesh', skip_material=True)
+
+
+    # query nearest triangles for a set of points
+    points = torch.randn(opt.N, 3, device='cuda', dtype=torch.float32)
+
+    # Ours
+    _t0 = time.time()
+    BVH = cubvh.cuBVH(mesh.vertices, mesh.faces)
+    torch.cuda.synchronize()
+    _t1 = time.time()
+    distances, face_id, uvw = BVH.unsigned_distance(points, return_uvw=True)
+    torch.cuda.synchronize()
+    _t2 = time.time()
+    print(f'[TIME] Ours total {_t2 - _t0:.6f}s = build {_t1 - _t0:.6f}s + query {_t2 - _t1:.6f}s')
+
+    # GT results by trimesh
+    _t0 = time.time()
+    gt_cpoint, gt_distances, gt_face_id = trimesh.proximity.closest_point(mesh, points.cpu().numpy())
+    _t1 = time.time()
+    print(f'[TIME] Trimesh total {_t1 - _t0:.6f}s')
+
+    # verify correctness
+    face_id = face_id.cpu().numpy().astype(np.int64)
+
+    # NOTE: if there are duplicated vertices, this will fail... but it won't affect later correctness.
+    # np.testing.assert_array_equal(
+    #     face_id,
+    #     gt_face_id.astype(np.int64)
+    # )
+    
+    distances = distances.cpu().numpy().astype(np.float32)
+    np.testing.assert_allclose(
+        distances,
+        gt_distances.astype(np.float32),
+        atol=1e-5
+    )
+    
+    # calc cpoint from uvw and face_id
+    uvw = uvw.cpu().numpy().astype(np.float32) # [M, 3]
+    trigs = mesh.faces[face_id] # [M, 3]
+
+    cpoint = mesh.vertices[trigs[:, 0]] * uvw[:, [0]] + \
+             mesh.vertices[trigs[:, 1]] * uvw[:, [1]] + \
+             mesh.vertices[trigs[:, 2]] * uvw[:, [2]]
+
+    np.testing.assert_allclose(
+        cpoint,
+        gt_cpoint.astype(np.float32),
+        atol=1e-5
+    )
+
+    
\ No newline at end of file