warp-lang 1.8.1__py3-none-macosx_10_13_universal2.whl → 1.9.1__py3-none-macosx_10_13_universal2.whl

warp/tests/test_fixedarray.py

@@ -0,0 +1,229 @@
+# 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 unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+devices = get_test_devices()
+
+
+@wp.kernel
+def test_zeros():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+@wp.func
+def test_func_arg_func(arr: wp.array(ndim=2, dtype=int)):
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+
+@wp.kernel
+def test_func_arg():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    test_func_arg_func(arr)
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+@wp.func
+def test_func_return_func():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    return arr
+
+
+@wp.kernel
+def test_func_return():
+    arr = test_func_return_func()
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+@wp.func
+def test_func_return_annotation_func() -> wp.fixedarray(shape=(2, 3), dtype=int):
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    return arr
+
+
+@wp.kernel
+def test_func_return_annotation():
+    arr = test_func_return_annotation_func()
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+def test_error_invalid_func_return_annotation(test, device):
+    @wp.func
+    def func() -> wp.array(ndim=2, dtype=int):
+        arr = wp.zeros(shape=(2, 3), dtype=int)
+        for i in range(arr.shape[0]):
+            for j in range(arr.shape[1]):
+                arr[i][j] = i * arr.shape[1] + j
+
+        return arr
+
+    @wp.kernel
+    def kernel():
+        arr = func()
+
+    with test.assertRaisesRegex(
+        wp.codegen.WarpCodegenError,
+        r"The function `func` returns a fixed-size array whereas it has its return type annotated as `Array\[int32\]`.$",
+    ):
+        wp.launch(kernel, 1, device=device)
+
+
+def test_error_runtime_shape(test, device):
+    @wp.kernel
+    def kernel():
+        tid = wp.tid()
+        wp.zeros(shape=(tid,), dtype=int)
+
+    with test.assertRaisesRegex(
+        RuntimeError,
+        r"the `shape` argument must be specified as a constant when zero-initializing an array$",
+    ):
+        wp.launch(kernel, 1, device=device)
+
+
+@wp.kernel
+def test_capture_if_kernel():
+    arr = wp.zeros(shape=(2, 3), dtype=int)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j] = i * arr.shape[1] + j
+
+    wp.expect_eq(arr[0][0], 0)
+    wp.expect_eq(arr[0][1], 1)
+    wp.expect_eq(arr[0][2], 2)
+    wp.expect_eq(arr[1][0], 3)
+    wp.expect_eq(arr[1][1], 4)
+    wp.expect_eq(arr[1][2], 5)
+
+
+def test_capture_if(test, device):
+    if (
+        not wp.get_device(device).is_cuda
+        or wp.context.runtime.toolkit_version < (12, 4)
+        or wp.context.runtime.driver_version < (12, 4)
+    ):
+        return
+
+    def foo():
+        wp.launch(test_capture_if_kernel, dim=512, block_dim=128, device=device)
+
+    cond = wp.ones(1, dtype=wp.int32, device=device)
+    with wp.ScopedCapture(device=device) as capture:
+        wp.capture_if(condition=cond, on_true=foo)
+
+    wp.capture_launch(capture.graph)
+
+
+@wp.struct
+class test_func_struct_MyStruct:
+    offset: int
+    dist: float
+
+
+@wp.func
+def test_func_struct_func():
+    arr = wp.zeros(shape=(2, 3), dtype=test_func_struct_MyStruct)
+    count = float(arr.shape[0] * arr.shape[1] - 1)
+    for i in range(arr.shape[0]):
+        for j in range(arr.shape[1]):
+            arr[i][j].offset = i * arr.shape[1] + j
+            arr[i][j].dist = float(arr[i][j].offset) / count
+
+    return arr
+
+
+@wp.kernel
+def test_func_struct():
+    arr = test_func_struct_func()
+
+    wp.expect_eq(arr[0][0].offset, 0)
+    wp.expect_near(arr[0][0].dist, 0.0)
+    wp.expect_eq(arr[0][1].offset, 1)
+    wp.expect_near(arr[0][1].dist, 0.2)
+    wp.expect_eq(arr[0][2].offset, 2)
+    wp.expect_near(arr[0][2].dist, 0.4)
+    wp.expect_eq(arr[1][0].offset, 3)
+    wp.expect_near(arr[1][0].dist, 0.6)
+    wp.expect_eq(arr[1][1].offset, 4)
+    wp.expect_near(arr[1][1].dist, 0.8)
+    wp.expect_eq(arr[1][2].offset, 5)
+    wp.expect_near(arr[1][2].dist, 1.0)
+
+
+class TestFixedArray(unittest.TestCase):
+    pass
+
+
+add_kernel_test(TestFixedArray, kernel=test_zeros, name="test_zeros", dim=1, devices=devices)
+add_kernel_test(TestFixedArray, kernel=test_func_arg, name="test_func_arg", dim=1, devices=devices)
+add_kernel_test(TestFixedArray, kernel=test_func_return, name="test_func_return", dim=1, devices=devices)
+add_kernel_test(
+    TestFixedArray, kernel=test_func_return_annotation, name="test_func_return_annotation", dim=1, devices=devices
+)
+add_function_test(
+    TestFixedArray,
+    "test_error_invalid_func_return_annotation",
+    test_error_invalid_func_return_annotation,
+    devices=devices,
+)
+add_function_test(TestFixedArray, "test_error_runtime_shape", test_error_runtime_shape, devices=devices)
+add_function_test(TestFixedArray, "test_capture_if", test_capture_if, devices=devices)
+add_kernel_test(TestFixedArray, kernel=test_func_struct, name="test_func_struct", dim=1, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_func.py

@@ -269,15 +269,14 @@ def normalize_vector(vec_a: wp.vec3):
     return wp.normalize(vec_a)
 
 
-# This pair is to test the situation where one overload throws an error, but a second one works.
 @wp.func
-def divide_by_zero_overload(x: wp.float32):
-    return x / 0
+def divide_float64(x: wp.float64):
+    return x / wp.float64(1.23)
 
 
 @wp.func
-def divide_by_zero_overload(x: wp.float64):
-    return wp.div(x, 0.0)
+def get_array_len(arr: wp.array(dtype=wp.float32)):
+    return len(arr)
 
 
 class TestFunc(unittest.TestCase):
@@ -444,26 +443,30 @@ class TestFunc(unittest.TestCase):
             a * b
 
     def test_cpython_call_user_function_with_error(self):
-        # Actually the following also includes a ZeroDivisionError in the message due to exception chaining,
-        # but I don't know how to test for that.
         with self.assertRaisesRegex(
-            RuntimeError,
-            "Error calling function 'divide_by_zero'. No version succeeded. "
-            "See above for the error from the last version that was tried.",
+            ZeroDivisionError,
+            "float division by zero",
         ):
             divide_by_zero(1.0)
 
-    def test_cpython_call_user_function_with_overloads(self):
-        self.assertEqual(divide_by_zero_overload(1.0), math.inf)
-
     def test_cpython_call_user_function_with_wrong_argument_types(self):
         with self.assertRaisesRegex(
             RuntimeError,
-            "Error calling function 'normalize_vector'. No version succeeded. "
-            "See above for the error from the last version that was tried.",
+            r"^Error calling function 'divide_float64', no overload found for arguments \(1.0,\)$",
+        ):
+            divide_float64(1.0)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"^Error calling function 'normalize_vector', no overload found for arguments \(1.0,\)$",
         ):
             normalize_vector(1.0)
 
