warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_snippet.py

@@ -0,0 +1,143 @@
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+def test_basic(test, device):
+    snippet = """
+    out[tid] = a * x[tid] + y[tid];
+    """
+    adj_snippet = """
+    adj_a = x[tid] * adj_out[tid];
+    adj_x[tid] = a * adj_out[tid];
+    adj_y[tid] = adj_out[tid];
+    """
+
+    @wp.func_native(snippet, adj_snippet)
+    def saxpy(
+        a: wp.float32,
+        x: wp.array(dtype=wp.float32),
+        y: wp.array(dtype=wp.float32),
+        out: wp.array(dtype=wp.float32),
+        tid: int,
+    ):
+        ...
+
+    @wp.kernel
+    def saxpy_cu(
+        a: wp.float32, x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32), out: wp.array(dtype=wp.float32)
+    ):
+        tid = wp.tid()
+        saxpy(a, x, y, out, tid)
+
+    @wp.kernel
+    def saxpy_py(
+        a: wp.float32, x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32), out: wp.array(dtype=wp.float32)
+    ):
+        tid = wp.tid()
+        out[tid] = a * x[tid] + y[tid]
+
+    N = 128
+
+    a1 = 2.0
+    x1 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device, requires_grad=True)
+    y1 = wp.zeros_like(x1)
+    out1 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device)
+    adj_out1 = wp.array(np.ones(N, dtype=np.float32), dtype=wp.float32, device=device)
+
+    a2 = 2.0
+    x2 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device, requires_grad=True)
+    y2 = wp.zeros_like(x2)
+    out2 = wp.array(np.arange(N, dtype=np.float32), dtype=wp.float32, device=device)
+    adj_out2 = wp.array(np.ones(N, dtype=np.float32), dtype=wp.float32, device=device)
+
+    tape = wp.Tape()
+
+    with tape:
+        wp.launch(kernel=saxpy_cu, dim=N, inputs=[a1, x1, y1], outputs=[out1], device=device)
+        wp.launch(kernel=saxpy_py, dim=N, inputs=[a2, x2, y2], outputs=[out2], device=device)
+
+    tape.backward(grads={out1: adj_out1, out2: adj_out2})
+
+    # test forward snippet
+    assert_np_equal(out1.numpy(), out2.numpy())
+
+    # test backward snippet
+    assert_np_equal(x1.grad.numpy(), a1 * np.ones(N, dtype=np.float32))
+    assert_np_equal(x1.grad.numpy(), x2.grad.numpy())
+
+    assert_np_equal(y1.grad.numpy(), np.ones(N, dtype=np.float32))
+    assert_np_equal(y1.grad.numpy(), y2.grad.numpy())
+
+
+def test_shared_memory(test, device):
+    snippet = """
+        __shared__ int s[128];
+
+        s[tid] = d[tid];
+        __syncthreads();
+        d[tid] = s[N - tid - 1];
+        """
+
+    @wp.func_native(snippet)
+    def reverse(d: wp.array(dtype=int), N: int, tid: int):
+        return
+
+    @wp.kernel
+    def reverse_kernel(d: wp.array(dtype=int), N: int):
+        tid = wp.tid()
+        reverse(d, N, tid)
+
+    N = 128
+    x = wp.array(np.arange(N, dtype=int), dtype=int, device=device)
+    y = np.arange(127, -1, -1, dtype=int)
+
+    wp.launch(kernel=reverse_kernel, dim=N, inputs=[x, N], device=device)
+
+    assert_np_equal(x.numpy(), y)
+
+
+def test_cpu_snippet(test, device):
+    snippet = """
+    int inc = 1;
+    out[tid] = x[tid] + inc;
+    """
+
+    @wp.func_native(snippet)
+    def increment_snippet(
+        x: wp.array(dtype=wp.int32),
+        out: wp.array(dtype=wp.int32),
+        tid: int,
+    ):
+        ...
+
+    @wp.kernel
+    def increment(x: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.int32)):
+        tid = wp.tid()
+        increment_snippet(x, out, tid)
+
+    N = 128
+    x = wp.array(np.arange(N, dtype=np.int32), dtype=wp.int32, device=device)
+    out = wp.zeros(N, dtype=wp.int32, device=device)
+
+    wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
+
+    assert_np_equal(out.numpy(), np.arange(1, N + 1, 1, dtype=np.int32))
+
+
+class TestSnippets(unittest.TestCase):
+    pass
+
+
+add_function_test(TestSnippets, "test_basic", test_basic, devices=get_unique_cuda_test_devices())
+add_function_test(TestSnippets, "test_shared_memory", test_shared_memory, devices=get_unique_cuda_test_devices())
+add_function_test(TestSnippets, "test_cpu_snippet", test_cpu_snippet, devices=["cpu"])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

