warp-lang 1.6.2__py3-none-win_amd64.whl → 1.7.0__py3-none-win_amd64.whl

warp/native/sparse.cpp

@@ -81,7 +81,8 @@ template <typename T> void bsr_dyn_block_transpose(const T* src, T* dest, int ro
 template <typename T>
 int bsr_matrix_from_triplets_host(const int rows_per_block, const int cols_per_block, const int row_count,
                                   const int nnz, const int* tpl_rows, const int* tpl_columns, const T* tpl_values,
-                                  const bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns, T* bsr_values)
+                                  const bool prune_numerical_zeros, const bool masked, int* bsr_offsets,
+                                  int* bsr_columns, T* bsr_values)
 {
 
     // get specialized accumulator for common block sizes (1,1), (1,2), (1,3),
@@ -124,14 +125,33 @@ int bsr_matrix_from_triplets_host(const int rows_per_block, const int cols_per_b
     std::iota(block_indices.begin(), block_indices.end(), 0);
 
     // remove zero blocks  and invalid row indices
-    block_indices.erase(std::remove_if(block_indices.begin(), block_indices.end(),
-                                       [&](int i)
-                                       {
-                                           return tpl_rows[i] < 0 || tpl_rows[i] >= row_count ||
-                                                  (prune_numerical_zeros && tpl_values &&
-                                                   block_is_zero_func(tpl_values + i * block_size, block_size));
-                                       }),
-                        block_indices.end());
+
+    auto discard_block = [&](int i)
+    {
+        const int row = tpl_rows[i];
+        if (row < 0 || row >= row_count)
+        {
+            return true;
+        }
+
+        if (prune_numerical_zeros && tpl_values && block_is_zero_func(tpl_values + i * block_size, block_size))
+        {
+            return true;
+        }
+
+        if (!masked)
+        {
+            return false;
+        }
+
+        const int* beg = bsr_columns + bsr_offsets[row];
+        const int* end = bsr_columns + bsr_offsets[row + 1];
+        const int col = tpl_columns[i];
+        const int* block = std::lower_bound(beg, end, col);
+        return block == end || *block != col;
+    };
+
+    block_indices.erase(std::remove_if(block_indices.begin(), block_indices.end(), discard_block), block_indices.end());
 
     // sort block indices according to lexico order
     std::sort(block_indices.begin(), block_indices.end(), [tpl_rows, tpl_columns](int i, int j) -> bool
@@ -281,12 +301,12 @@ void bsr_transpose_host(int rows_per_block, int cols_per_block, int row_count, i
 
 WP_API void bsr_matrix_from_triplets_float_host(int rows_per_block, int cols_per_block, int row_count, int nnz,
                                                 int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                                bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
-                                                void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
+                                                bool prune_numerical_zeros, bool masked, int* bsr_offsets,
+                                                int* bsr_columns, void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
 {
     bsr_matrix_from_triplets_host<float>(rows_per_block, cols_per_block, row_count, nnz, tpl_rows, tpl_columns,
-                                         static_cast<const float*>(tpl_values), prune_numerical_zeros, bsr_offsets,
-                                         bsr_columns, static_cast<float*>(bsr_values));
+                                         static_cast<const float*>(tpl_values), prune_numerical_zeros, masked,
+                                         bsr_offsets, bsr_columns, static_cast<float*>(bsr_values));
     if (bsr_nnz)
     {
         *bsr_nnz = bsr_offsets[row_count];
@@ -295,12 +315,12 @@ WP_API void bsr_matrix_from_triplets_float_host(int rows_per_block, int cols_per
 
 WP_API void bsr_matrix_from_triplets_double_host(int rows_per_block, int cols_per_block, int row_count, int nnz,
                                                  int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                                 bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
-                                                 void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
+                                                 bool prune_numerical_zeros, bool masked, int* bsr_offsets,
+                                                 int* bsr_columns, void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
 {
     bsr_matrix_from_triplets_host<double>(rows_per_block, cols_per_block, row_count, nnz, tpl_rows, tpl_columns,
-                                          static_cast<const double*>(tpl_values), prune_numerical_zeros, bsr_offsets,
-                                          bsr_columns, static_cast<double*>(bsr_values));
+                                          static_cast<const double*>(tpl_values), prune_numerical_zeros, masked,
+                                          bsr_offsets, bsr_columns, static_cast<double*>(bsr_values));
     if (bsr_nnz)
     {
         *bsr_nnz = bsr_offsets[row_count];
@@ -327,16 +347,17 @@ WP_API void bsr_transpose_double_host(int rows_per_block, int cols_per_block, in
 
 #if !WP_ENABLE_CUDA
 WP_API void bsr_matrix_from_triplets_float_device(int rows_per_block, int cols_per_block, int row_count, int nnz,
-                                                   int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                                   bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
-                                                   void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
+                                                  int* tpl_rows, int* tpl_columns, void* tpl_values,
+                                                  bool prune_numerical_zeros, bool masked, int* bsr_offsets,
+                                                  int* bsr_columns, void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
 {
 }
 
 WP_API void bsr_matrix_from_triplets_double_device(int rows_per_block, int cols_per_block, int row_count, int nnz,
                                                    int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                                   bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
-                                                   void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
+                                                   bool prune_numerical_zeros, bool masked, int* bsr_offsets,
+                                                   int* bsr_columns, void* bsr_values, int* bsr_nnz,
+                                                   void* bsr_nnz_event)
 {
 }
 

warp/native/sparse.cu

@@ -61,10 +61,41 @@ template <typename T> struct BsrBlockIsNotZero
     }
 };
 
+struct BsrBlockInMask
+{
+    const int* bsr_offsets;
+    const int* bsr_columns;
+
+    CUDA_CALLABLE_DEVICE bool operator()(int row, int col) const
+    {
+        if (bsr_offsets == nullptr)
+            return true;
+
+        int lower = bsr_offsets[row];
+        int upper = bsr_offsets[row + 1] - 1;
+
+        while (lower < upper)
+        {
+            const int mid = lower + (upper - lower) / 2;
+
+            if (bsr_columns[mid] < col)
+            {
+                lower = mid + 1;
+            }
+            else
+            {
+                upper = mid;
+            }
+        }
+
+        return lower == upper && (bsr_columns[lower] == col);
+    }
+};
+
 template <typename T>
 __global__ void bsr_fill_triplet_key_values(const int nnz, const int nrow, const int* tpl_rows, const int* tpl_columns,
-                                            const BsrBlockIsNotZero<T> nonZero, uint32_t* block_indices,
-                                            BsrRowCol* tpl_row_col)
+                                            const BsrBlockIsNotZero<T> nonZero, const BsrBlockInMask mask,
+                                            uint32_t* block_indices, BsrRowCol* tpl_row_col)
 {
     int block = blockIdx.x * blockDim.x + threadIdx.x;
     if (block >= nnz)
@@ -74,7 +105,8 @@ __global__ void bsr_fill_triplet_key_values(const int nnz, const int nrow, const
     const int col = tpl_columns[block];
     const bool is_valid = row >= 0 && row < nrow;
 
-    const BsrRowCol row_col = is_valid && nonZero(block) ? bsr_combine_row_col(row, col) : PRUNED_ROWCOL;
+    const BsrRowCol row_col =
+        is_valid && nonZero(block) && mask(row, col) ? bsr_combine_row_col(row, col) : PRUNED_ROWCOL;
     tpl_row_col[block] = row_col;
     block_indices[block] = block;
 }
@@ -122,7 +154,7 @@ __global__ void bsr_find_row_offsets(uint32_t row_count, const T* d_nnz, const B
 }
 
 template <typename T>
-__global__ void bsr_merge_blocks(const uint32_t* d_nnz, int block_size, const uint32_t* block_offsets,
+__global__ void bsr_merge_blocks(const int* d_nnz, int block_size, const uint32_t* block_offsets,
                                  const uint32_t* sorted_block_indices, const BsrRowCol* unique_row_cols,
                                  const T* tpl_values, int* bsr_cols, T* bsr_values)
 
@@ -163,8 +195,8 @@ __global__ void bsr_merge_blocks(const uint32_t* d_nnz, int block_size, const ui
 template <typename T>
 void bsr_matrix_from_triplets_device(const int rows_per_block, const int cols_per_block, const int row_count,
                                      const int nnz, const int* tpl_rows, const int* tpl_columns, const T* tpl_values,
-                                     const bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
-                                     T* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
+                                     const bool prune_numerical_zeros, const bool masked, int* bsr_offsets,
+                                     int* bsr_columns, T* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
 {
     const int block_size = rows_per_block * cols_per_block;
 
@@ -186,8 +218,9 @@ void bsr_matrix_from_triplets_device(const int rows_per_block, const int cols_pe
 
     // Combine rows and columns so we can sort on them both
     BsrBlockIsNotZero<T> isNotZero{block_size, prune_numerical_zeros ? tpl_values : nullptr};
+    BsrBlockInMask mask{masked ? bsr_offsets : nullptr, bsr_columns};
     wp_launch_device(WP_CURRENT_CONTEXT, bsr_fill_triplet_key_values, nnz,
-                     (nnz, row_count, tpl_rows, tpl_columns, isNotZero, d_keys.Current(), d_values.Current()));
+                     (nnz, row_count, tpl_rows, tpl_columns, isNotZero, mask, d_keys.Current(), d_values.Current()));
 
     // Sort
     {
@@ -214,7 +247,7 @@ void bsr_matrix_from_triplets_device(const int rows_per_block, const int cols_pe
 
     if (bsr_nnz)
     {
-        // Copy nnz to host, and record an event for the competed transfer if desired
+        // Copy nnz to host, and record an event for the completed transfer if desired
 
         memcpy_d2h(WP_CURRENT_CONTEXT, bsr_nnz, bsr_offsets + row_count, sizeof(int), stream);
 
@@ -236,7 +269,7 @@ void bsr_matrix_from_triplets_device(const int rows_per_block, const int cols_pe
 
     // Accumulate repeated blocks and set column indices
     wp_launch_device(WP_CURRENT_CONTEXT, bsr_merge_blocks, nnz,
-                     (unique_triplet_count, block_size, d_keys.Alternate(), d_keys.Current(), d_values.Alternate(),
+                     (bsr_offsets + row_count, block_size, d_keys.Alternate(), d_keys.Current(), d_values.Alternate(),
                       tpl_values, bsr_columns, bsr_values));
 }
 
@@ -452,22 +485,24 @@ void bsr_transpose_device(int rows_per_block, int cols_per_block, int row_count,
 
 void bsr_matrix_from_triplets_float_device(int rows_per_block, int cols_per_block, int row_count, int nnz,
                                            int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                           bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
+                                           bool prune_numerical_zeros, bool masked, int* bsr_offsets, int* bsr_columns,
                                            void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
 {
-    return bsr_matrix_from_triplets_device<float>(
-        rows_per_block, cols_per_block, row_count, nnz, tpl_rows, tpl_columns, static_cast<const float*>(tpl_values),
-        prune_numerical_zeros, bsr_offsets, bsr_columns, static_cast<float*>(bsr_values), bsr_nnz, bsr_nnz_event);
+    return bsr_matrix_from_triplets_device<float>(rows_per_block, cols_per_block, row_count, nnz, tpl_rows, tpl_columns,
+                                                  static_cast<const float*>(tpl_values), prune_numerical_zeros, masked,
+                                                  bsr_offsets, bsr_columns, static_cast<float*>(bsr_values), bsr_nnz,
+                                                  bsr_nnz_event);
 }
 
 void bsr_matrix_from_triplets_double_device(int rows_per_block, int cols_per_block, int row_count, int nnz,
                                             int* tpl_rows, int* tpl_columns, void* tpl_values,
-                                            bool prune_numerical_zeros, int* bsr_offsets, int* bsr_columns,
+                                            bool prune_numerical_zeros, bool masked, int* bsr_offsets, int* bsr_columns,
                                             void* bsr_values, int* bsr_nnz, void* bsr_nnz_event)
 {
-    return bsr_matrix_from_triplets_device<double>(
-        rows_per_block, cols_per_block, row_count, nnz, tpl_rows, tpl_columns, static_cast<const double*>(tpl_values),
-        prune_numerical_zeros, bsr_offsets, bsr_columns, static_cast<double*>(bsr_values), bsr_nnz, bsr_nnz_event);
+    return bsr_matrix_from_triplets_device<double>(rows_per_block, cols_per_block, row_count, nnz, tpl_rows,
+                                                   tpl_columns, static_cast<const double*>(tpl_values),
+                                                   prune_numerical_zeros, masked, bsr_offsets, bsr_columns,
+                                                   static_cast<double*>(bsr_values), bsr_nnz, bsr_nnz_event);
 }
 
 void bsr_transpose_float_device(int rows_per_block, int cols_per_block, int row_count, int col_count, int nnz,

warp/native/svd.h

@@ -432,6 +432,62 @@ void _svd(// input A
     );
 }
 
+
+template<typename Type>
+inline CUDA_CALLABLE
+void _svd_2(// input A
+        Type a11, Type a12,
+        Type a21, Type a22,
+        // output U
+        Type &u11, Type &u12,
+        Type &u21, Type &u22,
+        // output S
+        Type &s11, Type &s12,
+        Type &s21, Type &s22,
+        // output V
+        Type &v11, Type &v12,
+        Type &v21, Type &v22)
+{
+    // Step 1: Compute ATA
+    Type ATA11 = a11 * a11 + a21 * a21;
+    Type ATA12 = a11 * a12 + a21 * a22;
+    Type ATA22 = a12 * a12 + a22 * a22;
+
+    // Step 2: Eigenanalysis
+    Type trace = ATA11 + ATA22;
+    Type det = ATA11 * ATA22 - ATA12 * ATA12;
+    Type sqrt_term = sqrt(trace * trace - Type(4.0) * det);
+    Type lambda1 = (trace + sqrt_term) * Type(0.5);
+    Type lambda2 = (trace - sqrt_term) * Type(0.5);
+
+    // Step 3: Singular values
+    Type sigma1 = sqrt(lambda1);
+    Type sigma2 = sqrt(lambda2);
+
+    // Step 4: Eigenvectors (find V)
+    Type v1x = ATA12, v1y = lambda1 - ATA11; // For first eigenvector
+    Type v2x = ATA12, v2y = lambda2 - ATA11; // For second eigenvector
+    Type norm1 = sqrt(v1x * v1x + v1y * v1y);
+    Type norm2 = sqrt(v2x * v2x + v2y * v2y);
+
+    v11 = v1x / norm1; v12 = v2x / norm2;
+    v21 = v1y / norm1; v22 = v2y / norm2;
+
+    // Step 5: Compute U
+    Type inv_sigma1 = (sigma1 > Type(1e-6)) ? Type(1.0) / sigma1 : Type(0.0);
+    Type inv_sigma2 = (sigma2 > Type(1e-6)) ? Type(1.0) / sigma2 : Type(0.0);
+
+    u11 = (a11 * v11 + a12 * v21) * inv_sigma1;
+    u12 = (a11 * v12 + a12 * v22) * inv_sigma2;
+    u21 = (a21 * v11 + a22 * v21) * inv_sigma1;
+    u22 = (a21 * v12 + a22 * v22) * inv_sigma2;
+
+    // Step 6: Set S
+    s11 = sigma1; s12 = Type(0.0);
+    s21 = Type(0.0); s22 = sigma2;
+}
+
+
 template<typename Type>
 inline CUDA_CALLABLE void svd3(const mat_t<3,3,Type>& A, mat_t<3,3,Type>& U, vec_t<3,Type>& sigma, mat_t<3,3,Type>& V) {
   Type s12, s13, s21, s23, s31, s32;
@@ -492,6 +548,66 @@ inline CUDA_CALLABLE void adj_svd3(const mat_t<3,3,Type>& A,
   adj_A = adj_A + (u_term + v_term + sigma_term);
 }
 
+template<typename Type>
+inline CUDA_CALLABLE void svd2(const mat_t<2,2,Type>& A, mat_t<2,2,Type>& U, vec_t<2,Type>& sigma, mat_t<2,2,Type>& V) {
+  Type s12, s21;
+  _svd_2(A.data[0][0], A.data[0][1],
+       A.data[1][0], A.data[1][1],
+
+       U.data[0][0], U.data[0][1],
+       U.data[1][0], U.data[1][1],
+
+       sigma[0], s12,
+       s21, sigma[1],
+
+       V.data[0][0], V.data[0][1],
+       V.data[1][0], V.data[1][1]);
+}
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_svd2(const mat_t<2,2,Type>& A,
+                                   const mat_t<2,2,Type>& U,
+                                   const vec_t<2,Type>& sigma,
+                                   const mat_t<2,2,Type>& V,
+                                   mat_t<2,2,Type>& adj_A,
+                                   const mat_t<2,2,Type>& adj_U,
+                                   const vec_t<2,Type>& adj_sigma,
+                                   const mat_t<2,2,Type>& adj_V) {
+    Type s1_squared = sigma[0] * sigma[0];
+    Type s2_squared = sigma[1] * sigma[1];
+
+    // Compute inverse of (s1^2 - s2^2) if possible, use small epsilon to prevent division by zero
+    Type F01 = Type(1) / min(s2_squared - s1_squared, Type(-1e-6f));
+
+    // Construct the matrix F for the adjoint
+    mat_t<2,2,Type> F = mat_t<2,2,Type>(0.0, F01,
+                                        -F01, 0.0);
+
+    // Create a matrix to handle the adjoint of the singular values (diagonal matrix)
+    mat_t<2,2,Type> adj_sigma_mat = mat_t<2,2,Type>(adj_sigma[0], 0.0,
+                                                   0.0, adj_sigma[1]);
+
+    // Matrix for handling singular values (diagonal matrix with sigma values)
+    mat_t<2,2,Type> s_mat = mat_t<2,2,Type>(sigma[0], 0.0,
+                                            0.0, sigma[1]);
+
+    // Compute the transpose of U and V
+    mat_t<2,2,Type> UT = transpose(U);
+    mat_t<2,2,Type> VT = transpose(V);
+
+    // Compute the term for sigma (diagonal matrix of adjoint singular values)
+    mat_t<2,2,Type> sigma_term = mul(U, mul(adj_sigma_mat, VT));
+
+    // Compute the adjoint contributions for U (left singular vectors)
+    mat_t<2,2,Type> u_term = mul(mul(U, mul(cw_mul(F, (mul(UT, adj_U) - mul(transpose(adj_U), U))), s_mat)), VT);
+
+    // Compute the adjoint contributions for V (right singular vectors)
+    mat_t<2,2,Type> v_term = mul(U, mul(s_mat, mul(cw_mul(F, (mul(VT, adj_V) - mul(transpose(adj_V), V))), VT)));
+
+    // Combine the terms to compute the adjoint of A
+    adj_A = adj_A + (u_term + v_term + sigma_term);
+}
+
 
 template<typename Type>
 inline CUDA_CALLABLE void qr3(const mat_t<3,3,Type>& A, mat_t<3,3,Type>& Q, mat_t<3,3,Type>& R) {