warp-lang 1.8.1__py3-none-manylinux_2_34_aarch64.whl → 1.9.1__py3-none-manylinux_2_34_aarch64.whl

warp/native/spatial.h

@@ -34,7 +34,7 @@ CUDA_CALLABLE inline Type spatial_dot(const spatial_vector_t<Type>& a, const spa
 template<typename Type>
 CUDA_CALLABLE inline vec_t<3,Type> &w_vec( spatial_vector_t<Type>& a )
 {
-    return *(vec_t<3,Type>*)(&a);
+    return *reinterpret_cast<vec_t<3,Type>*>(&a);
 }
 
 template<typename Type>
@@ -46,14 +46,14 @@ CUDA_CALLABLE inline vec_t<3,Type> &v_vec( spatial_vector_t<Type>& a )
 template<typename Type>
 CUDA_CALLABLE inline const vec_t<3,Type> &w_vec( const spatial_vector_t<Type>& a )
 {
-    spatial_vector_t<Type> &non_const_vec = *(spatial_vector_t<Type>*)(const_cast<Type*>(&a.c[0]));
+    spatial_vector_t<Type> &non_const_vec = *reinterpret_cast<spatial_vector_t<Type>*>(const_cast<Type*>(&a.c[0]));
     return w_vec(non_const_vec);
 }
 
 template<typename Type>
 CUDA_CALLABLE inline const vec_t<3,Type> &v_vec( const spatial_vector_t<Type>& a )
 {
-    spatial_vector_t<Type> &non_const_vec = *(spatial_vector_t<Type>*)(const_cast<Type*>(&a.c[0]));
+    spatial_vector_t<Type> &non_const_vec = *reinterpret_cast<spatial_vector_t<Type>*>(const_cast<Type*>(&a.c[0]));
     return v_vec(non_const_vec);
 }
 
@@ -408,27 +408,64 @@ template<typename Type>
 inline CUDA_CALLABLE Type extract(const transform_t<Type>& t, int idx)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
-    
+
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     return t[idx];
 }
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE vec_t<SliceLength, Type> extract(const transform_t<Type> & t, slice_t slice)
+{
+    vec_t<SliceLength, Type> ret;
+
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        ret[ii] = t[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+    return ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE Type* index(transform_t<Type>& t, int idx)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     return &t[idx];
 }
 
@@ -436,13 +473,18 @@ template<typename Type>
 inline CUDA_CALLABLE Type* indexref(transform_t<Type>* t, int idx)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
-        printf("transformation store %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     return &((*t)[idx]);
 }
 
@@ -452,6 +494,34 @@ inline void CUDA_CALLABLE adj_extract(const transform_t<Type>& t, int idx, trans
     adj_t[idx] += adj_ret;
 }
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void adj_extract(
+    const transform_t<Type>& t, slice_t slice,
+    transform_t<Type>& adj_t, slice_t& adj_slice,
+    const vec_t<SliceLength, Type>& adj_ret
+)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        adj_t[i] += adj_ret[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_index(transform_t<Type>& t, int idx,
                                        transform_t<Type>& adj_t, int adj_idx, const Type& adj_value)
@@ -470,120 +540,325 @@ template<typename Type>
 inline CUDA_CALLABLE void add_inplace(transform_t<Type>& t, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     t[idx] += value;
 }
 
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void add_inplace(transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        t[i] += a[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_add_inplace(transform_t<Type>& t, int idx, Type value,
                                         transform_t<Type>& adj_t, int adj_idx, Type& adj_value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     adj_value += adj_t[idx];
 }
 
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void adj_add_inplace(
+    const transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a,
+    transform_t<Type>& adj_t, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a
+)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        adj_a[ii] += adj_t[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<typename Type>
 inline CUDA_CALLABLE void sub_inplace(transform_t<Type>& t, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     t[idx] -= value;
 }
 
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void sub_inplace(transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        t[i] -= a[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_sub_inplace(transform_t<Type>& t, int idx, Type value,
                                         transform_t<Type>& adj_t, int adj_idx, Type& adj_value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     adj_value -= adj_t[idx];
 }
 
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void adj_sub_inplace(
+    const transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a,
+    transform_t<Type>& adj_t, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a
+)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        adj_a[ii] -= adj_t[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<typename Type>
 inline CUDA_CALLABLE void assign_inplace(transform_t<Type>& t, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     t[idx] = value;
 }
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void assign_inplace(transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        t[i] = a[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_assign_inplace(transform_t<Type>& t, int idx, Type value, transform_t<Type>& adj_t, int& adj_idx, Type& adj_value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     adj_value += adj_t[idx];
 }
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void adj_assign_inplace(
+    const transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a,
+    transform_t<Type>& adj_t, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a
+)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (
+        int i = slice.start;
+        is_reversed ? (i > slice.stop) : (i < slice.stop);
+        i += slice.step
+    )
+    {
+        adj_a[ii] += adj_t[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
 
 template<typename Type>
 inline CUDA_CALLABLE transform_t<Type> assign_copy(transform_t<Type>& t, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     transform_t<Type> ret(t);
     ret[idx] = value;
     return ret;
 }
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE transform_t<Type> assign_copy(transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    transform_t<Type> ret(t);
+    assign_inplace<SliceLength>(ret, slice, a);
+    return ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_assign_copy(transform_t<Type>& t, int idx, Type value, transform_t<Type>& adj_t, int& adj_idx, Type& adj_value, const transform_t<Type>& adj_ret)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= 7)
+    if (idx < -7 || idx >= 7)
     {
         printf("transformation index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += 7;
+    }
+
     adj_value += adj_ret[idx];
     for(unsigned i=0; i < 7; ++i)
     {
@@ -592,6 +867,42 @@ inline CUDA_CALLABLE void adj_assign_copy(transform_t<Type>& t, int idx, Type va
     }
 }
 
+template<unsigned SliceLength, typename Type>
+inline CUDA_CALLABLE void adj_assign_copy(
+    transform_t<Type>& t, slice_t slice, const vec_t<SliceLength, Type> &a,
+    transform_t<Type>& adj_t, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a,
+    const transform_t<Type>& adj_ret
+)
+{
+    assert(slice.start >= 0 && slice.start <= 7);
+    assert(slice.stop >= -1 && slice.stop <= 7);
+    assert(slice.step != 0 && slice.step < 0 ? slice.start >= slice.stop : slice.start <= slice.stop);
+    assert(slice_get_length(slice) == SliceLength);
+
+    bool is_reversed = slice.step < 0;
+
+    int ii = 0;
+    for (int i = 0; i < 7; ++i)
+    {
+        bool in_slice = is_reversed
+            ? (i <= slice.start && i > slice.stop && (slice.start - i) % (-slice.step) == 0)
+            : (i >= slice.start && i < slice.stop && (i - slice.start) % slice.step == 0);
+
+        if (!in_slice)
+        {
+            adj_t[i] += adj_ret[i];
+        }
+        else
+        {
+            adj_a[ii] += adj_ret[i];
+            ++ii;
+        }
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 // adjoint methods
 template<typename Type>
 CUDA_CALLABLE inline void adj_add(const transform_t<Type>& a, const transform_t<Type>& b, transform_t<Type>& adj_a, transform_t<Type>& adj_b, const transform_t<Type>& adj_ret)
@@ -600,6 +911,25 @@ CUDA_CALLABLE inline void adj_add(const transform_t<Type>& a, const transform_t<
     adj_add(a.q, b.q, adj_a.q, adj_b.q, adj_ret.q);
 }
 
+template<typename Type>
+CUDA_CALLABLE inline void adj_add(
+    const transform_t<Type>& a, Type b,
+    transform_t<Type>& adj_a, Type& adj_b,
+    const transform_t<Type>& adj_ret
+)
+{
+    adj_a += adj_ret;
+
+    adj_b += adj_ret.p[0];
+    adj_b += adj_ret.p[1];
+    adj_b += adj_ret.p[2];
+
+    adj_b += adj_ret.q[0];
+    adj_b += adj_ret.q[1];
+    adj_b += adj_ret.q[2];
+    adj_b += adj_ret.q[3];
+}
+
 template<typename Type>
 CUDA_CALLABLE inline void adj_sub(const transform_t<Type>& a, const transform_t<Type>& b, transform_t<Type>& adj_a, transform_t<Type>& adj_b, const transform_t<Type>& adj_ret)
 {
@@ -607,6 +937,25 @@ CUDA_CALLABLE inline void adj_sub(const transform_t<Type>& a, const transform_t<
     adj_sub(a.q, b.q, adj_a.q, adj_b.q, adj_ret.q);
 }
 
+template<typename Type>
+CUDA_CALLABLE inline void adj_sub(
+    const transform_t<Type>& a, Type b,
+    transform_t<Type>& adj_a, Type& adj_b,
+    const transform_t<Type>& adj_ret
+)
+{
+    adj_a -= adj_ret;
+
+    adj_b -= adj_ret.p[0];
+    adj_b -= adj_ret.p[1];
+    adj_b -= adj_ret.p[2];
+
+    adj_b -= adj_ret.q[0];
+    adj_b -= adj_ret.q[1];
+    adj_b -= adj_ret.q[2];
+    adj_b -= adj_ret.q[3];
+}
+
 template<typename Type>
 CUDA_CALLABLE inline void adj_mul(const transform_t<Type>& a, Type s, transform_t<Type>& adj_a, Type& adj_s, const transform_t<Type>& adj_ret)
 {
@@ -942,4 +1291,4 @@ using spatial_matrixh = spatial_matrix_t<half>;
 using spatial_matrixf = spatial_matrix_t<float>;
 using spatial_matrixd = spatial_matrix_t<double>;
 
- } // namespace wp
+ } // namespace wp

warp/native/temp_buffer.h

@@ -26,13 +26,13 @@ template <typename T = char> struct ScopedTemporary
 {
 
     ScopedTemporary(void *context, size_t size)
-        : m_context(context), m_buffer(static_cast<T*>(alloc_device(m_context, size * sizeof(T))))
+        : m_context(context), m_buffer(static_cast<T*>(wp_alloc_device(m_context, size * sizeof(T))))
     {
     }
 
     ~ScopedTemporary()
     {
-        free_device(m_context, m_buffer);
+        wp_free_device(m_context, m_buffer);
     }
 
     T *buffer() const