warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/native/bvh.h

@@ -0,0 +1,555 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "builtin.h"
+#include "intersect.h"
+
+#ifdef __CUDA_ARCH__
+#define BVH_SHARED_STACK 1
+#else
+#define BVH_SHARED_STACK 0
+#endif
+
+#define SAH_NUM_BUCKETS (16)
+#define USE_LOAD4
+#define BVH_QUERY_STACK_SIZE (32)
+
+#define BVH_CONSTRUCTOR_SAH (0)
+#define BVH_CONSTRUCTOR_MEDIAN (1)
+#define BVH_CONSTRUCTOR_LBVH (2)
+
+namespace wp
+{
+
+struct bounds3
+{
+    CUDA_CALLABLE inline bounds3() : lower( FLT_MAX)
+                                   , upper(-FLT_MAX) {}
+
+    CUDA_CALLABLE inline bounds3(const vec3& lower, const vec3& upper) : lower(lower), upper(upper) {}
+
+    CUDA_CALLABLE inline vec3 center() const { return 0.5f*(lower+upper); }
+    CUDA_CALLABLE inline vec3 edges() const { return upper-lower; }
+
+    CUDA_CALLABLE inline void expand(float r)
+    {
+        lower -= vec3(r);
+        upper += vec3(r);
+    }
+
+    CUDA_CALLABLE inline void expand(const vec3& r)
+    {
+        lower -= r;
+        upper += r;
+    }
+
+    CUDA_CALLABLE inline bool empty() const { return lower[0] >= upper[0] || lower[1] >= upper[1] || lower[2] >= upper[2]; }
+
+    CUDA_CALLABLE inline bool overlaps(const vec3& p) const
+    {
+        if (p[0] < lower[0] ||
+            p[1] < lower[1] ||
+            p[2] < lower[2] ||
+            p[0] > upper[0] ||
+            p[1] > upper[1] ||
+            p[2] > upper[2])
+        {
+            return false;
+        }
+        else
+        {
+            return true;
+        }
+    }
+
+    CUDA_CALLABLE inline bool overlaps(const bounds3& b) const
+    {
+        if (lower[0] > b.upper[0] ||
+            lower[1] > b.upper[1] ||
+            lower[2] > b.upper[2] ||
+            upper[0] < b.lower[0] ||
+            upper[1] < b.lower[1] ||
+            upper[2] < b.lower[2])
+        {
+            return false;
+        }
+        else
+        {
+            return true;
+        }
+    }
+
+    CUDA_CALLABLE inline bool overlaps(const vec3& b_lower, const vec3& b_upper) const
+    {
+        if (lower[0] > b_upper[0] ||
+            lower[1] > b_upper[1] ||
+            lower[2] > b_upper[2] ||
+            upper[0] < b_lower[0] ||
+            upper[1] < b_lower[1] ||
+            upper[2] < b_lower[2])
+        {
+            return false;
+        }
+        else
+        {
+            return true;
+        }
+    }
+
+    CUDA_CALLABLE inline void add_point(const vec3& p)
+    {
+        lower = min(lower, p);
+        upper = max(upper, p);
+    }
+
+    CUDA_CALLABLE inline void add_bounds(const vec3& lower_other, const vec3& upper_other)
+    {
+        // lower_other will only impact the lower of the new bounds
+        // upper_other will only impact the upper of the new bounds
+        // this costs only half of the computation of adding lower_other and upper_other separately
+        lower = min(lower, lower_other);
+        upper = max(upper, upper_other);
+    }
+
+    CUDA_CALLABLE inline float area() const
+    {
+        vec3 e = upper-lower;
+        return 2.0f*(e[0]*e[1] + e[0]*e[2] + e[1]*e[2]);
+    }
+
+    vec3 lower;
+    vec3 upper;
+};
+
+CUDA_CALLABLE inline bounds3 bounds_union(const bounds3& a, const vec3& b) 
+{
+    return bounds3(min(a.lower, b), max(a.upper, b));
+}
+
+CUDA_CALLABLE inline bounds3 bounds_union(const bounds3& a, const bounds3& b) 
+{
+    return bounds3(min(a.lower, b.lower), max(a.upper, b.upper));
+}
+
+CUDA_CALLABLE inline bounds3 bounds_intersection(const bounds3& a, const bounds3& b)
+{
+    return bounds3(max(a.lower, b.lower), min(a.upper, b.upper));
+}
+
+struct BVHPackedNodeHalf
+{
+    float x;
+    float y;
+    float z;
+    // For non-leaf nodes:
+    // - 'lower.i' represents the index of the left child node.
+    // - 'upper.i' represents the index of the right child node.
+    //
+    // For leaf nodes:
+    // - 'lower.i' indicates the start index of the primitives in 'primitive_indices'.
+    // - 'upper.i' indicates the index just after the last primitive in 'primitive_indices'
+    unsigned int i : 31;
+    unsigned int b : 1;
+};
+
+struct BVH
+{		
+    BVHPackedNodeHalf* node_lowers;
+    BVHPackedNodeHalf* node_uppers;
+
+    // used for fast refits
+    int* node_parents;
+    int* node_counts;
+    // reordered primitive indices corresponds to the ordering of leaf nodes
+    int* primitive_indices;
+    
+    int max_depth;
+    int max_nodes;
+    int num_nodes;
+    // since we use packed leaf nodes, the number of them is no longer the number of items, but variable
+    int num_leaf_nodes;
+
+    // pointer (CPU or GPU) to a single integer index in node_lowers, node_uppers
+    // representing the root of the tree, this is not always the first node
+    // for bottom-up builders
+    int* root;
+
+    // item bounds are not owned by the BVH but by the caller
+    vec3* item_lowers;
+    vec3* item_uppers;
+    int num_items;
+
+    int leaf_size;
+
+    // cuda context
+    void* context;
+};
+
+CUDA_CALLABLE inline BVHPackedNodeHalf make_node(const vec3& bound, int child, bool leaf)
+{
+    BVHPackedNodeHalf n;
+    n.x = bound[0];
+    n.y = bound[1];
+    n.z = bound[2];
+    n.i = (unsigned int)child;
+    n.b = (unsigned int)(leaf?1:0);
+
+    return n;
+}
+
+// variation of make_node through volatile pointers used in build_hierarchy
+CUDA_CALLABLE inline void make_node(volatile BVHPackedNodeHalf* n, const vec3& bound, int child, bool leaf)
+{
+    n->x = bound[0];
+    n->y = bound[1];
+    n->z = bound[2];
+    n->i = (unsigned int)child;
+    n->b = (unsigned int)(leaf?1:0);
+}
+
+#ifdef __CUDA_ARCH__
+__device__ inline wp::BVHPackedNodeHalf bvh_load_node(const wp::BVHPackedNodeHalf* nodes, int index)
+{
+#ifdef USE_LOAD4
+    float4 f4 = __ldg((const float4*)(nodes)+index);
+    return  (const wp::BVHPackedNodeHalf&)f4;
+    //return  (const wp::BVHPackedNodeHalf&)(*((const float4*)(nodes)+index));
+#else
+    return  nodes[index];
+#endif // USE_LOAD4
+
+}
+#else
+inline wp::BVHPackedNodeHalf bvh_load_node(const wp::BVHPackedNodeHalf* nodes, int index)
+{
+    return  nodes[index];
+}
+#endif // __CUDACC__
+
+CUDA_CALLABLE inline int clz(int x)
+{
+    int n;
+    if (x == 0) return 32;
+    for (n = 0; ((x & 0x80000000) == 0); n++, x <<= 1);
+    return n;
+}
+
+CUDA_CALLABLE inline uint32_t part1by2(uint32_t n)
+{
+    n = (n ^ (n << 16)) & 0xff0000ff;
+    n = (n ^ (n <<  8)) & 0x0300f00f;
+    n = (n ^ (n <<  4)) & 0x030c30c3;
+    n = (n ^ (n <<  2)) & 0x09249249;
+
+    return n;
+}
+
+// Takes values in the range [0, 1] and assigns an index based Morton codes of length 3*lwp2(dim) bits 
+template <int dim>
+CUDA_CALLABLE inline uint32_t morton3(float x, float y, float z)
+{
+    uint32_t ux = clamp(int(x*dim), 0, dim-1);
+    uint32_t uy = clamp(int(y*dim), 0, dim-1);
+    uint32_t uz = clamp(int(z*dim), 0, dim-1);
+
+    return (part1by2(uz) << 2) | (part1by2(uy) << 1) | part1by2(ux);
+}
+
+// making the class accessible from python
+
+CUDA_CALLABLE inline BVH bvh_get(uint64_t id)
+{
+    return *(BVH*)(id);
+}
+
+CUDA_CALLABLE inline int bvh_get_num_bounds(uint64_t id)
+{
+    BVH bvh = bvh_get(id);
+    return bvh.num_items;
+}
+
+// represents a strided stack in shared memory
+// so each level of the stack is stored contiguously
+// across the block
+struct bvh_stack_t
+{
+    inline int operator[](int depth) const { return ptr[depth*WP_TILE_BLOCK_DIM]; }
+    inline int& operator[](int depth) { return ptr[depth*WP_TILE_BLOCK_DIM]; }
+
+    int* ptr;
+
+};
+
+
+// stores state required to traverse the BVH nodes that 
+// overlap with a query AABB.
+struct bvh_query_t
+{
+    CUDA_CALLABLE bvh_query_t()
+        : bvh(),
+          stack(),
+          count(0),
+          is_ray(false),
+          input_lower(),
+          input_upper(),
+          bounds_nr(0),
+          primitive_counter(-1)
+    {}
+
+    // Required for adjoint computations.
+    CUDA_CALLABLE inline bvh_query_t& operator+=(const bvh_query_t& other)
+    {
+        return *this;
+    }
+
+    BVH bvh;
+
+    // BVH traversal stack:
+#if BVH_SHARED_STACK
+    bvh_stack_t stack;
+#else
+    int stack[BVH_QUERY_STACK_SIZE];
+#endif
+
+    int count;
+
+    // >= 0 if currently in a packed leaf node
+    int primitive_counter;
+    
+    // inputs
+    wp::vec3 input_lower;	// start for ray
+    wp::vec3 input_upper;	// dir for ray
+
+    int bounds_nr;
+    bool is_ray;
+};
+
+CUDA_CALLABLE inline bool bvh_query_intersection_test(const bvh_query_t& query, const vec3& node_lower, const vec3& node_upper)
+{
+    if (query.is_ray)
+    {
+        float t = 0.0f;
+        return intersect_ray_aabb(query.input_lower, query.input_upper, node_lower, node_upper, t);
+    }
+    else
+    {
+        return intersect_aabb_aabb(query.input_lower, query.input_upper, node_lower, node_upper);
+    }
+}
+
+CUDA_CALLABLE inline bvh_query_t bvh_query(
+    uint64_t id, bool is_ray, const vec3& lower, const vec3& upper)
+{
+    // This routine traverses the BVH tree until it finds
+    // the first overlapping bound. 
+
+    // initialize empty
+    bvh_query_t query;
+    
+#if BVH_SHARED_STACK
+    __shared__ int stack[BVH_QUERY_STACK_SIZE*WP_TILE_BLOCK_DIM];
+    query.stack.ptr = &stack[threadIdx.x];
+#endif
+
+    query.bounds_nr = -1;
+
+    BVH bvh = bvh_get(id);
+
+    query.bvh = bvh;
+    query.is_ray = is_ray;
+
+    // optimization: make the latest	
+    query.stack[0] = *bvh.root;
+    query.count = 1;
+    query.input_lower = lower;
+    query.input_upper = upper;
+
+    // Navigate through the bvh, find the first overlapping leaf node.
+    while (query.count)
+    {
+        const int node_index = query.stack[--query.count];
+        BVHPackedNodeHalf node_lower = bvh_load_node(bvh.node_lowers, node_index);
+        BVHPackedNodeHalf node_upper = bvh_load_node(bvh.node_uppers, node_index);
+
+        if (!bvh_query_intersection_test(query, reinterpret_cast<vec3&>(node_lower), reinterpret_cast<vec3&>(node_upper)))
+        {
+            continue;
+        }
+
+        const int left_index = node_lower.i;
+        const int right_index = node_upper.i;
+        // Make bounds from this AABB
+        if (node_lower.b)
+        {
+            // Reached a leaf node, point to its first primitive
+            // Back up one level and return 
+            query.primitive_counter = 0;
+            query.stack[query.count++] = node_index;
+            return query;
+        }
+        else
+        {
+            query.stack[query.count++] = left_index;
+            query.stack[query.count++] = right_index;
+        }
+    }
+
+    return query;
+}
+
+CUDA_CALLABLE inline bvh_query_t bvh_query_aabb(
+    uint64_t id, const vec3& lower, const vec3& upper)
+{
+    return bvh_query(id, false, lower, upper);
+}
+
+CUDA_CALLABLE inline bvh_query_t bvh_query_ray(uint64_t id, const vec3& start, const vec3& dir)
+{
+    return bvh_query(id, true, start, 1.0f / dir);
+}
+
+//Stub
+CUDA_CALLABLE inline void adj_bvh_query_aabb(uint64_t id, const vec3& lower, const vec3& upper,
+                                               uint64_t, vec3&, vec3&, bvh_query_t&)
+{
+}
+
+
+CUDA_CALLABLE inline void adj_bvh_query_ray(uint64_t id, const vec3& start, const vec3& dir,
+                                               uint64_t, vec3&, vec3&, bvh_query_t&)
+{
+}
+
+
+CUDA_CALLABLE inline bool bvh_query_next(bvh_query_t& query, int& index)
+{
+    BVH bvh = query.bvh;
+
+    // Navigate through the bvh, find the first overlapping leaf node.
+    while (query.count)
+    {
+        const int node_index = query.stack[--query.count];
+
+        BVHPackedNodeHalf node_lower = bvh_load_node(bvh.node_lowers, node_index);
+        BVHPackedNodeHalf node_upper = bvh_load_node(bvh.node_uppers, node_index);
+
+        if (query.primitive_counter == 0) {
+            if (!bvh_query_intersection_test(query, reinterpret_cast<vec3&>(node_lower), reinterpret_cast<vec3&>(node_upper)))
+            {
+                continue;
+            }
+        }
+
+        const int left_index = node_lower.i;
+        const int right_index = node_upper.i;
+
+        if (node_lower.b)
+        {
+            // found leaf, loop through its content primitives
+            const int start = left_index;
+            const int end = right_index;
+
+            if (end - start <= 1)
+            {
+                int primitive_index = bvh.primitive_indices[start];
+                index = primitive_index;
+                query.bounds_nr = primitive_index;
+                return true;
+            }
+            else
+            {   
+                int primitive_index = bvh.primitive_indices[start + (query.primitive_counter++)];
+
+                // if already visited the last primitive in the leaf node
+                // move to the next node and reset the primitive counter to 0
+                if (start + query.primitive_counter == end)
+                {
+                    query.primitive_counter = 0;
+                }
+                // otherwise we need to keep this leaf node in stack for a future visit
+                else
+                {
+                    query.stack[query.count++] = node_index;
+                }
+                // return true;
+                if (bvh_query_intersection_test(query, bvh.item_lowers[primitive_index], bvh.item_uppers[primitive_index]))
+                {
+                    index = primitive_index;
+                    query.bounds_nr = primitive_index;
+
+                    return true;
+                }
+            }
+        }
+        else
+        {
+            // if it's not a leaf node we treat it as if we have visited the last primitive
+            query.primitive_counter = 0;
+            query.stack[query.count++] = left_index;
+            query.stack[query.count++] = right_index;
+        }
+    }
+    return false;
+}
+
+CUDA_CALLABLE inline int iter_next(bvh_query_t& query)
+{
+    return query.bounds_nr;
+}
+
+CUDA_CALLABLE inline bool iter_cmp(bvh_query_t& query)
+{
+    bool finished = bvh_query_next(query, query.bounds_nr);
+    return finished;
+}
+
+CUDA_CALLABLE inline bvh_query_t iter_reverse(const bvh_query_t& query)
+{
+    // can't reverse BVH queries, users should not rely on traversal ordering
+    return query;
+}
+
+CUDA_CALLABLE inline void adj_iter_reverse(const bvh_query_t& query, bvh_query_t& adj_query, bvh_query_t& adj_ret)
+{
+}
+
+
+// stub
+CUDA_CALLABLE inline void adj_bvh_query_next(bvh_query_t& query, int& index, bvh_query_t&, int&, bool&) 
+{
+
+}
+
+CUDA_CALLABLE bool bvh_get_descriptor(uint64_t id, BVH& bvh);
+CUDA_CALLABLE void bvh_add_descriptor(uint64_t id, const BVH& bvh);
+CUDA_CALLABLE void bvh_rem_descriptor(uint64_t id);
+
+void bvh_create_host(vec3* lowers, vec3* uppers, int num_items,  int constructor_type, BVH& bvh, int leaf_size);
+void bvh_destroy_host(wp::BVH& bvh);
+void bvh_refit_host(wp::BVH& bvh);
+
+#if WP_ENABLE_CUDA
+
+void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh_device_on_host, int leaf_size);
+void bvh_destroy_device(BVH& bvh);
+void bvh_refit_device(BVH& bvh);
+
+#endif // WP_ENABLE_CUDA
+
+} // namespace wp