warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/native/mesh.cpp

@@ -126,7 +126,7 @@ void bvh_refit_with_solid_angle_recursive_host(BVH& bvh, int index, Mesh& mesh)
         
         // write new BVH nodes
         reinterpret_cast<vec3&>(lower) = new_lower;
-        reinterpret_cast<vec3&>(upper) = new_upper; 
+        reinterpret_cast<vec3&>(upper) = new_upper;        
     }
 }
 
@@ -135,7 +135,7 @@ void bvh_refit_with_solid_angle_host(BVH& bvh, Mesh& mesh)
     bvh_refit_with_solid_angle_recursive_host(bvh, 0, mesh);
 }
 
-uint64_t wp_mesh_create_host(array_t<wp::vec3> points, array_t<wp::vec3> velocities, array_t<int> indices, int num_points, int num_tris, int support_winding_number, int constructor_type)
+uint64_t wp_mesh_create_host(array_t<wp::vec3> points, array_t<wp::vec3> velocities, array_t<int> indices, int num_points, int num_tris, int support_winding_number, int constructor_type, int bvh_leaf_size)
 {
     Mesh* m = new Mesh(points, velocities, indices, num_points, num_tris);
 
@@ -163,7 +163,7 @@ uint64_t wp_mesh_create_host(array_t<wp::vec3> points, array_t<wp::vec3> velocit
     }
     m->average_edge_length = sum / (num_tris*3);
 
-    wp::bvh_create_host(m->lowers, m->uppers, num_tris, constructor_type, m->bvh);
+    wp::bvh_create_host(m->lowers, m->uppers, num_tris, constructor_type, m->bvh, bvh_leaf_size);
     
     if (support_winding_number) 
     {
@@ -256,7 +256,7 @@ void wp_mesh_set_velocities_host(uint64_t id, wp::array_t<wp::vec3> velocities)
 #if !WP_ENABLE_CUDA
 
 
-WP_API uint64_t wp_mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number, int constructor_type) { return 0; }
+WP_API uint64_t wp_mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number, int constructor_type, int bvh_leaf_size) { return 0; }
 WP_API void wp_mesh_destroy_device(uint64_t id) {}
 WP_API void wp_mesh_refit_device(uint64_t id) {}
 WP_API void wp_mesh_set_points_device(uint64_t id, wp::array_t<wp::vec3> points) {};

warp/native/mesh.cu

@@ -21,6 +21,8 @@
 #include "bvh.h"
 #include "scan.h"
 
+extern CUcontext get_current_context();
+
 namespace wp
 {
 
@@ -245,7 +247,7 @@ void bvh_refit_with_solid_angle_device(BVH& bvh, Mesh& mesh)
 } // namespace wp
 
 
