warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/native/bvh.cpp

@@ -27,35 +27,34 @@ class MedianBVHBuilder
 {	
 public:
 
-    void build(BVH& bvh, const bounds3* items, int n);
+    void build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n);
 
 private:
 
-    bounds3 calc_bounds(const bounds3* bounds, const int* indices, int start, int end);
+    bounds3 calc_bounds(const vec3* lowers, const vec3* uppers, const int* indices, int start, int end);
 
-    int partition_median(const bounds3* bounds, int* indices, int start, int end, bounds3 range_bounds);
-    int partition_midpoint(const bounds3* bounds, int* indices, int start, int end, bounds3 range_bounds);
-    int partition_sah(const bounds3* bounds, int* indices, int start, int end, bounds3 range_bounds);
+    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);
 
-    int build_recursive(BVH& bvh, const bounds3* bounds, int* indices, int start, int end, int depth, int parent);
+    int build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int* indices, int start, int end, int depth, int parent);
 };
 
 //////////////////////////////////////////////////////////////////////
 
-void MedianBVHBuilder::build(BVH& bvh, const bounds3* items, int n)
+void MedianBVHBuilder::build(BVH& bvh, const vec3* lowers, const vec3* uppers, int n)
 {
-    memset(&bvh, 0, sizeof(BVH));
-
+    bvh.max_depth = 0;
     bvh.max_nodes = 2*n-1;
 
     bvh.node_lowers = new BVHPackedNodeHalf[bvh.max_nodes];
     bvh.node_uppers = new BVHPackedNodeHalf[bvh.max_nodes];
     bvh.node_parents = new int[bvh.max_nodes];
-
-    bvh.num_nodes = 0;
-
+    bvh.node_counts = NULL;
+    
     // root is always in first slot for top down builders
-    bvh.root = 0;
+    bvh.root = new int[1];
+    bvh.root[0] = 0;
 
     if (n == 0)
         return;
@@ -64,35 +63,42 @@ void MedianBVHBuilder::build(BVH& bvh, const bounds3* items, int n)
     for (int i=0; i < n; ++i)
         indices[i] = i;
 
-    build_recursive(bvh, items, &indices[0], 0, n, 0, -1);
+    build_recursive(bvh, lowers, uppers, &indices[0], 0, n, 0, -1);
 }
 
 
-bounds3 MedianBVHBuilder::calc_bounds(const bounds3* bounds, const int* indices, int start, int end)
+bounds3 MedianBVHBuilder::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 = bounds_union(u, bounds[indices[i]]);
+    {
+        u.add_point(lowers[indices[i]]);
+        u.add_point(uppers[indices[i]]);
+    }
 
     return u;
 }
 
 struct PartitionPredicateMedian
 {
-    PartitionPredicateMedian(const bounds3* bounds, int a) : bounds(bounds), axis(a) {}
+    PartitionPredicateMedian(const vec3* lowers, const vec3* uppers, int a) : lowers(lowers), uppers(uppers), axis(a) {}
 
     bool operator()(int a, int b) const
     {
-        return bounds[a].center()[axis] < bounds[b].center()[axis];
+        vec3 a_center = 0.5f*(lowers[a] + uppers[a]);
+        vec3 b_center = 0.5f*(lowers[b] + uppers[b]);
+        
+        return a_center[axis] < b_center[axis];
     }
 
-    const bounds3* bounds;
+    const vec3* lowers;
+    const vec3* uppers;
     int axis;
 };
 
 
