warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/test_math.py

@@ -0,0 +1,212 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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, NamedTuple
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import add_function_test, assert_np_equal, get_test_devices
+
+
+class ScalarFloatValues(NamedTuple):
+    degrees: wp.float32 = None
+    radians: wp.float32 = None
+
+
+@wp.kernel
+def scalar_float_kernel(
+    i: int,
+    x: wp.array(dtype=wp.float32),
+    out: wp.array(dtype=wp.float32),
+):
+    if i == 0:
+        out[0] = wp.degrees(x[0])
+    elif i == 1:
+        out[0] = wp.radians(x[0])
+
+
+def test_scalar_math(test, device):
+    float_values = ScalarFloatValues(degrees=(0.123,), radians=(123.0,))
+    float_results_expected = ScalarFloatValues(degrees=7.047381, radians=2.146755)
+    adj_float_results_expected = ScalarFloatValues(degrees=57.29578, radians=0.017453)
+    for i, values in enumerate(float_values):
+        x = wp.array([values[0]], dtype=wp.float32, requires_grad=True, device=device)
+        out = wp.array([0.0], dtype=wp.float32, requires_grad=True, device=device)
+
+        tape = wp.Tape()
+        with tape:
+            wp.launch(scalar_float_kernel, dim=1, inputs=[i, x, out], device=device)
+
+        assert_np_equal(out.numpy(), np.array([float_results_expected[i]]), tol=1e-6)
+
+        tape.backward(out)
+
+        assert_np_equal(tape.gradients[x].numpy(), np.array([adj_float_results_expected[i]]), tol=1e-6)
+
+
+@wp.kernel
+def erf_kernel(x: wp.array(dtype=Any), out: wp.array(dtype=Any)):
+    i = wp.tid()
+
+    if i == 0:
+        out[i] = wp.erf(x[i])
+    elif i == 1:
+        out[i] = wp.erfc(x[i])
+    elif i == 2:
+        out[i] = wp.erfinv(x[i])
+    elif i == 3:
+        out[i] = wp.erfcinv(x[i])
+
+
+def test_erf_math(test, device):
+    for type, tol in ((wp.float16, 1e-3), (wp.float32, 1e-6), (wp.float64, 1e-6)):
+        x = wp.full(4, value=0.123, dtype=type, requires_grad=True, device=device)
+        out = wp.zeros(4, dtype=type, requires_grad=True, device=device)
+
+        with wp.Tape() as tape:
+            wp.launch(erf_kernel, dim=4, inputs=[x], outputs=[out], device=device)
+
+        out.grad = wp.ones_like(out)
+
+        tape.backward()
+
+        out_true = np.array([0.13809388, 0.86190612, 0.10944129, 1.09057285])
+        adj_x_true = np.array([1.11143641, -1.11143641, 0.89690544, -2.91120449])
+
+        assert_np_equal(out.numpy(), out_true, tol=tol)
+        assert_np_equal(adj_x_true, x.grad.numpy(), tol=tol)
+
+
+@wp.kernel
+def test_vec_norm_kernel(vs: wp.array(dtype=Any), out: wp.array(dtype=float, ndim=2)):
+    tid = wp.tid()
+    out[tid, 0] = wp.norm_l1(vs[tid])
+    out[tid, 1] = wp.norm_l2(vs[tid])
+    out[tid, 2] = wp.norm_huber(vs[tid])
+    out[tid, 3] = wp.norm_pseudo_huber(vs[tid])
+
+
+def test_vec_norm(test, device):
+    # ground-truth implementations from SciPy
+    def huber(delta, x):
+        if x <= delta:
+            return 0.5 * x**2
+        else:
+            return delta * (x - 0.5 * delta)
+
+    def pseudo_huber(delta, x):
+        return delta**2 * (np.sqrt(1 + (x / delta) ** 2) - 1)
+
+    v0 = wp.vec3(-2.0, -1.0, -3.0)
+    v1 = wp.vec3(2.0, 1.0, 3.0)
+    v2 = wp.vec3(0.0, 0.0, 0.0)
+
+    xs = wp.array([v0, v1, v2], dtype=wp.vec3, requires_grad=True, device=device)
+    out = wp.empty((len(xs), 4), dtype=wp.float32, requires_grad=True, device=device)
+
+    wp.launch(test_vec_norm_kernel, dim=len(xs), inputs=[xs], outputs=[out], device=device)
+
+    for i, x in enumerate([v0, v1, v2]):
+        assert_np_equal(
+            out.numpy()[i],
+            np.array(
+                [
+                    np.linalg.norm(x, ord=1),
+                    np.linalg.norm(x, ord=2),
+                    huber(1.0, wp.length(x)),
+                    # note SciPy defines the Pseudo-Huber loss slightly differently
+                    pseudo_huber(1.0, wp.length(x)) + 1.0,
+                ]
+            ),
+            tol=1e-6,
+        )
+
+
+devices = get_test_devices()
+
+
+class TestMath(unittest.TestCase):
+    def test_vec_type(self):
+        vec5 = wp.vec(length=5, dtype=float)
+        v = vec5()
+        w = vec5()
+        a = vec5(1.0)
+        b = vec5(0.0, 0.0, 0.0, 0.0, 0.0)
+        c = vec5(0.0)
+
+        v[0] = 1.0
+        v.x = 0.0
+        v[1:] = [1.0, 1.0, 1.0, 1.0]
+
+        w[0] = 1.0
+        w[1:] = [0.0, 0.0, 0.0, 0.0]
+
+        self.assertEqual(v[0], w[1], "vec setter error")
+        self.assertEqual(v.x, w.y, "vec setter error")
+
+        for x in v[1:]:
+            self.assertEqual(x, 1.0, "vec slicing error")
+
+        self.assertEqual(b, c, "vec equality error")
+
+        self.assertEqual(str(v), "[0.0, 1.0, 1.0, 1.0, 1.0]", "vec to string error")
+
+    def test_mat_type(self):
+        mat55 = wp.mat(shape=(5, 5), dtype=float)
+        m1 = mat55()
+        m2 = mat55()
+
+        for i in range(5):
+            for j in range(5):
+                if i == j:
+                    m1[i, j] = 1.0
+                else:
+                    m1[i, j] = 0.0
+
+        for i in range(5):
+            m2[i] = [1.0, 1.0, 1.0, 1.0, 1.0]
+
+        a = mat55(1.0)
+        # fmt: off
+        b = mat55(
+            1.0, 0.0, 0.0, 0.0, 0.0,
+            0.0, 1.0, 0.0, 0.0, 0.0,
+            0.0, 0.0, 1.0, 0.0, 0.0,
+            0.0, 0.0, 0.0, 1.0, 0.0,
+            0.0, 0.0, 0.0, 0.0, 1.0,
+        )
+        # fmt: on
+
+        self.assertEqual(m1, b, "mat element setting error")
+        self.assertEqual(m2, a, "mat row setting error")
+        self.assertEqual(m1[0, 0], 1.0, "mat element getting error")
+        self.assertEqual(m2[0], [1.0, 1.0, 1.0, 1.0, 1.0], "mat row getting error")
+        self.assertEqual(
+            str(b),
+            "[[1.0, 0.0, 0.0, 0.0, 0.0],\n [0.0, 1.0, 0.0, 0.0, 0.0],\n [0.0, 0.0, 1.0, 0.0, 0.0],\n [0.0, 0.0, 0.0, 1.0, 0.0],\n [0.0, 0.0, 0.0, 0.0, 1.0]]",
+            "mat to string error",
+        )
+
+
+add_function_test(TestMath, "test_scalar_math", test_scalar_math, devices=devices)
+add_function_test(TestMath, "test_erf_math", test_erf_math, devices=devices)
+add_function_test(TestMath, "test_vec_norm", test_vec_norm, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_module_aot.py

@@ -0,0 +1,287 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 importlib
+import importlib.util
+import os
+import shutil
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+import warp as wp
+import warp.tests.aux_test_module_aot
+from warp.tests.unittest_utils import *
+
+ADD_KERNEL_START = """import warp as wp
+
+
+@wp.kernel
+def add_kernel(a: wp.array(dtype=wp.int32), b: wp.array(dtype=wp.int32), res: wp.array(dtype=wp.int32)):
+    pass
+"""
+
+ADD_KERNEL_FINAL = """import warp as wp
+
+
+@wp.kernel
+def add_kernel(a: wp.array(dtype=wp.int32), b: wp.array(dtype=wp.int32), res: wp.array(dtype=wp.int32)):
+    i = wp.tid()
+    res[i] = a[i] + b[i]
+"""
+
+
+def reload_module(module):
+    # Clearing the .pyc file associated with a module is a necessary workaround
+    # for `importlib.reload` to work as expected when run from within Kit.
+    cache_file = importlib.util.cache_from_source(module.__file__)
+    if os.path.exists(cache_file):
+        os.remove(cache_file)
+    importlib.reload(module)
+
+
+TEST_CACHE_DIR = Path(os.path.abspath(os.path.join(os.path.dirname(__file__), "test_module_aot_cache")))
+
+
+def test_disable_hashing(test, device):
+    """Test that module hashing can be disabled.
+
+    A module is run, modified, and run again. The second run should not trigger
+    a recompilation since the hash will not be used to detect changes.
+    """
+
+    try:
+        shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+        TEST_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        wp.set_module_options(
+            {"block_dim": 1 if device.is_cpu else 256},
+            warp.tests.aux_test_module_aot,
+        )
+
+        a = wp.ones(10, dtype=wp.int32, device=device)
+        b = wp.ones(10, dtype=wp.int32, device=device)
+        res = wp.zeros((10,), dtype=wp.int32, device=device)
+
+        # Write out the module and import it
+        with open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_module_aot.py")), "w") as f:
+            f.writelines(ADD_KERNEL_START)
+        reload_module(warp.tests.aux_test_module_aot)
+
+        # First launch, cold compile, expect res to be unchanged since kernel is empty
+        wp.compile_aot_module(warp.tests.aux_test_module_aot, device, module_dir=TEST_CACHE_DIR, strip_hash=True)
+        wp.load_aot_module(warp.tests.aux_test_module_aot, device, module_dir=TEST_CACHE_DIR, strip_hash=True)
+
+        wp.launch(
+            warp.tests.aux_test_module_aot.add_kernel,
+            dim=a.shape,
+            inputs=[a, b],
+            outputs=[res],
+            device=device,
+        )
+
+        assert_np_equal(res.numpy(), np.zeros((10,), dtype=np.int32))
+
+        res.zero_()
+
+        # Write out the modified module and import it
+        with open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_module_aot.py")), "w") as f:
+            f.writelines(ADD_KERNEL_FINAL)
+        reload_module(warp.tests.aux_test_module_aot)
+
+        # This time, the hash checks will be skipped so the previously compiled module will be loaded
+        wp.load_aot_module(warp.tests.aux_test_module_aot, device, module_dir=TEST_CACHE_DIR, strip_hash=True)
+
+        # Kernel is executed with the ADD_KERNEL_START code, not the ADD_KERNEL_FINAL code
+        wp.launch(
+            warp.tests.aux_test_module_aot.add_kernel,
+            dim=a.shape,
+            inputs=[a, b],
+            outputs=[res],
+            device=device,
+        )
+
+        assert_np_equal(res.numpy(), np.zeros((10,), dtype=np.int32))
+    finally:
+        # Clear the cache directory
+        shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+        # Revert the module default options and auxiliary file to the original states
+        wp.set_module_options({"cuda_output": None, "strip_hash": False}, warp.tests.aux_test_module_aot)
+
+        with open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_module_aot.py")), "w") as f:
+            f.writelines(ADD_KERNEL_FINAL)
+
+
+def test_enable_hashing(test, device):
+    """Ensure that the logic of test_disable_hashing is sound.
+
+    This test sets "strip_hash" to False, so normal module hashing rules
+    should be in effect.
+    """
+
+    try:
+        shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+        TEST_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        wp.set_module_options(
+            {"block_dim": 1 if device.is_cpu else 256},
+            warp.tests.aux_test_module_aot,
+        )
+
+        a = wp.ones(10, dtype=wp.int32, device=device)
+        b = wp.ones(10, dtype=wp.int32, device=device)
+        res = wp.zeros((10,), dtype=wp.int32, device=device)
+
+        # Write out the module and import it
+        with open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_module_aot.py")), "w") as f:
+            f.writelines(ADD_KERNEL_START)
+        reload_module(warp.tests.aux_test_module_aot)
+
+        # First launch, cold compile, expect no-op result
+        wp.compile_aot_module(warp.tests.aux_test_module_aot, device, module_dir=TEST_CACHE_DIR, strip_hash=False)
+        wp.load_aot_module(warp.tests.aux_test_module_aot, device, module_dir=TEST_CACHE_DIR, strip_hash=False)
+        wp.launch(
+            warp.tests.aux_test_module_aot.add_kernel,
+            dim=a.shape,
+            inputs=[a, b],
+            outputs=[res],
+            device=device,
+        )
+
+        assert_np_equal(res.numpy(), np.zeros((10,), dtype=np.int32))
+
+        # Write out the modified module (results in a different hash) and import it
+        with open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_module_aot.py")), "w") as f:
+            f.writelines(ADD_KERNEL_FINAL)
+        reload_module(warp.tests.aux_test_module_aot)
+
+        # Trying to load the module should fail since a compiled module with the expected hash does not exist
+        with test.assertRaises(FileNotFoundError):
+            wp.load_aot_module("warp.tests.aux_test_module_aot", device, module_dir=TEST_CACHE_DIR, strip_hash=False)
+    finally:
+        # Clear the cache directory
+        shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+        # Revert the module default options and auxiliary file to the original states
+        wp.set_module_options({"cuda_output": None, "strip_hash": False}, warp.tests.aux_test_module_aot)
+
+        with open(os.path.abspath(os.path.join(os.path.dirname(__file__), "aux_test_module_aot.py")), "w") as f:
+            f.writelines(ADD_KERNEL_FINAL)
+
+
+def test_module_load_resolution(test, device):
+    """Test various ways to resolving a module when loading and compiling."""
+
+    wp.set_module_options(
+        {"block_dim": 1 if device.is_cpu else 256},
+        warp.tests.aux_test_module_aot,
+    )
+
+    a = wp.ones(10, dtype=wp.int32, device=device)
+    b = wp.ones(10, dtype=wp.int32, device=device)
+    res = wp.zeros((10,), dtype=wp.int32, device=device)
+
+    reload_module(warp.tests.aux_test_module_aot)
+    wp.compile_aot_module(warp.tests.aux_test_module_aot, device)
+    wp.load_aot_module(warp.tests.aux_test_module_aot, device)
+
+    wp.launch(
+        warp.tests.aux_test_module_aot.add_kernel,
+        dim=a.shape,
+        inputs=[a, b],
+        outputs=[res],
+        device=device,
+    )
+    assert_np_equal(res.numpy(), np.full((10,), 2, dtype=np.int32))
+
+    reload_module(warp.tests.aux_test_module_aot)
+    res.zero_()
+    wp.compile_aot_module("warp.tests.aux_test_module_aot", device)
+    wp.load_aot_module("warp.tests.aux_test_module_aot", device)
+
+    wp.launch(
+        warp.tests.aux_test_module_aot.add_kernel,
+        dim=a.shape,
+        inputs=[a, b],
+        outputs=[res],
+        device=device,
+    )
+    assert_np_equal(res.numpy(), np.full((10,), 2, dtype=np.int32))
+
+
+class TestModuleAOT(unittest.TestCase):
+    def test_module_compile_specified_arch_ptx(self):
+        """Test that a module can be compiled for a specific architecture or architectures (PTX)."""
+
+        if wp.get_cuda_device_count() == 0:
+            self.skipTest("No CUDA devices found")
+
+        if len(wp._src.context.runtime.nvrtc_supported_archs) < 2:
+            self.skipTest("NVRTC must support at least two architectures to run this test")
+
+        try:
+            shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+            TEST_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+            archs = list(wp._src.context.runtime.nvrtc_supported_archs)[:2]
+
+            wp.compile_aot_module(warp.tests.aux_test_module_aot, arch=archs, module_dir=TEST_CACHE_DIR, use_ptx=True)
+
+            # Make sure the expected files exist
+            module_identifier = wp.get_module("warp.tests.aux_test_module_aot").get_module_identifier()
+            for arch in archs:
+                expected_filename = f"{module_identifier}.sm{arch}.ptx"
+                expected_path = TEST_CACHE_DIR / expected_filename
+                self.assertTrue(expected_path.exists(), f"Expected compiled PTX file not found: {expected_path}")
+
+        finally:
+            # Clear the cache directory
+            shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+
+    def test_module_compile_specified_arch_cubin(self):
+        """Test that a module can be compiled for a specific architecture or architectures (CUBIN)."""
+
+        if wp.get_cuda_device_count() == 0:
+            self.skipTest("No CUDA devices found")
+
+        if len(wp._src.context.runtime.nvrtc_supported_archs) < 2:
+            self.skipTest("NVRTC must support at least two architectures to run this test")
+
+        try:
+            shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+            TEST_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+            archs = list(wp._src.context.runtime.nvrtc_supported_archs)[:2]
+
+            wp.compile_aot_module(warp.tests.aux_test_module_aot, arch=archs, module_dir=TEST_CACHE_DIR, use_ptx=False)
+
+            # Make sure the expected files exist
+            module_identifier = wp.get_module("warp.tests.aux_test_module_aot").get_module_identifier()
+            for arch in archs:
+                expected_filename = f"{module_identifier}.sm{arch}.cubin"
+                expected_path = TEST_CACHE_DIR / expected_filename
+                self.assertTrue(expected_path.exists(), f"Expected compiled CUBIN file not found: {expected_path}")
+
+        finally:
+            # Clear the cache directory
+            shutil.rmtree(TEST_CACHE_DIR, ignore_errors=True)
+
+
+devices = get_test_devices()
+add_function_test(TestModuleAOT, "test_disable_hashing", test_disable_hashing, devices=devices)
+add_function_test(TestModuleAOT, "test_enable_hashing", test_enable_hashing, devices=devices)
+add_function_test(TestModuleAOT, "test_module_load_resolution", test_module_load_resolution, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)