-uint64_t wp_mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> indices, int num_points, int num_tris, int support_winding_number, int constructor_type)
+uint64_t wp_mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> indices, int num_points, int num_tris, int support_winding_number, int constructor_type, int bvh_leaf_size)
 {
     ContextGuard guard(context);
 
@@ -280,7 +282,7 @@ uint64_t wp_mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::
 
     // compute triangle bound and construct BVH
     wp_launch_device(WP_CURRENT_CONTEXT, wp::compute_triangle_bounds, mesh.num_tris, (mesh.num_tris, mesh.points, mesh.indices, mesh.lowers, mesh.uppers));
-    wp::bvh_create_device(mesh.context, mesh.lowers, mesh.uppers, num_tris, constructor_type, mesh.bvh);
+    wp::bvh_create_device(mesh.context, mesh.lowers, mesh.uppers, num_tris, constructor_type, mesh.bvh, bvh_leaf_size);
 
     // we need to overwrite mesh.bvh because it is not initialized when we construct it on device
     wp_memcpy_h2d(WP_CURRENT_CONTEXT, &(mesh_device->bvh), &mesh.bvh, sizeof(wp::BVH));

warp/native/mesh.h

@@ -1374,6 +1374,7 @@ CUDA_CALLABLE inline bool mesh_query_ray(uint64_t id, const vec3& start, const v
     Mesh mesh = mesh_get(id);
 
     int stack[BVH_QUERY_STACK_SIZE];
+
     stack[0] = *mesh.bvh.root;
     int count = 1;
 
@@ -1463,6 +1464,139 @@ CUDA_CALLABLE inline bool mesh_query_ray(uint64_t id, const vec3& start, const v
     
 }
 
+template <typename T>
+CUDA_CALLABLE inline void _swap(T& a, T& b) 
+{
+    T t = a; a = b; b = t;
+}
+
+CUDA_CALLABLE inline bool mesh_query_ray_ordered(uint64_t id, const vec3& start, const vec3& dir, float max_t, float& t, float& u, float& v, float& sign, vec3& normal, int& face)
+{
+    Mesh mesh = mesh_get(id);
+
+    int stack[BVH_QUERY_STACK_SIZE];
+    float stack_dist[BVH_QUERY_STACK_SIZE];
+
+    stack[0] = *mesh.bvh.root;
+    stack_dist[0] = -FLT_MAX;
+
+    int count = 1;
+
+    vec3 rcp_dir = vec3(1.0f/dir[0], 1.0f/dir[1], 1.0f/dir[2]);
+
+    float min_t = max_t;
+    int min_face;
+    float min_u;
+    float min_v;
+    float min_sign = 1.0f;
+    vec3 min_normal;
+
+    while (count)
+    {
+        count -= 1;
+
+        const int nodeIndex = stack[count];
+        const float nodeDist = stack_dist[count];
+
+        if (nodeDist < min_t)
+        {
+            int left_index = mesh.bvh.node_lowers[nodeIndex].i;
+            int right_index = mesh.bvh.node_uppers[nodeIndex].i;
+            bool leaf = mesh.bvh.node_lowers[nodeIndex].b;
+
+            if (leaf)
+            {	
+                const int start_index = left_index;
+                const int end_index = right_index;
+                // loops through primitives in the leaf
+                for (int primitive_counter = start_index; primitive_counter < end_index ; primitive_counter++)
+                {
+                    int primitive_index = mesh.bvh.primitive_indices[primitive_counter];
+                    int i = mesh.indices[primitive_index * 3 + 0];
+                    int j = mesh.indices[primitive_index * 3 + 1];
+                    int k = mesh.indices[primitive_index * 3 + 2];
+
+                    vec3 p = mesh.points[i];
+                    vec3 q = mesh.points[j];
+                    vec3 r = mesh.points[k];
+
+                    float t, u, v, w, sign;
+                    vec3 n;
+
+                    if (intersect_ray_tri_rtcd(start, dir, p, q, r, t, u, v, w, sign, &n))
+                    {
+                        if (t < min_t && t >= 0.0f)
+                        {
+                            min_t = t;
+                            min_face = primitive_index;
+                            min_u = u;
+                            min_v = v;
+                            min_sign = sign;
+                            min_normal = n;
+                        }
+                    }
+                }
+            }
+            else
+            {
+                const float eps = 1.e-3f;
+
+                BVHPackedNodeHalf left_lower = bvh_load_node(mesh.bvh.node_lowers, left_index);
+                BVHPackedNodeHalf left_upper = bvh_load_node(mesh.bvh.node_uppers, left_index);
+        
+                BVHPackedNodeHalf right_lower = bvh_load_node(mesh.bvh.node_lowers, right_index);
+                BVHPackedNodeHalf right_upper = bvh_load_node(mesh.bvh.node_uppers, right_index);
+
+                float left_dist = FLT_MAX;
+                bool left_hit = intersect_ray_aabb(start, rcp_dir, vec3(left_lower.x-eps, left_lower.y-eps, left_lower.z-eps), vec3(left_upper.x+eps, left_upper.y+eps, left_upper.z+eps), left_dist);
+                       
+                float right_dist = FLT_MAX;
+                bool right_hit = intersect_ray_aabb(start, rcp_dir, vec3(right_lower.x-eps, right_lower.y-eps, right_lower.z-eps), vec3(right_upper.x+eps, right_upper.y+eps, right_upper.z+eps), right_dist);
+                
+
+                if (left_dist < right_dist)
+                {
+                    _swap(left_index, right_index);
+                    _swap(left_dist, right_dist);
+                    _swap(left_hit, right_hit);
+                }
+
+                if (left_hit && left_dist < min_t)
+                {
+                    stack[count] = left_index;
+                    stack_dist[count] = left_dist;
+                    count += 1;
+                }
+
+                if (right_hit && right_dist < min_t)
+                {
+                    stack[count] = right_index;
+                    stack_dist[count] = right_dist;
+                    count += 1;
+                }
+            }
+        }
+    }
+
+    if (min_t < max_t)
+    {
+        // write outputs
+        u = min_u;
+        v = min_v;
+        sign = min_sign;
+        t = min_t;
+        normal = normalize(min_normal);
+        face = min_face;
+
+        return true;
+    }
+    else
+    {
+        return false;
+    }
+    
+}
+
 
 CUDA_CALLABLE inline void adj_mesh_query_ray(
     uint64_t id, const vec3& start, const vec3& dir, float max_t, float t, float u, float v, float sign, const vec3& n, int face,
@@ -1589,7 +1723,12 @@ struct mesh_query_aabb_t
     // Mesh Id
     Mesh mesh;
     // BVH traversal stack:
+#if BVH_SHARED_STACK
+    bvh_stack_t stack;
+#else
     int stack[BVH_QUERY_STACK_SIZE];
+#endif
+
     int count;
 
     // inputs
@@ -1617,13 +1756,16 @@ CUDA_CALLABLE inline mesh_query_aabb_t mesh_query_aabb(
 
     Mesh mesh = mesh_get(id);
     query.mesh = mesh;
+
+#if BVH_SHARED_STACK
+    __shared__ int stack[BVH_QUERY_STACK_SIZE * WP_TILE_BLOCK_DIM];
+    query.stack.ptr = &stack[threadIdx.x];
+#endif
     
     query.stack[0] = *mesh.bvh.root;
     query.count = 1;
     query.input_lower = lower;
     query.input_upper = upper;
-
-    wp::bounds3 input_bounds(query.input_lower, query.input_upper);
     
     // Navigate through the bvh, find the first overlapping leaf node.
     while (query.count)
@@ -1632,10 +1774,13 @@ CUDA_CALLABLE inline mesh_query_aabb_t mesh_query_aabb(
         BVHPackedNodeHalf node_lower = bvh_load_node(mesh.bvh.node_lowers, nodeIndex);
         BVHPackedNodeHalf node_upper = bvh_load_node(mesh.bvh.node_uppers, nodeIndex);
 
-        if (!input_bounds.overlaps(reinterpret_cast<vec3&>(node_lower), reinterpret_cast<vec3&>(node_upper)))
+        if (query.primitive_counter == 0)
         {
-            // Skip this box, it doesn't overlap with our target box.
-            continue;
+            if (!intersect_aabb_aabb(query.input_lower, query.input_upper, reinterpret_cast<vec3&>(node_lower), reinterpret_cast<vec3&>(node_upper)))
+            {
+                // Skip this box, it doesn't overlap with our target box.
+                continue;
+            }
         }
 
         const int left_index = node_lower.i;
@@ -1646,7 +1791,7 @@ CUDA_CALLABLE inline mesh_query_aabb_t mesh_query_aabb(
         {
             // Reached a leaf node, point to its first primitive
             // Back up one level and return 
-            query.primitive_counter = left_index;
+            query.primitive_counter = 0;
             query.stack[query.count++] = nodeIndex;
             return query;
         }
@@ -1671,45 +1816,6 @@ CUDA_CALLABLE inline bool mesh_query_aabb_next(mesh_query_aabb_t& query, int& in
 {
     Mesh mesh = query.mesh;
 
-    wp::bounds3 input_bounds(query.input_lower, query.input_upper);
-
-    if (query.primitive_counter != -1)
-        // currently in a leaf node which is the last node in the stack
-    {
-        const int node_index = query.stack[query.count - 1];
-        BVHPackedNodeHalf node_lower = bvh_load_node(mesh.bvh.node_lowers, node_index);
-        BVHPackedNodeHalf node_upper = bvh_load_node(mesh.bvh.node_uppers, node_index);
-
-        const int end = node_upper.i;
-        for (int primitive_counter = query.primitive_counter; primitive_counter < end; primitive_counter++)
-        {
-            int primitive_index = mesh.bvh.primitive_indices[primitive_counter];
-            if (input_bounds.overlaps(mesh.lowers[primitive_index], mesh.uppers[primitive_index]))
-            {
-                if (primitive_counter < end - 1)
-                    // still need to come back to this leaf node for the leftover primitives
-                {
-                    query.primitive_counter = primitive_counter + 1;
-                }
-                else
-                    // no need to come back to this leaf node
-                {
-                    query.count--;
-                    query.primitive_counter = -1;
-                }
-                index = primitive_index;
-                query.face = primitive_index;
-
-                return true;
-            }
-        }
-        // if we reach here it means we have finished the current leaf node without finding intersections
-        query.primitive_counter = -1;
-        // remove the leaf node from the back of the stack because it is finished
-        // and continue the bvh traversal
-        query.count--;
-    }
-    
     // Navigate through the bvh, find the first overlapping leaf node.
     while (query.count)
     {
@@ -1717,7 +1823,7 @@ CUDA_CALLABLE inline bool mesh_query_aabb_next(mesh_query_aabb_t& query, int& in
         BVHPackedNodeHalf node_lower = bvh_load_node(mesh.bvh.node_lowers, node_index);
         BVHPackedNodeHalf node_upper = bvh_load_node(mesh.bvh.node_uppers, node_index);
 
-        if (!input_bounds.overlaps(reinterpret_cast<vec3&>(node_lower), reinterpret_cast<vec3&>(node_upper)))
+        if (!intersect_aabb_aabb(query.input_lower, query.input_upper, reinterpret_cast<vec3&>(node_lower), reinterpret_cast<vec3&>(node_upper)))
         {
             // Skip this box, it doesn't overlap with our target box.
             continue;
@@ -1731,24 +1837,32 @@ CUDA_CALLABLE inline bool mesh_query_aabb_next(mesh_query_aabb_t& query, int& in
         {
             // found leaf, loop through its content primitives
             const int start = left_index;
-            const int end = right_index;
 
-            for (int primitive_counter = start; primitive_counter < end; primitive_counter++)
+            if (mesh.bvh.leaf_size == 1)
             {
-                int primitive_index = mesh.bvh.primitive_indices[primitive_counter];
-                if (input_bounds.overlaps(mesh.lowers[primitive_index], mesh.uppers[primitive_index]))
+                int primitive_index = mesh.bvh.primitive_indices[start];
+                index = primitive_index;
+                query.face = primitive_index;
+                return true;
+            }
+            else
+            {
+                const int end = right_index;
+                int primitive_index = mesh.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.count++;
+                }
+
+                if (intersect_aabb_aabb(query.input_lower, query.input_upper, mesh.lowers[primitive_index], mesh.uppers[primitive_index]))
                 {
-                    if (primitive_counter < end - 1)
-                        // still need to come back to this leaf node for the leftover primitives
-                    {
-                        query.primitive_counter = primitive_counter + 1;
-                        query.stack[query.count++] = node_index;
-                    }
-                    else
-                        // no need to come back to this leaf node
-                    {
-                        query.primitive_counter = -1;
-                    }
                     index = primitive_index;
                     query.face = primitive_index;
 
@@ -1758,6 +1872,7 @@ CUDA_CALLABLE inline bool mesh_query_aabb_next(mesh_query_aabb_t& query, int& in
         }
         else
         {
+            query.primitive_counter = 0;
             query.stack[query.count++] = left_index;
             query.stack[query.count++] = right_index;
         }

warp/native/quat.h

@@ -1621,58 +1621,6 @@ inline CUDA_CALLABLE void adj_quat_from_matrix(const mat_t<Rows,Cols,Type>& m, m
     adj_m.data[2][2] += dot(dq_dm22, adj_q);
 }
 
-template<typename Type>
-inline CUDA_CALLABLE void adj_mat_t(const vec_t<3,Type>& pos, const quat_t<Type>& rot, const vec_t<3,Type>& scale,
-                                          vec_t<3,Type>& adj_pos, quat_t<Type>& adj_rot, vec_t<3,Type>& adj_scale, const mat_t<4,4,Type>& adj_ret)
-{
-    mat_t<3,3,Type> R = quat_to_matrix(rot);
-    mat_t<3,3,Type> adj_R(0);
-
-    adj_pos[0] += adj_ret.data[0][3];
-    adj_pos[1] += adj_ret.data[1][3];
-    adj_pos[2] += adj_ret.data[2][3];
-
-    adj_mul(R.data[0][0], scale[0], adj_R.data[0][0], adj_scale[0], adj_ret.data[0][0]);
-    adj_mul(R.data[1][0], scale[0], adj_R.data[1][0], adj_scale[0], adj_ret.data[1][0]);
-    adj_mul(R.data[2][0], scale[0], adj_R.data[2][0], adj_scale[0], adj_ret.data[2][0]);
-
-    adj_mul(R.data[0][1], scale[1], adj_R.data[0][1], adj_scale[1], adj_ret.data[0][1]);
-    adj_mul(R.data[1][1], scale[1], adj_R.data[1][1], adj_scale[1], adj_ret.data[1][1]);
-    adj_mul(R.data[2][1], scale[1], adj_R.data[2][1], adj_scale[1], adj_ret.data[2][1]);
-
-    adj_mul(R.data[0][2], scale[2], adj_R.data[0][2], adj_scale[2], adj_ret.data[0][2]);
-    adj_mul(R.data[1][2], scale[2], adj_R.data[1][2], adj_scale[2], adj_ret.data[1][2]);
-    adj_mul(R.data[2][2], scale[2], adj_R.data[2][2], adj_scale[2], adj_ret.data[2][2]);
-
-    adj_quat_to_matrix(rot, adj_rot, adj_R);
-}
-
-template<unsigned Rows, unsigned Cols, typename Type>                
-inline CUDA_CALLABLE mat_t<Rows,Cols,Type>::mat_t(const vec_t<3,Type>& pos, const quat_t<Type>& rot, const vec_t<3,Type>& scale)
-{
-    mat_t<3,3,Type> R = quat_to_matrix(rot);
-
-    data[0][0] = R.data[0][0]*scale[0];
-    data[1][0] = R.data[1][0]*scale[0];
-    data[2][0] = R.data[2][0]*scale[0];
-    data[3][0] = Type(0);
-
-    data[0][1] = R.data[0][1]*scale[1];
-    data[1][1] = R.data[1][1]*scale[1];
-    data[2][1] = R.data[2][1]*scale[1];
-    data[3][1] = Type(0);
-
-    data[0][2] = R.data[0][2]*scale[2];
-    data[1][2] = R.data[1][2]*scale[2];
-    data[2][2] = R.data[2][2]*scale[2];
-    data[3][2] = Type(0);
-
-    data[0][3] = pos[0];
-    data[1][3] = pos[1];
-    data[2][3] = pos[2];
-    data[3][3] = Type(1);
-}
-
 template<typename Type=float32>
 inline CUDA_CALLABLE quat_t<Type> quat_identity()
 {

warp/native/scan.cu

@@ -18,6 +18,8 @@
 #include "warp.h"
 #include "scan.h"
 
+#include "cuda_util.h"
+
 #define THRUST_IGNORE_CUB_VERSION_CHECK
 
 #include <cub/device/device_scan.cuh>

warp/native/sort.cu

@@ -23,7 +23,7 @@
 
 #include <cub/cub.cuh>
 
-#include <map>
+#include <unordered_map>
 
 // temporary buffer for radix sort
 struct RadixSortTemp
@@ -32,8 +32,8 @@ struct RadixSortTemp
     size_t size = 0;
 };
 
-// map temp buffers to CUDA contexts
-static std::map<void*, RadixSortTemp> g_radix_sort_temp_map;
+// use unique temp buffers per CUDA stream to avoid race conditions
+static std::unordered_map<void*, RadixSortTemp> g_radix_sort_temp_map;
 
 
 template <typename KeyType>
@@ -44,6 +44,8 @@ void radix_sort_reserve_internal(void* context, int n, void** mem_out, size_t* s
     cub::DoubleBuffer<KeyType> d_keys;
 	cub::DoubleBuffer<int> d_values;
 
+    CUstream stream = static_cast<CUstream>(wp_cuda_stream_get_current());
+
     // compute temporary memory required
 	size_t sort_temp_size;
     check_cuda(cub::DeviceRadixSort::SortPairs(
@@ -52,12 +54,9 @@ void radix_sort_reserve_internal(void* context, int n, void** mem_out, size_t* s
         d_keys,
         d_values,
         n, 0, sizeof(KeyType)*8,
-        (cudaStream_t)wp_cuda_stream_get_current()));
-
-    if (!context)
-        context = wp_cuda_context_get_current();
+        stream));
 
-    RadixSortTemp& temp = g_radix_sort_temp_map[context];
+    RadixSortTemp& temp = g_radix_sort_temp_map[stream];
 
     if (sort_temp_size > temp.size)
     {
@@ -77,6 +76,17 @@ void radix_sort_reserve(void* context, int n, void** mem_out, size_t* size_out)
     radix_sort_reserve_internal<int>(context, n, mem_out, size_out);
 }
 
+void radix_sort_release(void* context, void* stream)
+{
+    // release temporary buffer for the given stream, if it exists
+    auto it = g_radix_sort_temp_map.find(stream);
+    if (it != g_radix_sort_temp_map.end())
+    {
+        wp_free_device(context, it->second.mem);
+        g_radix_sort_temp_map.erase(it);
+    }
+}
+
 template <typename KeyType>
 void radix_sort_pairs_device(void* context, KeyType* keys, int* values, int n)
 {
@@ -153,6 +163,8 @@ void segmented_sort_reserve(void* context, int n, int num_segments, void** mem_o
     int* start_indices = NULL;
     int* end_indices = NULL;
 
+    CUstream stream = static_cast<CUstream>(wp_cuda_stream_get_current());
+
     // compute temporary memory required
 	size_t sort_temp_size;
     check_cuda(cub::DeviceSegmentedRadixSort::SortPairs(
@@ -166,12 +178,9 @@ void segmented_sort_reserve(void* context, int n, int num_segments, void** mem_o
         end_indices,
         0,
         32,
-        (cudaStream_t)wp_cuda_stream_get_current()));
-
-    if (!context)
-        context = wp_cuda_context_get_current();
+        stream));
 
-    RadixSortTemp& temp = g_radix_sort_temp_map[context];
+    RadixSortTemp& temp = g_radix_sort_temp_map[stream];
 
     if (sort_temp_size > temp.size)
     {

warp/native/sort.h

@@ -20,6 +20,8 @@
 #include <stddef.h>
 
 void radix_sort_reserve(void* context, int n, void** mem_out=NULL, size_t* size_out=NULL);
+void radix_sort_release(void* context, void* stream);
+
 void radix_sort_pairs_host(int* keys, int* values, int n);
 void radix_sort_pairs_host(float* keys, int* values, int n);
 void radix_sort_pairs_host(int64_t* keys, int* values, int n);

warp/native/sparse.cu

@@ -16,8 +16,9 @@
  */
 
 #include "cuda_util.h"
+#include "temp_buffer.h"
 #include "warp.h"
-#include "stdint.h"
+
 #include <cstdint>
 
 #define THRUST_IGNORE_CUB_VERSION_CHECK
@@ -26,6 +27,8 @@
 #include <cub/device/device_run_length_encode.cuh>
 #include <cub/device/device_scan.cuh>
 
+extern CUcontext get_current_context();
+
 namespace
 {
 
@@ -361,7 +364,8 @@ WP_API void wp_bsr_matrix_from_triplets_device(
 
         if (bsr_nnz_event)
         {
-            wp_cuda_event_record(bsr_nnz_event, stream);
+            const bool external = true;
+            wp_cuda_event_record(bsr_nnz_event, stream, external);
         }
     }
 
@@ -416,7 +420,7 @@ WP_API void wp_bsr_transpose_device(int row_count, int col_count, int nnz,
         // Ensures the sorted keys are available in summed_block_indices if needed
         if(d_keys.Current() != src_block_indices)
         {
-            check_cuda(cudaMemcpy(src_block_indices, src_block_indices+nnz, size_t(nnz) * sizeof(int), cudaMemcpyDeviceToDevice));
+            check_cuda(cudaMemcpyAsync(src_block_indices, src_block_indices+nnz, size_t(nnz) * sizeof(int), cudaMemcpyDeviceToDevice, stream));
         }
     }
 

warp/native/spatial.h

@@ -992,6 +992,18 @@ CUDA_CALLABLE inline void adj_transform_t(const vec_t<3,Type>& p, const quat_t<T
     adj_q += adj_ret.q;
 }
 
+template<typename Type>
+CUDA_CALLABLE inline void adj_transform_t(const initializer_array<7, Type> &l, const initializer_array<7, Type*>& adj_l, const transform_t<Type>& adj_ret)
+{
+    *adj_l[0] += adj_ret.p[0];
+    *adj_l[1] += adj_ret.p[1];
+    *adj_l[2] += adj_ret.p[2];
+    *adj_l[3] += adj_ret.q[0];
+    *adj_l[4] += adj_ret.q[1];
+    *adj_l[5] += adj_ret.q[2];
+    *adj_l[6] += adj_ret.q[3];
+}
+
 template<typename Type>
 CUDA_CALLABLE inline void adj_transform_inverse(const transform_t<Type>& t, transform_t<Type>& adj_t, const transform_t<Type>& adj_ret)
 {