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

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)

warp/tests/walkthrough_debug.py

@@ -68,7 +68,9 @@ wp.init()
 wp.config.mode = "debug"
 
 # Make sure Warp was built with `build_lib.py --mode=debug`
-assert wp.context.runtime.core.wp_is_debug_enabled(), "Warp must be built in debug mode to enable debugging kernels"
+assert wp._src.context.runtime.core.wp_is_debug_enabled(), (
+    "Warp must be built in debug mode to enable debugging kernels"
+)
 
 
 @wp.kernel