warp-lang 1.5.0__py3-none-manylinux2014_aarch64.whl → 1.6.0__py3-none-manylinux2014_aarch64.whl

warp/native/bvh.cu

@@ -18,31 +18,49 @@
 #include <cuda_runtime_api.h>
 
 #define THRUST_IGNORE_CUB_VERSION_CHECK
+#define REORDER_HOST_TREE
 
 #include <cub/cub.cuh>
 
 
 namespace wp
 {
+    void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh);
+    void bvh_destroy_host(BVH& bvh);
 
-__global__ void bvh_refit_kernel(int n, const int* __restrict__ parents, int* __restrict__ child_count, BVHPackedNodeHalf* __restrict__ node_lowers, BVHPackedNodeHalf* __restrict__ node_uppers, const vec3* item_lowers, const vec3* item_uppers)
+// 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
+// for a compact tree such as those from SAH or Median constructor, there is no muted leaf nodes
+__global__ void bvh_refit_kernel(int n, const int* __restrict__ parents, int* __restrict__ child_count, int* __restrict__ primitive_indices, BVHPackedNodeHalf* __restrict__ node_lowers, BVHPackedNodeHalf* __restrict__ node_uppers, const vec3* item_lowers, const vec3* item_uppers)
 {
     int index = blockDim.x*blockIdx.x + threadIdx.x;
 
     if (index < n)
     {
         bool leaf = node_lowers[index].b;
+        int parent = parents[index];
 
         if (leaf)
         {
+            BVHPackedNodeHalf& lower = node_lowers[index];
+            BVHPackedNodeHalf& upper = node_uppers[index];
             // update the leaf node
-            const int leaf_index = node_lowers[index].i;
 
-            vec3 lower = item_lowers[leaf_index];
-            vec3 upper = item_uppers[leaf_index];
-
-            make_node(node_lowers+index, lower, leaf_index, true);
-            make_node(node_uppers+index, upper, 0, false);
+            // only need to compute bound when this is a valid leaf node
+            if (!node_lowers[parent].b)
+            {
+                const int start = lower.i;
+                const int end = upper.i;
+
+                bounds3 bound;
+                for (int primitive_counter = start; primitive_counter < end; primitive_counter++)
+                {
+                    const int primitive = primitive_indices[primitive_counter];
+                    bound.add_bounds(item_lowers[primitive], item_uppers[primitive]);
+                }
+                (vec3&)lower = bound.lower;
+                (vec3&)upper = bound.upper;
+            }
         }
         else
         {
@@ -53,8 +71,7 @@ __global__ void bvh_refit_kernel(int n, const int* __restrict__ parents, int* __
         // update hierarchy
         for (;;)
         {
-            int parent = parents[index];
-            
+            parent = parents[index];
             // reached root
             if (parent == -1)
                 return;
@@ -68,34 +85,49 @@ __global__ void bvh_refit_kernel(int n, const int* __restrict__ parents, int* __
             // then update its bounds and move onto the next parent in the hierarchy
             if (finished == 1)
             {
-                const int left_child = node_lowers[parent].i;
-                const int right_child = node_uppers[parent].i;
-
-                vec3 left_lower = vec3(node_lowers[left_child].x,
-                                       node_lowers[left_child].y, 
-                                       node_lowers[left_child].z);
-
-                vec3 left_upper = vec3(node_uppers[left_child].x,
-                                       node_uppers[left_child].y, 
-                                       node_uppers[left_child].z);
-
-                vec3 right_lower = vec3(node_lowers[right_child].x,
-                                        node_lowers[right_child].y,
-                                        node_lowers[right_child].z);
-
-
-                vec3 right_upper = vec3(node_uppers[right_child].x, 
-                                        node_uppers[right_child].y, 
-                                        node_uppers[right_child].z);
-
-                // union of child bounds
-                vec3 lower = min(left_lower, right_lower);
-                vec3 upper = max(left_upper, right_upper);
-                
-                // write new BVH nodes
-                make_node(node_lowers+parent, lower, left_child, false);
-                make_node(node_uppers+parent, upper, right_child, false);
-
+                BVHPackedNodeHalf& parent_lower = node_lowers[parent];
+                BVHPackedNodeHalf& parent_upper = node_uppers[parent];
+                if (parent_lower.b)
+                    // a packed leaf node can still be a parent in LBVH, we need to recompute its bounds
+                    // since we've lost its left and right child node index in the muting process
+                {
+                    // update the leaf node
+                    int parent_parent = parents[parent];;
+
+                    // only need to compute bound when this is a valid leaf node
+                    if (!node_lowers[parent_parent].b)
+                    {
+                        const int start = parent_lower.i;
+                        const int end = parent_upper.i;
+                        bounds3 bound;
+                        for (int primitive_counter = start; primitive_counter < end; primitive_counter++)
+                        {
+                            const int primitive = primitive_indices[primitive_counter];
+                            bound.add_bounds(item_lowers[primitive], item_uppers[primitive]);
+                        }
+
+                        (vec3&)parent_lower = bound.lower;
+                        (vec3&)parent_upper = bound.upper;
+                    }
+                }
+                else
+                {
+                    const int left_child = parent_lower.i;
+                    const int right_child = parent_upper.i;
+
+                    vec3 left_lower = (vec3&)(node_lowers[left_child]);
+                    vec3 left_upper = (vec3&)(node_uppers[left_child]);
+                    vec3 right_lower = (vec3&)(node_lowers[right_child]);
+                    vec3 right_upper = (vec3&)(node_uppers[right_child]);
+
+                    // union of child bounds
+                    vec3 lower = min(left_lower, right_lower);
+                    vec3 upper = max(left_upper, right_upper);
+
+                    // write new BVH nodes
+                    (vec3&)parent_lower = lower;
+                    (vec3&)parent_upper = upper;
+                }
                 // move onto processing the parent
                 index = parent;
             }
@@ -114,9 +146,8 @@ void bvh_refit_device(BVH& bvh)
     ContextGuard guard(bvh.context);
 
     // clear child counters
-    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_items, (bvh.num_items, bvh.node_parents, bvh.node_counts, bvh.node_lowers, bvh.node_uppers, bvh.item_lowers, bvh.item_uppers));
+    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));
 }
 
 
@@ -316,6 +347,39 @@ __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
+* 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 
+* mark_packed_leaf_nodes .
+*/
+__global__ void mark_packed_leaf_nodes(int n, volatile int* __restrict__ range_lefts, volatile int* __restrict__ range_rights,
+    volatile BVHPackedNodeHalf* __restrict__ lowers, volatile BVHPackedNodeHalf* __restrict__ uppers)
+{
+    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
+        // this will forever mute its child nodes so that they will never be accessed
+
+        int left = range_lefts[node_index];
+        // 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;
+        // printf("node %d (left %d right %d)", node_index, left, right);
+        if (right - left <= BVH_LEAF_SIZE)
+        {
+            lowers[node_index].b = 1;
+            lowers[node_index].i = left;
+            uppers[node_index].i = right;
+
+            // printf("node %d (left %d right %d) is set to child\n", node_index, left, right);
+        }
+    }
+}
+
+
 CUDA_CALLABLE inline vec3 Vec3Max(const vec3& a, const vec3& b) { return wp::max(a, b); }
 CUDA_CALLABLE inline vec3 Vec3Min(const vec3& a, const vec3& b) { return wp::min(a, b); }
 
@@ -392,7 +456,7 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
     // allocate temporary memory used during  building
     indices = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items*2); 	// *2 for radix sort
     keys = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items*2);	    // *2 for radix sort