warp/tests/test_sparse.py

@@ -1,8 +1,20 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+
 import numpy as np
+
 import warp as wp
+from warp.sparse import bsr_zeros, bsr_set_from_triplets, bsr_get_diag, bsr_diag, bsr_identity, bsr_copy, bsr_scale
+from warp.sparse import bsr_set_transpose, bsr_transposed
+from warp.sparse import bsr_axpy, bsr_mm, bsr_axpy_work_arrays, bsr_mm_work_arrays, bsr_mv
+from warp.tests.unittest_utils import *
 
-from warp.sparse import bsr_zeros, bsr_set_from_triplets, bsr_get_diag, bsr_diag, bsr_set_transpose, bsr_axpy, bsr_mm
-from warp.tests.test_base import *
 
 wp.init()
 
@@ -46,45 +58,62 @@ def _bsr_to_dense(bsr):
 
 
 def test_csr_from_triplets(test, device):
+    rng = np.random.default_rng(123)
+
     shape = (8, 6)
     n = 100
 
-    rows = wp.array(np.random.randint(0, shape[0], n, dtype=int), dtype=int, device=device)
-    cols = wp.array(np.random.randint(0, shape[1], n, dtype=int), dtype=int, device=device)
-    vals = wp.array(np.random.rand(n), dtype=float, device=device)
+    rows = wp.array(rng.integers(0, high=shape[0], size=n, dtype=int), dtype=int, device=device)
+    cols = wp.array(rng.integers(0, high=shape[1], size=n, dtype=int), dtype=int, device=device)
+    vals = wp.array(rng.random(size=n), dtype=float, device=device)
 
     ref = _triplets_to_dense(shape, rows, cols, vals)
 
     csr = bsr_zeros(shape[0], shape[1], float, device=device)
     bsr_set_from_triplets(csr, rows, cols, vals)
+    test.assertEqual(csr.block_size, 1)
 
     res = _bsr_to_dense(csr)
 
-    assert_np_equal(ref, res, 0.0001)
+    assert_np_equal(res, ref, 0.0001)
 
 
 def test_bsr_from_triplets(test, device):
+    rng = np.random.default_rng(123)
+
     block_shape = (3, 2)
     nrow = 4
     ncol = 9
     shape = (block_shape[0] * nrow, block_shape[1] * ncol)
     n = 50
 
-    rows = wp.array(np.random.randint(0, nrow, n, dtype=int), dtype=int, device=device)
-    cols = wp.array(np.random.randint(0, ncol, n, dtype=int), dtype=int, device=device)
-    vals = wp.array(np.random.rand(n, block_shape[0], block_shape[1]), dtype=float, device=device)
+    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)
 
     ref = _triplets_to_dense(shape, rows, cols, vals)
 
     bsr = bsr_zeros(nrow, ncol, wp.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])
 
     res = _bsr_to_dense(bsr)
 
-    assert_np_equal(ref, res, 0.0001)
+    assert_np_equal(res, ref, 0.0001)
+
+    # test zero-length inputs
+    bsr_set_from_triplets(
+        bsr,
+        wp.array([], dtype=int, device=device),
+        wp.array([], dtype=int, device=device),
+        wp.array([], shape=(0, block_shape[0], block_shape[1]), dtype=float, device=device),
+    )
+    test.assertEqual(bsr.nnz, 0)
+
 
+def test_bsr_get_set_diag(test, device):
+    rng = np.random.default_rng(123)
 
-def test_bsr_get_diag(test, device):
     block_shape = (3, 3)
     nrow = 4
     ncol = 4
@@ -92,7 +121,7 @@ def test_bsr_get_diag(test, device):
 
     rows = wp.array([0, 1, 2, 3, 2, 1], dtype=int, device=device)
     cols = wp.array([1, 1, 1, 3, 2, 2], dtype=int, device=device)
