warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/native/bvh.cu

@@ -31,11 +31,22 @@
 
 #include <cub/cub.cuh>
 
+extern CUcontext get_current_context();
 
 namespace wp
 {
-    void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh);
-    void bvh_destroy_host(BVH& bvh);
+void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh, int leaf_size);
+void bvh_destroy_host(BVH& bvh);
+
+__global__ void memset_kernel(int* dest, int value, size_t n)
+{
+    const size_t tid = static_cast<size_t>(blockDim.x) * static_cast<size_t>(blockIdx.x) + static_cast<size_t>(threadIdx.x);
+    
+    if (tid < n)
+    {
+        dest[tid] = value;
+    }
+}
 
 // for LBVH: this will start with some muted leaf nodes, but that is okay, we can still trace up because there parents information is still valid
 // the only thing worth mentioning is that when the parent leaf node is also a leaf node, we need to recompute its bounds, since their child information are lost
@@ -149,17 +160,6 @@ __global__ void bvh_refit_kernel(int n, const int* __restrict__ parents, int* __
     }
 }
 
-
-void bvh_refit_device(BVH& bvh)
-{
-    ContextGuard guard(bvh.context);
-
-    // clear child counters
-    wp_memset_device(WP_CURRENT_CONTEXT, bvh.node_counts, 0, sizeof(int) * bvh.max_nodes);
-    wp_launch_device(WP_CURRENT_CONTEXT, bvh_refit_kernel, bvh.num_leaf_nodes, (bvh.num_leaf_nodes, bvh.node_parents, bvh.node_counts, bvh.primitive_indices, bvh.node_lowers, bvh.node_uppers, bvh.item_lowers, bvh.item_uppers));
-}
-
-
 /////////////////////////////////////////////////////////////////////////////////////////////
 
 // Create a linear BVH as described in Fast and Simple Agglomerative LBVH construction