-    deltas = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items);    	// highest differenting bit between keys for item i and i+1
+    deltas = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*num_items);    	// highest differentiating bit between keys for item i and i+1
     range_lefts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
     range_rights = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
     num_children = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh.max_nodes);
@@ -431,6 +495,7 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
     
     // sort items based on Morton key (note the 32-bit sort key corresponds to the template parameter to morton3, i.e. 3x9 bit keys combined)
     radix_sort_pairs_device(WP_CURRENT_CONTEXT, keys, indices, num_items);
+    memcpy_d2d(WP_CURRENT_CONTEXT, bvh.primitive_indices, indices, sizeof(int) * num_items);
 
     // calculate deltas between adjacent keys
     wp_launch_device(WP_CURRENT_CONTEXT, compute_key_deltas, num_items, (keys, deltas, num_items-1));
@@ -443,6 +508,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, 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_lowers, bvh.node_uppers));
 
     // free temporary memory
     free_device(WP_CURRENT_CONTEXT, indices);
@@ -455,27 +521,171 @@ void LinearBVHBuilderGPU::build(BVH& bvh, const vec3* item_lowers, const vec3* i
 
 }
 
+// buffer_size is the number of T, not the number of bytes
+template<typename T>
+T* make_device_buffer_of(void* context, T* host_buffer, size_t buffer_size)
+{
+    T* device_buffer = (T*)alloc_device(context, sizeof(T) * buffer_size);;
+    memcpy_h2d(context, device_buffer, host_buffer, sizeof(T) * buffer_size);
+
+    return device_buffer;
+}
+
+void copy_host_tree_to_device(void* context, BVH& bvh_host, BVH& bvh_device_on_host)
+{
+#ifdef REORDER_HOST_TREE
+
+
+    // reorder bvh_host such that its nodes are in the front
+    // this is essential for the device refit 
+    BVHPackedNodeHalf* node_lowers_reordered = new BVHPackedNodeHalf[bvh_host.max_nodes];
+    BVHPackedNodeHalf* node_uppers_reordered = new BVHPackedNodeHalf[bvh_host.max_nodes];
+
+    int* node_parents_reordered = new int[bvh_host.max_nodes];
+
+    std::vector<int> old_to_new(bvh_host.max_nodes, -1);
+
+    // We will place nodes in this order:
+    //   Pass 1: leaf nodes (except if it's the root index)
+    //   Pass 2: non-leaf, non-root
+    //   Pass 3: root node
+    int next_pos = 0;
+
+    const int root_index = *bvh_host.root;
+    // Pass 1: place leaf nodes at the front 
+    for (int i = 0; i < bvh_host.num_nodes; ++i)
+    {
+        if (bvh_host.node_lowers[i].b)
+        {
+            node_lowers_reordered[next_pos] = bvh_host.node_lowers[i];
+            node_uppers_reordered[next_pos] = bvh_host.node_uppers[i];
+            old_to_new[i] = next_pos;
+            next_pos++;
+        }
+    }
+
+    // Pass 2: place non-leaf, non-root nodes
+    for (int i = 0; i < bvh_host.num_nodes; ++i)
+    {
+        if (i == root_index)
+        {
+            if (bvh_host.node_lowers[i].b)
+                // if root node is leaf node, there must be only be one node
+            {
+                *bvh_host.root = 0;
+            }
+            else
+            {
+                *bvh_host.root = next_pos;
+            }
+        }
+        if (!bvh_host.node_lowers[i].b)
+        {
+            node_lowers_reordered[next_pos] = bvh_host.node_lowers[i];
+            node_uppers_reordered[next_pos] = bvh_host.node_uppers[i];
+            old_to_new[i] = next_pos;
+            next_pos++;
+        }
+    }
+
+    // We can do that by enumerating all old->new pairs:
+    for (int old_index = 0; old_index < bvh_host.num_nodes; ++old_index) {
+        int new_index = old_to_new[old_index];  // new index
+
+        int old_parent = bvh_host.node_parents[old_index];
+        if (old_parent != -1)
+        {
+            node_parents_reordered[new_index] = old_to_new[old_parent];
+        }
+        else
+        {
+            node_parents_reordered[new_index] = -1;
+        }
+
+        // only need to modify the child index of non-leaf nodes
+        if (!bvh_host.node_lowers[old_index].b)
+        {
+            node_lowers_reordered[new_index].i = old_to_new[bvh_host.node_lowers[old_index].i];
+            node_uppers_reordered[new_index].i = old_to_new[bvh_host.node_uppers[old_index].i];
+        }
+    }
+
+    delete[] bvh_host.node_lowers;
+    delete[] bvh_host.node_uppers;
+    delete[] bvh_host.node_parents;
+
+    bvh_host.node_lowers = node_lowers_reordered;
+    bvh_host.node_uppers = node_uppers_reordered;
+    bvh_host.node_parents = node_parents_reordered;
+#endif // REORDER_HOST_TREE
+
+    bvh_device_on_host.num_nodes = bvh_host.num_nodes;
+    bvh_device_on_host.num_leaf_nodes = bvh_host.num_leaf_nodes;
+    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.root = (int*)alloc_device(context, sizeof(int));
+    memcpy_h2d(context, bvh_device_on_host.root, bvh_host.root, sizeof(int));
+    bvh_device_on_host.context = context;
+
+    bvh_device_on_host.node_lowers = make_device_buffer_of(context, bvh_host.node_lowers, bvh_host.max_nodes);
+    bvh_device_on_host.node_uppers = make_device_buffer_of(context, bvh_host.node_uppers, bvh_host.max_nodes);
+    bvh_device_on_host.node_parents = make_device_buffer_of(context, bvh_host.node_parents, bvh_host.max_nodes);
+    bvh_device_on_host.primitive_indices = make_device_buffer_of(context, bvh_host.primitive_indices, bvh_host.num_items);
+}
+
 // create in-place given existing descriptor
