warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/tests/tile/test_tile_shared_memory.py

@@ -28,7 +28,7 @@ def test_tile_shared_mem_size(test, device):
 
     BLOCK_DIM = 256
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(out: wp.array2d(dtype=float)):
         a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
         b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
@@ -64,7 +64,7 @@ def test_tile_shared_mem_large(test, device):
     BLOCK_DIM = 256
 
     # we disable backward kernel gen since 128k is not supported on most architectures
-    @wp.kernel(enable_backward=False)
+    @wp.kernel(enable_backward=False, module="unique")
     def compute(out: wp.array2d(dtype=float)):
         a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
         b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
@@ -100,7 +100,7 @@ def test_tile_shared_mem_graph(test, device):
 
     BLOCK_DIM = 256
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(out: wp.array2d(dtype=float)):
         a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
         b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
@@ -110,7 +110,7 @@ def test_tile_shared_mem_graph(test, device):
 
     out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
 
-    wp.load_module(device=device)
+    compute.module.load(device)
 
     wp.capture_begin(device, force_module_load=False)
     wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
@@ -157,7 +157,7 @@ def test_tile_shared_mem_func(test, device):
 
         return a + b
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(out: wp.array2d(dtype=float)):
         s = add_tile_small()
         b = add_tile_big()
@@ -197,7 +197,7 @@ def test_tile_shared_non_aligned(test, device):
         b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 3.0
         return a + b
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(out: wp.array2d(dtype=float)):
         # This test the logic in the stack allocator, which should increment and
         # decrement the stack pointer each time foo() is called
@@ -224,6 +224,121 @@ def test_tile_shared_non_aligned(test, device):
     assert hooks.backward_smem_bytes == expected_required_shared * 2
 
 
