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

warp/tests/test_utils.py

@@ -87,7 +87,7 @@ def test_array_scan_error_unsupported_dtype(test, device):
 
 
 def test_radix_sort_pairs(test, device):
-    keyTypes = [int, wp.float32]
+    keyTypes = [int, wp.float32, wp.int64]
 
     for keyType in keyTypes:
         keys = wp.array((7, 2, 8, 4, 1, 6, 5, 3, 0, 0, 0, 0, 0, 0, 0, 0), dtype=keyType, device=device)
@@ -97,18 +97,46 @@ def test_radix_sort_pairs(test, device):
         assert_np_equal(values.numpy()[:8], np.array((5, 2, 8, 4, 7, 6, 1, 3)))
 
 
-def test_radix_sort_pairs_empty(test, device):
+def test_segmented_sort_pairs(test, device):
     keyTypes = [int, wp.float32]
 
+    for keyType in keyTypes:
+        keys = wp.array((7, 2, 8, 4, 1, 6, 5, 3, 0, 0, 0, 0, 0, 0, 0, 0), dtype=keyType, device=device)
+        values = wp.array((1, 2, 3, 4, 5, 6, 7, 8, 0, 0, 0, 0, 0, 0, 0, 0), dtype=int, device=device)
+        wp.utils.segmented_sort_pairs(
+            keys,
+            values,
+            8,
+            wp.array((0, 4), dtype=int, device=device),
+            wp.array((4, 8), dtype=int, device=device),
+        )
+        assert_np_equal(keys.numpy()[:8], np.array((2, 4, 7, 8, 1, 3, 5, 6)))
+        assert_np_equal(values.numpy()[:8], np.array((2, 4, 1, 3, 5, 8, 7, 6)))
+
+
+def test_radix_sort_pairs_empty(test, device):
+    keyTypes = [int, wp.float32, wp.int64]
+
     for keyType in keyTypes:
         keys = wp.array((), dtype=keyType, device=device)
         values = wp.array((), dtype=int, device=device)
         wp.utils.radix_sort_pairs(keys, values, 0)
 
 
-def test_radix_sort_pairs_error_insufficient_storage(test, device):
+def test_segmented_sort_pairs_empty(test, device):
     keyTypes = [int, wp.float32]
 
+    for keyType in keyTypes:
+        keys = wp.array((), dtype=keyType, device=device)
+        values = wp.array((), dtype=int, device=device)
+        wp.utils.segmented_sort_pairs(
+            keys, values, 0, wp.array((), dtype=int, device=device), wp.array((), dtype=int, device=device)
+        )
+
+
+def test_radix_sort_pairs_error_insufficient_storage(test, device):
+    keyTypes = [int, wp.float32, wp.int64]
+
     for keyType in keyTypes:
         keys = wp.array((1, 2, 3), dtype=keyType, device=device)
         values = wp.array((1, 2, 3), dtype=int, device=device)
@@ -119,9 +147,28 @@ def test_radix_sort_pairs_error_insufficient_storage(test, device):
             wp.utils.radix_sort_pairs(keys, values, 3)
 
 
-def test_radix_sort_pairs_error_unsupported_dtype(test, device):
+def test_segmented_sort_pairs_error_insufficient_storage(test, device):
     keyTypes = [int, wp.float32]
 
+    for keyType in keyTypes:
+        keys = wp.array((1, 2, 3), dtype=keyType, device=device)
+        values = wp.array((1, 2, 3), dtype=int, device=device)
+        with test.assertRaisesRegex(
+            RuntimeError,
+            r"Array storage must be large enough to contain 2\*count elements$",
+        ):
+            wp.utils.segmented_sort_pairs(
+                keys,
+                values,
+                3,
+                wp.array((0,), dtype=int, device=device),
+                wp.array((3,), dtype=int, device=device),
+            )
+
+
+def test_radix_sort_pairs_error_unsupported_dtype(test, device):
+    keyTypes = [int, wp.float32, wp.int64]
+
     for keyType in keyTypes:
         keys = wp.array((1.0, 2.0, 3.0), dtype=keyType, device=device)
         values = wp.array((1.0, 2.0, 3.0), dtype=float, device=device)
@@ -132,6 +179,25 @@ def test_radix_sort_pairs_error_unsupported_dtype(test, device):
             wp.utils.radix_sort_pairs(keys, values, 1)
 
 
