warp-lang 1.8.0__py3-none-manylinux_2_34_aarch64.whl → 1.9.0__py3-none-manylinux_2_34_aarch64.whl

warp/native/vec.h

@@ -25,7 +25,7 @@ namespace wp
 template<unsigned Length, typename Type>
 struct vec_t
 {
-    Type c[Length];
+    Type c[Length < 1 ? 1 : Length];
 
     inline CUDA_CALLABLE vec_t()
         : c()
@@ -343,6 +343,17 @@ inline CUDA_CALLABLE vec_t<Length, Type> add(vec_t<Length, Type> a, vec_t<Length
     return ret;
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> add(Type a, vec_t<Length, Type> b)
+{
+    vec_t<Length, Type> ret;
+    for( unsigned i=0; i < Length; ++i )
+    {
+        ret[i] = a + b[i];
+    }
+    return ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE vec_t<2, Type> add(vec_t<2, Type> a, vec_t<2, Type> b)
 {
@@ -367,6 +378,18 @@ inline CUDA_CALLABLE vec_t<Length, Type> sub(vec_t<Length, Type> a, vec_t<Length
     return ret;
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> sub(Type a, vec_t<Length, Type> b)
+{
+    vec_t<Length, Type> ret;
+    for (unsigned i=0; i < Length; ++i)
+    {
+        ret[i] = Type(a - b[i]);
+    }
+
+    return ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE vec_t<2, Type> sub(vec_t<2, Type> a, vec_t<2, Type> b)
 {
@@ -440,27 +463,64 @@ template<unsigned Length, typename Type>
 inline CUDA_CALLABLE Type extract(const vec_t<Length, Type> & a, int idx)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     return a[idx];        
 }
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<SliceLength, Type> extract(const vec_t<Length, Type> & a, slice_t slice)
+{
+    vec_t<SliceLength, Type> ret;
+
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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] = a[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+    return ret;
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE Type* index(vec_t<Length, Type>& v, int idx)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     return &v[idx];
 }
 
@@ -468,13 +528,18 @@ template<unsigned Length, typename Type>
 inline CUDA_CALLABLE Type* indexref(vec_t<Length, Type>* v, int idx)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec store %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     return &((*v)[idx]);
 }
 
@@ -498,120 +563,325 @@ template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void add_inplace(vec_t<Length, Type>& v, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     v[idx] += value;
 }
 
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void add_inplace(vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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
+    )
+    {
+        v[i] += a[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_add_inplace(vec_t<Length, Type>& v, int idx, Type value,
                                         vec_t<Length, Type>& adj_v, int adj_idx, Type& adj_value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     adj_value += adj_v[idx];
 }
 
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_add_inplace(
+    const vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a,
+    vec_t<Length, Type>& adj_v, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a
+)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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_v[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void sub_inplace(vec_t<Length, Type>& v, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     v[idx] -= value;
 }
 
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void sub_inplace(vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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
+    )
+    {
+        v[i] -= a[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_sub_inplace(vec_t<Length, Type>& v, int idx, Type value,
                                         vec_t<Length, Type>& adj_v, int adj_idx, Type& adj_value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     adj_value -= adj_v[idx];
 }
 
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_sub_inplace(
+    const vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a,
+    vec_t<Length, Type>& adj_v, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a
+)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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_v[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void assign_inplace(vec_t<Length, Type>& v, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     v[idx] = value;
 }
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void assign_inplace(vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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
+    )
+    {
+        v[i] = a[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_assign_inplace(vec_t<Length, Type>& v, int idx, Type value, vec_t<Length, Type>& adj_v, int& adj_idx, Type& adj_value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     adj_value += adj_v[idx];
 }
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_assign_inplace(
+    const vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a,
+    vec_t<Length, Type>& adj_v, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a
+)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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_v[i];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
 
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE vec_t<Length, Type> assign_copy(vec_t<Length, Type>& v, int idx, Type value)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx >= Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     vec_t<Length, Type> ret(v);
     ret[idx] = value;
     return ret;
 }
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> assign_copy(vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a)
+{
+    vec_t<Length, Type> ret(v);
+    assign_inplace<SliceLength>(ret, slice, a);
+    return ret;
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_assign_copy(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)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
         printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     adj_value += adj_ret[idx];
     for(unsigned i=0; i < Length; ++i)
     {
@@ -620,6 +890,40 @@ inline CUDA_CALLABLE void adj_assign_copy(vec_t<Length, Type>& v, int idx, Type
     }
 }
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_assign_copy(
+    vec_t<Length, Type>& v, slice_t slice, const vec_t<SliceLength, Type> &a,
+    vec_t<Length, Type>& adj_v, slice_t& adj_slice, vec_t<SliceLength, Type>& adj_a,
+    const vec_t<Length, Type>& adj_ret)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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 < Length; ++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_v[i] += adj_ret[i];
+        }
+        else
+        {
+            adj_a[ii] += adj_ret[i];
+            ++ii;
+        }
+    }
+
+    assert(ii == SliceLength);
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE Type length(vec_t<Length, Type> a)
 {
@@ -969,11 +1273,11 @@ template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_div(Type s, vec_t<Length, Type> a, Type& adj_s, vec_t<Length, Type>& adj_a, const vec_t<Length, Type>& adj_ret)
 {
 
-    adj_s -= dot(a , adj_ret)/ (s * s); // - a / s^2
-
-    for( unsigned i=0; i < Length; ++i )
+    for (unsigned i=0; i < Length; ++i)
     {
-        adj_a[i] += s / adj_ret[i];
+        Type inv = Type(1) / a[i];
+        adj_a[i] -= s * adj_ret[i] * inv * inv;
+        adj_s += adj_ret[i] * inv;
     }
 
 #if FP_CHECK
@@ -999,6 +1303,21 @@ inline CUDA_CALLABLE void adj_add(vec_t<Length, Type> a, vec_t<Length, Type> b,
     adj_b += adj_ret;
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_add(
+    Type a, vec_t<Length, Type> b,
+    Type& adj_a, vec_t<Length, Type>& adj_b,
+    const vec_t<Length, Type>& adj_ret
+)
+{
+    for (unsigned i = 0; i < Length; ++i)
+    {
+        adj_a += adj_ret.c[i];
+    }
+
+    adj_b += adj_ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_add(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)
 {
@@ -1026,6 +1345,21 @@ inline CUDA_CALLABLE void adj_sub(vec_t<Length, Type> a, vec_t<Length, Type> b,
     adj_b -= adj_ret;
 }
 
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_sub(
+    Type a, vec_t<Length, Type> b,
+    Type& adj_a, vec_t<Length, Type>& adj_b,
+    const vec_t<Length, Type>& adj_ret
+)
+{
+    for (unsigned i = 0; i < Length; ++i)
+    {
+        adj_a += adj_ret.c[i];
+    }
+
+    adj_b -= adj_ret;
+}
+
 template<typename Type>
 inline CUDA_CALLABLE void adj_sub(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)
 {
@@ -1106,16 +1440,49 @@ template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_extract(const vec_t<Length, Type> & a, int idx, vec_t<Length, Type> & adj_a, int & adj_idx, Type & adj_ret)
 {
 #ifndef NDEBUG
-    if (idx < 0 || idx > Length)
+    if (idx < -(int)Length || idx >= (int)Length)
     {
-        printf("Tvec2<Scalar> index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
         assert(0);
     }
 #endif
 
+    if (idx < 0)
+    {
+        idx += Length;
+    }
+
     adj_a[idx] += adj_ret;
 }
 
+template<unsigned SliceLength, unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_extract(
+    const vec_t<Length, Type>& a, slice_t slice,
+    vec_t<Length, Type>& adj_a, slice_t& adj_slice,
+    const vec_t<SliceLength, Type>& adj_ret
+)
+{
+    assert(slice.start >= 0 && slice.start <= (int)Length);
+    assert(slice.stop >= -1 && slice.stop <= (int)Length);
+    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[i] += adj_ret[ii];
+        ++ii;
+    }
+
+    assert(ii == SliceLength);
+}
+
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_length(vec_t<Length, Type> a, Type ret, vec_t<Length, Type>& adj_a, const Type adj_ret)
 {