-    vals_np = np.random.rand(nnz, block_shape[0], block_shape[1])
+    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)
@@ -106,14 +135,46 @@ def test_bsr_get_diag(test, device):
     assert_np_equal(diag_np[2], vals_np[4], tol=0.00001)
     assert_np_equal(diag_np[3], vals_np[3], tol=0.00001)
 
-    # Test round-trip
+    # Test set_diag/get_diag round-trips with various block types
+
+    # Array of blocks
     diag_bsr = bsr_diag(diag)
-    diag = bsr_get_diag(diag_bsr)
+    bsr_get_diag(diag_bsr, out=diag)
     assert_np_equal(diag_np, diag.numpy())
 
+    diag_scalar_np = rng.random(size=nrow)
+    diag_scalar = wp.array(diag_scalar_np, device=device)
+    diag_bsr = bsr_diag(diag_scalar)
+    diag = bsr_get_diag(diag_bsr)
+    assert_np_equal(diag_scalar_np, diag.numpy(), tol=0.000001)
+
+    # Uniform block diagonal
+
+    with test.assertRaisesRegex(ValueError, "BsrMatrix block type must be either warp matrix or scalar"):
+        # 1d block type -- invalid
+        diag_bsr = bsr_diag(diag=vals_np[0, 0], rows_of_blocks=nrow, cols_of_blocks=nrow + 1)
+
+    diag_bsr = bsr_diag(diag=vals_np[0], rows_of_blocks=nrow, cols_of_blocks=nrow + 1)
+    assert diag_bsr.values.shape[0] == nrow
+    assert_np_equal(diag_bsr.values.numpy(), np.broadcast_to(vals_np[0], shape=(nrow, *block_shape)), tol=0.000001)
+
+    diag_bsr = bsr_diag(diag=float(diag_scalar_np[0]), rows_of_blocks=nrow, cols_of_blocks=nrow + 1)
+    assert diag_bsr.values.shape[0] == nrow
+    assert_np_equal(diag_bsr.values.numpy(), np.full(nrow, diag_scalar_np[0]), tol=0.000001)
+
+    # Identity matrix
+    diag_bsr = bsr_identity(nrow, block_type=wp.mat44, device=device)
+    assert diag_bsr.values.shape[0] == nrow
+    assert_np_equal(diag_bsr.values.numpy(), np.broadcast_to(np.eye(4), shape=(nrow, 4, 4)), tol=0.000001)
+
+    diag_csr = bsr_identity(nrow, block_type=wp.float64, device=device)
+    assert np.all(diag_csr.values.numpy() == np.ones(nrow, dtype=float))
+
 
 def make_test_bsr_transpose(block_shape, scalar_type):
     def test_bsr_transpose(test, device):
+        rng = np.random.default_rng(123)
+
         nrow = 4
         ncol = 5
         nnz = 6
@@ -121,7 +182,7 @@ def make_test_bsr_transpose(block_shape, scalar_type):
         rows = wp.array([0, 1, 2, 3, 2, 1], dtype=int, device=device)
         cols = wp.array([1, 4, 1, 3, 0, 2], dtype=int, device=device)
 
-        vals_np = np.random.rand(nnz, block_shape[0], block_shape[1])
+        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)
@@ -134,49 +195,92 @@ def make_test_bsr_transpose(block_shape, scalar_type):
         bsr_set_transpose(dest=bsr_transposed, src=bsr)
 
         res = _bsr_to_dense(bsr_transposed)
+        assert_np_equal(res, ref, 0.0001)
 
-        assert_np_equal(ref, res, 0.0001)
+        if block_shape[0] != block_shape[-1]:
+            # test incompatible block shape
+            with test.assertRaisesRegex(ValueError, "Destination block shape must be"):
+                bsr_set_transpose(dest=bsr, src=bsr)
 
     return test_bsr_transpose
 
 
