warp-lang 1.7.2__py3-none-win_amd64.whl → 1.8.0__py3-none-win_amd64.whl

warp/tests/tile/test_tile_reduce.py

@@ -106,6 +106,39 @@ def test_tile_reduce_min(test, device):
         test.assertAlmostEqual(min_wp[i], min_np, places=4)
 
 
+@wp.kernel
+def tile_argmin_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=int)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input[i], shape=TILE_DIM)
+    m = wp.tile_argmin(a)
+
+    wp.tile_store(output, m, offset=i)
+
+
+def test_tile_reduce_argmin(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    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:
+        wp.launch_tiled(
+            tile_argmin_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    argmin_wp = output_wp.numpy()
+    for i in range(batch_count):
+        argmin_np = np.argmin(input[i])
+        test.assertAlmostEqual(argmin_wp[i], argmin_np, places=4)
+
+
 @wp.kernel
 def tile_max_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
     # output tile index
@@ -139,6 +172,39 @@ def test_tile_reduce_max(test, device):
         test.assertAlmostEqual(max_wp[i], max_np, places=4)
 
 
+@wp.kernel
+def tile_argmax_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=int)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input[i], shape=TILE_DIM)
+    m = wp.tile_argmax(a)
+
+    wp.tile_store(output, m, offset=i)
+
+
+def test_tile_reduce_argmax(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    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:
+        wp.launch_tiled(
+            tile_argmax_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    argmax_wp = output_wp.numpy()
+    for i in range(batch_count):
+        argmax_np = np.argmax(input[i])
+        test.assertAlmostEqual(argmax_wp[i], argmax_np, places=4)
+
+
 @wp.kernel
 def tile_reduce_custom_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
     # output tile index
@@ -176,6 +242,79 @@ def test_tile_reduce_custom(test, device):
         test.assertAlmostEqual(prod_wp[i], prod_np, places=4)
 
 
+def create_tile_scan_inclusive_kernel(tile_dim: int):
+    @wp.kernel
+    def tile_scan_inclusive_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+        i = wp.tid()
+        t = wp.tile_load(input[i], shape=tile_dim)
+        t = wp.tile_scan_inclusive(t)
+        wp.tile_store(output[i], t)
+
+    return tile_scan_inclusive_kernel
+
+
+def test_tile_scan_inclusive(test, device):
+    batch_count = 56
+    N = 1234
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    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:
+        wp.launch_tiled(
+            create_tile_scan_inclusive_kernel(N),
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    scan_wp = output_wp.numpy()
+    for i in range(batch_count):
+        scan_np = np.cumsum(input[i])
+        np.testing.assert_allclose(scan_wp[i], scan_np, rtol=1e-5, atol=1e-6)
+
+
+def create_tile_scan_exclusive_kernel(tile_dim: int):
+    @wp.kernel
+    def tile_scan_exclusive_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+        i = wp.tid()
+        t = wp.tile_load(input[i], shape=tile_dim)
+        t = wp.tile_scan_exclusive(t)
+        wp.tile_store(output[i], t)
+
+    return tile_scan_exclusive_kernel
+
+
+def test_tile_scan_exclusive(test, device):
+    batch_count = 56
+    N = 1234
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    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:
+        wp.launch_tiled(
+            create_tile_scan_exclusive_kernel(N),
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    scan_wp = output_wp.numpy()
+    for i in range(batch_count):
+        scan_np = np.zeros(N, dtype=np.float32)
+        scan_np[1:] = np.cumsum(input[i][:-1])
+        np.testing.assert_allclose(scan_wp[i], scan_np, rtol=1e-5, atol=1e-6)
+
+
 @wp.struct
 class KeyValue:
     key: wp.int32
@@ -423,7 +562,75 @@ def test_tile_arange(test, device):
     assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
 
 
+@wp.kernel
+def tile_strided_loop_kernel(arr: wp.array(dtype=float), max_val: wp.array(dtype=float)):
+    tid, lane = wp.tid()
+
+    num_threads = wp.block_dim()
+
+    thread_max = wp.float32(-wp.inf)
+
+    length = arr.shape[0]
+    upper = ((length + num_threads - 1) // num_threads) * num_threads
+    for el_id in range(lane, upper, num_threads):
+        if el_id < length:
+            val = arr[el_id]
+        else:
+            val = wp.float32(-wp.inf)
+
+        t = wp.tile(val)
+        local_max = wp.tile_max(t)
+
+        thread_max = wp.max(thread_max, local_max[0])
+
+    if lane == 0:
+        max_val[0] = thread_max
+
+
+def test_tile_strided_loop(test, device):
+    N = 5  # Length of array
+
+    rng = np.random.default_rng(42)
+    input = rng.random(N, dtype=np.float32)
+
+    input_wp = wp.array(input, device=device)
+    output_wp = wp.zeros(1, dtype=wp.float32, device=device)
+
+    wp.launch_tiled(
+        tile_strided_loop_kernel,
+        dim=[1],
+        inputs=[input_wp, output_wp],
+        device=device,
+        block_dim=128,
+    )
+
+    max_wp = output_wp.numpy()
+    max_np = np.max(input)
+    test.assertAlmostEqual(max_wp[0], max_np, places=4)
+
+
+@wp.kernel
+def test_tile_reduce_matrix_kernel(y: wp.array(dtype=wp.mat33)):
+    i = wp.tid()
+    I = wp.identity(3, dtype=wp.float32)
+    m = wp.float32(i) * I
+
+    t = wp.tile(m, preserve_type=True)
+    sum = wp.tile_reduce(wp.add, t)
+
+    wp.tile_store(y, sum)
+
+
+def test_tile_reduce_matrix(test, device):
+    y = wp.zeros(shape=1, dtype=wp.mat33, device=device)
+
+    wp.launch(test_tile_reduce_matrix_kernel, dim=TILE_DIM, inputs=[], outputs=[y], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(y.numpy().squeeze(), np.eye(3, dtype=np.float32) * 2016.0)
+
+
 devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
 
 
 class TestTileReduce(unittest.TestCase):
@@ -433,6 +640,8 @@ class TestTileReduce(unittest.TestCase):
 add_function_test(TestTileReduce, "test_tile_reduce_sum", test_tile_reduce_sum, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_min", test_tile_reduce_min, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_max", test_tile_reduce_max, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_argmin", test_tile_reduce_argmin, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_argmax", test_tile_reduce_argmax, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_custom", test_tile_reduce_custom, devices=devices)
 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_sum, devices=devices)
@@ -441,6 +650,10 @@ add_function_test(TestTileReduce, "test_tile_ones", test_tile_ones, devices=devi
 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)
 add_function_test(TestTileReduce, "test_tile_untile_vector", test_tile_untile_vector, devices=devices)
+add_function_test(TestTileReduce, "test_tile_strided_loop", test_tile_strided_loop, devices=devices)
+add_function_test(TestTileReduce, "test_tile_scan_inclusive", test_tile_scan_inclusive, devices=devices)
+add_function_test(TestTileReduce, "test_tile_scan_exclusive", test_tile_scan_exclusive, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_matrix", test_tile_reduce_matrix, devices=cuda_devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/tile/test_tile_shared_memory.py

@@ -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 = 64
+
+    @wp.kernel
+    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 = 64
+
+    @wp.kernel
+    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 = 64
+
+    @wp.kernel
+    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,117 @@
+# 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):
+    for dtype in [int, float]:  # Loop over int and float keys
+        for j in range(5, 10):
+            TILE_DIM = 2**j
+            for i in range(0, 11):  # Start from 1 to avoid zero-length cases
+                length = 2**i + 1
+
+                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)
+
+                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
+                kernel = create_sort_kernel(dtype, length)
+                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
+                assert keys_match, f"Key sorting mismatch for dtype={dtype}!"
+                assert 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
@@ -208,10 +208,12 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestArrayReduce,
         TestAsync,
         TestAtomic,
+        TestAtomicCAS,
         TestBool,
         TestBuiltinsResolution,
         TestBvh,
         TestClosestPointEdgeEdgeMethods,
+        TestCloth,
         TestCodeGen,
         TestCodeGenInstancing,
         TestCollision,
@@ -262,7 +264,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
-        TestMLP,
         TestModel,
         TestModuleHashing,
         TestModuleLite,
@@ -300,7 +301,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestTriangleClosestPoint,
         TestTypes,
         TestUtils,
-        TestVbd,
         TestVec,
         TestVecLite,
         TestVecScalarOps,
@@ -350,7 +350,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 +396,6 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
-        TestMLP,
         TestModuleHashing,
         TestModuleLite,
         TestNoise,