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

warp/tests/tile/test_tile_matmul.py

@@ -0,0 +1,179 @@
+# 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
+from typing import Any
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+TILE_M = wp.constant(8)
+TILE_N = wp.constant(4)
+TILE_K = wp.constant(8)
+
+# num threads per-tile
+TILE_DIM = 64
+
+
+@wp.kernel
+def tile_grouped_gemm(A: wp.array3d(dtype=float), B: wp.array3d(dtype=float), C: wp.array3d(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(A[i], shape=(TILE_M, TILE_K))
+    b = wp.tile_load(B[i], shape=(TILE_K, TILE_N))
+
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32)
+
+    wp.tile_matmul(a, b, sum)
+
+    wp.tile_store(C[i], sum)
+
+
+def test_tile_grouped_gemm(test, device):
+    batch_count = 56
+
+    M = TILE_M
+    N = TILE_N
+    K = TILE_K
+
+    rng = np.random.default_rng(42)
+    A = rng.random((batch_count, M, K), dtype=np.float32)
+    B = rng.random((batch_count, K, N), dtype=np.float32)
+    C = A @ B
+
+    A_wp = wp.array(A, requires_grad=True, device=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:
+        wp.launch_tiled(
+            tile_grouped_gemm, dim=[batch_count], inputs=[A_wp, B_wp, C_wp], block_dim=TILE_DIM, device=device
+        )
+
+    # TODO: 32 mismatched elements
+    assert_np_equal(C_wp.numpy(), C, 1e-6)
+
+
+@wp.kernel
+def tile_gemm(A: wp.array2d(dtype=Any), B: wp.array2d(dtype=Any), C: wp.array2d(dtype=Any)):
+    # output tile index
+    i, j = wp.tid()
+
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=A.dtype)
+
+    M = A.shape[0]
+    N = B.shape[1]
+    K = A.shape[1]
+
+    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))
+
+        # sum += a*b
+        wp.tile_matmul(a, b, sum)
+
+    wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
+
+
+wp.overload(
+    tile_gemm, {"A": wp.array2d(dtype=wp.float16), "B": wp.array2d(dtype=wp.float16), "C": wp.array2d(dtype=wp.float16)}
+)
+wp.overload(
+    tile_gemm, {"A": wp.array2d(dtype=wp.float32), "B": wp.array2d(dtype=wp.float32), "C": wp.array2d(dtype=wp.float32)}
+)
+wp.overload(
+    tile_gemm, {"A": wp.array2d(dtype=wp.float64), "B": wp.array2d(dtype=wp.float64), "C": wp.array2d(dtype=wp.float64)}
+)
+
+
+def test_tile_gemm(dtype):
+    def test(test, device):
+        M = TILE_M * 7
+        K = TILE_K * 6
+        N = TILE_N * 5
+
+        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)
+
+        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-1)
+
+        adj_C = np.ones_like(C)
+
+        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-1)
+        assert_np_equal(B_wp.grad.numpy(), A.T @ adj_C, 1.0e-1)
+
+    return test
+
+
+@wp.kernel
+def test_tile_transpose_matmul_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+    x = wp.tile_load(input, shape=(TILE_M, TILE_N))
+    y = wp.tile_transpose(x)
+
+    z = wp.tile_zeros(dtype=float, shape=(TILE_N, TILE_N))
+    wp.tile_matmul(y, x, z)
+
+    wp.tile_store(output, z)
+
+
+def test_tile_transpose_matmul(test, device):
+    rng = np.random.default_rng(42)
+    input = wp.array(rng.random((TILE_M, TILE_N), dtype=np.float32), device=device)
+    output = wp.zeros((TILE_N, TILE_N), dtype=float, device=device)
+
+    wp.launch_tiled(
+        test_tile_transpose_matmul_kernel, dim=[1], inputs=[input, output], block_dim=TILE_DIM, device=device
+    )
+
+    assert_np_equal(output.numpy(), input.numpy().T @ input.numpy(), 1e-6)
+
+
+class TestTileMatmul(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestTileMatmul, "test_tile_gemm_fp16", test_tile_gemm(wp.float16), devices=devices)
+add_function_test(TestTileMatmul, "test_tile_gemm_fp32", test_tile_gemm(wp.float32), devices=devices)
+add_function_test(TestTileMatmul, "test_tile_gemm_fp64", test_tile_gemm(wp.float64), devices=devices)
+add_function_test(TestTileMatmul, "test_tile_grouped_gemm", test_tile_grouped_gemm, devices=devices)
+add_function_test(TestTileMatmul, "test_tile_transpose_matmul", test_tile_transpose_matmul, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

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.is_mathdx_enabled():
+    if device.is_cuda and not wp.context.runtime.core.wp_is_mathdx_enabled():
         test.skipTest("Skipping test on CUDA device without MathDx (tolerance)")
 
     NUM_FREQ = wp.constant(8)

warp/tests/tile/test_tile_reduce.py

@@ -73,6 +73,46 @@ def test_tile_reduce_sum(test, device):
     assert_np_equal(input_wp.grad.numpy(), np.ones_like(input) * 0.5, tol=1.0e-4)
 
 
+@wp.kernel
+def tile_sum_to_shared_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
+    i, lane = wp.tid()
+
+    a = wp.tile_load(input[i], shape=TILE_DIM)
+    s = wp.tile_sum(a)
+    v = s[0]  # force shared storage for s
+    wp.tile_store(output, s * 0.5, offset=i)
+
+
+def test_tile_sum_to_shared(test, device):
+    batch_count = 1
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, TILE_DIM), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device, dtype=float)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device, dtype=float)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_sum_to_shared_kernel,
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    sum_wp = output_wp.numpy()
+    for i in range(batch_count):
+        sum_np = np.sum(input[i], axis=0) * 0.5
+        assert_np_equal(sum_wp[i], sum_np, tol=0.0001)
+
+    output_wp.grad.fill_(1.0)
+
+    tape.backward()
+
+    assert_np_equal(input_wp.grad.numpy(), np.ones_like(input) * 0.5, tol=1.0e-4)
+
+
 @wp.kernel
 def tile_min_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
     # output tile index