+def test_segmented_sort_pairs_error_unsupported_dtype(test, device):
+    keyTypes = [int, wp.float32]
+
+    for keyType in keyTypes:
+        keys = wp.array((1.0, 2.0, 3.0), dtype=keyType, device=device)
+        values = wp.array((1.0, 2.0, 3.0), dtype=float, device=device)
+        with test.assertRaisesRegex(
+            RuntimeError,
+            r"Unsupported data type$",
+        ):
+            wp.utils.segmented_sort_pairs(
+                keys,
+                values,
+                1,
+                wp.array((0,), dtype=int, device=device),
+                wp.array((3,), dtype=int, device=device),
+            )
+
+
 def test_array_sum(test, device):
     for dtype in (wp.float32, wp.float64):
         with test.subTest(dtype=dtype):
@@ -468,6 +534,20 @@ add_function_test(
     test_radix_sort_pairs_error_unsupported_dtype,
     devices=devices,
 )
+add_function_test(TestUtils, "test_segmented_sort_pairs", test_segmented_sort_pairs, devices=devices)
+add_function_test(TestUtils, "test_segmented_sort_pairs_empty", test_segmented_sort_pairs, devices=devices)
+add_function_test(
+    TestUtils,
+    "test_segmented_sort_pairs_error_insufficient_storage",
+    test_segmented_sort_pairs_error_insufficient_storage,
+    devices=devices,
+)
+add_function_test(
+    TestUtils,
+    "test_segmented_sort_pairs_error_unsupported_dtype",
+    test_segmented_sort_pairs_error_unsupported_dtype,
+    devices=devices,
+)
 add_function_test(TestUtils, "test_array_sum", test_array_sum, devices=devices)
 add_function_test(
     TestUtils, "test_array_sum_error_out_dtype_mismatch", test_array_sum_error_out_dtype_mismatch, devices=devices

warp/tests/test_vec.py

@@ -1044,7 +1044,7 @@ def test_casting_constructors(test, device, dtype, register_kernels=False):
     assert_np_equal(out, a_grad.numpy())
 
 
-def test_vec_assign(test, device, dtype, register_kernels=False):
+def test_vector_assign_inplace(test, device, dtype, register_kernels=False):
     np_type = np.dtype(dtype)
     wp_type = wp.types.np_dtype_to_warp_type[np_type]
 
@@ -1085,16 +1085,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
         g = a_vec[0] + a_vec[1]
         x[tid] = g
 
-    def vectest_in_register_overwrite(x: wp.array(dtype=vec3), a: wp.array(dtype=vec3)):
-        tid = wp.tid()
-
-        f = vec3(wp_type(0.0))
-        a_vec = a[tid]
-        f = a_vec
-        f[1] = wp_type(3.0)
-
-        x[tid] = f
-
     def vectest_component(x: wp.array(dtype=vec3), y: wp.array(dtype=wp_type)):
         i = wp.tid()
 
@@ -1106,7 +1096,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
 
     kernel_read_write_store = getkernel(vectest_read_write_store, suffix=dtype.__name__)
     kernel_in_register = getkernel(vectest_in_register, suffix=dtype.__name__)
-    kernel_in_register_overwrite = getkernel(vectest_in_register_overwrite, suffix=dtype.__name__)
     kernel_component = getkernel(vectest_component, suffix=dtype.__name__)
 
     if register_kernels:
@@ -1156,7 +1145,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
     x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
     y = wp.ones(1, dtype=wp_type, device=device, requires_grad=True)
 
-    tape = wp.Tape()
     with tape:
         wp.launch(kernel_component, dim=1, inputs=[x, y], device=device)
 
@@ -1165,20 +1153,6 @@ def test_vec_assign(test, device, dtype, register_kernels=False):
     assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=np_type))
     assert_np_equal(y.grad.numpy(), np.array([6.0], dtype=np_type))
 
-    tape.reset()
-
-    x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
-    a = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
-
-    tape = wp.Tape()
-    with tape:
-        wp.launch(kernel_in_register_overwrite, dim=1, inputs=[x, a], device=device)
-
-    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
-
-    assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=np_type))
-    assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
-
 
 @wp.kernel
 def test_vector_constructor_value_func():
