warp-lang 1.3.3__py3-none-manylinux2014_x86_64.whl → 1.4.1__py3-none-manylinux2014_x86_64.whl

warp/native/mesh.cu

@@ -101,7 +101,7 @@ __global__ void bvh_refit_with_solid_angle_kernel(int n, const int* __restrict__
             int finished = atomicAdd(&child_count[parent], 1);
 
             // if we have are the last thread (such that the parent node is now complete)
-            // then update its bounds and move onto the the next parent in the hierarchy
+            // then update its bounds and move onto the next parent in the hierarchy
             if (finished == 1)
             {
                 //printf("Compute non-leaf at %d\n", index);
@@ -290,3 +290,54 @@ void mesh_refit_device(uint64_t id)
     }
 }
 
+void mesh_set_points_device(uint64_t id, wp::array_t<wp::vec3> points)
+{
+    wp::Mesh m;
+    if (mesh_get_descriptor(id, m))
+    {
+        if (points.ndim != 1 || points.shape[0] != m.points.shape[0])
+        {
+            fprintf(stderr, "The new points input for mesh_set_points_device does not match the shape of the original points!\n");
+            return;
+        }
+
+        m.points = points;
+
+        wp::Mesh* mesh_device = (wp::Mesh*)id;
+        memcpy_h2d(WP_CURRENT_CONTEXT, mesh_device, &m, sizeof(wp::Mesh));
+
+        // update the cpu copy as well
+        mesh_set_descriptor(id, m);
+
+        mesh_refit_device(id);
+    }
+    else 
+    {
+        fprintf(stderr, "The mesh id provided to mesh_set_points_device is not valid!\n");
+        return;
+    }
+}
+
+void mesh_set_velocities_device(uint64_t id, wp::array_t<wp::vec3> velocities)
+{
+    wp::Mesh m;
+    if (mesh_get_descriptor(id, m))
+    {
+        if (velocities.ndim != 1 || velocities.shape[0] != m.velocities.shape[0])
+        {
+            fprintf(stderr, "The new velocities input for mesh_set_velocities_device does not match the shape of the original velocities\n");
+            return;
+        }
+
+        m.velocities = velocities;
+
+        wp::Mesh* mesh_device = (wp::Mesh*)id;
+        memcpy_h2d(WP_CURRENT_CONTEXT, mesh_device, &m, sizeof(wp::Mesh));
+        mesh_set_descriptor(id, m);
+    }
+    else 
+    {
+        fprintf(stderr, "The mesh id provided to mesh_set_velocities_device is not valid!\n");
+        return;
+    }
+}

warp/native/mesh.h

@@ -1881,6 +1881,7 @@ CUDA_CALLABLE inline void adj_mesh_get_index(uint64_t id, int index,
 }
 
 CUDA_CALLABLE bool mesh_get_descriptor(uint64_t id, Mesh& mesh);
+CUDA_CALLABLE bool mesh_set_descriptor(uint64_t id, const Mesh& mesh);
 CUDA_CALLABLE void mesh_add_descriptor(uint64_t id, const Mesh& mesh);
 CUDA_CALLABLE void mesh_rem_descriptor(uint64_t id);
 

warp/native/quat.h

@@ -29,6 +29,14 @@ struct quat_t
         w = static_cast<Type>(other.w);
     }
 
+    inline CUDA_CALLABLE quat_t(const initializer_array<4, Type> &l)
+    {
+        x = l[0];
+        y = l[1];
+        z = l[2];
+        w = l[3];
+    }
+
     // imaginary part
     Type x;
     Type y;
@@ -479,6 +487,41 @@ inline CUDA_CALLABLE void adj_indexref(quat_t<Type>* q, int idx,
     // nop
 }
 
