warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/tests/tile/test_tile_cholesky.py

@@ -20,7 +20,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.wp_is_mathdx_enabled()
+wp.init()  # For wp._src.context.runtime.core.wp_is_mathdx_enabled()
 
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
@@ -490,7 +490,7 @@ def test_tile_upper_solve(L: wp.array2d(dtype=float), y: wp.array(dtype=float),
 
 
 def test_tile_cholesky_singular_matrices(test, device):
-    if not wp.context.runtime.core.wp_is_mathdx_enabled():
+    if not wp._src.context.runtime.core.wp_is_mathdx_enabled():
         test.skipTest("MathDx is not enabled")
 
     rng = np.random.default_rng(42)
@@ -527,8 +527,11 @@ cuda_devices = get_cuda_test_devices()
 
 
 @unittest.skipUnless(
-    not wp.context.runtime.core.wp_is_mathdx_enabled()
-    or (wp.context.runtime.core.wp_is_mathdx_enabled() and wp.context.runtime.core.wp_cuda_toolkit_version() >= 12060),
+    not wp._src.context.runtime.core.wp_is_mathdx_enabled()
+    or (
+        wp._src.context.runtime.core.wp_is_mathdx_enabled()
+        and wp._src.context.runtime.core.wp_cuda_toolkit_version() >= 12060
+    ),
     "MathDx is not enabled or is enabled but CUDA toolkit version is less than 12.6",
 )
 class TestTileCholesky(unittest.TestCase):

warp/tests/tile/test_tile_load.py

@@ -40,6 +40,7 @@ def tile_load_1d_kernel(
     input: wp.array1d(dtype=float),
     out_full: wp.array1d(dtype=float),
     out_padded: wp.array1d(dtype=float),
+    out_sliced: wp.array1d(dtype=float),
     out_offset: wp.array1d(dtype=float),
 ):
     full0 = wp.tile_load(input, TILE_M)
@@ -50,8 +51,13 @@ def tile_load_1d_kernel(
     padded1 = wp.tile_load(input, shape=TILE_M, offset=TILE_OFFSET)
     padded2 = wp.tile_load(input, shape=(TILE_M,), offset=(TILE_OFFSET,))
 
+    sliced0 = wp.tile_load(input[::2], TILE_M)
+    sliced1 = wp.tile_load(input[::2], shape=TILE_M)
+    sliced2 = wp.tile_load(input[::2], shape=(TILE_M,))
+
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET,))
 
 
@@ -60,13 +66,16 @@ def tile_load_2d_kernel(
     input: wp.array2d(dtype=float),
     out_full: wp.array2d(dtype=float),
     out_padded: wp.array2d(dtype=float),
+    out_sliced: wp.array2d(dtype=float),
     out_offset: wp.array2d(dtype=float),
 ):
     full0 = wp.tile_load(input, shape=(TILE_M, TILE_N))
     padded0 = wp.tile_load(input, shape=(TILE_M, TILE_N), offset=(TILE_OFFSET, TILE_OFFSET))
+    sliced0 = wp.tile_load(input[::2, ::2], shape=(TILE_M, TILE_N))
 
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET))
 
 
@@ -75,13 +84,16 @@ def tile_load_3d_kernel(
     input: wp.array3d(dtype=float),
     out_full: wp.array3d(dtype=float),
     out_padded: wp.array3d(dtype=float),
+    out_sliced: wp.array3d(dtype=float),
     out_offset: wp.array3d(dtype=float),
 ):
     full0 = wp.tile_load(input, shape=(TILE_M, TILE_N, TILE_O))
     padded0 = wp.tile_load(input, shape=(TILE_M, TILE_N, TILE_O), offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
+    sliced0 = wp.tile_load(input[::2, ::2, ::2], shape=(TILE_M, TILE_N, TILE_O))
 
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
 
 
@@ -90,15 +102,18 @@ def tile_load_4d_kernel(
     input: wp.array4d(dtype=float),
     out_full: wp.array4d(dtype=float),
     out_padded: wp.array4d(dtype=float),
+    out_sliced: wp.array4d(dtype=float),
     out_offset: wp.array4d(dtype=float),
 ):
     full0 = wp.tile_load(input, shape=(TILE_M, TILE_N, TILE_O, TILE_P))
     padded0 = wp.tile_load(
         input, shape=(TILE_M, TILE_N, TILE_O, TILE_P), offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET, TILE_OFFSET)
     )