+def test_bsr_copy_scale(test, device):
+    nrow = 6
+    bsize = 2
+
+    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)
+
+    test.assertTrue(wp.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)
+    ref = np.eye(nrow * bsize)
+    assert_np_equal(res, ref, 0.0001)
+
+    bsr_scale(x=diag_copy, alpha=0.0)
+    test.assertEqual(diag_copy.nrow, nrow)
+    test.assertEqual(diag_copy.ncol, nrow)
+    test.assertEqual(diag_copy.nnz, 0)
+
+
 def make_test_bsr_axpy(block_shape, scalar_type):
     def test_bsr_axpy(test, device):
+        rng = np.random.default_rng(123)
+
         nrow = 2
         ncol = 3
         nnz = 6
 
-        alpha = -1.0
-        beta = 2.0
+        alphas = [-1.0, 0.0, 1.0]
+        betas = [2.0, -1.0, 0.0]
 
-        x_rows = wp.array(np.random.randint(0, nrow, nnz, dtype=int), dtype=int, device=device)
-        x_cols = wp.array(np.random.randint(0, ncol, nnz, dtype=int), dtype=int, device=device)
-        x_vals = wp.array(np.random.rand(nnz, block_shape[0], block_shape[1]), dtype=scalar_type, device=device)
+        x_rows = wp.array(rng.integers(0, high=nrow, size=nnz, dtype=int), dtype=int, device=device)
+        x_cols = wp.array(rng.integers(0, high=ncol, size=nnz, dtype=int), dtype=int, device=device)
+        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)
         bsr_set_from_triplets(x, x_rows, x_cols, x_vals)
 
-        y_rows = wp.array(np.random.randint(0, nrow, nnz, dtype=int), dtype=int, device=device)
-        y_cols = wp.array(np.random.randint(0, ncol, nnz, dtype=int), dtype=int, device=device)
-        y_vals = wp.array(np.random.rand(nnz, block_shape[0], block_shape[1]), dtype=scalar_type, device=device)
+        y_rows = wp.array(rng.integers(0, high=nrow, size=nnz, dtype=int), dtype=int, device=device)
+        y_cols = wp.array(rng.integers(0, high=ncol, size=nnz, dtype=int), dtype=int, device=device)
+        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)
         bsr_set_from_triplets(y, y_rows, y_cols, y_vals)
 
-        ref = alpha * _bsr_to_dense(x) + beta * _bsr_to_dense(y)
+        work_arrays = bsr_axpy_work_arrays()
+        for alpha, beta in zip(alphas, betas):
+            ref = alpha * _bsr_to_dense(x) + beta * _bsr_to_dense(y)
+            if beta == 0.0:
+                y = bsr_axpy(x, alpha=alpha, beta=beta, work_arrays=work_arrays)
+            else:
+                bsr_axpy(x, y, alpha, beta, work_arrays=work_arrays)
 
-        bsr_axpy(x, y, alpha, beta)
+            res = _bsr_to_dense(y)
+            assert_np_equal(res, ref, 0.0001)
 
+        # test aliasing
+        ref = 3.0 * _bsr_to_dense(y)
+        bsr_axpy(y, y, alpha=1.0, beta=2.0)
         res = _bsr_to_dense(y)
-        assert_np_equal(ref, res, 0.0001)
+        assert_np_equal(res, ref, 0.0001)
+
+        # test incompatible shapes
+        y.ncol = y.ncol + 1
+        with test.assertRaisesRegex(ValueError, "Matrices must have the same number of rows and columns"):
+            bsr_axpy(x, y)
 
     return test_bsr_axpy
 
 
 def make_test_bsr_mm(block_shape, scalar_type):
     def test_bsr_mm(test, device):
+        rng = np.random.default_rng(123)
+
         x_nrow = 3
         x_ncol = 2
         x_block_shape = block_shape
@@ -191,72 +295,166 @@ def make_test_bsr_mm(block_shape, scalar_type):
 
         nnz = 6
 
-        alpha = -1.0
-        beta = 2.0
+        alphas = [-1.0, 0.0, 1.0]
+        betas = [2.0, -1.0, 0.0]
 
-        x_rows = wp.array(np.random.randint(0, x_nrow, nnz, dtype=int), dtype=int, device=device)
-        x_cols = wp.array(np.random.randint(0, x_ncol, nnz, dtype=int), dtype=int, device=device)
-        x_vals = wp.array(np.random.rand(nnz, x_block_shape[0], x_block_shape[1]), dtype=scalar_type, device=device)
+        x_rows = wp.array(rng.integers(0, high=x_nrow, size=nnz, dtype=int), dtype=int, device=device)
+        x_cols = wp.array(rng.integers(0, high=x_ncol, size=nnz, dtype=int), dtype=int, device=device)
+        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)
         bsr_set_from_triplets(x, x_rows, x_cols, x_vals)
 