+    def test_cpython_call_user_function_with_array_type(self):
+        arr = wp.array((1, 2, 3, 4, 5, 6, 7, 8), dtype=wp.float32)
+        length = get_array_len(arr)
+        assert length == 8
+
 
 devices = get_test_devices()
 

warp/tests/test_future_annotations.py

@@ -77,20 +77,22 @@ def test_future_annotations(test, device):
     foo_data.x = 1.23
     foo_data.y = 2.34
 
-    out = wp.empty(1, dtype=float)
+    out = wp.empty(1, dtype=float, device=device)
 
     kernel_3 = create_kernel_3(foo)
 
-    wp.launch(kernel_1, dim=out.shape, outputs=(out,))
-    wp.launch(kernel_2, dim=out.shape, outputs=(out,))
-    wp.launch(kernel_3, dim=out.shape, inputs=(foo_data,), outputs=(out,))
+    wp.launch(kernel_1, dim=out.shape, outputs=(out,), device=device)
+    wp.launch(kernel_2, dim=out.shape, outputs=(out,), device=device)
+    wp.launch(kernel_3, dim=out.shape, inputs=(foo_data,), outputs=(out,), device=device)
 
 
 class TestFutureAnnotations(unittest.TestCase):
     pass
 
 
-add_function_test(TestFutureAnnotations, "test_future_annotations", test_future_annotations)
+devices = get_test_devices()
+
+add_function_test(TestFutureAnnotations, "test_future_annotations", test_future_annotations, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_linear_solvers.py

@@ -18,10 +18,19 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp.context import assert_conditional_graph_support
 from warp.optim.linear import bicgstab, cg, cr, gmres, preconditioner
 from warp.tests.unittest_utils import *
 
 
+def check_conditional_graph_support():
+    try:
+        assert_conditional_graph_support()
+    except RuntimeError:
+        return False
+    return True
+
+
 def _check_linear_solve(test, A, b, func, *args, **kwargs):
     # test from zero
     x = wp.zeros_like(b)
@@ -30,10 +39,30 @@ def _check_linear_solve(test, A, b, func, *args, **kwargs):
 
     test.assertLessEqual(err, atol)
 
+    # Test with capturable graph
+    if A.device.is_cuda and check_conditional_graph_support():
+        x.zero_()
+        with wp.ScopedDevice(A.device):
+            with wp.ScopedCapture() as capture:
+                niter, err, atol = func(A, b, x, *args, use_cuda_graph=True, check_every=0, **kwargs)
+
+            wp.capture_launch(capture.graph)
+
+        niter = niter.numpy()[0]
+        err = np.sqrt(err.numpy()[0])
+        atol = np.sqrt(atol.numpy()[0])
+
+        test.assertLessEqual(err, atol)
+
     # test with warm start
     with wp.ScopedDevice(A.device):
         niter_warm, err, atol = func(A, b, x, *args, use_cuda_graph=False, **kwargs)
 
+    if isinstance(niter_warm, wp.array):
+        niter_warm = niter_warm.numpy()[0]
+        err = np.sqrt(err.numpy()[0])
+        atol = np.sqrt(atol.numpy()[0])
+
     test.assertLessEqual(err, atol)
 
     if func in [cr, gmres]:
@@ -45,6 +74,7 @@ def _check_linear_solve(test, A, b, func, *args, **kwargs):
     # This can lead to accumulated inaccuracies over iterations, esp in float32
     residual = A.numpy() @ x.numpy() - b.numpy()
     err_np = np.linalg.norm(residual)
+
     if A.dtype == wp.float64:
         test.assertLessEqual(err_np, 2.0 * atol)
     else:

warp/tests/test_map.py

@@ -318,7 +318,7 @@ def test_input_validity(test, device):
 
     with test.assertRaisesRegex(
         TypeError,
-        'Incorrect input provided for argument "i": received array of dtype float32, expected int$',
+        "Function test_input_validity__locals__int_function does not support the provided argument types float32",
     ):
         wp.map(int_function, a1)
 
@@ -476,6 +476,20 @@ add_function_test(TestMap, "test_graph_capture", test_graph_capture, devices=cud
 add_function_test(TestMap, "test_renamed_warp_module", test_renamed_warp_module, devices=devices)
 
 
+class TestMapDebug(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.config.mode
+        wp.config.mode = "debug"
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.config.mode = cls._saved_mode
+
+
+add_function_test(TestMapDebug, "test_mixed_inputs", test_mixed_inputs, devices=devices)
+add_function_test(TestMapDebug, "test_kernel_creation", test_kernel_creation, devices=devices)
+
 if __name__ == "__main__":
     wp.clear_kernel_cache()
     unittest.main(verbosity=2)