warp-lang 1.6.1__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/tests/{test_tile.py → tile/test_tile.py}

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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
 
@@ -12,8 +20,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.is_mathdx_enabled()
-
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
 TILE_K = wp.constant(8)
@@ -208,7 +214,6 @@ def test_tile_binary_map(test, device):
     assert_np_equal(B_wp.grad.numpy(), B_grad)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_grouped_gemm(test, device):
     @wp.kernel
     def tile_grouped_gemm(A: wp.array3d(dtype=float), B: wp.array3d(dtype=float), C: wp.array3d(dtype=float)):
@@ -248,60 +253,62 @@ def test_tile_grouped_gemm(test, device):
     assert_np_equal(C_wp.numpy(), C, 1e-6)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
-def test_tile_gemm(test, device):
-    @wp.kernel
-    def tile_gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)):
-        # output tile index
-        i, j = wp.tid()
+def test_tile_gemm(dtype):
+    def test(test, device):
+        @wp.kernel
+        def tile_gemm(A: wp.array2d(dtype=dtype), B: wp.array2d(dtype=dtype), C: wp.array2d(dtype=dtype)):
+            # output tile index
+            i, j = wp.tid()
 
-        sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32)
+            sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=dtype)
 
-        M = A.shape[0]
-        N = B.shape[1]
-        K = A.shape[1]
+            M = A.shape[0]
+            N = B.shape[1]
+            K = A.shape[1]
 
-        count = int(K / TILE_K)
+            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))
+            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)
+                # sum += a*b
+                wp.tile_matmul(a, b, sum)
 
-        wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
+            wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
 
-    M = TILE_M * 7
-    K = TILE_K * 6
-    N = TILE_N * 5
+        M = TILE_M * 7
+        K = TILE_K * 6
+        N = TILE_N * 5
 
-    rng = np.random.default_rng(42)
-    A = rng.random((M, K), dtype=np.float32)
-    B = rng.random((K, N), dtype=np.float32)
-    C = np.zeros((M, N), dtype=np.float32)
+        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)
+        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,
-        )
+        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-5)
+        assert_np_equal(C_wp.numpy(), A @ B, tol=1.0e-1)
 
-    adj_C = np.ones_like(C)
+        adj_C = np.ones_like(C)
 
-    tape.backward(grads={C_wp: wp.array(adj_C, device=device)})
+        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-5)
-    assert_np_equal(B_wp.grad.numpy(), A.T @ adj_C, 1.0e-5)
+        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
@@ -542,7 +549,6 @@ def test_tile_transpose(test, device):
     assert_np_equal(output.numpy(), input.numpy().T)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_transpose_matmul(test, device):
     @wp.kernel
     def test_tile_transpose_matmul_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
@@ -564,9 +570,36 @@ def test_tile_transpose_matmul(test, device):
 
 
 @wp.kernel
