warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_sparse.py

@@ -18,6 +18,7 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp._src.sparse import bsr_set_zero
 from warp.sparse import (
     bsr_assign,
     bsr_axpy,
@@ -59,6 +60,17 @@ def _triplets_to_dense(shape, rows, cols, values):
     return mat
 
 
+def _bsr_pruned(bsr):
+    return bsr_from_triplets(
+        rows_of_blocks=bsr.nrow,
+        cols_of_blocks=bsr.ncol,
+        rows=bsr.uncompress_rows(),
+        columns=bsr.columns,
+        values=bsr.values,
+        prune_numerical_zeros=True,
+    )
+
+
 def _bsr_to_dense(bsr):
     mat = np.zeros(bsr.shape)
 
@@ -113,7 +125,7 @@ def test_bsr_from_triplets(test, device):
 
     ref = _triplets_to_dense(shape, rows, cols, vals)
 
-    bsr = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=float), device=device)
+    bsr = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=float), device=device)
     bsr_set_from_triplets(bsr, rows, cols, vals)
     test.assertEqual(bsr.block_size, block_shape[0] * block_shape[1])
 
@@ -218,7 +230,7 @@ def test_bsr_get_set_diag(test, device):
     vals_np = rng.random(size=(nnz, block_shape[0], block_shape[1]))
     vals = wp.array(vals_np, dtype=float, device=device)
 
-    bsr = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=float), device=device)
+    bsr = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=float), device=device)
     bsr_set_from_triplets(bsr, rows, cols, vals)
 
     diag = bsr_get_diag(bsr)
@@ -274,14 +286,13 @@ def test_bsr_split_merge(test, device):
     block_shape = (4, 2)
     nrow = 4
     ncol = 8
-    shape = (block_shape[0] * nrow, block_shape[1] * ncol)
     n = 20
 
     rows = wp.array(rng.integers(0, high=nrow, size=n, dtype=int), dtype=int, device=device)
     cols = wp.array(rng.integers(0, high=ncol, size=n, dtype=int), dtype=int, device=device)
     vals = wp.array(rng.random(size=(n, block_shape[0], block_shape[1])), dtype=float, device=device)
 
-    bsr = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=float), device=device)
+    bsr = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=float), device=device)
     bsr_set_from_triplets(bsr, rows, cols, vals)
     ref = _bsr_to_dense(bsr)
 
@@ -359,13 +370,13 @@ def make_test_bsr_transpose(block_shape, scalar_type):
         vals_np = rng.random(size=(nnz, block_shape[0], block_shape[1]))
         vals = wp.array(vals_np, dtype=scalar_type, device=device).reshape((nnz, block_shape[0], block_shape[1]))
 
-        bsr = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
+        bsr = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(bsr, rows, cols, vals)
         ref = 2.0 * np.transpose(_bsr_to_dense(bsr))
 
-        bsr_transposed = (2.0 * bsr).transpose()
+        bsr_transposed = (2.0 * bsr).transpose().eval()
 
-        res = _bsr_to_dense(bsr_transposed.eval())
+        res = _bsr_to_dense(bsr_transposed)
         assert_np_equal(res, ref, 0.0001)
 
         if block_shape[0] != block_shape[-1]:
@@ -373,6 +384,22 @@ def make_test_bsr_transpose(block_shape, scalar_type):
             with test.assertRaisesRegex(ValueError, "Destination block shape must be"):
                 bsr_set_transpose(dest=bsr, src=bsr)
 
+        # test masked transpose
+        # remove some non zeros from src and dest matrices
+        bsr_set_from_triplets(bsr, rows[:3], cols[:3], vals[:3])
+        bsr_transposed = bsr_from_triplets(
+            bsr_transposed.nrow,
+            bsr_transposed.ncol,
+            bsr_transposed.uncompress_rows()[:3],
+            bsr_transposed.columns[:3],
+            bsr_transposed.values[:3],
+        )
+
+        assert_np_equal(bsr_transposed.uncompress_rows().numpy()[:3], [0, 1, 1])
+        assert_np_equal(bsr_transposed.columns.numpy()[:3], [2, 0, 2])
+        bsr_set_transpose(bsr_transposed, bsr, masked=True)
+        assert _bsr_pruned(bsr_transposed).nnz_sync() == 2
+
     return test_bsr_transpose
 
 
@@ -392,7 +419,7 @@ def make_test_bsr_axpy(block_shape, scalar_type):
         x_vals = wp.array(rng.random(size=(nnz, block_shape[0], block_shape[1])), dtype=scalar_type, device=device)
         x_vals = x_vals.reshape((nnz, block_shape[0], block_shape[1]))
 