-        y_rows = wp.array(np.random.randint(0, y_nrow, nnz, dtype=int), dtype=int, device=device)
-        y_cols = wp.array(np.random.randint(0, y_ncol, nnz, dtype=int), dtype=int, device=device)
-        y_vals = wp.array(np.random.rand(nnz, y_block_shape[0], y_block_shape[1]), dtype=scalar_type, device=device)
+        y_rows = wp.array(rng.integers(0, high=y_nrow, size=nnz, dtype=int), dtype=int, device=device)
+        y_cols = wp.array(rng.integers(0, high=y_ncol, size=nnz, dtype=int), dtype=int, device=device)
+        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)
         bsr_set_from_triplets(y, y_rows, y_cols, y_vals)
 
-        z_rows = wp.array(np.random.randint(0, z_nrow, nnz, dtype=int), dtype=int, device=device)
-        z_cols = wp.array(np.random.randint(0, z_ncol, nnz, dtype=int), dtype=int, device=device)
-        z_vals = wp.array(np.random.rand(nnz, z_block_shape[0], z_block_shape[1]), dtype=scalar_type, device=device)
+        z_rows = wp.array(rng.integers(0, high=z_nrow, size=nnz, dtype=int), dtype=int, device=device)
+        z_cols = wp.array(rng.integers(0, high=z_ncol, size=nnz, dtype=int), dtype=int, device=device)
+        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)
         bsr_set_from_triplets(z, z_rows, z_cols, z_vals)
 
-        ref = alpha * (_bsr_to_dense(x) @ _bsr_to_dense(y)) + beta * _bsr_to_dense(z)
+        work_arrays = bsr_mm_work_arrays()
+        for alpha, beta in zip(alphas, betas):
+            ref = alpha * (_bsr_to_dense(x) @ _bsr_to_dense(y)) + beta * _bsr_to_dense(z)
+
+            bsr_mm(x, y, z, alpha, beta, work_arrays=work_arrays)
+
+            res = _bsr_to_dense(z)
+            assert_np_equal(res, ref, 0.0001)
+
+        # test aliasing of matrix arguments
+        # x = alpha * z * x + beta * x
+        alpha, beta = alphas[0], betas[0]
+        ref = alpha * (_bsr_to_dense(z) @ _bsr_to_dense(x)) + beta * _bsr_to_dense(x)
+        bsr_mm(z, x, x, alpha, beta)
 
-        bsr_mm(x, y, z, alpha, beta)
+        res = _bsr_to_dense(x)
+        assert_np_equal(res, ref, 0.0001)
+
+        # z = alpha * z * z + beta * z
+        ref = alpha * (_bsr_to_dense(z) @ _bsr_to_dense(z)) + beta * _bsr_to_dense(z)
+        bsr_mm(z, z, z, alpha, beta)
 
         res = _bsr_to_dense(z)
-        assert_np_equal(ref, res, 0.0001)
+        assert_np_equal(res, ref, 0.0001)
+
+        # test incompatible shapes
+        if block_shape[0] != block_shape[-1]:
+            with test.assertRaisesRegex(ValueError, "Incompatible block sizes"):
+                bsr_mm(z, y)
+
+        y.ncol = y.ncol * 2
+        with test.assertRaisesRegex(ValueError, "Incompatible number of rows/columns"):
+            bsr_mm(y, z)
 
     return test_bsr_mm
 
 
-def register(parent):
-    devices = get_test_devices()
+def make_test_bsr_mv(block_shape, scalar_type):
+    def test_bsr_mv(test, device):
+        rng = np.random.default_rng(123)
 
-    class TestSparse(parent):
-        pass
+        nrow = 2
+        ncol = 3
+        nnz = 6
 
-    add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
-    add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
-    add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_diag, devices=devices)
+        alphas = [-1.0, 0.0, 1.0]
+        betas = [2.0, -1.0, 0.0]
+        A_rows = wp.array(rng.integers(0, high=nrow, size=nnz, dtype=int), dtype=int, device=device)
+        A_cols = wp.array(rng.integers(0, high=ncol, size=nnz, dtype=int), dtype=int, device=device)
+        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]))
 