+def test_tile_shared_vec_accumulation(test, device):
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(indices: wp.array(dtype=int), vecs: wp.array(dtype=wp.vec3), output: wp.array2d(dtype=float)):
+        i, j = wp.tid()
+
+        idx_tile = wp.tile_load(indices, shape=BLOCK_DIM, offset=i * BLOCK_DIM)
+        idx = idx_tile[j]
+
+        s = wp.tile_zeros(shape=(1, 3), dtype=float)
+
+        s[0, 0] += vecs[idx].x
+        s[0, 1] += vecs[idx].y
+        s[0, 2] += vecs[idx].z
+
+        wp.tile_store(output, s, offset=(i, 0))
+
+    N = BLOCK_DIM * 3
+
+    basis_vecs = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32)
+    vecs = wp.array(basis_vecs, dtype=wp.vec3, requires_grad=True, device=device)
+
+    rng = np.random.default_rng(42)
+    indices_np = rng.integers(0, 3, size=N)
+
+    indices = wp.array(indices_np, dtype=int, requires_grad=True, device=device)
+
+    output = wp.zeros(shape=(3, 3), dtype=float, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch_tiled(compute, dim=3, inputs=[indices, vecs, output], block_dim=BLOCK_DIM, device=device)
+
+    output.grad = wp.ones_like(output)
+
+    tape.backward()
+
+    n0 = np.count_nonzero(indices_np == 0)
+    n1 = np.count_nonzero(indices_np == 1)
+    n2 = np.count_nonzero(indices_np == 2)
+    true_grads = np.array([[n0, n0, n0], [n1, n1, n1], [n2, n2, n2]])
+
+    indices_np = indices_np.reshape((3, BLOCK_DIM))
+
+    def compute_row(idx):
+        n0 = np.count_nonzero(indices_np[idx, :] == 0)
+        n1 = np.count_nonzero(indices_np[idx, :] == 1)
+        n2 = np.count_nonzero(indices_np[idx, :] == 2)
+        return np.array([1, 0, 0]) * n0 + np.array([0, 1, 0]) * n1 + np.array([0, 0, 1]) * n2
+
+    row_0 = compute_row(0)
+    row_1 = compute_row(1)
+    row_2 = compute_row(2)
+
+    true_vecs = np.stack([row_0, row_1, row_2])
+
+    assert_np_equal(output.numpy(), true_vecs)
+    assert_np_equal(vecs.grad.numpy(), true_grads)
+
+
+def test_tile_shared_simple_reduction_add(test, device):
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+        i, j = wp.tid()
+
+        t = wp.tile_load(x, shape=BLOCK_DIM, offset=BLOCK_DIM * i)
+
+        k = BLOCK_DIM // 2
+        while k > 0:
+            if j < k:
+                t[j] += t[j + k]
+            k //= 2
+
+        wp.tile_store(y, wp.tile_view(t, offset=(0,), shape=(1,)), i)
+
+    N = BLOCK_DIM * 4
+    x_np = np.arange(N, dtype=np.float32)
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(4, dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=4, inputs=[x], outputs=[y], block_dim=BLOCK_DIM, device=device)
+
+    assert_np_equal(np.sum(y.numpy()), np.sum(x_np))
+
+
+def test_tile_shared_simple_reduction_sub(test, device):
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+        i, j = wp.tid()
+
+        t = wp.tile_load(x, shape=BLOCK_DIM, offset=BLOCK_DIM * i)
+
+        k = BLOCK_DIM // 2
+        while k > 0:
+            if j < k:
+                t[j] -= t[j + k]
+            k //= 2
+
+        wp.tile_store(y, wp.tile_view(t, offset=(0,), shape=(1,)), i)
+
+    N = BLOCK_DIM * 4
+    x_np = np.arange(N, dtype=np.float32)
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(4, dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=4, inputs=[x], outputs=[y], block_dim=BLOCK_DIM, device=device)
+
+    assert_np_equal(np.sum(y.numpy()), 0.0)
+
+
 devices = get_cuda_test_devices()
 
 
@@ -240,7 +355,21 @@ add_function_test(
 add_function_test(TestTileSharedMemory, "test_tile_shared_mem_graph", test_tile_shared_mem_graph, devices=devices)
 add_function_test(TestTileSharedMemory, "test_tile_shared_mem_func", test_tile_shared_mem_func, devices=devices)
 add_function_test(TestTileSharedMemory, "test_tile_shared_non_aligned", test_tile_shared_non_aligned, devices=devices)
-
+add_function_test(
+    TestTileSharedMemory, "test_tile_shared_vec_accumulation", test_tile_shared_vec_accumulation, devices=devices
+)
+add_function_test(
+    TestTileSharedMemory,
+    "test_tile_shared_simple_reduction_add",
+    test_tile_shared_simple_reduction_add,
+    devices=devices,
+)
+add_function_test(
+    TestTileSharedMemory,
+    "test_tile_shared_simple_reduction_sub",
+    test_tile_shared_simple_reduction_sub,
+    devices=devices,
+)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/tile/test_tile_sort.py

@@ -0,0 +1,121 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+def create_sort_kernel(KEY_TYPE, MAX_SORT_LENGTH):
+    @wp.kernel
+    def tile_sort_kernel(
+        input_keys: wp.array(dtype=KEY_TYPE),
+        input_values: wp.array(dtype=wp.int32),
+        output_keys: wp.array(dtype=KEY_TYPE),
+        output_values: wp.array(dtype=wp.int32),
+    ):
+        # Load input into shared memory
+        keys = wp.tile_load(input_keys, shape=MAX_SORT_LENGTH, storage="shared")
+        values = wp.tile_load(input_values, shape=MAX_SORT_LENGTH, storage="shared")
+
+        # Perform in-place sorting
+        wp.tile_sort(keys, values)
+
+        # Store sorted shared memory into output arrays
+        wp.tile_store(output_keys, keys)
+        wp.tile_store(output_values, values)
+
+    return tile_sort_kernel
+
+
+def test_tile_sort(test, device):
+    # Forward-declare kernels for more efficient compilation
+    kernels = {}
+    for dtype in [int, float]:
+        for i in range(0, 11):
+            length = 2**i + 1
+            kernels[(dtype, length)] = create_sort_kernel(dtype, length)
+
+    for (dtype, length), kernel in kernels.items():
+        for j in range(5, 10):
+            TILE_DIM = 2**j
+
+            rng = np.random.default_rng(42)  # Create a random generator instance
+
+            if dtype == int:
+                np_keys = rng.choice(1000000000, size=length, replace=False)
+            else:  # dtype == float
+                np_keys = rng.uniform(0, 1000000000, size=length).astype(dtype)
+
+            np_values = np.arange(length)
+
+            # Generate random keys and iota indexer
+            input_keys = wp.array(np_keys, dtype=dtype, device=device)
+            input_values = wp.array(np_values, dtype=int, device=device)
+            output_keys = wp.zeros_like(input_keys, device=device)
+            output_values = wp.zeros_like(input_values, device=device)
+
+            # Execute sorting kernel
+            wp.launch_tiled(
+                kernel,
+                dim=1,
+                inputs=[input_keys, input_values, output_keys, output_values],
+                block_dim=TILE_DIM,
+                device=device,
+            )
+            wp.synchronize()
+
+            # Sort using NumPy for validation
+            sorted_indices = np.argsort(np_keys)
+            np_sorted_keys = np_keys[sorted_indices]
+            np_sorted_values = np_values[sorted_indices]
+
+            if dtype == int:
+                keys_match = np.array_equal(output_keys.numpy(), np_sorted_keys)
+            else:  # dtype == float
+                keys_match = np.allclose(output_keys.numpy(), np_sorted_keys, atol=1e-6)  # Use tolerance for floats
+
+            values_match = np.array_equal(output_values.numpy(), np_sorted_values)
+
+            if not keys_match or not values_match:
+                print(f"Test failed for dtype={dtype}, TILE_DIM={TILE_DIM}, length={length}")
+                print("")
+                print(output_keys.numpy())
+                print(np_sorted_keys)
+                print("")
+                print(output_values.numpy())
+                print(np_sorted_values)
+                print("")
+
+            # Validate results
+            test.assertTrue(keys_match, f"Key sorting mismatch for dtype={dtype}!")
+            test.assertTrue(values_match, f"Value sorting mismatch for dtype={dtype}!")
+
+
+devices = get_test_devices()
+
+
+class TestTileSort(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTileSort, "test_tile_sort", test_tile_sort, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/unittest_suites.py

@@ -113,17 +113,18 @@ 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.sim.test_vbd import TestVbd
     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_cas import TestAtomicCAS
     from warp.tests.test_bool import TestBool
     from warp.tests.test_builtins_resolution import TestBuiltinsResolution
     from warp.tests.test_closest_point_edge_edge import TestClosestPointEdgeEdgeMethods
@@ -166,7 +167,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_mat_lite import TestMatLite
     from warp.tests.test_mat_scalar_ops import TestMatScalarOps
     from warp.tests.test_math import TestMath
-    from warp.tests.test_mlp import TestMLP
     from warp.tests.test_module_hashing import TestModuleHashing
     from warp.tests.test_modules_lite import TestModuleLite
     from warp.tests.test_noise import TestNoise
@@ -193,13 +193,18 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_types import TestTypes
     from warp.tests.test_utils import TestUtils
     from warp.tests.test_vec import TestVec
+    from warp.tests.test_vec_constructors import TestVecConstructors
     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.tile.test_tile import TestTile
+    from warp.tests.tile.test_tile_load import TestTileLoad
     from warp.tests.tile.test_tile_mathdx import TestTileMathDx
+    from warp.tests.tile.test_tile_matmul import TestTileMatmul
     from warp.tests.tile.test_tile_reduce import TestTileReduce
     from warp.tests.tile.test_tile_shared_memory import TestTileSharedMemory
+    from warp.tests.tile.test_tile_sort import TestTileSort
+    from warp.tests.tile.test_tile_view import TestTileView
 
     test_classes = [
         TestAdam,
@@ -208,10 +213,12 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestArrayReduce,
         TestAsync,
         TestAtomic,
+        TestAtomicCAS,
         TestBool,
         TestBuiltinsResolution,
         TestBvh,
         TestClosestPointEdgeEdgeMethods,
+        TestCloth,
         TestCodeGen,
         TestCodeGenInstancing,
         TestCollision,
@@ -262,7 +269,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
-        TestMLP,
         TestModel,
         TestModuleHashing,
         TestModuleLite,
@@ -292,16 +298,20 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestStruct,
         TestTape,
         TestTile,
+        TestTileLoad,
         TestTileMathDx,
+        TestTileMatmul,
         TestTileReduce,
         TestTileSharedMemory,
+        TestTileSort,
+        TestTileView,
         TestTorch,
         TestTransientModule,
         TestTriangleClosestPoint,
         TestTypes,
         TestUtils,
-        TestVbd,
         TestVec,
+        TestVecConstructors,
         TestVecLite,
         TestVecScalarOps,
         TestVerifyFP,
@@ -350,7 +360,6 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     from warp.tests.test_lvalue import TestLValue
     from warp.tests.test_mat_lite import TestMatLite
     from warp.tests.test_math import TestMath
-    from warp.tests.test_mlp import TestMLP
     from warp.tests.test_module_hashing import TestModuleHashing
     from warp.tests.test_modules_lite import TestModuleLite
     from warp.tests.test_noise import TestNoise
@@ -397,7 +406,6 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
-        TestMLP,
         TestModuleHashing,
         TestModuleLite,
         TestNoise,