@@ -1325,15 +1299,15 @@ def vector_augassign_kernel(
 def test_vector_augassign(test, device):
     N = 3
 
-    a = wp.zeros(N, dtype=wp.vec3, requires_grad=True)
-    b = wp.ones(N, dtype=wp.vec3, requires_grad=True)
+    a = wp.zeros(N, dtype=wp.vec3, requires_grad=True, device=device)
+    b = wp.ones(N, dtype=wp.vec3, requires_grad=True, device=device)
 
-    c = wp.zeros(N, dtype=wp.vec3, requires_grad=True)
-    d = wp.ones(N, dtype=wp.vec3, requires_grad=True)
+    c = wp.zeros(N, dtype=wp.vec3, requires_grad=True, device=device)
+    d = wp.ones(N, dtype=wp.vec3, requires_grad=True, device=device)
 
     tape = wp.Tape()
     with tape:
-        wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d])
+        wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d], device=device)
 
     tape.backward(grads={a: wp.ones_like(a), c: wp.ones_like(c)})
 
@@ -1346,6 +1320,38 @@ def test_vector_augassign(test, device):
     assert_np_equal(d.grad.numpy(), -wp.ones_like(d).numpy())
 
 
+def test_vector_assign_copy(test, device):
+    saved_enable_vector_component_overwrites_setting = wp.config.enable_vector_component_overwrites
+    try:
+        wp.config.enable_vector_component_overwrites = True
+
+        @wp.kernel
+        def vec_in_register_overwrite(x: wp.array(dtype=wp.vec3), a: wp.array(dtype=wp.vec3)):
+            tid = wp.tid()
+
+            f = wp.vec3(0.0)
+            a_vec = a[tid]
+            f = a_vec
+            f[1] = 3.0
+
+            x[tid] = f
+
+        x = wp.zeros(1, dtype=wp.vec3, device=device, requires_grad=True)
+        a = wp.ones(1, dtype=wp.vec3, device=device, requires_grad=True)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(vec_in_register_overwrite, dim=1, inputs=[x, a], device=device)
+
+        tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+        assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=float))
+        assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=float))
+
+    finally:
+        wp.config.enable_vector_component_overwrites = saved_enable_vector_component_overwrites_setting
+
+
 devices = get_test_devices()
 
 
@@ -1414,8 +1420,8 @@ for dtype in np_float_types:
     )
     add_function_test_register_kernel(
         TestVec,
-        f"test_vec_assign_{dtype.__name__}",
-        test_vec_assign,
+        f"test_vector_assign_inplace_{dtype.__name__}",
+        test_vector_assign_inplace,
         devices=devices,
         dtype=dtype,
     )
@@ -1468,6 +1474,12 @@ add_function_test(
     test_vector_augassign,
     devices=devices,
 )
+add_function_test(
+    TestVec,
+    "test_vector_assign_copy",
+    test_vector_assign_copy,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":

warp/tests/{test_tile.py → tile/test_tile.py}

@@ -20,8 +20,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.is_mathdx_enabled()
-
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
 TILE_K = wp.constant(8)
@@ -216,7 +214,6 @@ def test_tile_binary_map(test, device):
     assert_np_equal(B_wp.grad.numpy(), B_grad)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_grouped_gemm(test, device):
     @wp.kernel
     def tile_grouped_gemm(A: wp.array3d(dtype=float), B: wp.array3d(dtype=float), C: wp.array3d(dtype=float)):
@@ -256,60 +253,62 @@ def test_tile_grouped_gemm(test, device):
     assert_np_equal(C_wp.numpy(), C, 1e-6)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
-def test_tile_gemm(test, device):
-    @wp.kernel
-    def tile_gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)):
-        # output tile index
-        i, j = wp.tid()
+def test_tile_gemm(dtype):
+    def test(test, device):
+        @wp.kernel
+        def tile_gemm(A: wp.array2d(dtype=dtype), B: wp.array2d(dtype=dtype), C: wp.array2d(dtype=dtype)):
+            # output tile index
+            i, j = wp.tid()
 
-        sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32)
+            sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=dtype)
 
-        M = A.shape[0]
-        N = B.shape[1]
-        K = A.shape[1]
+            M = A.shape[0]
+            N = B.shape[1]
+            K = A.shape[1]
 