@@ -84,6 +124,13 @@ def tile_min_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float
     wp.tile_store(output, m, offset=i)
 
 
+@wp.kernel
+def tile_min_kernel_edge_case(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    t = wp.tile_load(x, shape=(3, 3))
+    min = wp.tile_min(t)
+    wp.tile_store(y, min)
+
+
 def test_tile_reduce_min(test, device):
     batch_count = 56
 
@@ -105,6 +152,14 @@ def test_tile_reduce_min(test, device):
         min_np = np.min(input[i])
         test.assertAlmostEqual(min_wp[i], min_np, places=4)
 
+    # test edge case: tile is multiple warps in size but at least one is empty
+    x = wp.array(np.array([[2.0, 2.0, 3.0], [4.0, 1.0, 6.0], [7.0, 3.0, 9.0]]), dtype=float, device=device)
+    y = wp.zeros(1, dtype=float, device=device)
+
+    wp.launch_tiled(tile_min_kernel_edge_case, dim=1, inputs=[x, y], block_dim=64, device=device)
+
+    assert_np_equal(y.numpy(), np.array([1.0]))
+
 
 @wp.kernel
 def tile_argmin_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=int)):
@@ -117,6 +172,13 @@ def tile_argmin_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=in
     wp.tile_store(output, m, offset=i)
 
 
+@wp.kernel
+def tile_argmin_kernel_edge_case(x: wp.array2d(dtype=float), y: wp.array(dtype=int)):
+    t = wp.tile_load(x, shape=(3, 3))
+    min = wp.tile_argmin(t)
+    wp.tile_store(y, min)
+
+
 def test_tile_reduce_argmin(test, device):
     batch_count = 56
 
@@ -138,6 +200,14 @@ def test_tile_reduce_argmin(test, device):
         argmin_np = np.argmin(input[i])
         test.assertAlmostEqual(argmin_wp[i], argmin_np, places=4)
 
+    # test edge case: tile is multiple warps in size but at least one is empty
+    x = wp.array(np.array([[2.0, 2.0, 3.0], [4.0, 1.0, 6.0], [7.0, 3.0, 9.0]]), dtype=float, device=device)
+    y = wp.zeros(1, dtype=int, device=device)
+
+    wp.launch_tiled(tile_argmin_kernel_edge_case, dim=1, inputs=[x, y], block_dim=64, device=device)
+
+    assert_np_equal(y.numpy(), np.array([4]))
+
 
 @wp.kernel
 def tile_max_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
@@ -243,7 +313,7 @@ def test_tile_reduce_custom(test, device):
 
 
 def create_tile_scan_inclusive_kernel(tile_dim: int):
-    @wp.kernel
+    @wp.kernel(module="unique")
     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)
@@ -279,7 +349,7 @@ def test_tile_scan_inclusive(test, device):
 
 
 def create_tile_scan_exclusive_kernel(tile_dim: int):
-    @wp.kernel
+    @wp.kernel(module="unique")
     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)
@@ -398,7 +468,7 @@ def test_tile_reduce_grouped_sum(test, device):
 
     with wp.Tape() as tape:
         wp.launch_tiled(
-            tile_sum_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+            tile_grouped_sum_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
         )
 
     sum_wp = output_wp.numpy()