-    add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
-    add_function_test(
-        TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices
-    )
-    add_function_test(
-        TestSparse, "test_bsr_transpose_3_3", make_test_bsr_transpose((3, 3), wp.float64), devices=devices
-    )
+        A = bsr_zeros(nrow, ncol, wp.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:
+            x = wp.array(rng.random(size=ncol), dtype=scalar_type, device=device)
+        else:
+            x = wp.array(
+                rng.random(size=(ncol, block_shape[1])),
+                dtype=wp.vec(length=block_shape[1], dtype=scalar_type),
+                device=device,
+            )
+
+        if block_shape[0] == 1:
+            y = wp.array(rng.random(size=nrow), dtype=scalar_type, device=device)
+        else:
+            y = wp.array(
+                rng.random(size=(nrow, block_shape[0])),
+                dtype=wp.vec(length=block_shape[0], dtype=scalar_type),
+                device=device,
+            )
+
+        work_buffer = wp.empty_like(y)
+        for alpha, beta in zip(alphas, betas):
+            ref = alpha * _bsr_to_dense(A) @ x.numpy().flatten() + beta * y.numpy().flatten()
+            if beta == 0.0:
+                y = bsr_mv(A, x, alpha=alpha, beta=beta, work_buffer=work_buffer)
+            else:
+                bsr_mv(A, x, y, alpha, beta, work_buffer=work_buffer)
+
+            res = y.numpy().flatten()
+            assert_np_equal(res, ref, 0.0001)
+
+        # test aliasing
+        alpha, beta = alphas[0], betas[0]
+        AAt = bsr_mm(A, bsr_transposed(A))
+        ref = alpha * _bsr_to_dense(AAt) @ y.numpy().flatten() + beta * y.numpy().flatten()
+        bsr_mv(AAt, y, y, alpha, beta)
+        res = y.numpy().flatten()
+        assert_np_equal(res, ref, 0.0001)
+
+        A.ncol = A.ncol + 1
+        with test.assertRaisesRegex(ValueError, "Number of columns"):
+            bsr_mv(A, x, y)
+
+        A.ncol = A.ncol - 1
+        A.nrow = A.nrow - 1
+        with test.assertRaisesRegex(ValueError, "Number of rows"):
+            bsr_mv(A, x, y)
+
+    return test_bsr_mv
+
+
+devices = get_test_devices()
+
+
+class TestSparse(unittest.TestCase):
+    pass
+
+
+add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
+add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
+add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
+add_function_test(TestSparse, "test_bsr_copy_scale", test_bsr_copy_scale, devices=devices)
+
+add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_transpose_3_3", make_test_bsr_transpose((3, 3), wp.float64), devices=devices)
 
-    add_function_test(TestSparse, "test_csr_axpy", make_test_bsr_axpy((1, 1), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_axpy_1_3", make_test_bsr_axpy((1, 3), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_axpy_3_3", make_test_bsr_axpy((3, 3), wp.float64), devices=devices)
+add_function_test(TestSparse, "test_csr_axpy", make_test_bsr_axpy((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_axpy_1_3", make_test_bsr_axpy((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_axpy_3_3", make_test_bsr_axpy((3, 3), wp.float64), devices=devices)
 
-    add_function_test(TestSparse, "test_csr_mm", make_test_bsr_mm((1, 1), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_mm_1_3", make_test_bsr_mm((1, 3), wp.float32), devices=devices)
-    add_function_test(TestSparse, "test_bsr_mm_3_3", make_test_bsr_mm((3, 3), wp.float64), devices=devices)
+add_function_test(TestSparse, "test_csr_mm", make_test_bsr_mm((1, 1), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_mm_1_3", make_test_bsr_mm((1, 3), wp.float32), devices=devices)
+add_function_test(TestSparse, "test_bsr_mm_3_3", make_test_bsr_mm((3, 3), wp.float64), devices=devices)
 
-    return TestSparse
+add_function_test(TestSparse, "test_csr_mv", make_test_bsr_mv((1, 1), wp.float32), devices=devices)
+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)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)