warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/native/mat.h

@@ -389,23 +389,17 @@ inline CUDA_CALLABLE bool operator==(const mat_t<Rows,Cols,Type>& a, const mat_t
     return true;
 }
 
-
-// negation:
 template<unsigned Rows, unsigned Cols, typename Type>
-inline CUDA_CALLABLE mat_t<Rows,Cols,Type> operator - (mat_t<Rows,Cols,Type> a)
+inline CUDA_CALLABLE mat_t<Rows,Cols,Type> operator - (const mat_t<Rows,Cols,Type>& x)
 {
-    // NB: this constructor will initialize all ret's components to 0, which is
-    // unnecessary... 
     mat_t<Rows,Cols,Type> ret;
     for (unsigned i=0; i < Rows; ++i)
         for (unsigned j=0; j < Cols; ++j)
-            ret.data[i][j] = -a.data[i][j];
+            ret.data[i][j] = -x.data[i][j];
 
-    // Wonder if this does a load of copying when it returns... hopefully not as it's inlined?
     return ret;
 }
 
-
 template<unsigned Rows, unsigned Cols, typename Type>
 CUDA_CALLABLE inline mat_t<Rows,Cols,Type> pos(const mat_t<Rows,Cols,Type>& x)
 {
@@ -1539,13 +1533,13 @@ inline CUDA_CALLABLE void adj_div(const mat_t<Rows,Cols,Type>& a, Type s, mat_t<
 template<unsigned Rows, unsigned Cols, typename Type>
 inline CUDA_CALLABLE void adj_div(Type s, const mat_t<Rows,Cols,Type>& a, Type& adj_s, mat_t<Rows,Cols,Type>& adj_a, const mat_t<Rows,Cols,Type>& adj_ret)
 {
-    adj_s -= tensordot(a , adj_ret)/ (s * s); // - a / s^2
-
     for (unsigned i=0; i < Rows; ++i)
     {
         for (unsigned j=0; j < Cols; ++j)
         {
-            adj_a.data[i][j] += s / adj_ret.data[i][j];
+            Type inv = Type(1) / a.data[i][j];
+            adj_a.data[i][j] -= s * adj_ret.data[i][j] * inv * inv;
+            adj_s += adj_ret.data[i][j] * inv;
         }
     }
 }
@@ -2206,8 +2200,9 @@ inline CUDA_CALLABLE void expect_near(const mat_t<Rows,Cols,Type>& actual, const
     if (diff > tolerance)
     {
         printf("Error, expect_near() failed with tolerance "); print(tolerance);
-        printf("\t Expected: "); print(expected);
-        printf("\t Actual: "); print(actual);
+        printf("    Expected: "); print(expected);
+        printf("    Actual: "); print(actual);
+        printf("    Max absolute difference: "); print(diff);
     }
 }
 

warp/native/mathdx.cpp

@@ -26,7 +26,8 @@ extern "C"
 WP_API
 bool cuda_compile_fft(
                       const char* ltoir_output_path,
-                      const char* symbol_name, int num_include_dirs,
+                      const char* symbol_name,
+                      int num_include_dirs,
                       const char** include_dirs,
                       const char* mathdx_include_dir,
                       int arch,
@@ -41,7 +42,6 @@ bool cuda_compile_fft(
 }
 
 WP_API bool cuda_compile_dot(
-                             const char* fatbin_output_path,
                              const char* ltoir_output_path,
                              const char* symbol_name,
                              int num_include_dirs,
@@ -55,9 +55,9 @@ WP_API bool cuda_compile_dot(
                              int precision_B,
                              int precision_C,
                              int type,
-                             int a_arrangement,
-                             int b_arrangement,
-                             int c_arrangement,
+                             int arrangement_A,
+                             int arrangement_B,
+                             int arrangement_C,
                              int num_threads)
 {
     printf("CUDA is disabled and/or Warp was not compiled with MathDx support.\n");
@@ -65,6 +65,7 @@ WP_API bool cuda_compile_dot(
 }
 
 WP_API bool cuda_compile_solver(
+                                const char* fatbin_output_path,
                                 const char* ltoir_output_path,
                                 const char* symbol_name,
                                 int num_include_dirs,
@@ -73,8 +74,13 @@ WP_API bool cuda_compile_solver(
                                 int arch,
                                 int M,
                                 int N,
+                                int NRHS,
                                 int function,
+                                int side,
+                                int diag,
                                 int precision,
+                                int arrangement_A,
+                                int arrangement_B,
                                 int fill_mode,
                                 int num_threads)
 {

warp/native/matnn.h

@@ -218,126 +218,4 @@ CUDA_CALLABLE inline void adj_dense_solve(int n,
     }
 }
 
-
-template <typename F>
-CUDA_CALLABLE inline void mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int index, const array_t<float>& x, array_t<float>& out)
-{
-    const int m = weights.shape[0];
-    const int n = weights.shape[1];
-    const int b = x.shape[1];
-
-    for (int i=0; i < m; ++i)
-    {
-        float tmp = bias.data[i];
-
-        for(int j=0; j < n; ++j)
-        {
-            tmp += weights.data[i*n + j]*x.data[index + b*j];
-        }
-
-        out.data[index + b*i] = activation(tmp);
-    }
-}
-
-template <typename F, typename AdjF>
-CUDA_CALLABLE inline void adj_mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int index, const array_t<float>& x, array_t<float>& out,
-                                  array_t<float>& adj_weights, array_t<float>& adj_bias, AdjF adj_activation, int adj_index, array_t<float>& adj_x, array_t<float>& adj_out)
-{
-    const int m = weights.shape[0];
-    const int n = weights.shape[1];
-    const int b = x.shape[1];
-
-    for (int i=0; i < m; ++i)
-    {
-        // recompute forward pass so we don't have to store pre-activation outputs
-        float tmp = bias.data[i];
-
-        for(int j=0; j < n; ++j)
-        {
-            tmp += weights.data[i*n + j]*x.data[index + b*j];
-        }
-
-        // adjoint w.r.t to activation
-        float adj_f = 0.0f;
-    
-        if (adj_out.data)
-            adj_activation(tmp, adj_f, adj_out.data[index + b*i]);
-
-        for (int j=0; j < n; ++j)
-        {
-            // adjoint w.r.t M_i
-            if (adj_weights.data)
-                atomic_add(&adj_weights.data[i*n + j], x.data[index + b*j]*adj_f);    // todo: reduce these atomic stores using warp/block level reductions
-
-            // adjoint w.r.t x
-            if (adj_x.data)
-                atomic_add(&adj_x.data[index + b*j], weights.data[i*n + j]*adj_f);
-        }
-
-        // adjoint w.r.t b
-        if (adj_bias.data)
-            atomic_add(&adj_bias.data[i], adj_f);
-
-    }
-}
-
-
-// template <typename F>
-// CUDA_CALLABLE inline void mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int m, int n, int b, int index, const array_t<float>& x, array_t<float>& out)
-// {
-//     x += index*n;
-//     out += index*m;
-
-
-//     for (int i=0; i < m; ++i)
-//     {
-//         float tmp = bias[i];
-
-//         for(int j=0; j < n; ++j)
-//         {
-//             tmp += weights[i*n + j]*x[j];
-//         }
-
-//         out[i] = activation(tmp);
-//     }
-// }
-
-// template <typename F, typename AdjF>
-// CUDA_CALLABLE inline void adj_mlp(const array_t<float>& weights, const array_t<float>& bias, F activation, int m, int n, int b, int index, const array_t<float>& x, const array_t<float>& out,
-//                                   array_t<float>& adj_weights, array_t<float>& adj_bias, AdjF adj_activation, int adj_m, int adj_n, int adj_b, int adj_index, array_t<float>& adj_x, array_t<float>& adj_out)
-// {
-//     x += index*n;
-//     out += index*m;
-
-//     adj_x += index*n;
-//     adj_out += index*m;
-
-//     for (int i=0; i < m; ++i)
-//     {
-//         // recompute forward pass so we don't have to store pre-activation outputs
-//         float tmp = bias[i];
-
-//         for(int j=0; j < n; ++j)
-//         {
-//             tmp += weights[i*n + j]*x[index + b*j];            
-//         }
-
-//         // adjoint w.r.t to activation
-//         float adj_f = 0.0f;
-//         adj_activation(tmp, adj_f, adj_out[index + b*i]);
-
-//         for (int j=0; j < n; ++j)
-//         {
-//             // adjoint w.r.t M_i
-//             adj_weights[i*n + j] += x[j]*adj_f;
-
-//             // adjoint w.r.t x
-//             adj_x[index + b*j] += weights[i*n + j]*adj_f;
-//         }
-
-//         // adjoint w.r.t b
-//         adj_bias[i] += adj_f;
-//     }
-// }
-
-} // namespace wp
+} // namespace wp

warp/native/mesh.h

@@ -1357,7 +1357,7 @@ CUDA_CALLABLE inline void adj_mesh_query_point_sign_normal(uint64_t id, const ve
                                                            uint64_t adj_id, vec3& adj_point, float& adj_max_dist, float& adj_epsilon, mesh_query_point_t& adj_ret)
 {
     adj_mesh_query_point_sign_normal(id, point, max_dist, ret.sign, ret.face, ret.u, ret.v, epsilon,
-                                     adj_id, adj_point, adj_max_dist, adj_ret.sign, adj_ret.face, adj_ret.u, adj_ret.v, epsilon, adj_ret.result);
+                                     adj_id, adj_point, adj_max_dist, adj_ret.sign, adj_ret.face, adj_ret.u, adj_ret.v, adj_epsilon, adj_ret.result);
 }
 
 CUDA_CALLABLE inline void adj_mesh_query_point_sign_winding_number(uint64_t id, const vec3& point, float max_dist, float accuracy, float winding_number_threshold, const mesh_query_point_t& ret,

warp/native/quat.h

@@ -274,8 +274,32 @@ inline CUDA_CALLABLE quat_t<Type> add(const quat_t<Type>& a, const quat_t<Type>&
 template<typename Type>
 inline CUDA_CALLABLE quat_t<Type> sub(const quat_t<Type>& a, const quat_t<Type>& b)
 {
-    return quat_t<Type>(a.x-b.x, a.y-b.y, a.z-b.z, a.w-b.w);}
+    return quat_t<Type>(a.x-b.x, a.y-b.y, a.z-b.z, a.w-b.w);
+}
+
+template<typename Type>
+inline CUDA_CALLABLE quat_t<Type> operator - (const quat_t<Type>& q)
+{
+    return quat_t<Type>(-q.x, -q.y, -q.z, -q.w);
+}
 
+template<typename Type>
+CUDA_CALLABLE inline quat_t<Type> pos(const quat_t<Type>& q)
+{
+    return q;
+}
+
+template<typename Type>
+CUDA_CALLABLE inline quat_t<Type> neg(const quat_t<Type>& q)
+{
+    return -q;
+}
+
+template<typename Type>
+CUDA_CALLABLE inline void adj_neg(const quat_t<Type>& q, quat_t<Type>& adj_q, const quat_t<Type>& adj_ret)
+{
+    adj_q -= adj_ret;
+}
 
 template<typename Type>
 inline CUDA_CALLABLE quat_t<Type> mul(const quat_t<Type>& a, const quat_t<Type>& b)
@@ -298,7 +322,6 @@ inline CUDA_CALLABLE quat_t<Type> mul(Type s, const quat_t<Type>& a)
     return mul(a, s);
 }
 
-// division
 template<typename Type>
 inline CUDA_CALLABLE quat_t<Type> div(quat_t<Type> q, Type s)
 {
@@ -881,8 +904,12 @@ inline CUDA_CALLABLE void adj_div(quat_t<Type> a, Type s, quat_t<Type>& adj_a, T
 template<typename Type>
 inline CUDA_CALLABLE void adj_div(Type s, quat_t<Type> a, Type& adj_s, quat_t<Type>& adj_a, const quat_t<Type>& adj_ret)
 {
-    adj_s -= dot(a, adj_ret)/ (s * s); // - a / s^2
-    adj_a += s / adj_ret;
+    for (unsigned i=0; i < 4; ++i)
+    {
+        Type inv = Type(1) / a[i];
+        adj_a[i] -= s * adj_ret[i] * inv * inv;
+        adj_s += adj_ret[i] * inv;
+    }
 }
 
 template<typename Type>
@@ -1357,8 +1384,9 @@ inline CUDA_CALLABLE void expect_near(const quat_t<Type>& actual, const quat_t<T
     if (diff > tolerance)
     {
         printf("Error, expect_near() failed with tolerance "); print(tolerance);
-        printf("\t Expected: "); print(expected);
-        printf("\t Actual: "); print(actual);
+        printf("    Expected: "); print(expected);
+        printf("    Actual: "); print(actual);
+        printf("    Max absolute difference: "); print(diff);
     }
 }
 

warp/native/rand.h

@@ -71,14 +71,14 @@ inline CUDA_CALLABLE float randf(uint32& state, float min, float max) { return (
 // Box-Muller method
 inline CUDA_CALLABLE float randn(uint32& state) { return sqrt(-2.f * log(randf(state) + RANDN_EPSILON)) * cos(2.f * M_PI_F * randf(state)); }
 
-inline CUDA_CALLABLE void adj_rand_init(int seed, int& adj_seed, float adj_ret) {}
-inline CUDA_CALLABLE void adj_rand_init(int seed, int offset, int& adj_seed, int& adj_offset, float adj_ret) {}
+inline CUDA_CALLABLE void adj_rand_init(int seed, int& adj_seed, uint32 adj_ret) {}
+inline CUDA_CALLABLE void adj_rand_init(int seed, int offset, int& adj_seed, int& adj_offset, uint32 adj_ret) {}
 
-inline CUDA_CALLABLE void adj_randi(uint32& state, uint32& adj_state, float adj_ret) {}
-inline CUDA_CALLABLE void adj_randi(uint32& state, int min, int max, uint32& adj_state, int& adj_min, int& adj_max, float adj_ret) {}
+inline CUDA_CALLABLE void adj_randi(uint32& state, uint32& adj_state, int adj_ret) {}
+inline CUDA_CALLABLE void adj_randi(uint32& state, int min, int max, uint32& adj_state, int& adj_min, int& adj_max, int adj_ret) {}
 
-inline CUDA_CALLABLE void adj_randu(uint32& state, uint32& adj_state, float adj_ret) {}
-inline CUDA_CALLABLE void adj_randu(uint32& state, uint32 min, uint32 max, uint32& adj_state, uint32& adj_min, uint32& adj_max, float adj_ret) {}
+inline CUDA_CALLABLE void adj_randu(uint32& state, uint32& adj_state, uint32 adj_ret) {}
+inline CUDA_CALLABLE void adj_randu(uint32& state, uint32 min, uint32 max, uint32& adj_state, uint32& adj_min, uint32& adj_max, uint32 adj_ret) {}
 
 inline CUDA_CALLABLE void adj_randf(uint32& state, uint32& adj_state, float adj_ret) {}
 inline CUDA_CALLABLE void adj_randf(uint32& state, float min, float max, uint32& adj_state, float& adj_min, float& adj_max, float adj_ret) {}
@@ -195,7 +195,7 @@ inline CUDA_CALLABLE void adj_sample_unit_hemisphere_surface(uint32& state, uint
 inline CUDA_CALLABLE void adj_sample_unit_hemisphere(uint32& state, uint32& adj_state, const vec3& adj_ret) {}
 inline CUDA_CALLABLE void adj_sample_unit_square(uint32& state, uint32& adj_state, const vec2& adj_ret) {}
 inline CUDA_CALLABLE void adj_sample_unit_cube(uint32& state, uint32& adj_state, const vec3& adj_ret) {}
-inline CUDA_CALLABLE void adj_sample_unit_hypercube(uint32& state, uint32& adj_state, const vec3& adj_ret) {}
+inline CUDA_CALLABLE void adj_sample_unit_hypercube(uint32& state, uint32& adj_state, const vec4& adj_ret) {}
 
 /*
  * log-gamma function to support some of these distributions. The