-        count = int(K / TILE_K)
+            count = int(K / TILE_K)
 
-        for k in range(0, count):
-            a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i * TILE_M, k * TILE_K))
-            b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k * TILE_K, j * TILE_N))
+            for k in range(0, count):
+                a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i * TILE_M, k * TILE_K))
+                b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k * TILE_K, j * TILE_N))
 
-            # sum += a*b
-            wp.tile_matmul(a, b, sum)
+                # sum += a*b
+                wp.tile_matmul(a, b, sum)
 
-        wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
+            wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
 
-    M = TILE_M * 7
-    K = TILE_K * 6
-    N = TILE_N * 5
+        M = TILE_M * 7
+        K = TILE_K * 6
+        N = TILE_N * 5
 
-    rng = np.random.default_rng(42)
-    A = rng.random((M, K), dtype=np.float32)
-    B = rng.random((K, N), dtype=np.float32)
-    C = np.zeros((M, N), dtype=np.float32)
+        rng = np.random.default_rng(42)
+        A = rng.random((M, K), dtype=float).astype(wp.dtype_to_numpy(dtype))
+        B = rng.random((K, N), dtype=float).astype(wp.dtype_to_numpy(dtype))
+        C = np.zeros((M, N), dtype=float).astype(wp.dtype_to_numpy(dtype))
 