+    sliced0 = wp.tile_load(input[::2, ::2, ::2, ::2], shape=(TILE_M, TILE_N, TILE_O, TILE_P))
 
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
 
 
@@ -112,13 +127,14 @@ def test_tile_load(kernel, ndim):
         input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
         output_full = wp.zeros(shape, dtype=float, device=device)
         output_padded = wp.zeros(shape, dtype=float, device=device)
+        output_sliced = wp.zeros(shape, dtype=float, device=device)
         output_offset = wp.zeros(shape, dtype=float, device=device)
 
         with wp.Tape() as tape:
             wp.launch_tiled(
                 kernel,
                 dim=[1],
-                inputs=[input, output_full, output_padded, output_offset],
+                inputs=[input, output_full, output_padded, output_sliced, output_offset],
                 block_dim=TILE_DIM,
                 device=device,
             )
@@ -134,8 +150,16 @@ def test_tile_load(kernel, ndim):
         ref_offset = np.zeros_like(ref_full)
         ref_offset[src_slice] = ref_full[dest_slice]
 
+        # construct a slice for the source/dest sliced arrays
+        src_slice = tuple(slice(0, dim, 2) for dim in shape)
+        dest_slice = tuple(slice(0, (dim + 1) // 2) for dim in shape)
+
+        ref_sliced = np.zeros_like(ref_full)
+        ref_sliced[dest_slice] = ref_full[src_slice]
+
         assert_np_equal(output_full.numpy(), ref_full)
         assert_np_equal(output_padded.numpy(), ref_padded)
+        assert_np_equal(output_sliced.numpy(), ref_sliced)
         assert_np_equal(output_offset.numpy(), ref_offset)
 
         output_full.grad = wp.ones_like(output_full)
@@ -570,7 +594,7 @@ def test_tile_assign(kernel, ndim):
         input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
         output = wp.zeros_like(input)
 
-        with wp.Tape() as tape:
+        with wp.Tape():
             wp.launch(
                 kernel,
                 dim=shape,

warp/tests/tile/test_tile_mathdx.py

@@ -21,7 +21,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.wp_is_mathdx_enabled()
+wp.init()  # For wp._src.context.runtime.core.wp_is_mathdx_enabled()
 
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
@@ -92,7 +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.wp_is_mathdx_enabled(), "Warp was not built with MathDx support")
+@unittest.skipUnless(wp._src.context.runtime.core.wp_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]
@@ -113,7 +113,7 @@ def test_tile_math_fft(test, device, wp_dtype):
     X_c64 = X.view(np_cplx_dtype).reshape(fft_size, fft_size)
     Y_c64 = np.fft.fft(X_c64, axis=-1)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(kernel, dim=[1, 1], inputs=[X_wp, Y_wp], block_dim=TILE_DIM, device=device)
 
     Y_wp_c64 = Y_wp.numpy().view(np_cplx_dtype).reshape(fft_size, fft_size)

warp/tests/tile/test_tile_matmul.py

@@ -60,7 +60,7 @@ def test_tile_grouped_gemm(test, device):
     B_wp = wp.array(B, requires_grad=True, device=device)
     C_wp = wp.zeros((batch_count, TILE_M, TILE_N), requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_grouped_gemm, dim=[batch_count], inputs=[A_wp, B_wp, C_wp], block_dim=TILE_DIM, device=device
         )

warp/tests/tile/test_tile_mlp.py

@@ -43,7 +43,7 @@ def create_array(rng, dim_in, dim_hid, dtype=float):
 def test_multi_layer_nn(test, device):
     import torch as tc
 
-    if device.is_cuda and not wp.context.runtime.core.wp_is_mathdx_enabled():
+    if device.is_cuda and not wp._src.context.runtime.core.wp_is_mathdx_enabled():
         test.skipTest("Skipping test on CUDA device without MathDx (tolerance)")
 
     NUM_FREQ = wp.constant(8)
@@ -63,7 +63,7 @@ def test_multi_layer_nn(test, device):
         NUM_THREADS = 32
 
     dtype = wp.float16
-    npdtype = wp.types.warp_type_to_np_dtype[dtype]
+    npdtype = wp._src.types.warp_type_to_np_dtype[dtype]
 
     @wp.func
     def relu(x: dtype):
@@ -188,7 +188,6 @@ def test_multi_layer_nn(test, device):
         optimizer_inputs = [p.flatten() for p in params]
         optimizer = warp.optim.Adam(optimizer_inputs, lr=0.01)
 
-        num_batches = int((IMG_WIDTH * IMG_HEIGHT) / BATCH_SIZE)
         max_epochs = 30
 
         # create randomized batch indices
@@ -288,7 +287,6 @@ def test_single_layer_nn(test, device):
     import torch as tc
 
     DIM_IN = 8
-    DIM_HID = 32
     DIM_OUT = 16
 
     NUM_BLOCKS = 56

warp/tests/tile/test_tile_reduce.py

@@ -79,7 +79,7 @@ def tile_sum_to_shared_kernel(input: wp.array2d(dtype=float), output: wp.array(d
 
     a = wp.tile_load(input[i], shape=TILE_DIM)
     s = wp.tile_sum(a)
-    v = s[0]  # force shared storage for s
+    v = s[0]
     wp.tile_store(output, s * 0.5, offset=i)
 
 
@@ -142,7 +142,7 @@ def test_tile_reduce_min(test, device):
     input_wp = wp.array(input, requires_grad=True, device=device)
     output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_min_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
         )
@@ -190,7 +190,7 @@ def test_tile_reduce_argmin(test, device):
     input_wp = wp.array(input, requires_grad=True, device=device)
     output_wp = wp.zeros(batch_count, dtype=wp.int32, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_argmin_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
         )
@@ -231,7 +231,7 @@ def test_tile_reduce_max(test, device):
     input_wp = wp.array(input, requires_grad=True, device=device)
     output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_max_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
         )
@@ -264,7 +264,7 @@ def test_tile_reduce_argmax(test, device):
     input_wp = wp.array(input, requires_grad=True, device=device)
     output_wp = wp.zeros(batch_count, dtype=wp.int32, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_argmax_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
         )
@@ -297,7 +297,7 @@ def test_tile_reduce_custom(test, device):
     input_wp = wp.array(input, requires_grad=True, device=device)
     output_wp = wp.zeros(batch_count, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_reduce_custom_kernel,
             dim=[batch_count],
@@ -333,7 +333,7 @@ def test_tile_scan_inclusive(test, device):
     input_wp = wp.array2d(input, requires_grad=True, device=device)
     output_wp = wp.zeros_like(input_wp, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             create_tile_scan_inclusive_kernel(N),
             dim=[batch_count],
@@ -369,7 +369,7 @@ def test_tile_scan_exclusive(test, device):
     input_wp = wp.array2d(input, requires_grad=True, device=device)
     output_wp = wp.zeros_like(input_wp, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             create_tile_scan_exclusive_kernel(N),
             dim=[batch_count],
@@ -501,12 +501,210 @@ def test_tile_reduce_simt(test, device):
 
     output = wp.zeros(shape=1, dtype=int, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch(tile_reduce_simt_kernel, dim=N, inputs=[output], block_dim=TILE_DIM, device=device)
 
     test.assertEqual(output.numpy()[0], np.sum(np.arange(N)))
 
 
+# Tier 1: axis size <= 32
+@wp.kernel
+def tile_reduce_axis_tier1_sum_axis0_kernel(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    a = wp.tile_load(x, shape=(32, 64), storage="shared")
+    b = wp.tile_sum(a, axis=0)
+    wp.tile_store(y, b)
+
+
+@wp.kernel
+def tile_reduce_axis_tier1_prod_axis1_kernel(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    a = wp.tile_load(x, shape=(32, 8), storage="shared")
+    b = wp.tile_reduce(wp.mul, a, axis=1)
+    wp.tile_store(y, b)
+
+
+@wp.kernel
+def tile_reduce_axis_tier1_sum_axis2_kernel(x: wp.array3d(dtype=float), y: wp.array2d(dtype=float)):
+    a = wp.tile_load(x, shape=(8, 8, 16), storage="shared")
+    b = wp.tile_sum(a, axis=2)
+    wp.tile_store(y, b)
+
+
+# Tier 2: 32 < axis size <= 256
+@wp.kernel
+def tile_reduce_axis_tier2_sum_axis0_kernel(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    a = wp.tile_load(x, shape=(200, 32), storage="shared")
+    b = wp.tile_sum(a, axis=0)
+    wp.tile_store(y, b)
+
+
+@wp.kernel
+def tile_reduce_axis_tier2_prod_axis1_kernel(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    a = wp.tile_load(x, shape=(16, 64), storage="shared")
+    b = wp.tile_reduce(wp.mul, a, axis=1)
+    wp.tile_store(y, b)
+
+
+@wp.kernel
+def tile_reduce_axis_tier2_sum_axis2_kernel(x: wp.array3d(dtype=float), y: wp.array2d(dtype=float)):
+    a = wp.tile_load(x, shape=(8, 8, 128), storage="shared")
+    b = wp.tile_sum(a, axis=2)
+    wp.tile_store(y, b)
+
+
+# Tier 3: axis size > 256
+@wp.kernel
+def tile_reduce_axis_tier3_sum_axis0_kernel(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    a = wp.tile_load(x, shape=(400, 16), storage="shared")
+    b = wp.tile_sum(a, axis=0)
+    wp.tile_store(y, b)
+
+
+@wp.kernel
+def tile_reduce_axis_tier3_prod_axis1_kernel(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    a = wp.tile_load(x, shape=(8, 300), storage="shared")
+    b = wp.tile_reduce(wp.mul, a, axis=1)
+    wp.tile_store(y, b)
+
+
+@wp.kernel
+def tile_reduce_axis_tier3_sum_axis2_kernel(x: wp.array3d(dtype=float), y: wp.array2d(dtype=float)):
+    a = wp.tile_load(x, shape=(4, 4, 384), storage="shared")
+    b = wp.tile_sum(a, axis=2)
+    wp.tile_store(y, b)
+
+
+def test_tile_reduce_axis_tier1(test, device):
+    # 2D sum: axis=0, size 32 (forward and backward)
+    x = wp.ones((32, 64), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros(64, dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_axis_tier1_sum_axis0_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+        )
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.ones(64, dtype=float) * 32.0)
+    assert_np_equal(x.grad.numpy(), np.ones((32, 64), dtype=float))
+
+    # 2D product: axis=1, size 8
+    rng = np.random.default_rng(42)
+    x_np = rng.random((32, 8), dtype=np.float32) * 0.1 + 1.0
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(32, dtype=float, device=device)
+
+    wp.launch_tiled(
+        tile_reduce_axis_tier1_prod_axis1_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+    )
+
+    assert_np_equal(y.numpy(), np.prod(x_np, axis=1), tol=1e-3)
+
+    # 3D sum: axis=2, size 16 (forward and backward)
+    x = wp.ones((8, 8, 16), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros((8, 8), dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_axis_tier1_sum_axis2_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+        )
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.ones((8, 8), dtype=float) * 16.0)
+    assert_np_equal(x.grad.numpy(), np.ones((8, 8, 16), dtype=float))
+
+
+def test_tile_reduce_axis_tier2(test, device):
+    # 2D sum: axis=0, size 200 (forward and backward)
+    x = wp.ones((200, 32), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros(32, dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_axis_tier2_sum_axis0_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+        )
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.ones(32, dtype=float) * 200.0)
+    assert_np_equal(x.grad.numpy(), np.ones((200, 32), dtype=float))
+
+    # 2D product: axis=1, size 64
+    rng = np.random.default_rng(42)
+    x_np = rng.random((16, 64), dtype=np.float32) * 0.05 + 1.0
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(16, dtype=float, device=device)
+
+    wp.launch_tiled(
+        tile_reduce_axis_tier2_prod_axis1_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+    )
+
+    assert_np_equal(y.numpy(), np.prod(x_np, axis=1), tol=1e-2)
+
+    # 3D sum: axis=2, size 128 (forward and backward)
+    x = wp.ones((8, 8, 128), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros((8, 8), dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_axis_tier2_sum_axis2_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+        )
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.ones((8, 8), dtype=float) * 128.0)
+    assert_np_equal(x.grad.numpy(), np.ones((8, 8, 128), dtype=float))
+
+
+def test_tile_reduce_axis_tier3(test, device):
+    # 2D sum: axis=0, size 400 (forward and backward)
+    x = wp.ones((400, 16), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros(16, dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_axis_tier3_sum_axis0_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+        )
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.ones(16, dtype=float) * 400.0, tol=2e-4)
+    assert_np_equal(x.grad.numpy(), np.ones((400, 16), dtype=float))
+
+    # 2D product: axis=1, size 300
+    rng = np.random.default_rng(42)
+    x_np = rng.random((8, 300), dtype=np.float32) * 0.01 + 1.0
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(8, dtype=float, device=device)
+
+    wp.launch_tiled(
+        tile_reduce_axis_tier3_prod_axis1_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+    )
+
+    assert_np_equal(y.numpy(), np.prod(x_np, axis=1), tol=1e-1)
+
+    # 3D sum: axis=2, size 384 (forward and backward)
+    x = wp.ones((4, 4, 384), dtype=float, requires_grad=True, device=device)
+    y = wp.zeros((4, 4), dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_reduce_axis_tier3_sum_axis2_kernel, dim=[1], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device
+        )
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    assert_np_equal(y.numpy(), np.ones((4, 4), dtype=float) * 384.0, tol=2e-4)
+    assert_np_equal(x.grad.numpy(), np.ones((4, 4, 384), dtype=float))
+
+
 @wp.kernel
 def tile_untile_kernel(output: wp.array(dtype=int)):
     # thread index
@@ -525,7 +723,7 @@ def test_tile_untile(test, device):
 
     output = wp.zeros(shape=N, dtype=int, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch(tile_untile_kernel, dim=N, inputs=[output], block_dim=TILE_DIM, device=device)
 
     assert_np_equal(output.numpy(), np.arange(N) * 2)
@@ -549,7 +747,7 @@ def test_tile_untile_scalar(test, device):
 
     output = wp.zeros(shape=N, dtype=int, requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch(tile_untile_kernel, dim=N, inputs=[output], block_dim=TILE_DIM, device=device)
 
     assert_np_equal(output.numpy(), np.arange(N) * 2)
@@ -594,7 +792,7 @@ def tile_ones_kernel(out: wp.array(dtype=float)):
 def test_tile_ones(test, device):
     output = wp.zeros(1, dtype=float, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(tile_ones_kernel, dim=[1], inputs=[output], block_dim=TILE_DIM, device=device)
 
     test.assertAlmostEqual(output.numpy()[0], 256.0)
@@ -622,7 +820,7 @@ def test_tile_arange(test, device):
 
     output = wp.zeros(shape=(5, N), dtype=int, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(tile_arange_kernel, dim=[1], inputs=[output], block_dim=TILE_DIM, device=device)
 
     assert_np_equal(output.numpy()[0], np.arange(17))
@@ -734,6 +932,9 @@ add_function_test(TestTileReduce, "test_tile_reduce_custom", test_tile_reduce_cu
 add_function_test(TestTileReduce, "test_tile_reduce_custom_struct", test_tile_reduce_custom_struct, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_grouped_sum", test_tile_reduce_grouped_sum, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_simt", test_tile_reduce_simt, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_axis_tier1", test_tile_reduce_axis_tier1, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_axis_tier2", test_tile_reduce_axis_tier2, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_axis_tier3", test_tile_reduce_axis_tier3, devices=devices)
 add_function_test(TestTileReduce, "test_tile_ones", test_tile_ones, devices=devices)
 add_function_test(TestTileReduce, "test_tile_arange", test_tile_arange, devices=devices)
 add_function_test(TestTileReduce, "test_tile_untile_scalar", test_tile_untile_scalar, devices=devices)

warp/tests/unittest_suites.py

@@ -113,17 +113,12 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.interop.test_dlpack import TestDLPack
     from warp.tests.interop.test_jax import TestJax
     from warp.tests.interop.test_torch import TestTorch
-    from warp.tests.sim.test_cloth import TestCloth
-    from warp.tests.sim.test_collision import TestCollision
-    from warp.tests.sim.test_coloring import TestColoring
-    from warp.tests.sim.test_model import TestModel
-    from warp.tests.sim.test_sim_grad import TestSimGradients
-    from warp.tests.sim.test_sim_kinematics import TestSimKinematics
     from warp.tests.test_adam import TestAdam
     from warp.tests.test_arithmetic import TestArithmetic
     from warp.tests.test_array import TestArray
     from warp.tests.test_array_reduce import TestArrayReduce
     from warp.tests.test_atomic import TestAtomic
+    from warp.tests.test_atomic_bitwise import TestAtomicBitwise
     from warp.tests.test_atomic_cas import TestAtomicCAS
     from warp.tests.test_bool import TestBool
     from warp.tests.test_builtins_resolution import TestBuiltinsResolution
@@ -142,7 +137,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestFemDiffusionExamples,
         TestFemExamples,
         TestOptimExamples,
-        TestSimExamples,
     )
     from warp.tests.test_fabricarray import TestFabricArray
     from warp.tests.test_fast_math import TestFastMath
@@ -198,7 +192,9 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_vec_lite import TestVecLite
     from warp.tests.test_vec_scalar_ops import TestVecScalarOps
     from warp.tests.test_verify_fp import TestVerifyFP
+    from warp.tests.test_version import TestVersion
     from warp.tests.tile.test_tile import TestTile
+    from warp.tests.tile.test_tile_atomic_bitwise import TestTileAtomicBitwise
     from warp.tests.tile.test_tile_cholesky import TestTileCholesky
     from warp.tests.tile.test_tile_load import TestTileLoad
     from warp.tests.tile.test_tile_mathdx import TestTileMathDx
@@ -215,16 +211,14 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestArrayReduce,
         TestAsync,
         TestAtomic,
+        TestAtomicBitwise,
         TestAtomicCAS,
         TestBool,
         TestBuiltinsResolution,
         TestBvh,
         TestClosestPointEdgeEdgeMethods,
-        TestCloth,
         TestCodeGen,
         TestCodeGenInstancing,
-        TestCollision,
-        TestColoring,
         TestConditional,
         TestConstants,
         TestContext,
@@ -237,7 +231,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestFemDiffusionExamples,
         TestFemExamples,
         TestOptimExamples,
-        TestSimExamples,
         TestFabricArray,
         TestFastMath,
         TestFem,
@@ -272,7 +265,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
-        TestModel,
         TestModuleHashing,
         TestModuleLite,
         TestMultiGPU,
@@ -289,8 +281,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestRounding,
         TestRunlengthEncode,
         TestScalarOps,
-        TestSimGradients,
-        TestSimKinematics,
         TestSmoothstep,
         TestSnippets,
         TestSparse,
@@ -301,6 +291,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestStruct,
         TestTape,
         TestTile,
+        TestTileAtomicBitwise,
         TestTileCholesky,
         TestTileLoad,
         TestTileMathDx,
@@ -319,6 +310,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestVecLite,
         TestVecScalarOps,
         TestVerifyFP,
+        TestVersion,
         TestVolume,
         TestVolumeWrite,
     ]

warp/tests/unittest_utils.py

@@ -376,24 +376,25 @@ def write_junit_results(
     )
 
     for test_data in test_records:
-        test = test_data[0]
-        test_duration = test_data[1]
-        test_status = test_data[2]
+        test_classname = test_data[0]
+        test_methodname = test_data[1]
+        test_duration = test_data[2]
+        test_status = test_data[3]
 
         test_case = ET.SubElement(
-            root, "testcase", classname=test.__class__.__name__, name=test._testMethodName, time=f"{test_duration:.3f}"
+            root, "testcase", classname=test_classname, name=test_methodname, time=f"{test_duration:.3f}"
         )
 
         if test_status == "FAIL":
-            failure = ET.SubElement(test_case, "failure", message=str(test_data[3]))
-            failure.text = str(test_data[4])  # Stacktrace
+            failure = ET.SubElement(test_case, "failure", message=str(test_data[4]))
+            failure.text = str(test_data[5])  # Stacktrace
         elif test_status == "ERROR":
             error = ET.SubElement(test_case, "error")
-            error.text = str(test_data[4])  # Stacktrace
+            error.text = str(test_data[5])  # Stacktrace
         elif test_status == "SKIP":
             skip = ET.SubElement(test_case, "skipped")
             # Set the skip reason
-            skip.set("message", str(test_data[3]))
+            skip.set("message", str(test_data[4]))
 
     tree = ET.ElementTree(root)
 
@@ -425,7 +426,7 @@ class ParallelJunitTestResult(unittest.TextTestResult):
 
     def _record_test(self, test, code, message=None, details=None):
         duration = round((time.perf_counter_ns() - self.start_time) * 1e-9, 3)  # [s]
-        self.test_record.append((test, duration, code, message, details))
+        self.test_record.append((test.__class__.__name__, test._testMethodName, duration, code, message, details))
 
     def addSuccess(self, test):
         super(unittest.TextTestResult, self).addSuccess(test)