warp-lang 1.5.1__py3-none-win_amd64.whl → 1.6.0__py3-none-win_amd64.whl

warp/native/bvh.cpp

@@ -23,11 +23,11 @@ namespace wp
 
 /////////////////////////////////////////////////////////////////////////////////////////////
 
-class MedianBVHBuilder
+class TopDownBVHBuilder
 {	
 public:
 
-    void build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n);
+    void build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n, int in_constructor_type);
 
 private:
 
@@ -35,15 +35,26 @@ private:
 
     int partition_median(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds);
     int partition_midpoint(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds);
-    int partition_sah(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds);
+    float partition_sah(BVH& bvh, const vec3* lowers, const vec3* uppers,
+       int start, int end, bounds3 range_bounds, int& split_axis);
 
-    int build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int* indices, int start, int end, int depth, int parent);
+    int build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int start, int end, int depth, int parent);
+
+    int constructor_type = -1;
 };
 
 //////////////////////////////////////////////////////////////////////
 
-void MedianBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n)
+void TopDownBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n, int in_constructor_type)
 {
+    constructor_type = in_constructor_type;
+    if (constructor_type != BVH_CONSTRUCTOR_SAH && constructor_type != BVH_CONSTRUCTOR_MEDIAN)
+    {
+        printf("Unrecognized Constructor type: %d! For CPU constructor it should be either SAH (%d) or Median (%d)!\n",
+            constructor_type, BVH_CONSTRUCTOR_SAH, BVH_CONSTRUCTOR_MEDIAN);
+        return;
+    }
+
     bvh.max_depth = 0;
     bvh.max_nodes = 2*n-1;
 
@@ -51,7 +62,7 @@ void MedianBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers, i
     bvh.node_uppers = new BVHPackedNodeHalf[bvh.max_nodes];
     bvh.node_parents = new int[bvh.max_nodes];
     bvh.node_counts = NULL;
-    
+
     // root is always in first slot for top down builders
     bvh.root = new int[1];
     bvh.root[0] = 0;
@@ -59,22 +70,21 @@ void MedianBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers, i
     if (n == 0)
         return;
     
-    std::vector<int> indices(n);
-    for (int i=0; i < n; ++i)
-        indices[i] = i;
+    bvh.primitive_indices = new int[n];
+    for (int i = 0; i < n; ++i)
+        bvh.primitive_indices[i] = i;
 
-    build_recursive(bvh, lowers, uppers, &indices[0], 0, n, 0, -1);
+    build_recursive(bvh, lowers, uppers,  0, n, 0, -1);
 }
 
 
-bounds3 MedianBVHBuilder::calc_bounds(const vec3* lowers, const vec3* uppers, const int* indices, int start, int end)
+bounds3 TopDownBVHBuilder::calc_bounds(const vec3* lowers, const vec3* uppers, const int* indices, int start, int end)
 {
     bounds3 u;
 
     for (int i=start; i < end; ++i)
     {
-        u.add_point(lowers[indices[i]]);
-        u.add_point(uppers[indices[i]]);
+        u.add_bounds(lowers[indices[i]], uppers[indices[i]]);
     }
 
     return u;
@@ -98,7 +108,7 @@ struct PartitionPredicateMedian
 };
 
 