-def test_tile_broadcast_add_kernel(
+def test_tile_broadcast_add_1d_kernel(
+    input_a: wp.array(dtype=float), input_b: wp.array(dtype=float), output: wp.array(dtype=float)
+):
+    a = wp.tile_load(input_a, shape=(10,))
+    b = wp.tile_load(input_b, shape=(1,))
+
+    c = wp.tile_broadcast(b, shape=(10,))
+    d = a + c
+
+    wp.tile_store(output, d)
+
+
+def test_tile_broadcast_add_1d(test, device):
+    N = 10
+
+    # implicit 1-dim ([1], 1)
+    a = wp.array(np.arange(0, N, dtype=np.float32), device=device)
+    b = wp.array(np.ones(1, dtype=np.float32), device=device)
+    out = wp.zeros((N,), dtype=float, device=device)
+
+    wp.launch_tiled(test_tile_broadcast_add_1d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+
+    assert_np_equal(out.numpy(), a.numpy() + b.numpy())
+
+
+@wp.kernel
+def test_tile_broadcast_add_2d_kernel(
     input_a: wp.array2d(dtype=float), input_b: wp.array(dtype=float), output: wp.array2d(dtype=float)
 ):
+    # implicit 1-dim ([1], 10)
     a = wp.tile_load(input_a, shape=(10, 10))
     b = wp.tile_load(input_b, shape=10)
 
@@ -576,7 +609,7 @@ def test_tile_broadcast_add_kernel(
     wp.tile_store(output, d)
 
 
-def test_tile_broadcast_add(test, device):
+def test_tile_broadcast_add_2d(test, device):
     M = 10
     N = 10
 
@@ -584,7 +617,62 @@ def test_tile_broadcast_add(test, device):
     b = wp.array(np.arange(0, N, dtype=np.float32), device=device)
     out = wp.zeros((M, N), dtype=float, device=device)
 
-    wp.launch_tiled(test_tile_broadcast_add_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+    wp.launch_tiled(test_tile_broadcast_add_2d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+
+    assert_np_equal(out.numpy(), a.numpy() + b.numpy())
+
+
+@wp.kernel
+def test_tile_broadcast_add_3d_kernel(
+    input_a: wp.array3d(dtype=float), input_b: wp.array3d(dtype=float), output: wp.array3d(dtype=float)
+):
+    a = wp.tile_load(input_a, shape=(4, 10, 12))
+    b = wp.tile_load(input_b, shape=(4, 10, 1))
+
+    c = wp.tile_broadcast(b, shape=(4, 10, 12))
+    d = a + c
+
+    wp.tile_store(output, d)
+
+
+def test_tile_broadcast_add_3d(test, device):
+    M = 4
+    N = 10
+    O = 12
+
+    # explicit 1-dim (M, N, 1) to (M, N, O)
+    a = wp.array(np.ones((M, N, O), dtype=np.float32), device=device)
+    b = wp.array(np.arange(0, M * N, dtype=np.float32).reshape((M, N, 1)), device=device)
+    out = wp.zeros((M, N, O), dtype=float, device=device)
+
+    wp.launch_tiled(test_tile_broadcast_add_3d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
+    assert_np_equal(out.numpy(), a.numpy() + b.numpy())
+
+
+@wp.kernel
+def test_tile_broadcast_add_4d_kernel(
+    input_a: wp.array4d(dtype=float), input_b: wp.array4d(dtype=float), output: wp.array4d(dtype=float)
+):
+    a = wp.tile_load(input_a, shape=(4, 10, 5, 6))
+    b = wp.tile_load(input_b, shape=(4, 1, 5, 1))
+    c = wp.tile_broadcast(b, shape=(4, 10, 5, 6))
+    d = a + c
+
+    wp.tile_store(output, d)
+
+
+def test_tile_broadcast_add_4d(test, device):
+    M = 4
+    N = 10
+    O = 5
+    P = 6
+
+    # explicit 1-dims (M, 1, O, 1) to (M, N, O, P)
+    a = wp.array(np.ones((M, N, O, P), dtype=np.float32), device=device)
+    b = wp.array(np.arange(0, M * O, dtype=np.float32).reshape((M, 1, O, 1)), device=device)
+    out = wp.zeros((M, N, O, P), dtype=float, device=device)
+
+    wp.launch_tiled(test_tile_broadcast_add_4d_kernel, dim=[1], inputs=[a, b, out], block_dim=32, device=device)
 
     assert_np_equal(out.numpy(), a.numpy() + b.numpy())
 
@@ -657,7 +745,7 @@ def test_tile_print(test, device):
     wp.synchronize()
 
 
-devices = get_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestTile(unittest.TestCase):
@@ -669,15 +757,20 @@ add_function_test(TestTile, "test_tile_copy_2d", test_tile_copy_2d, devices=devi
 add_function_test(TestTile, "test_tile_unary_map", test_tile_unary_map, devices=devices)
 add_function_test(TestTile, "test_tile_binary_map", test_tile_binary_map, devices=devices)
 add_function_test(TestTile, "test_tile_grouped_gemm", test_tile_grouped_gemm, devices=devices)
-add_function_test(TestTile, "test_tile_gemm", test_tile_gemm, devices=devices)
+add_function_test(TestTile, "test_tile_gemm_fp16", test_tile_gemm(wp.float16), devices=devices)
+add_function_test(TestTile, "test_tile_gemm_fp32", test_tile_gemm(wp.float32), devices=devices)
+add_function_test(TestTile, "test_tile_gemm_fp64", test_tile_gemm(wp.float64), devices=devices)
 add_function_test(TestTile, "test_tile_transpose", test_tile_transpose, devices=devices)
 add_function_test(TestTile, "test_tile_transpose_matmul", test_tile_transpose_matmul, devices=devices)
 add_function_test(TestTile, "test_tile_operators", test_tile_operators, devices=devices)
-add_function_test(TestTile, "test_tile_sum", test_tile_sum, devices=devices)
+add_function_test(TestTile, "test_tile_sum", test_tile_sum, devices=devices, check_output=False)
 add_function_test(TestTile, "test_tile_sum_launch", test_tile_sum_launch, devices=devices)
 add_function_test(TestTile, "test_tile_extract", test_tile_extract, devices=devices)
 add_function_test(TestTile, "test_tile_extract_repeated", test_tile_extract_repeated, devices=devices)
-add_function_test(TestTile, "test_tile_broadcast_add", test_tile_broadcast_add, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_1d", test_tile_broadcast_add_1d, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_2d", test_tile_broadcast_add_2d, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_3d", test_tile_broadcast_add_3d, devices=devices)
+add_function_test(TestTile, "test_tile_broadcast_add_4d", test_tile_broadcast_add_4d, devices=devices)
 add_function_test(TestTile, "test_tile_broadcast_grad", test_tile_broadcast_grad, devices=devices)
 add_function_test(TestTile, "test_tile_len", test_tile_len, devices=devices)
 add_function_test(TestTile, "test_tile_print", test_tile_print, devices=devices, check_output=False)

warp/tests/{test_tile_load.py → tile/test_tile_load.py}

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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
 
@@ -368,7 +376,7 @@ def test_tile_load_fortran(test, device):
     assert_array_equal(B_wp.grad, A_wp.grad)
 
 
-devices = get_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestTileLoad(unittest.TestCase):

warp/tests/{test_tile_mathdx.py → tile/test_tile_mathdx.py}

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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 functools
 import unittest
@@ -84,6 +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.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]
@@ -164,31 +173,33 @@ def test_tile_math_cholesky(test, device):
     # TODO: implement and test backward pass
 
 
-devices = get_cuda_test_devices()
+all_devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 class TestTileMathDx(unittest.TestCase):
     pass
 
 
 # check_output=False so we can enable libmathdx's logging without failing the tests
-add_function_test(TestTileMathDx, "test_tile_math_matmul", test_tile_math_matmul, devices=devices, check_output=False)
 add_function_test(
-    TestTileMathDx, "test_tile_math_cholesky", test_tile_math_cholesky, devices=devices, check_output=False
+    TestTileMathDx, "test_tile_math_matmul", test_tile_math_matmul, devices=all_devices, check_output=False
+)
+add_function_test(
+    TestTileMathDx, "test_tile_math_cholesky", test_tile_math_cholesky, devices=all_devices, check_output=False
 )
 add_function_test(
     TestTileMathDx,
     "test_tile_math_fft_vec2f",
     functools.partial(test_tile_math_fft, wp_dtype=wp.vec2f),
-    devices=devices,
+    devices=cuda_devices,
     check_output=False,
 )
 add_function_test(
     TestTileMathDx,
     "test_tile_math_fft_vec2d",
     functools.partial(test_tile_math_fft, wp_dtype=wp.vec2d),
-    devices=devices,
+    devices=cuda_devices,
     check_output=False,
 )
 

warp/tests/{test_tile_mlp.py → tile/test_tile_mlp.py}

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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 os
 
@@ -14,11 +22,6 @@ import warp.examples
 import warp.optim
 from warp.tests.unittest_utils import *
 
-wp.init()
-
-# needs to be constant for the whole module
-NUM_THREADS = 32
-
 
 def create_layer(rng, dim_in, dim_hid, dtype=float):
     w = rng.uniform(-1.0 / np.sqrt(dim_in), 1.0 / np.sqrt(dim_in), (dim_hid, dim_in))
@@ -37,10 +40,12 @@ def create_array(rng, dim_in, dim_hid, dtype=float):
     return a
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_multi_layer_nn(test, device):
     import torch as tc
 
+    if device.is_cuda and not wp.context.runtime.core.is_mathdx_enabled():
+        test.skipTest("Skipping test on CUDA device without MathDx (tolerance)")
+
     NUM_FREQ = wp.constant(8)
 
     DIM_IN = wp.constant(4 * NUM_FREQ)  # sin,cos for both x,y at each frequency
@@ -52,7 +57,13 @@ def test_multi_layer_nn(test, device):
 
     BATCH_SIZE = min(512, int((IMG_WIDTH * IMG_HEIGHT) / 8))
 
+    if device.is_cpu:
+        NUM_THREADS = 1
+    else:
+        NUM_THREADS = 32
+
     dtype = wp.float16
+    npdtype = wp.types.warp_type_to_np_dtype[dtype]
 
     @wp.func
     def relu(x: dtype):
@@ -66,7 +77,7 @@ def test_multi_layer_nn(test, device):
     def zero(loss: wp.array(dtype=float)):
         loss[0] = 0.0
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(
         batches: wp.array(dtype=int),
         input: wp.array2d(dtype=dtype),
@@ -162,7 +173,9 @@ def test_multi_layer_nn(test, device):
         input = create_array(rng, IMG_WIDTH * IMG_HEIGHT, DIM_IN, dtype=dtype)
         output = create_array(rng, IMG_WIDTH * IMG_HEIGHT, DIM_OUT)
 
-        reference_np = np.load(os.path.join(os.path.dirname(__file__), "assets/pixel.npy"), allow_pickle=True) / 255.0
+        reference_np = (
+            np.load(os.path.join(os.path.dirname(__file__), "..", "assets", "pixel.npy"), allow_pickle=True) / 255.0
+        )
         reference = wp.array(reference_np, dtype=float)
 
         assert reference.shape[1] == IMG_WIDTH * IMG_HEIGHT
@@ -224,7 +237,7 @@ def test_multi_layer_nn(test, device):
                         z_np = np.maximum(weights_3.numpy() @ z_np + bias_3.numpy(), 0.0)
 
                         # test numpy forward
-                        assert_np_equal(output.numpy()[:, indices], z_np, tol=1.0e-2)
+                        assert_np_equal(output.numpy()[:, indices].astype(npdtype), z_np, tol=1.0e-2)
 
                         # torch
                         input_tc = tc.tensor(input.numpy()[:, indices], requires_grad=True, device=torch_device)
@@ -252,7 +265,9 @@ def test_multi_layer_nn(test, device):
                         l_tc.backward()
 
                         # test torch
-                        assert_np_equal(z_tc.cpu().detach().numpy(), output.numpy()[:, indices], tol=1.0e-2)
+                        assert_np_equal(
+                            z_tc.cpu().detach().numpy(), output.numpy()[:, indices].astype(npdtype), tol=1.0e-2
+                        )
                         assert_np_equal(weights_0.grad.numpy(), weights_0_tc.grad.cpu().detach().numpy(), tol=1.0e-2)
                         assert_np_equal(bias_0.grad.numpy(), bias_0_tc.grad.cpu().detach().numpy(), tol=1.0e-2)
                         assert_np_equal(weights_1.grad.numpy(), weights_1_tc.grad.cpu().detach().numpy(), tol=1.0e-2)
@@ -269,7 +284,6 @@ def test_multi_layer_nn(test, device):
         test.assertLess(loss.numpy()[0], 0.002)
 
 
-@unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_single_layer_nn(test, device):
     import torch as tc
 
@@ -279,11 +293,16 @@ def test_single_layer_nn(test, device):
 
     NUM_BLOCKS = 56
 
+    if device.is_cpu:
+        NUM_THREADS = 1
+    else:
+        NUM_THREADS = 32
+
     @wp.func
     def relu(x: float):
         return wp.max(x, 0.0)
 
-    @wp.kernel
+    @wp.kernel(module="unique")
     def compute(
         input: wp.array2d(dtype=float),
         weights: wp.array2d(dtype=float),
@@ -345,7 +364,6 @@ try:
     import torch
 
     # check which Warp devices work with Torch
-    # CUDA devices may fail if Torch was not compiled with CUDA support
     torch_compatible_devices = []
     torch_compatible_cuda_devices = []
 
@@ -364,7 +382,7 @@ try:
         "test_single_layer_nn",
         test_single_layer_nn,
         check_output=False,
-        devices=torch_compatible_cuda_devices,
+        devices=torch_compatible_devices,
     )
     add_function_test(
         TestTileMLP,
@@ -380,4 +398,5 @@ except Exception as e:
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()
+    wp.clear_lto_cache()
     unittest.main(verbosity=2, failfast=True)

warp/tests/{test_tile_reduce.py → tile/test_tile_reduce.py}

@@ -1,9 +1,17 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# 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
 
@@ -168,6 +176,64 @@ def test_tile_reduce_custom(test, device):
         test.assertAlmostEqual(prod_wp[i], prod_np, places=4)
 
 
+@wp.struct
+class KeyValue:
+    key: wp.int32
+    value: wp.float32
+
+
+@wp.func
+def kv_max(a: KeyValue, b: KeyValue) -> KeyValue:
+    return wp.where(a.value < b.value, b, a)
+
+
+@wp.kernel
+def initialize_key_value(values: wp.array2d(dtype=wp.float32), keyvalues: wp.array2d(dtype=KeyValue)):
+    batch, idx = wp.tid()
+    keyvalues[batch, idx] = KeyValue(idx, values[batch, idx])
+
+
+@wp.kernel(enable_backward=False)
+def tile_reduce_custom_struct_kernel(values: wp.array2d(dtype=KeyValue), res: wp.array(dtype=KeyValue)):
+    # output tile index
+    i = wp.tid()
+
+    t = wp.tile_load(values, shape=(1, TILE_DIM), offset=(i, 0))
+
+    max_el = wp.tile_reduce(kv_max, t)
+    wp.tile_store(res, max_el, offset=i)
+
+
+def test_tile_reduce_custom_struct(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, dtype=wp.float32, device=device)
+    keyvalues_wp = wp.empty(input_wp.shape, dtype=KeyValue, device=device)
+
+    wp.launch(initialize_key_value, dim=[batch_count, N], inputs=[input_wp], outputs=[keyvalues_wp], device=device)
+
+    output_wp = wp.empty(batch_count, dtype=KeyValue, device=device)
+
+    wp.launch_tiled(
+        tile_reduce_custom_struct_kernel,
+        dim=[batch_count],
+        inputs=[keyvalues_wp],
+        outputs=[output_wp],
+        block_dim=TILE_DIM,
+        device=device,
+    )
+
+    prod_wp = np.array([k for k, v in output_wp.numpy()])
+    expected = np.argmax(input, axis=1)
+
+    assert_np_equal(prod_wp, expected)
+
+
 @wp.kernel
 def tile_grouped_sum_kernel(input: wp.array3d(dtype=float), output: wp.array(dtype=float)):
     # output tile index
@@ -357,7 +423,7 @@ def test_tile_arange(test, device):
     assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
 
 
-devices = get_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestTileReduce(unittest.TestCase):
@@ -368,6 +434,7 @@ add_function_test(TestTileReduce, "test_tile_reduce_sum", test_tile_reduce_sum,
 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_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_simt", test_tile_reduce_simt, devices=devices)
 add_function_test(TestTileReduce, "test_tile_ones", test_tile_ones, devices=devices)