+template<typename Type>
+inline CUDA_CALLABLE quat_t<Type> assign(quat_t<Type>& q, int idx, Type value)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    quat_t<Type> ret(q);
+    ret[idx] = value;
+    return ret;
+}
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_assign(quat_t<Type>& q, int idx, Type value, quat_t<Type>& adj_q, int& adj_idx, Type& adj_value, const quat_t<Type>& adj_ret)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    adj_value += adj_ret[idx];
+    for(unsigned i=0; i < 4; ++i)
+    {
+        if(i != idx)
+            adj_q[i] += adj_ret[i];
+    }
+}
+
 template<typename Type>
 CUDA_CALLABLE inline quat_t<Type> lerp(const quat_t<Type>& a, const quat_t<Type>& b, Type t)
 {

warp/native/range.h

@@ -97,8 +97,17 @@ CUDA_CALLABLE inline range_t iter_reverse(const range_t& r)
 {
     // generates a reverse range, equivalent to reversed(range())
     range_t rev;
-    rev.start = r.end-1;
-    rev.end = r.start-1;
+
+    if (r.step > 0)
+    {
+        rev.start = r.start + int((r.end - r.start - 1) / r.step) * r.step;
+    }
+    else
+    {
+        rev.start = r.start + int((r.end - r.start + 1) / r.step) * r.step;
+    }
+
+    rev.end = r.start - r.step;
     rev.step = -r.step;
 
     rev.i = rev.start;

warp/native/spatial.h

@@ -127,6 +127,12 @@ struct transform_t
     CUDA_CALLABLE inline transform_t(vec_t<3,Type> p=vec_t<3,Type>(), quat_t<Type> q=quat_t<Type>()) : p(p), q(q) {}
     CUDA_CALLABLE inline transform_t(Type)  {}  // helps uniform initialization
 
+    CUDA_CALLABLE inline transform_t(const initializer_array<7, Type> &l)
+    {
+        p = vec_t<3,Type>(l[0], l[1], l[2]);
+        q = quat_t<Type>(l[3], l[4], l[5], l[6]);
+    }
+
     CUDA_CALLABLE inline Type operator[](int index) const
     {
         assert(index < 7);

warp/native/vec.h

@@ -381,6 +381,31 @@ inline CUDA_CALLABLE vec_t<3, Type> sub(vec_t<3, Type> a, vec_t<3, Type> b)
     return vec_t<3, Type>( a.c[0] - b.c[0], a.c[1] - b.c[1], a.c[2] - b.c[2]);
 }
 
+// modulo
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> mod(vec_t<Length, Type> a, vec_t<Length, Type> b)
+{
+    vec_t<Length, Type> ret;
+    for (unsigned i=0; i < Length; ++i)
+    {
+        ret[i] = mod(a[i], b[i]);
+    }
+
+    return ret;
+}
+
+template<typename Type>
+inline CUDA_CALLABLE vec_t<2, Type> mod(vec_t<2, Type> a, vec_t<2, Type> b)
+{
+    return vec_t<2, Type>(mod(a.c[0], b.c[0]), mod(a.c[1], b.c[1]));
+}
+
+template<typename Type>
+inline CUDA_CALLABLE vec_t<3, Type> mod(vec_t<3, Type> a, vec_t<3, Type> b)
+{
+    return vec_t<3, Type>(mod(a.c[0], b.c[0]), mod(a.c[1], b.c[1]), mod(a.c[2], b.c[2]));
+}
+
 // dot product:
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE Type dot(vec_t<Length, Type> a, vec_t<Length, Type> b)
@@ -470,6 +495,40 @@ inline CUDA_CALLABLE void adj_indexref(vec_t<Length, Type>* v, int idx,
     // nop
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> assign(vec_t<Length, Type>& v, int idx, Type value)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    vec_t<Length, Type> ret(v);
+    ret[idx] = value;
+    return ret;
+}
+
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_assign(vec_t<Length, Type>& v, int idx, Type value, vec_t<Length, Type>& adj_v, int& adj_idx, Type& adj_value, const vec_t<Length, Type>& adj_ret)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    adj_value += adj_ret[idx];
+    for(unsigned i=0; i < Length; ++i)
+    {
+        if (i != idx)
+            adj_v[i] += adj_ret[i];
+    }
+}
 
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE Type length(vec_t<Length, Type> a)
@@ -896,6 +955,21 @@ inline CUDA_CALLABLE void adj_sub(vec_t<3, Type> a, vec_t<3, Type> b, vec_t<3, T
     adj_b.c[2] -= adj_ret.c[2];
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_mod(vec_t<Length, Type> a, vec_t<Length, Type> b, vec_t<Length, Type>& adj_a, vec_t<Length, Type>& adj_b, const vec_t<Length, Type>& adj_ret)
+{
+}
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_mod(vec_t<2, Type> a, vec_t<2, Type> b, vec_t<2, Type>& adj_a, vec_t<2, Type>& adj_b, const vec_t<2, Type>& adj_ret)
+{
+}
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_mod(vec_t<3, Type> a, vec_t<3, Type> b, vec_t<3, Type>& adj_a, vec_t<3, Type>& adj_b, const vec_t<3, Type>& adj_ret)
+{
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_dot(vec_t<Length, Type> a, vec_t<Length, Type> b, vec_t<Length, Type>& adj_a, vec_t<Length, Type>& adj_b, const Type adj_ret)
 {

warp/native/warp.cpp

@@ -1010,7 +1010,7 @@ WP_API int cuda_set_peer_access_enabled(void* target_context, void* peer_context
 WP_API int cuda_is_mempool_access_enabled(int target_ordinal, int peer_ordinal) { return 0; }
 WP_API int cuda_set_mempool_access_enabled(int target_ordinal, int peer_ordinal, int enable) { return 0; }
 
-WP_API void* cuda_stream_create(void* context) { return NULL; }
+WP_API void* cuda_stream_create(void* context, int priority) { return NULL; }
 WP_API void cuda_stream_destroy(void* context, void* stream) {}
 WP_API void cuda_stream_register(void* context, void* stream) {}
 WP_API void cuda_stream_unregister(void* context, void* stream) {}
@@ -1020,6 +1020,7 @@ WP_API void cuda_stream_wait_event(void* stream, void* event) {}
 WP_API void cuda_stream_wait_stream(void* stream, void* other_stream, void* event) {}
 WP_API int cuda_stream_is_capturing(void* stream) { return 0; }
 WP_API uint64_t cuda_stream_get_capture_id(void* stream) { return 0; }
+WP_API int cuda_stream_get_priority(void* stream) { return 0; }
 
 WP_API void* cuda_event_create(void* context, unsigned flags) { return NULL; }
 WP_API void cuda_event_destroy(void* event) {}

warp/native/warp.cu

@@ -1963,7 +1963,6 @@ void cuda_context_set_stream(void* context, void* stream, int sync)
     }
 }
 
-
 int cuda_is_peer_access_supported(int target_ordinal, int peer_ordinal)
 {
     int num_devices = int(g_devices.size());
@@ -2174,12 +2173,12 @@ int cuda_set_mempool_access_enabled(int target_ordinal, int peer_ordinal, int en
 }
 
 
-void* cuda_stream_create(void* context)
+void* cuda_stream_create(void* context, int priority)
 {
     ContextGuard guard(context, true);
 
     CUstream stream;
-    if (check_cu(cuStreamCreate_f(&stream, CU_STREAM_DEFAULT)))
+    if (check_cu(cuStreamCreateWithPriority_f(&stream, CU_STREAM_DEFAULT, priority)))
     {
         cuda_stream_register(WP_CURRENT_CONTEXT, stream);
         return stream;
@@ -2268,6 +2267,14 @@ uint64_t cuda_stream_get_capture_id(void* stream)
     return get_capture_id(static_cast<CUstream>(stream));
 }
 
+int cuda_stream_get_priority(void* stream)
+{
+    int priority = 0;
+    check_cuda(cuStreamGetPriority_f(static_cast<CUstream>(stream), &priority));
+
+    return priority;
+}
+
 void* cuda_event_create(void* context, unsigned flags)
 {
     ContextGuard guard(context, true);

warp/native/warp.h

@@ -83,6 +83,12 @@ extern "C"
 	WP_API void mesh_destroy_device(uint64_t id);
     WP_API void mesh_refit_device(uint64_t id);
 
+    WP_API void mesh_set_points_host(uint64_t id, wp::array_t<wp::vec3> points);
+    WP_API void mesh_set_points_device(uint64_t id, wp::array_t<wp::vec3> points);
+
+    WP_API void mesh_set_velocities_host(uint64_t id, wp::array_t<wp::vec3> velocities);
+    WP_API void mesh_set_velocities_device(uint64_t id, wp::array_t<wp::vec3> velocities);
+
     WP_API uint64_t hash_grid_create_host(int dim_x, int dim_y, int dim_z);
     WP_API void hash_grid_reserve_host(uint64_t id, int num_points);
     WP_API void hash_grid_destroy_host(uint64_t id);
@@ -286,7 +292,7 @@ extern "C"
     WP_API int cuda_is_mempool_access_enabled(int target_ordinal, int peer_ordinal);
     WP_API int cuda_set_mempool_access_enabled(int target_ordinal, int peer_ordinal, int enable);
 
-    WP_API void* cuda_stream_create(void* context);
+    WP_API void* cuda_stream_create(void* context, int priority);
     WP_API void cuda_stream_destroy(void* context, void* stream);
     WP_API void cuda_stream_register(void* context, void* stream);
     WP_API void cuda_stream_unregister(void* context, void* stream);
@@ -296,6 +302,7 @@ extern "C"
     WP_API void cuda_stream_wait_stream(void* stream, void* other_stream, void* event);
     WP_API int cuda_stream_is_capturing(void* stream);
     WP_API uint64_t cuda_stream_get_capture_id(void* stream);
+    WP_API int cuda_stream_get_priority(void* stream);
 
     WP_API void* cuda_event_create(void* context, unsigned flags);
     WP_API void cuda_event_destroy(void* event);