-        x = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
+        x = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(x, x_rows, x_cols, x_vals)
 
         y_rows = wp.array(rng.integers(0, high=nrow, size=nnz, dtype=int), dtype=int, device=device)
@@ -400,7 +427,7 @@ def make_test_bsr_axpy(block_shape, scalar_type):
         y_vals = wp.array(rng.random(size=(nnz, block_shape[0], block_shape[1])), dtype=scalar_type, device=device)
         y_vals = y_vals.reshape((nnz, block_shape[0], block_shape[1]))
 
-        y = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
+        y = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(y, y_rows, y_cols, y_vals)
 
         work_arrays = bsr_axpy_work_arrays()
@@ -457,7 +484,7 @@ def make_test_bsr_mm(block_shape, scalar_type):
         x_vals = wp.array(rng.random(size=(nnz, x_block_shape[0], x_block_shape[1])), dtype=scalar_type, device=device)
         x_vals = x_vals.reshape((nnz, x_block_shape[0], x_block_shape[1]))
 
-        x = bsr_zeros(x_nrow, x_ncol, wp.types.matrix(shape=x_block_shape, dtype=scalar_type), device=device)
+        x = bsr_zeros(x_nrow, x_ncol, wp._src.types.matrix(shape=x_block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(x, x_rows, x_cols, x_vals)
 
         y_rows = wp.array(rng.integers(0, high=y_nrow, size=nnz, dtype=int), dtype=int, device=device)
@@ -465,7 +492,7 @@ def make_test_bsr_mm(block_shape, scalar_type):
         y_vals = wp.array(rng.random(size=(nnz, y_block_shape[0], y_block_shape[1])), dtype=scalar_type, device=device)
         y_vals = y_vals.reshape((nnz, y_block_shape[0], y_block_shape[1]))
 
-        y = bsr_zeros(y_nrow, y_ncol, wp.types.matrix(shape=y_block_shape, dtype=scalar_type), device=device)
+        y = bsr_zeros(y_nrow, y_ncol, wp._src.types.matrix(shape=y_block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(y, y_rows, y_cols, y_vals)
 
         z_rows = wp.array(rng.integers(0, high=z_nrow, size=nnz, dtype=int), dtype=int, device=device)
@@ -473,7 +500,7 @@ def make_test_bsr_mm(block_shape, scalar_type):
         z_vals = wp.array(rng.random(size=(nnz, z_block_shape[0], z_block_shape[1])), dtype=scalar_type, device=device)
         z_vals = z_vals.reshape((nnz, z_block_shape[0], z_block_shape[1]))
 
-        z = bsr_zeros(z_nrow, z_ncol, wp.types.matrix(shape=z_block_shape, dtype=scalar_type), device=device)
+        z = bsr_zeros(z_nrow, z_ncol, wp._src.types.matrix(shape=z_block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(z, z_rows, z_cols, z_vals)
 
         work_arrays = bsr_mm_work_arrays()
@@ -544,7 +571,7 @@ def make_test_bsr_mv(block_shape, scalar_type):
         A_vals = wp.array(rng.random(size=(nnz, block_shape[0], block_shape[1])), dtype=scalar_type, device=device)
         A_vals = A_vals.reshape((nnz, block_shape[0], block_shape[1]))
 
-        A = bsr_zeros(nrow, ncol, wp.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
+        A = bsr_zeros(nrow, ncol, wp._src.types.matrix(shape=block_shape, dtype=scalar_type), device=device)
         bsr_set_from_triplets(A, A_rows, A_cols, A_vals)
 
         if block_shape[1] == 1:
@@ -664,6 +691,83 @@ def make_test_bsr_multiply_deep(block_shape, scalar_type):
     return test_bsr_multiply_deep
 
 
+def test_bsr_mm_max_new_nnz(test, device):
+    """Test that BSR matrix multiplication with max_new_nnz works"""
+    A = bsr_from_triplets(
+        2,
+        2,
+        wp.array([0, 0, 1, 1], dtype=int, device=device),
+        wp.array([0, 1, 0, 1], dtype=int, device=device),
+        wp.array([1.0, 2.0, 3.0, 4.0], dtype=wp.float32, device=device),
+    )
+    B = bsr_from_triplets(
+        2,
+        2,
+        wp.array([0, 0, 1, 1], dtype=int, device=device),
+        wp.array([0, 1, 0, 1], dtype=int, device=device),
+        wp.array([1.0, 2.0, 3.0, 4.0], dtype=wp.float32, device=device),
+    )
+    C = bsr_zeros(2, 2, wp.float32, device=device)
+
+    # max_new_nnz big enough
+    bsr_mm(A, B, C, max_new_nnz=4)
+    test.assertEqual(C.nnz_sync(), 4)
+
+    bsr_set_zero(C)
+    test.assertEqual(C.nnz_sync(), 0)
+
+    # max_new_nnz too small, check warning
+    capture = StdOutCapture()
+    capture.begin()
+    bsr_mm(A, B, C, max_new_nnz=2)
+    test.assertEqual(C.nnz_sync(), 2)
+    output = capture.end()
+
+    # Check that the output contains warnings about "max_new_nnz" being exceeded.
+    # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+    if output != "" or sys.platform != "win32":
+        test.assertRegex(output, r"exceeded")
+
+
+def test_capturability(test, device):
+    """Test that BSR operations are graph-capturable"""
+
+    N = 5
+    M = 3
+
+    C = bsr_diag(wp.zeros(N, dtype=wp.mat33, device=device))
+
+    rows = wp.array([3, 4, 2, 0, 1], dtype=int, device=device)
+    columns = wp.array([2, 0, 1, 2, 1], dtype=int, device=device)
+    values = wp.ones(5, dtype=wp.mat33, device=device)
+
+    def test_body():
+        A = bsr_from_triplets(
+            N,
+            M,
+            rows=rows,
+            columns=columns,
+            values=values,
+        )
+        B = A + bsr_copy(A * 2.0)
+        bsr_mm(A, bsr_transposed(B), C, max_new_nnz=N * N)
+
+    # ensure necessary modules are loaded and reset result
+    test_body()
+    bsr_set_zero(C)
+    test.assertEqual(C.nnz_sync(), 0)
+
+    with wp.ScopedDevice(device):
+        with wp.ScopedCapture(force_module_load=False) as capture:
+            test_body()
+
+    assert_array_equal(bsr_get_diag(C), wp.zeros(N, dtype=wp.mat33, device=device))
+
+    wp.capture_launch(capture.graph)
+    test.assertEqual(C.nnz_sync(), 9)
+    assert_array_equal(bsr_get_diag(C), wp.full(N, value=wp.mat33(9.0), dtype=wp.mat33, device=device))
+
+
 devices = get_test_devices()
 cuda_test_devices = get_selected_cuda_test_devices()
 
@@ -676,7 +780,9 @@ class TestSparse(unittest.TestCase):
         diag_bsr = bsr_diag(diag=np.eye(bsize, dtype=float) * 2.0, rows_of_blocks=nrow)
         diag_copy = bsr_copy(diag_bsr, scalar_type=wp.float64)
 
-        self.assertTrue(wp.types.types_equal(diag_copy.values.dtype, wp.mat(shape=(bsize, bsize), dtype=wp.float64)))
+        self.assertTrue(
+            wp._src.types.types_equal(diag_copy.values.dtype, wp.mat(shape=(bsize, bsize), dtype=wp.float64))
+        )
         bsr_scale(x=diag_copy, alpha=0.5)
 
         res = _bsr_to_dense(diag_copy)
@@ -686,7 +792,10 @@ class TestSparse(unittest.TestCase):
         bsr_scale(x=diag_copy, alpha=0.0)
         self.assertEqual(diag_copy.nrow, nrow)
         self.assertEqual(diag_copy.ncol, nrow)
-        self.assertEqual(diag_copy.nnz, 0)
+        self.assertEqual(diag_copy.nnz, diag_bsr.nnz)
+
+        diag_pruned = _bsr_pruned(diag_copy)
+        self.assertEqual(diag_pruned.nnz_sync(), 0)
 
 
 add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
@@ -728,6 +837,8 @@ add_function_test(TestSparse, "test_csr_mv", make_test_bsr_mv((1, 1), wp.float32
 add_function_test(TestSparse, "test_bsr_mv_1_3", make_test_bsr_mv((1, 3), wp.float32), devices=devices)
 add_function_test(TestSparse, "test_bsr_mv_3_3", make_test_bsr_mv((3, 3), wp.float64), devices=devices)
 
+add_function_test(TestSparse, "test_capturability", test_capturability, devices=cuda_test_devices)
+add_function_test(TestSparse, "test_bsr_mm_max_new_nnz", test_bsr_mm_max_new_nnz, devices=devices, check_output=False)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()