@@ -379,18 +379,18 @@ __global__ void build_hierarchy(int n, int* root, const int* __restrict__ deltas
 
 /*
 * LBVH uses a bottom-up constructor which makes variable-sized leaf nodes more challenging to achieve. 
-* Simply splitting the ordered primitives into uniform groups of size BVH_LEAF_SIZE will result in poor
+* Simply splitting the ordered primitives into uniform groups of size leaf_size will result in poor
 * quality. Instead, after the hierarchy is built, we convert any intermediate node whose size is 
-* <= BVH_LEAF_SIZE into a new leaf node. This process is done using the new kernel function called 
+* <= leaf_size into a new leaf node. This process is done using the new kernel function called 
 * mark_packed_leaf_nodes .
 */
 __global__ void mark_packed_leaf_nodes(int n, const int* __restrict__ range_lefts, const int* __restrict__ range_rights, const int* __restrict__ parents,
-    BVHPackedNodeHalf* __restrict__ lowers, BVHPackedNodeHalf* __restrict__ uppers)
+    BVHPackedNodeHalf* __restrict__ lowers, BVHPackedNodeHalf* __restrict__ uppers, const int leaf_size)
 {
     int node_index = blockDim.x * blockIdx.x + threadIdx.x;
     if (node_index < n)
     {
-        // mark the node as leaf if its range is less than LEAF_SIZE_LBVH or it is deeper than BVH_QUERY_STACK_SIZE
+        // mark the node as leaf if its range is less than leaf_size or it is deeper than BVH_QUERY_STACK_SIZE
         // this will forever mute its child nodes so that they will never be accessed
 
         // calculate depth
@@ -407,7 +407,7 @@ __global__ void mark_packed_leaf_nodes(int n, const int* __restrict__ range_left
         // the LBVH constructor's range is defined as left <= i <= right
         // we need to convert it to our convention: left <= i < right
         int right = range_rights[node_index] + 1;
-        if (right - left <= BVH_LEAF_SIZE || depth >= BVH_QUERY_STACK_SIZE)
+        if (right - left <= leaf_size || depth >= BVH_QUERY_STACK_SIZE)
         {
             lowers[node_index].b = 1;
             lowers[node_index].i = left;
@@ -514,11 +514,15 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
     }
     else
     {
-        static vec3 upper(-FLT_MAX);
-        static vec3 lower(FLT_MAX);
+        // IEEE-754 bit patterns for +/- FLT_MAX
+        constexpr int FLT_MAX_BITS = 0x7f7fffff;
+        constexpr int NEG_FLT_MAX_BITS = 0xff7fffff;
+
+        // total_lower := ( +FLT_MAX, +FLT_MAX, +FLT_MAX )
+        wp_launch_device(WP_CURRENT_CONTEXT, memset_kernel, sizeof(vec3) / 4, ((int*)total_lower, FLT_MAX_BITS, sizeof(vec3) / 4));
 
-        wp_memcpy_h2d(WP_CURRENT_CONTEXT, total_lower, &lower, sizeof(lower));
-        wp_memcpy_h2d(WP_CURRENT_CONTEXT, total_upper, &upper, sizeof(upper));
+        // total_upper := ( -FLT_MAX, -FLT_MAX, -FLT_MAX )
+        wp_launch_device(WP_CURRENT_CONTEXT, memset_kernel, sizeof(vec3) / 4, ((int*)total_upper, NEG_FLT_MAX_BITS, sizeof(vec3) / 4));
 
         // compute the total bounds on the GPU
         wp_launch_device(WP_CURRENT_CONTEXT, compute_total_bounds, num_items, (item_lowers, item_uppers, total_lower, total_upper, num_items));
@@ -545,7 +549,7 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
 
     // build the tree and internal node bounds
     wp_launch_device(WP_CURRENT_CONTEXT, build_hierarchy, num_items, (num_items, bvh.root, deltas, num_children, bvh.primitive_indices, range_lefts, range_rights, bvh.node_parents, bvh.node_lowers, bvh.node_uppers));
-    wp_launch_device(WP_CURRENT_CONTEXT, mark_packed_leaf_nodes, bvh.max_nodes, (bvh.max_nodes, range_lefts, range_rights, bvh.node_parents, bvh.node_lowers, bvh.node_uppers));
+    wp_launch_device(WP_CURRENT_CONTEXT, mark_packed_leaf_nodes, bvh.max_nodes, (bvh.max_nodes, range_lefts, range_rights, bvh.node_parents, bvh.node_lowers, bvh.node_uppers, bvh.leaf_size));
 
     // free temporary memory
     wp_free_device(WP_CURRENT_CONTEXT, indices);
@@ -661,6 +665,7 @@ void copy_host_tree_to_device(void* context, BVH& bvh_host, BVH& bvh_device_on_h
     bvh_device_on_host.max_nodes = bvh_host.max_nodes;
     bvh_device_on_host.num_items = bvh_host.num_items;
     bvh_device_on_host.max_depth = bvh_host.max_depth;
+    bvh_device_on_host.leaf_size = bvh_host.leaf_size;
 
     bvh_device_on_host.root = (int*)wp_alloc_device(context, sizeof(int));
     wp_memcpy_h2d(context, bvh_device_on_host.root, bvh_host.root, sizeof(int));
@@ -673,7 +678,7 @@ void copy_host_tree_to_device(void* context, BVH& bvh_host, BVH& bvh_device_on_h
 }
 
 // create in-place given existing descriptor
-void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh_device_on_host)
+void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh_device_on_host, int leaf_size)
 {
     ContextGuard guard(context);
     if (constructor_type == BVH_CONSTRUCTOR_SAH || constructor_type == BVH_CONSTRUCTOR_MEDIAN)
@@ -687,7 +692,7 @@ void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items,
 
         // run CPU based constructor
         wp::BVH bvh_host;
-        wp::bvh_create_host(lowers_host.data(), uppers_host.data(), num_items, constructor_type, bvh_host);
+        wp::bvh_create_host(lowers_host.data(), uppers_host.data(), num_items, constructor_type, bvh_host, leaf_size);
 
         // copy host tree to device
         wp::copy_host_tree_to_device(WP_CURRENT_CONTEXT, bvh_host, bvh_device_on_host);
@@ -700,6 +705,7 @@ void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items,
     }
     else if (constructor_type == BVH_CONSTRUCTOR_LBVH)
     {
+        bvh_device_on_host.leaf_size = leaf_size;
         bvh_device_on_host.num_items = num_items;
         bvh_device_on_host.max_nodes = 2 * num_items - 1;
         bvh_device_on_host.num_leaf_nodes = num_items;
@@ -737,6 +743,24 @@ void bvh_destroy_device(BVH& bvh)
     wp_free_device(WP_CURRENT_CONTEXT, bvh.root); bvh.root = NULL;
 }
 
+void bvh_refit_device(BVH& bvh)
+{
+    ContextGuard guard(bvh.context);
+
+    // clear child counters
+    wp_memset_device(WP_CURRENT_CONTEXT, bvh.node_counts, 0, sizeof(int) * bvh.max_nodes);
+    wp_launch_device(WP_CURRENT_CONTEXT, bvh_refit_kernel, bvh.num_leaf_nodes, (bvh.num_leaf_nodes, bvh.node_parents, bvh.node_counts, bvh.primitive_indices, bvh.node_lowers, bvh.node_uppers, bvh.item_lowers, bvh.item_uppers));
+}
+
+void bvh_rebuild_device(BVH& bvh)
+{
+    ContextGuard guard(bvh.context);
+
+    LinearBVHBuilderGPU builder;
+    builder.build(bvh, bvh.item_lowers, bvh.item_uppers, bvh.num_items, NULL);
+}
+
+
 
 } // namespace wp
 
@@ -752,6 +776,17 @@ void wp_bvh_refit_device(uint64_t id)
     }
 }
 
+void wp_bvh_rebuild_device(uint64_t id)
+{
+    wp::BVH bvh;
+    if (bvh_get_descriptor(id, bvh))
+    {
+        ContextGuard guard(bvh.context);
+
+        wp::bvh_rebuild_device(bvh);
+    }
+}
+
 /*
 * Since we don't even know the number of true leaf nodes, never mention where they are, we will launch
 * the num_items threads, which are identical to the number of leaf nodes in the original tree. The 
@@ -759,13 +794,13 @@ void wp_bvh_refit_device(uint64_t id)
 * muted. However, the muted leaf nodes will still have the pointer to their parents, thus the up-tracing
 * can still work. We will only compute the bounding box of a leaf node if its parent is not a leaf node.
 */
-uint64_t wp_bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type)
+uint64_t wp_bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type, int leaf_size)
 {
     ContextGuard guard(context);
     wp::BVH bvh_device_on_host;
     wp::BVH* bvh_device_ptr = nullptr;
     
-    wp::bvh_create_device(WP_CURRENT_CONTEXT, lowers, uppers, num_items, constructor_type, bvh_device_on_host);
+    wp::bvh_create_device(WP_CURRENT_CONTEXT, lowers, uppers, num_items, constructor_type, bvh_device_on_host, leaf_size);
 
     // create device-side BVH descriptor
     bvh_device_ptr = (wp::BVH*)wp_alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::BVH));