-void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items, BVH& bvh_host)
+void bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh_device_on_host)
 {
     ContextGuard guard(context);
-
-    bvh_host.num_items = num_items;
-    bvh_host.max_nodes = 2*num_items;
-    bvh_host.node_lowers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    memset_device(WP_CURRENT_CONTEXT, bvh_host.node_lowers, 0, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    bvh_host.node_uppers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    memset_device(WP_CURRENT_CONTEXT, bvh_host.node_uppers, 0, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    bvh_host.node_parents = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
-    bvh_host.node_counts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
-    bvh_host.root = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int));
-    bvh_host.item_lowers = lowers;
-    bvh_host.item_uppers = uppers;
-
-    bvh_host.context = context ? context : cuda_context_get_current();
-
-    LinearBVHBuilderGPU builder;
-    builder.build(bvh_host, lowers, uppers, num_items, NULL);
+    if (constructor_type == BVH_CONSTRUCTOR_SAH || constructor_type == BVH_CONSTRUCTOR_MEDIAN)
+        // CPU based constructors
+    {
+        // copy bounds back to CPU
+        std::vector<vec3> lowers_host(num_items);
+        std::vector<vec3> uppers_host(num_items);
+        memcpy_d2h(WP_CURRENT_CONTEXT, lowers_host.data(), lowers, sizeof(vec3) * num_items);
+        memcpy_d2h(WP_CURRENT_CONTEXT, uppers_host.data(), uppers, sizeof(vec3) * num_items);
+
+        // run CPU based constructor
+        wp::BVH bvh_host;
+        bvh_create_host(lowers_host.data(), uppers_host.data(), num_items, constructor_type, bvh_host);
+
+        // copy host tree to device
+        wp::copy_host_tree_to_device(WP_CURRENT_CONTEXT, bvh_host, bvh_device_on_host);
+        // replace host bounds with device bounds
+        bvh_device_on_host.item_lowers = lowers;
+        bvh_device_on_host.item_uppers = uppers;
+        // node_counts is not allocated for host tree
+        bvh_device_on_host.node_counts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int) * bvh_device_on_host.max_nodes);
+        bvh_destroy_host(bvh_host);
+    }
+    else if (constructor_type == BVH_CONSTRUCTOR_LBVH)
+    {
+        bvh_device_on_host.num_items = num_items;
+        bvh_device_on_host.max_nodes = 2 * num_items;
+        bvh_device_on_host.num_leaf_nodes = num_items;
+        bvh_device_on_host.node_lowers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf) * bvh_device_on_host.max_nodes);
+        memset_device(WP_CURRENT_CONTEXT, bvh_device_on_host.node_lowers, 0, sizeof(BVHPackedNodeHalf) * bvh_device_on_host.max_nodes);
+        bvh_device_on_host.node_uppers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf) * bvh_device_on_host.max_nodes);
+        memset_device(WP_CURRENT_CONTEXT, bvh_device_on_host.node_uppers, 0, sizeof(BVHPackedNodeHalf) * bvh_device_on_host.max_nodes);
+        bvh_device_on_host.node_parents = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int) * bvh_device_on_host.max_nodes);
+        bvh_device_on_host.node_counts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int) * bvh_device_on_host.max_nodes);
+        bvh_device_on_host.root = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int));
+        bvh_device_on_host.primitive_indices = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int) * num_items);
+        bvh_device_on_host.item_lowers = lowers;
+        bvh_device_on_host.item_uppers = uppers;
+
+        bvh_device_on_host.context = context ? context : cuda_context_get_current();
+
+        LinearBVHBuilderGPU builder;
+        builder.build(bvh_device_on_host, lowers, uppers, num_items, NULL);
+    }
+    else
+    {
+        printf("Unrecognized Constructor type: %d! For GPU constructor it should be SAH (0), Median (1), or LBVH (2)!\n", constructor_type);
+    }
 }
 
 void bvh_destroy_device(BVH& bvh)