-    A_wp = wp.array(A, requires_grad=True, device=device)
-    B_wp = wp.array(B, requires_grad=True, device=device)
-    C_wp = wp.array(C, requires_grad=True, device=device)
+        A_wp = wp.array(A, requires_grad=True, device=device)
+        B_wp = wp.array(B, requires_grad=True, device=device)
+        C_wp = wp.array(C, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
-        wp.launch_tiled(
-            tile_gemm,
-            dim=(int(M / TILE_M), int(N / TILE_N)),
-            inputs=[A_wp, B_wp, C_wp],
-            block_dim=TILE_DIM,
-            device=device,
-        )
+        with wp.Tape() as tape:
+            wp.launch_tiled(
+                tile_gemm,
+                dim=(int(M / TILE_M), int(N / TILE_N)),
+                inputs=[A_wp, B_wp, C_wp],
+                block_dim=TILE_DIM,
+                device=device,
+            )
 
-    assert_np_equal(C_wp.numpy(), A @ B, tol=1.0e-5)
+        assert_np_equal(C_wp.numpy(), A @ B, tol=1.0e-1)
 
-    adj_C = np.ones_like(C)
+        adj_C = np.ones_like(C)
 
-    tape.backward(grads={C_wp: wp.array(adj_C, device=device)})
+        tape.backward(grads={C_wp: wp.array(adj_C, device=device)})
 
-    assert_np_equal(A_wp.grad.numpy(), adj_C @ B.T, tol=1.0e-5)
-    assert_np_equal(B_wp.grad.numpy(), A.T @ adj_C, 1.0e-5)
+        assert_np_equal(A_wp.grad.numpy(), adj_C @ B.T, tol=1.0e-1)
+        assert_np_equal(B_wp.grad.numpy(), A.T @ adj_C, 1.0e-1)
+
+    return test
 
 
 @wp.kernel
@@ -550,7 +549,6 @@ def test_tile_transpose(test, device):
     assert_np_equal(output.numpy(), input.numpy().T)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_transpose_matmul(test, device):
     @wp.kernel
     def test_tile_transpose_matmul_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
@@ -572,9 +570,36 @@ def test_tile_transpose_matmul(test, device):
 
 
 @wp.kernel
-def test_tile_broadcast_add_kernel(
+def test_tile_broadcast_add_1d_kernel(
+    input_a: wp.array(dtype=float), input_b: wp.array(dtype=float), output: wp.array(dtype=float)
+):
+    a = wp.tile_load(input_a, shape=(10,))
+    b = wp.tile_load(input_b, shape=(1,))
+
+    c = wp.tile_broadcast(b, shape=(10,))
+    d = a + c
+
+    wp.tile_store(output, d)
+
+
+def test_tile_broadcast_add_1d(test, device):
+    N = 10
+
+    # implicit 1-dim ([1], 1)
+    a = wp.array(np.arange(0, N, dtype=np.float32), device=device)
+    b = wp.array(np.ones(1, dtype=np.float32), device=device)
+    out = wp.zeros((N,), dtype=float, device=device)
+
+    wp.launch_tiled(test_tile_broadcast_add_1d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+
+    assert_np_equal(out.numpy(), a.numpy() + b.numpy())
+
+
+@wp.kernel
+def test_tile_broadcast_add_2d_kernel(
     input_a: wp.array2d(dtype=float), input_b: wp.array(dtype=float), output: wp.array2d(dtype=float)
 ):
+    # implicit 1-dim ([1], 10)
     a = wp.tile_load(input_a, shape=(10, 10))
     b = wp.tile_load(input_b, shape=10)
 
@@ -584,7 +609,7 @@ def test_tile_broadcast_add_kernel(
     wp.tile_store(output, d)
 
 
-def test_tile_broadcast_add(test, device):
+def test_tile_broadcast_add_2d(test, device):
     M = 10
     N = 10
 
@@ -592,7 +617,62 @@ def test_tile_broadcast_add(test, device):
     b = wp.array(np.arange(0, N, dtype=np.float32), device=device)
     out = wp.zeros((M, N), dtype=float, device=device)
 
-    wp.launch_tiled(test_tile_broadcast_add_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+    wp.launch_tiled(test_tile_broadcast_add_2d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+
+    assert_np_equal(out.numpy(), a.numpy() + b.numpy())
+
+
+@wp.kernel
+def test_tile_broadcast_add_3d_kernel(
+    input_a: wp.array3d(dtype=float), input_b: wp.array3d(dtype=float), output: wp.array3d(dtype=float)
+):
+    a = wp.tile_load(input_a, shape=(4, 10, 12))
+    b = wp.tile_load(input_b, shape=(4, 10, 1))
+
+    c = wp.tile_broadcast(b, shape=(4, 10, 12))
+    d = a + c
+
+    wp.tile_store(output, d)
+
+
+def test_tile_broadcast_add_3d(test, device):
+    M = 4
+    N = 10
+    O = 12
+
+    # explicit 1-dim (M, N, 1) to (M, N, O)
+    a = wp.array(np.ones((M, N, O), dtype=np.float32), device=device)
+    b = wp.array(np.arange(0, M * N, dtype=np.float32).reshape((M, N, 1)), device=device)
+    out = wp.zeros((M, N, O), dtype=float, device=device)
+
+    wp.launch_tiled(test_tile_broadcast_add_3d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+    assert_np_equal(out.numpy(), a.numpy() + b.numpy())
+
+
+@wp.kernel
+def test_tile_broadcast_add_4d_kernel(
+    input_a: wp.array4d(dtype=float), input_b: wp.array4d(dtype=float), output: wp.array4d(dtype=float)
+):
+    a = wp.tile_load(input_a, shape=(4, 10, 5, 6))
+    b = wp.tile_load(input_b, shape=(4, 1, 5, 1))
+    c = wp.tile_broadcast(b, shape=(4, 10, 5, 6))
+    d = a + c
+
+    wp.tile_store(output, d)
+
+
+def test_tile_broadcast_add_4d(test, device):
+    M = 4
+    N = 10
+    O = 5
+    P = 6
+
+    # explicit 1-dims (M, 1, O, 1) to (M, N, O, P)
+    a = wp.array(np.ones((M, N, O, P), dtype=np.float32), device=device)
+    b = wp.array(np.arange(0, M * O, dtype=np.float32).reshape((M, 1, O, 1)), device=device)
+    out = wp.zeros((M, N, O, P), dtype=float, device=device)
+
+    wp.launch_tiled(test_tile_broadcast_add_4d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
 
     assert_np_equal(out.numpy(), a.numpy() + b.numpy())
 
@@ -665,7 +745,7 @@ def test_tile_print(test, device):
     wp.synchronize()
 
 
-devices = get_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestTile(unittest.TestCase):
@@ -677,15 +757,20 @@ add_function_test(TestTile, "test_tile_copy_2d", test_tile_copy_2d, devices=devi
 add_function_test(TestTile, "test_tile_unary_map", test_tile_unary_map, devices=devices)
 add_function_test(TestTile, "test_tile_binary_map", test_tile_binary_map, devices=devices)
 add_function_test(TestTile, "test_tile_grouped_gemm", test_tile_grouped_gemm, devices=devices)
-add_function_test(TestTile, "test_tile_gemm", test_tile_gemm, devices=devices)
+add_function_test(TestTile, "test_tile_gemm_fp16", test_tile_gemm(wp.float16), devices=devices)
+add_function_test(TestTile, "test_tile_gemm_fp32", test_tile_gemm(wp.float32), devices=devices)
+add_function_test(TestTile, "test_tile_gemm_fp64", test_tile_gemm(wp.float64), devices=devices)
 add_function_test(TestTile, "test_tile_transpose", test_tile_transpose, devices=devices)
 add_function_test(TestTile, "test_tile_transpose_matmul", test_tile_transpose_matmul, devices=devices)
 add_function_test(TestTile, "test_tile_operators", test_tile_operators, devices=devices)
-add_function_test(TestTile, "test_tile_sum", test_tile_sum, devices=devices)
+add_function_test(TestTile, "test_tile_sum", test_tile_sum, devices=devices, check_output=False)
 add_function_test(TestTile, "test_tile_sum_launch", test_tile_sum_launch, devices=devices)
 add_function_test(TestTile, "test_tile_extract", test_tile_extract, devices=devices)
 add_function_test(TestTile, "test_tile_extract_repeated", test_tile_extract_repeated, devices=devices)
-add_function_test(TestTile, "test_tile_broadcast_add", test_tile_broadcast_add, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_1d", test_tile_broadcast_add_1d, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_2d", test_tile_broadcast_add_2d, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_3d", test_tile_broadcast_add_3d, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_4d", test_tile_broadcast_add_4d, devices=devices)
 add_function_test(TestTile, "test_tile_broadcast_grad", test_tile_broadcast_grad, devices=devices)
 add_function_test(TestTile, "test_tile_len", test_tile_len, devices=devices)
 add_function_test(TestTile, "test_tile_print", test_tile_print, devices=devices, check_output=False)

warp/tests/{test_tile_load.py → tile/test_tile_load.py}

@@ -376,7 +376,7 @@ def test_tile_load_fortran(test, device):
     assert_array_equal(B_wp.grad, A_wp.grad)
 
 
-devices = get_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestTileLoad(unittest.TestCase):

warp/tests/{test_tile_mathdx.py → tile/test_tile_mathdx.py}

@@ -92,6 +92,7 @@ def tile_math_fft_kernel_vec2d(gx: wp.array2d(dtype=wp.vec2d), gy: wp.array2d(dt
     wp.tile_store(gy, xy)
 
 
+@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_math_fft(test, device, wp_dtype):
     np_real_dtype = {wp.vec2f: np.float32, wp.vec2d: np.float64}[wp_dtype]
     np_cplx_dtype = {wp.vec2f: np.complex64, wp.vec2d: np.complex128}[wp_dtype]
@@ -172,31 +173,33 @@ def test_tile_math_cholesky(test, device):
     # TODO: implement and test backward pass
 
 
-devices = get_cuda_test_devices()
+all_devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 class TestTileMathDx(unittest.TestCase):
     pass
 
 
 # check_output=False so we can enable libmathdx's logging without failing the tests
-add_function_test(TestTileMathDx, "test_tile_math_matmul", test_tile_math_matmul, devices=devices, check_output=False)
 add_function_test(
-    TestTileMathDx, "test_tile_math_cholesky", test_tile_math_cholesky, devices=devices, check_output=False
+    TestTileMathDx, "test_tile_math_matmul", test_tile_math_matmul, devices=all_devices, check_output=False
+)
+add_function_test(
+    TestTileMathDx, "test_tile_math_cholesky", test_tile_math_cholesky, devices=all_devices, check_output=False
 )
 add_function_test(
     TestTileMathDx,
     "test_tile_math_fft_vec2f",
     functools.partial(test_tile_math_fft, wp_dtype=wp.vec2f),
-    devices=devices,
+    devices=cuda_devices,
     check_output=False,
 )
 add_function_test(
     TestTileMathDx,
     "test_tile_math_fft_vec2d",
     functools.partial(test_tile_math_fft, wp_dtype=wp.vec2d),
-    devices=devices,
+    devices=cuda_devices,
     check_output=False,
 )