@@ -498,17 +568,17 @@ def test_untile_vector_kernel(input: wp.array(dtype=wp.vec3), output: wp.array(d
 
 
 def test_tile_untile_vector(test, device):
-    input = wp.full(16, wp.vec3(1.0, 2.0, 3.0), requires_grad=True, device=device)
+    input = wp.full(TILE_DIM, wp.vec3(1.0, 2.0, 3.0), requires_grad=True, device=device)
     output = wp.zeros_like(input, device=device)
 
     with wp.Tape() as tape:
-        wp.launch(test_untile_vector_kernel, dim=16, inputs=[input, output], block_dim=16, device=device)
+        wp.launch(test_untile_vector_kernel, dim=TILE_DIM, inputs=[input, output], block_dim=TILE_DIM, device=device)
 
     output.grad = wp.ones_like(output, device=device)
     tape.backward()
 
     assert_np_equal(output.numpy(), input.numpy())
-    assert_np_equal(input.grad.numpy(), np.ones((16, 3)))
+    assert_np_equal(input.grad.numpy(), np.ones((TILE_DIM, 3)))
 
 
 @wp.kernel
@@ -562,7 +632,7 @@ def test_tile_arange(test, device):
     assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
 
 
-@wp.kernel
+@wp.kernel(module="unique")
 def tile_strided_loop_kernel(arr: wp.array(dtype=float), max_val: wp.array(dtype=float)):
     tid, lane = wp.tid()
 
@@ -618,7 +688,7 @@ def test_tile_reduce_matrix_kernel(y: wp.array(dtype=wp.mat33)):
     t = wp.tile(m, preserve_type=True)
     sum = wp.tile_reduce(wp.add, t)
 
-    wp.tile_store(y, sum)
+    wp.tile_atomic_add(y, sum)
 
 
 def test_tile_reduce_matrix(test, device):
@@ -629,8 +699,25 @@ def test_tile_reduce_matrix(test, device):
     assert_np_equal(y.numpy().squeeze(), np.eye(3, dtype=np.float32) * 2016.0)
 
 
+@wp.kernel
+def test_tile_reduce_vector_kernel(out: wp.array(dtype=wp.vec3)):
+    v = wp.vec3f(1.0)
+    v_tile = wp.tile(v, preserve_type=True)
+
+    sum = wp.tile_reduce(wp.add, v_tile)
+
+    wp.tile_atomic_add(out, sum)
+
+
+def test_tile_reduce_vector(test, device):
+    out = wp.zeros(1, dtype=wp.vec3, device=device)
+
+    wp.launch(kernel=test_tile_reduce_vector_kernel, dim=8, inputs=[], outputs=[out], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(out.numpy(), np.array([[8.0, 8.0, 8.0]]))
+
+
 devices = get_test_devices()
-cuda_devices = get_cuda_test_devices()
 
 
 class TestTileReduce(unittest.TestCase):
@@ -638,13 +725,14 @@ class TestTileReduce(unittest.TestCase):
 
 
 add_function_test(TestTileReduce, "test_tile_reduce_sum", test_tile_reduce_sum, devices=devices)
+add_function_test(TestTileReduce, "test_tile_sum_to_shared", test_tile_sum_to_shared, 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)
+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_ones", test_tile_ones, devices=devices)
 add_function_test(TestTileReduce, "test_tile_arange", test_tile_arange, devices=devices)
@@ -653,7 +741,8 @@ add_function_test(TestTileReduce, "test_tile_untile_vector", test_tile_untile_ve
 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)
+add_function_test(TestTileReduce, "test_tile_reduce_matrix", test_tile_reduce_matrix, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_vector", test_tile_reduce_vector, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

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,13 +110,13 @@ def test_tile_shared_mem_graph(test, device):
 
     out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
 
-    wp.load_module(device=device)
+    # preload the unique module
+    wp.load_module(compute.module, device=device, block_dim=BLOCK_DIM)
 
-    wp.capture_begin(device, force_module_load=False)
-    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
-    graph = wp.capture_end(device)
+    with wp.ScopedCapture(device, force_module_load=False) as capture:
+        wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
 
-    wp.capture_launch(graph)
+    wp.capture_launch(capture.graph)
 
     # check output
     assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
@@ -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
@@ -225,9 +225,9 @@ def test_tile_shared_non_aligned(test, device):
 
 
 def test_tile_shared_vec_accumulation(test, device):
-    BLOCK_DIM = 64
+    BLOCK_DIM = 256
 
-    @wp.kernel
+    @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()
 
@@ -286,9 +286,9 @@ def test_tile_shared_vec_accumulation(test, device):
 
 
 def test_tile_shared_simple_reduction_add(test, device):
-    BLOCK_DIM = 64
+    BLOCK_DIM = 256
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(x: wp.array(dtype=float), y: wp.array(dtype=float)):
         i, j = wp.tid()
 
@@ -313,9 +313,9 @@ def test_tile_shared_simple_reduction_add(test, device):
 
 
 def test_tile_shared_simple_reduction_sub(test, device):
-    BLOCK_DIM = 64
+    BLOCK_DIM = 256
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(x: wp.array(dtype=float), y: wp.array(dtype=float)):
         i, j = wp.tid()
 

warp/tests/tile/test_tile_sort.py

@@ -44,63 +44,67 @@ def create_sort_kernel(KEY_TYPE, MAX_SORT_LENGTH):
 
 
 def test_tile_sort(test, device):
-    for dtype in [int, float]:  # Loop over int and float keys
+    # 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
-            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}!"
+
+            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()