-int MedianBVHBuilder::partition_median(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds)
+int TopDownBVHBuilder::partition_median(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds)
 {
     assert(end-start >= 2);
 
@@ -113,9 +123,9 @@ int MedianBVHBuilder::partition_median(const vec3* lowers, const vec3* uppers, i
     return k;
 }	
     
-struct PartitionPredictateMidPoint
+struct PartitionPredicateMidPoint
 {
-    PartitionPredictateMidPoint(const vec3* lowers, const vec3* uppers, int a, float m) : lowers(lowers), uppers(uppers), axis(a), mid(m) {}
+    PartitionPredicateMidPoint(const vec3* lowers, const vec3* uppers, int a, float m) : lowers(lowers), uppers(uppers), axis(a), mid(m) {}
 
     bool operator()(int index) const 
     {
@@ -132,7 +142,7 @@ struct PartitionPredictateMidPoint
 };
 
 
-int MedianBVHBuilder::partition_midpoint(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds)
+int TopDownBVHBuilder::partition_midpoint(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds)
 {
     assert(end-start >= 2);
 
@@ -142,7 +152,7 @@ int MedianBVHBuilder::partition_midpoint(const vec3* lowers, const vec3* uppers,
     int axis = longest_axis(edges);
     float mid = center[axis];
 
-    int* upper = std::partition(indices+start, indices+end, PartitionPredictateMidPoint(lowers, uppers, axis, mid));
+    int* upper = std::partition(indices+start, indices+end, PartitionPredicateMidPoint(lowers, uppers, axis, mid));
 
     int k = upper-indices;
 
@@ -153,50 +163,88 @@ int MedianBVHBuilder::partition_midpoint(const vec3* lowers, const vec3* uppers,
     return k;
 }
 
-// disable std::sort workaround for macOS error
-#if 0
-int MedianBVHBuilder::partition_sah(const bounds3* bounds, int* indices, int start, int end, bounds3 range_bounds)
+float TopDownBVHBuilder::partition_sah(BVH& bvh, const vec3* lowers, const vec3* uppers, int start, int end, bounds3 range_bounds, int& split_axis)
 {
-    assert(end-start >= 2);
+    int buckets_counts[SAH_NUM_BUCKETS];
+    bounds3 buckets[SAH_NUM_BUCKETS];
+    float left_areas[SAH_NUM_BUCKETS - 1];
+    float right_areas[SAH_NUM_BUCKETS - 1];
+
+    assert(end - start >= 2);
 
-    int n = end-start;
+    int n = end - start;
     vec3 edges = range_bounds.edges();
 
-    int longestAxis = longest_axis(edges);
+    bounds3 b = calc_bounds(lowers, uppers, bvh.primitive_indices, start, end);
 
-    // sort along longest axis
-    std::sort(&indices[0]+start, &indices[0]+end, PartitionPredicateMedian(&bounds[0], longestAxis));
+    split_axis = longest_axis(edges);
 
-    // total area for range from [0, split]
-    std::vector<float> left_areas(n);
-    // total area for range from (split, end]
-    std::vector<float> right_areas(n);
+    // compute each bucket
+    float range_start = b.lower[split_axis];
+    float range_end = b.upper[split_axis];
+
+    std::fill(buckets_counts, buckets_counts + SAH_NUM_BUCKETS, 0);
+    for (int item_idx = start; item_idx < end; item_idx++)
+    {
+        vec3 item_center = 0.5f * (lowers[bvh.primitive_indices[item_idx]] + uppers[bvh.primitive_indices[item_idx]]);
+        int bucket_idx = SAH_NUM_BUCKETS * (item_center[split_axis] - range_start) / (range_end - range_start);
+        assert(bucket_idx >= 0 && bucket_idx <= SAH_NUM_BUCKETS);
+        // one of them will have the range_end, we put it into the last bucket
+        bucket_idx = bucket_idx < SAH_NUM_BUCKETS ? bucket_idx : SAH_NUM_BUCKETS - 1;
+
+        bounds3 item_bound(lowers[bvh.primitive_indices[item_idx]], uppers[bvh.primitive_indices[item_idx]]);
+
+        if (buckets_counts[bucket_idx])
+        {
+            buckets[bucket_idx] = bounds_union(item_bound, buckets[bucket_idx]);
+        }
+        else
+        {
+            buckets[bucket_idx] = item_bound;
+        }
+
+        buckets_counts[bucket_idx]++;
+    }
 
     bounds3 left;
     bounds3 right;
 
+    // n - 1 division points for n buckets
+    int counts_l[SAH_NUM_BUCKETS - 1];
+    int counts_r[SAH_NUM_BUCKETS - 1];
+
+    int count_l = 0;
+    int count_r = 0;
     // build cumulative bounds and area from left and right
-    for (int i=0; i < n; ++i)
+    for (int i = 0; i < SAH_NUM_BUCKETS - 1; ++i)
     {
-        left = bounds_union(left, bounds[indices[start+i]]);
-        right = bounds_union(right, bounds[indices[end-i-1]]);
+        bounds3 bound_start = buckets[i];
+        bounds3 bound_end = buckets[SAH_NUM_BUCKETS - i - 1];
+
+        left = bounds_union(left, bound_start);
+        right = bounds_union(right, bound_end);
 
         left_areas[i] = left.area();
-        right_areas[n-i-1] = right.area();
+        right_areas[SAH_NUM_BUCKETS - i - 2] = right.area();
+
+        count_l += buckets_counts[i];
+        count_r += buckets_counts[SAH_NUM_BUCKETS - i - 1];
+
+        counts_l[i] = count_l;
+        counts_r[SAH_NUM_BUCKETS - i - 2] = count_r;
     }
 
-    float invTotalArea = 1.0f/range_bounds.area();
+    float invTotalArea = 1.0f / range_bounds.area();
 
-    // find split point i that minimizes area(left[i]) + area(right[i])
+    // find split point i that minimizes area(left[i]) * count[left[i]] + area(right[i]) * count[right[i]]
     int minSplit = 0;
     float minCost = FLT_MAX;
-
-    for (int i=0; i < n; ++i)
+    for (int i = 0; i < SAH_NUM_BUCKETS - 1; ++i)
     {
-        float pBelow = left_areas[i]*invTotalArea;
-        float pAbove = right_areas[i]*invTotalArea;
+        float pBelow = left_areas[i] * invTotalArea;
+        float pAbove = right_areas[i] * invTotalArea;
 
-        float cost = pBelow*i + pAbove*(n-i);
+        float cost = pBelow * counts_l[i] + pAbove * counts_r[i];
 
         if (cost < minCost)
         {
@@ -205,15 +253,20 @@ int MedianBVHBuilder::partition_sah(const bounds3* bounds, int* indices, int sta
         }
     }
 
-    return start + minSplit + 1;
+    // return the dividing 
+    assert(minSplit >= 0 && minSplit < SAH_NUM_BUCKETS - 1);
+    float split_point = range_start + (minSplit + 1) * (range_end - range_start) / SAH_NUM_BUCKETS;
+
+    return split_point;
 }
-#endif
 
-int MedianBVHBuilder::build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int* indices, int start, int end, int depth, int parent)
+int TopDownBVHBuilder::build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int start, int end, int depth, int parent)
 {
     assert(start < end);
 
-    const int n = end-start;
+    // printf("start %d end %d\n", start, end);
+
+    const int n = end - start;
     const int node_index = bvh.num_nodes++;
 
     assert(node_index < bvh.max_nodes);
@@ -221,31 +274,50 @@ int MedianBVHBuilder::build_recursive(BVH& bvh, const vec3* lowers, const vec3*
     if (depth > bvh.max_depth)
         bvh.max_depth = depth;
 
-    bounds3 b = calc_bounds(lowers, uppers, indices, start, end);
-    
-    const int kMaxItemsPerLeaf = 1;
+    bounds3 b = calc_bounds(lowers, uppers, bvh.primitive_indices, start, end);
 
-    if (n <= kMaxItemsPerLeaf)
+    if (n <= BVH_LEAF_SIZE)
     {
-        bvh.node_lowers[node_index] = make_node(b.lower, indices[start], true);
-        bvh.node_uppers[node_index] = make_node(b.upper, indices[start], false);
+        bvh.node_lowers[node_index] = make_node(b.lower, start, true);
+        bvh.node_uppers[node_index] = make_node(b.upper, end, false);
         bvh.node_parents[node_index] = parent;
+        bvh.num_leaf_nodes++;
     }
-    else    
+    else
     {
-        //int split = partition_midpoint(bounds, indices, start, end, b);
-        int split = partition_median(lowers, uppers, indices, start, end, b);
-        //int split = partition_sah(bounds, indices, start, end, b);
+        int split = -1;
+        if (constructor_type == BVH_CONSTRUCTOR_SAH)
+            // SAH constructor
+        {
+            int split_axis = -1;
+            float split_point = partition_sah(bvh, lowers, uppers, start, end, b, split_axis);
+            auto boundary = std::partition(bvh.primitive_indices + start, bvh.primitive_indices + end,
+                [&](int i) {
+                    return 0.5f * (lowers[i] + uppers[i])[split_axis] < split_point;
+                });
+
+            split = std::distance(bvh.primitive_indices + start, boundary) + start;
+        }
+        else if (constructor_type == BVH_CONSTRUCTOR_MEDIAN)
+            // Median constructor
+        {
+            split = partition_median(lowers, uppers, bvh.primitive_indices, start, end, b);
+        }
+        else
+        {
+            printf("Unknown type of BVH constructor: %d!\n", constructor_type);
+            return -1;
+        }
 
         if (split == start || split == end)
         {
             // partitioning failed, split down the middle
-            split = (start+end)/2;
+            split = (start + end) / 2;
         }
-    
-        int left_child = build_recursive(bvh, lowers, uppers, indices, start, split, depth+1, node_index);
-        int right_child = build_recursive(bvh, lowers, uppers, indices, split, end, depth+1, node_index);
-        
+
+        int left_child = build_recursive(bvh, lowers, uppers, start, split, depth + 1, node_index);
+        int right_child = build_recursive(bvh, lowers, uppers, split, end, depth + 1, node_index);
+
         bvh.node_lowers[node_index] = make_node(b.lower, left_child, false);
         bvh.node_uppers[node_index] = make_node(b.upper, right_child, false);
         bvh.node_parents[node_index] = parent;
@@ -262,11 +334,16 @@ void bvh_refit_recursive(BVH& bvh, int index)
 
     if (lower.b)
     {
-        const int leaf_index = lower.i;
-
         // update leaf from items
-        (vec3&)lower = bvh.item_lowers[leaf_index];
-        (vec3&)upper = bvh.item_uppers[leaf_index];
+        bounds3 bound;
+        for (int item_counter = lower.i; item_counter < upper.i; item_counter++)
+        {
+            const int item = bvh.primitive_indices[item_counter];
+            bound.add_bounds(bvh.item_lowers[item], bvh.item_uppers[item]);
+        }
+
+        (vec3&)lower = bound.lower;
+        (vec3&)upper = bound.upper;
     }
     else
     {
@@ -340,7 +417,7 @@ void bvh_rem_descriptor(uint64_t id)
 
 
 // create in-place given existing descriptor
-void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, BVH& bvh)
+void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type, BVH& bvh)
 {
     memset(&bvh, 0, sizeof(BVH));
 
@@ -348,8 +425,8 @@ void bvh_create_host(vec3* lowers, vec3* uppers, int num_items, BVH& bvh)
     bvh.item_uppers = uppers;
     bvh.num_items = num_items;
 
-    MedianBVHBuilder builder;
-    builder.build(bvh, lowers, uppers, num_items);
+    TopDownBVHBuilder builder;
+    builder.build(bvh, lowers, uppers, num_items, constructor_type);
 }
 
 void bvh_destroy_host(BVH& bvh)
@@ -357,11 +434,13 @@ void bvh_destroy_host(BVH& bvh)
     delete[] bvh.node_lowers;
     delete[] bvh.node_uppers;
     delete[] bvh.node_parents;
+    delete[] bvh.primitive_indices;
     delete[] bvh.root;
 
     bvh.node_lowers = NULL;
     bvh.node_uppers = NULL;
     bvh.node_parents = NULL;
+    bvh.primitive_indices = NULL;
     bvh.root = NULL;
 
     bvh.max_nodes = 0;
@@ -370,10 +449,10 @@ void bvh_destroy_host(BVH& bvh)
 
 } // namespace wp
 
-uint64_t bvh_create_host(vec3* lowers, vec3* uppers, int num_items)
+uint64_t bvh_create_host(vec3* lowers, vec3* uppers, int num_items, int constructor_type)
 {
     BVH* bvh = new BVH();
-    wp::bvh_create_host(lowers, uppers, num_items, *bvh);
+    wp::bvh_create_host(lowers, uppers, num_items, constructor_type, *bvh);
 
     return (uint64_t)bvh;
 }
@@ -395,7 +474,7 @@ void bvh_destroy_host(uint64_t id)
 // stubs for non-CUDA platforms
 #if !WP_ENABLE_CUDA
 
-uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items) { return 0; }
+uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type) { return 0; }
 void bvh_refit_device(uint64_t id) {}
 void bvh_destroy_device(uint64_t id) {}