-int MedianBVHBuilder::partition_median(const bounds3* bounds, int* indices, int start, int end, bounds3 range_bounds)
+int MedianBVHBuilder::partition_median(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds)
 {
     assert(end-start >= 2);
 
@@ -102,27 +108,31 @@ int MedianBVHBuilder::partition_median(const bounds3* bounds, int* indices, int
 
     const int k = (start+end)/2;
 
-    std::nth_element(&indices[start], &indices[k], &indices[end], PartitionPredicateMedian(&bounds[0], axis));
+    std::nth_element(&indices[start], &indices[k], &indices[end], PartitionPredicateMedian(lowers, uppers, axis));
 
     return k;
 }	
     
 struct PartitionPredictateMidPoint
 {
-    PartitionPredictateMidPoint(const bounds3* bounds, int a, float m) : bounds(bounds), axis(a), mid(m) {}
+    PartitionPredictateMidPoint(const vec3* lowers, const vec3* uppers, int a, float m) : lowers(lowers), uppers(uppers), axis(a), mid(m) {}
 
     bool operator()(int index) const 
     {
-        return bounds[index].center()[axis] <= mid;
+        vec3 center = 0.5f*(lowers[index] + uppers[index]);
+
+        return center[axis] <= mid;
     }
 
-    const bounds3* bounds;
+    const vec3* lowers;
+    const vec3* uppers;
+    
     int axis;
     float mid;
 };
 
 
-int MedianBVHBuilder::partition_midpoint(const bounds3* bounds, int* indices, int start, int end, bounds3 range_bounds)
+int MedianBVHBuilder::partition_midpoint(const vec3* lowers, const vec3* uppers, int* indices, int start, int end, bounds3 range_bounds)
 {
     assert(end-start >= 2);
 
@@ -132,7 +142,7 @@ int MedianBVHBuilder::partition_midpoint(const bounds3* bounds, int* indices, in
     int axis = longest_axis(edges);
     float mid = center[axis];
 
-    int* upper = std::partition(indices+start, indices+end, PartitionPredictateMidPoint(&bounds[0], axis, mid));
+    int* upper = std::partition(indices+start, indices+end, PartitionPredictateMidPoint(lowers, uppers, axis, mid));
 
     int k = upper-indices;
 
@@ -140,7 +150,6 @@ int MedianBVHBuilder::partition_midpoint(const bounds3* bounds, int* indices, in
     if (k == start || k == end)
         k = (start+end)/2;
 
-
     return k;
 }
 
@@ -200,7 +209,7 @@ int MedianBVHBuilder::partition_sah(const bounds3* bounds, int* indices, int sta
 }
 #endif
 
-int MedianBVHBuilder::build_recursive(BVH& bvh, const bounds3* bounds, int* indices, int start, int end, int depth, int parent)
+int MedianBVHBuilder::build_recursive(BVH& bvh, const vec3* lowers, const vec3* uppers, int* indices, int start, int end, int depth, int parent)
 {
     assert(start < end);
 
@@ -212,7 +221,7 @@ int MedianBVHBuilder::build_recursive(BVH& bvh, const bounds3* bounds, int* indi
     if (depth > bvh.max_depth)
         bvh.max_depth = depth;
 
-    bounds3 b = calc_bounds(bounds, indices, start, end);
+    bounds3 b = calc_bounds(lowers, uppers, indices, start, end);
     
     const int kMaxItemsPerLeaf = 1;
 
@@ -225,7 +234,7 @@ int MedianBVHBuilder::build_recursive(BVH& bvh, const bounds3* bounds, int* indi
     else    
     {
         //int split = partition_midpoint(bounds, indices, start, end, b);
-        int split = partition_median(bounds, indices, start, end, b);
+        int split = partition_median(lowers, uppers, indices, start, end, b);
         //int split = partition_sah(bounds, indices, start, end, b);
 
         if (split == start || split == end)
@@ -234,8 +243,8 @@ int MedianBVHBuilder::build_recursive(BVH& bvh, const bounds3* bounds, int* indi
             split = (start+end)/2;
         }
     
-        int left_child = build_recursive(bvh, bounds, indices, start, split, depth+1, node_index);
-        int right_child = build_recursive(bvh, bounds, indices, split, end, depth+1, node_index);
+        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);
         
         bvh.node_lowers[node_index] = make_node(b.lower, left_child, false);
         bvh.node_uppers[node_index] = make_node(b.upper, right_child, false);
@@ -245,218 +254,8 @@ int MedianBVHBuilder::build_recursive(BVH& bvh, const bounds3* bounds, int* indi
     return node_index;
 }
 
-class LinearBVHBuilderCPU
-{
-public:
-    
-    void build(BVH& bvh, const bounds3* items, int n);
-
-private:
-
-    // calculate Morton codes
-    struct KeyIndexPair
-    {
-        uint32_t key;
-        int index;
-
-        inline bool operator < (const KeyIndexPair& rhs) const { return key < rhs.key; }
-    };	
-
-    bounds3 calc_bounds(const bounds3* bounds, const KeyIndexPair* keys, int start, int end);
-    int find_split(const KeyIndexPair* pairs, int start, int end);
-    int build_recursive(BVH& bvh, const KeyIndexPair* keys, const bounds3* bounds, int start, int end, int depth);
-
-};
-
-
-// disable std::sort workaround for macOS error
-#if 0
-void LinearBVHBuilderCPU::build(BVH& bvh, const bounds3* items, int n)
-{
-	memset(&bvh, 0, sizeof(BVH));
-
-	bvh.max_nodes = 2*n-1;
-
-	bvh.node_lowers = new BVHPackedNodeHalf[bvh.max_nodes];
-	bvh.node_uppers = new BVHPackedNodeHalf[bvh.max_nodes];
-	bvh.num_nodes = 0;
-
-	// root is always in first slot for top down builders
-	bvh.root = 0;
-
-	std::vector<KeyIndexPair> keys;
-	keys.reserve(n);
-
-	bounds3 totalbounds3;
-	for (int i=0; i < n; ++i)
-		totalbounds3 = bounds_union(totalbounds3, items[i]);
-
-	// ensure non-zero edge length in all dimensions
-	totalbounds3.expand(0.001f);
-
-	vec3 edges = totalbounds3.edges();
-	vec3 invEdges = cw_div(vec3(1.0f), edges);
-
-	for (int i=0; i < n; ++i)
-	{
-		vec3 center = items[i].center();
-		vec3 local = cw_mul(center-totalbounds3.lower, invEdges);
 
-		KeyIndexPair l;
-		l.key = morton3<1024>(local.x, local.y, local.z);
-		l.index = i;
-
-		keys.push_back(l);
-	}
-
-	// sort by key
-	std::sort(keys.begin(), keys.end());
-
-	build_recursive(bvh, &keys[0], items,  0, n, 0);
-
-	printf("Created BVH for %d items with %d nodes, max depth of %d\n", n, bvh.num_nodes, bvh.max_depth);
-}
-#endif
-
-inline bounds3 LinearBVHBuilderCPU::calc_bounds(const bounds3* bounds, const KeyIndexPair* keys, int start, int end)
-{
-	bounds3 u;
-
-	for (int i=start; i < end; ++i)
-		u = bounds_union(u, bounds[keys[i].index]);
-
-	return u;
-}
-
-inline int LinearBVHBuilderCPU::find_split(const KeyIndexPair* pairs, int start, int end)
-{
-	if (pairs[start].key == pairs[end-1].key)
-		return (start+end)/2;
-
-	// find split point between keys, xor here means all bits 
-	// of the result are zero up until the first differing bit
-	int common_prefix = clz(pairs[start].key ^ pairs[end-1].key);
-
-	// use binary search to find the point at which this bit changes
-	// from zero to a 1		
-	const int mask = 1 << (31-common_prefix);
-
-	while (end-start > 0)
-	{
-		int index = (start+end)/2;
-
-		if (pairs[index].key&mask)
-		{
-			end = index;
-		}
-		else
-			start = index+1;
-	}
-
-	assert(start == end);
-
-	return start;
-}
-
-int LinearBVHBuilderCPU::build_recursive(BVH& bvh, const KeyIndexPair* keys, const bounds3* bounds, int start, int end, int depth)
-{
-	assert(start < end);
-
-	const int n = end-start;
-	const int nodeIndex = bvh.num_nodes++;
-
-	assert(nodeIndex < bvh.max_nodes);
-
-	if (depth > bvh.max_depth)
-		bvh.max_depth = depth;
-
-	bounds3 b = calc_bounds(bounds, keys, start, end);
-		
-	const int kMaxItemsPerLeaf = 1;
-
-	if (n <= kMaxItemsPerLeaf)
-	{
-		bvh.node_lowers[nodeIndex] = make_node(b.lower, keys[start].index, true);
-		bvh.node_uppers[nodeIndex] = make_node(b.upper, keys[start].index, false);
-	}
-	else
-	{
-		int split = find_split(keys, start, end);
-		
-		int leftChild = build_recursive(bvh, keys, bounds, start, split, depth+1);
-		int rightChild = build_recursive(bvh, keys, bounds, split, end, depth+1);
-			
-		bvh.node_lowers[nodeIndex] = make_node(b.lower, leftChild, false);
-		bvh.node_uppers[nodeIndex] = make_node(b.upper, rightChild, false);		
-	}
-
-	return nodeIndex;
-}
-
-
-
-// create only happens on host currently, use bvh_clone() to transfer BVH To device
-BVH bvh_create(const bounds3* bounds, int num_bounds)
-{
-    BVH bvh;
-    memset(&bvh, 0, sizeof(bvh));
-
-    MedianBVHBuilder builder;
-    //LinearBVHBuilderCPU builder;
-    builder.build(bvh, bounds, num_bounds);
-
-    return bvh;
-}
-
-void bvh_destroy_host(BVH& bvh)
-{
-    delete[] bvh.node_lowers;
-    delete[] bvh.node_uppers;
-    delete[] bvh.node_parents;
-	delete[] bvh.bounds;
-
-    bvh.node_lowers = NULL;
-    bvh.node_uppers = NULL;
-    bvh.max_nodes = 0;
-    bvh.num_nodes = 0;
-    bvh.num_bounds = 0;
-}
-
-void bvh_destroy_device(BVH& bvh)
-{
-    ContextGuard guard(bvh.context);
-
-    free_device(WP_CURRENT_CONTEXT, bvh.node_lowers); bvh.node_lowers = NULL;
-    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.bounds); bvh.bounds = NULL;
-}
-
-BVH bvh_clone(void* context, const BVH& bvh_host)
-{
-    ContextGuard guard(context);
-
-    BVH bvh_device = bvh_host;
-
-    bvh_device.context = context ? context : cuda_context_get_current();
-
-    bvh_device.node_lowers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    bvh_device.node_uppers = (BVHPackedNodeHalf*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    bvh_device.node_parents = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
-    bvh_device.node_counts = (int*)alloc_device(WP_CURRENT_CONTEXT, sizeof(int)*bvh_host.max_nodes);
-	bvh_device.bounds = (bounds3*)alloc_device(WP_CURRENT_CONTEXT, sizeof(bounds3)*bvh_host.num_bounds);
-
-    // copy host data to device
-    memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device.node_lowers, bvh_host.node_lowers, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device.node_uppers, bvh_host.node_uppers, sizeof(BVHPackedNodeHalf)*bvh_host.max_nodes);
-    memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device.node_parents, bvh_host.node_parents, sizeof(int)*bvh_host.max_nodes);
-	memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device.bounds, bvh_host.bounds, sizeof(bounds3)*bvh_host.num_bounds);
-
-    return bvh_device;
-}
-
-void bvh_refit_recursive(BVH& bvh, int index, const bounds3* bounds)
+void bvh_refit_recursive(BVH& bvh, int index)
 {
     BVHPackedNodeHalf& lower = bvh.node_lowers[index];
     BVHPackedNodeHalf& upper = bvh.node_uppers[index];
@@ -465,16 +264,17 @@ void bvh_refit_recursive(BVH& bvh, int index, const bounds3* bounds)
     {
         const int leaf_index = lower.i;
 
-        (vec3&)lower = bounds[leaf_index].lower;
-        (vec3&)upper = bounds[leaf_index].upper;
+        // update leaf from items
+        (vec3&)lower = bvh.item_lowers[leaf_index];
+        (vec3&)upper = bvh.item_uppers[leaf_index];
     }
     else
     {
         int left_index = lower.i;
         int right_index = upper.i;
 
-        bvh_refit_recursive(bvh, left_index, bounds);
-        bvh_refit_recursive(bvh, right_index, bounds);
+        bvh_refit_recursive(bvh, left_index);
+        bvh_refit_recursive(bvh, right_index);
 
         // compute union of children
         const vec3& left_lower = (vec3&)bvh.node_lowers[left_index];
@@ -493,9 +293,9 @@ void bvh_refit_recursive(BVH& bvh, int index, const bounds3* bounds)
     }
 }
 
-void bvh_refit_host(BVH& bvh, const bounds3* b)
+void bvh_refit_host(BVH& bvh)
 {
-    bvh_refit_recursive(bvh, 0, b);
+    bvh_refit_recursive(bvh, 0);
 }
 
 
@@ -538,87 +338,46 @@ void bvh_rem_descriptor(uint64_t id)
 
 }
 
+
+void bvh_destroy_host(BVH& bvh)
+{
+    delete[] bvh.node_lowers;
+    delete[] bvh.node_uppers;
+    delete[] bvh.node_parents;
+    delete[] bvh.root;
+
+    bvh.node_lowers = NULL;
+    bvh.node_uppers = NULL;
+    bvh.node_parents = NULL;
+    bvh.root = NULL;
+
+    bvh.max_nodes = 0;
+    bvh.num_items = 0;
+}
+
 } // namespace wp
 
-uint64_t bvh_create_host(vec3* lowers, vec3* uppers, int num_bounds)
+uint64_t bvh_create_host(vec3* lowers, vec3* uppers, int num_items)
 {
     BVH* bvh = new BVH();
     memset(bvh, 0, sizeof(BVH));
 
     bvh->context = NULL;
 
-    bvh->lowers = lowers;
-    bvh->uppers = uppers;
-    bvh->num_bounds = num_bounds;
-
-    bvh->bounds = new bounds3[num_bounds];  
-
-    for (int i=0; i < num_bounds; ++i)
-    {
-        bvh->bounds[i].lower = lowers[i];
-        bvh->bounds[i].upper = uppers[i];
-    }
+    bvh->item_lowers = lowers;
+    bvh->item_uppers = uppers;
+    bvh->num_items = num_items;
 
     MedianBVHBuilder builder;
-    builder.build(*bvh, bvh->bounds, num_bounds);
+    builder.build(*bvh, lowers, uppers, num_items);
 
     return (uint64_t)bvh;
 }
 
-uint64_t bvh_create_device(void* context, vec3* lowers, vec3* uppers, int num_bounds)
-{
-    ContextGuard guard(context);
-
-    // todo: BVH creation only on CPU at the moment so temporarily bring all the data back to host
-    vec3* lowers_host = (vec3*)alloc_host(sizeof(vec3)*num_bounds);
-    vec3* uppers_host = (vec3*)alloc_host(sizeof(vec3)*num_bounds);
-    bounds3* bounds_host = (bounds3*)alloc_host(sizeof(bounds3)*num_bounds);
-
-    memcpy_d2h(WP_CURRENT_CONTEXT, lowers_host, lowers, sizeof(vec3)*num_bounds);
-    memcpy_d2h(WP_CURRENT_CONTEXT, uppers_host, uppers, sizeof(vec3)*num_bounds);
-    cuda_context_synchronize(WP_CURRENT_CONTEXT);
-
-    for (int i=0; i < num_bounds; ++i)
-    {
-        bounds_host[i] = bounds3();
-        bounds_host[i].lower = lowers_host[i];
-        bounds_host[i].upper = uppers_host[i];
-    }
-
-    BVH bvh_host = bvh_create(bounds_host, num_bounds);
-    bvh_host.context = context ? context : cuda_context_get_current();
-	bvh_host.bounds = bounds_host;
-    bvh_host.num_bounds = num_bounds;        
-    BVH bvh_device_clone = bvh_clone(WP_CURRENT_CONTEXT, bvh_host);
-
-	bvh_device_clone.lowers = lowers;		// managed by the user
-	bvh_device_clone.uppers = uppers;		// managed by the user
-
-    BVH* bvh_device = (BVH*)alloc_device(WP_CURRENT_CONTEXT, sizeof(BVH));
-    memcpy_h2d(WP_CURRENT_CONTEXT, bvh_device, &bvh_device_clone, sizeof(BVH));
-
-    bvh_destroy_host(bvh_host);
-    free_host(lowers_host);
-    free_host(uppers_host);
-
-    uint64_t bvh_id = (uint64_t)bvh_device;
-    bvh_add_descriptor(bvh_id, bvh_device_clone);
-
-    return bvh_id;
-}
-
 void bvh_refit_host(uint64_t id)
 {
     BVH* bvh = (BVH*)(id);
-
-    for (int i=0; i < bvh->num_bounds; ++i)
-    {
-        bvh->bounds[i] = bounds3();
-        bvh->bounds[i].lower = bvh->lowers[i];
-        bvh->bounds[i].upper = bvh->uppers[i];
-    }
-
-    bvh_refit_host(*bvh, bvh->bounds);
+    bvh_refit_host(*bvh);
 }
 
 void bvh_destroy_host(uint64_t id)
@@ -629,23 +388,11 @@ void bvh_destroy_host(uint64_t id)
 }
 
 
-void bvh_destroy_device(uint64_t id)
-{
-    BVH bvh;
-    if (bvh_get_descriptor(id, bvh))
-    {
-        bvh_destroy_device(bvh);
-        mesh_rem_descriptor(id);
-    }
-}
-
 // stubs for non-CUDA platforms
 #if !WP_ENABLE_CUDA
 
-void bvh_refit_device(uint64_t id)
-{
-}
-
-
+uint64_t bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items) { return 0; }
+void bvh_refit_device(uint64_t id) {}
+void bvh_destroy_device(uint64_t id) {}
 
 #endif // !WP_ENABLE_CUDA