@@ -486,9 +696,11 @@ void bvh_destroy_device(BVH& bvh)
     free_device(WP_CURRENT_CONTEXT, bvh.node_uppers); bvh.node_uppers = NULL;
     free_device(WP_CURRENT_CONTEXT, bvh.node_parents); bvh.node_parents = NULL;
     free_device(WP_CURRENT_CONTEXT, bvh.node_counts); bvh.node_counts = NULL;
+    free_device(WP_CURRENT_CONTEXT, bvh.primitive_indices); bvh.primitive_indices = NULL;
     free_device(WP_CURRENT_CONTEXT, bvh.root); bvh.root = NULL;
 }
 
+
 } // namespace wp
 
 
@@ -503,20 +715,27 @@ void bvh_refit_device(uint64_t id)
     }
 }
 
-uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items)
+/*
+* 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 
+* refitting threads will start from the nodes corresponding to the original leaf nodes, which might be 
+* 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 bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type)
 {
     ContextGuard guard(context);
-
-    wp::BVH bvh_host;
-    bvh_create_device(WP_CURRENT_CONTEXT, lowers, uppers, num_items, bvh_host);
+    wp::BVH bvh_device_on_host;
+    wp::BVH* bvh_device_ptr = nullptr;
+    
+    bvh_create_device(WP_CURRENT_CONTEXT, lowers, uppers, num_items, constructor_type, bvh_device_on_host);
 
     // create device-side BVH descriptor
-    wp::BVH* bvh_device = (wp::BVH*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::BVH));
-    memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device, &bvh_host, sizeof(wp::BVH));
-        
-    uint64_t bvh_id = (uint64_t)bvh_device;
-    wp::bvh_add_descriptor(bvh_id, bvh_host);
+    bvh_device_ptr = (wp::BVH*)alloc_device(WP_CURRENT_CONTEXT, sizeof(wp::BVH));
+    memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device_ptr, &bvh_device_on_host, sizeof(wp::BVH));
 
+    uint64_t bvh_id = (uint64_t)bvh_device_ptr;
+    wp::bvh_add_descriptor(bvh_id, bvh_device_on_host);
     return bvh_id;
 }