warp-lang 1.8.1__py3-none-manylinux_2_34_aarch64.whl → 1.9.0__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_enum.py

@@ -0,0 +1,136 @@
+# 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 enum
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+class MyIntEnum(enum.IntEnum):
+    A = 1
+    B = 2
+    C = 3
+
+
+class MyIntFlag(enum.IntFlag):
+    A = 1
+    B = enum.auto()
+    C = enum.auto()
+
+
+def test_intenum_ints(test, device):
+    @wp.kernel
+    def expect_intenum_ints():
+        wp.expect_eq(MyIntEnum.A, 1)
+        wp.expect_eq(MyIntEnum.B, 2)
+        wp.expect_eq(MyIntEnum.C, 3)
+        wp.expect_eq(MyIntEnum.A + MyIntEnum.B, MyIntEnum.C)
+
+    wp.launch(expect_intenum_ints, dim=1, device=device)
+
+
+def test_intflag_ints(test, device):
+    @wp.kernel
+    def expect_intflag_ints():
+        wp.expect_eq(MyIntFlag.A, 1)
+        wp.expect_eq(MyIntFlag.B, 2)
+        wp.expect_eq(MyIntFlag.C, 4)
+        wp.expect_eq(MyIntFlag.A | MyIntFlag.B, 3)
+        wp.expect_eq(MyIntFlag.A | MyIntFlag.B | MyIntFlag.C, 7)
+
+    wp.launch(expect_intflag_ints, dim=1, device=device)
+
+
+def test_alternative_accessors(test, device):
+    @wp.kernel
+    def expect_alternative_accessors():
+        wp.expect_eq(int(MyIntEnum.A), 1)
+        wp.expect_eq(int(MyIntEnum.B.value), 2)
+        wp.expect_eq(MyIntEnum.C.value, 3)
+        wp.expect_eq(MyIntEnum.A + int(MyIntEnum.B) + 0, MyIntEnum.C)
+        wp.expect_eq(int(MyIntFlag.A), 1)
+        wp.expect_eq(int(MyIntFlag.B.value), 2)
+        wp.expect_eq(MyIntFlag.C.value, 4)
+        wp.expect_eq(MyIntFlag.A | int(MyIntFlag.B), 3)
+        wp.expect_eq(MyIntFlag.A | MyIntFlag.B.value | MyIntFlag.C, 7)
+
+    wp.launch(expect_alternative_accessors, dim=1, device=device)
+
+
+def test_static_accessors(test, device):
+    @wp.kernel
+    def expect_static_accessors():
+        wp.expect_eq(wp.static(MyIntEnum.A), 1)
+        wp.expect_eq(wp.static(int(MyIntEnum.A)), 1)
+        wp.expect_eq(wp.static(MyIntEnum.A.value), 1)
+        wp.expect_eq(wp.static(MyIntFlag.A), 1)
+        wp.expect_eq(wp.static(int(MyIntFlag.A)), 1)
+        wp.expect_eq(wp.static(MyIntFlag.A.value), 1)
+
+    wp.launch(expect_static_accessors, dim=1, device=device)
+
+
+def test_intflag_compare(test, device):
+    @wp.kernel
+    def compute_intflag_compare(ins: wp.array(dtype=wp.int32), outs: wp.array(dtype=wp.int32)):
+        tid = wp.tid()
+        if ins[tid] & MyIntFlag.A:
+            outs[tid] += MyIntFlag.A
+        if ins[tid] & MyIntFlag.B:
+            outs[tid] += MyIntFlag.B
+        if ins[tid] & MyIntFlag.C:
+            outs[tid] += MyIntFlag.C
+
+    with wp.ScopedDevice(device):
+        ins = wp.array(
+            [
+                0,
+                MyIntFlag.A,
+                MyIntFlag.B,
+                MyIntFlag.C,
+                MyIntFlag.A | MyIntFlag.B,
+                MyIntFlag.A | MyIntFlag.B | MyIntFlag.C,
+            ],
+            dtype=wp.int32,
+        )
+        outs = wp.zeros(len(ins), dtype=wp.int32)
+        wp.launch(compute_intflag_compare, dim=len(ins), inputs=[ins], outputs=[outs])
+        outs = outs.numpy()
+        test.assertEqual(outs[0], 0)
+        test.assertEqual(outs[1], 1)
+        test.assertEqual(outs[2], 2)
+        test.assertEqual(outs[3], 4)
+        test.assertEqual(outs[4], 3)
+        test.assertEqual(outs[5], 7)
+
+
+class TestEnum(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestEnum, "test_intenum_ints", test_intenum_ints, devices=devices)
+add_function_test(TestEnum, "test_intflag_ints", test_intflag_ints, devices=devices)
+add_function_test(TestEnum, "test_intflag_compare", test_intflag_compare, devices=devices)
+add_function_test(TestEnum, "test_alternative_accessors", test_alternative_accessors, devices=devices)
+add_function_test(TestEnum, "test_static_accessors", test_static_accessors, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_examples.py

@@ -52,7 +52,7 @@ from warp.tests.unittest_utils import (
 )
 from warp.utils import check_p2p
 
-wp.init()  # For wp.context.runtime.core.is_debug_enabled()
+wp.init()  # For wp.context.runtime.core.wp_is_debug_enabled()
 
 
 def _build_command_line_options(test_options: Dict[str, Any]) -> list:
@@ -331,7 +331,7 @@ add_example_test(
     name="optim.example_softbody_properties",
     devices=test_devices,
     test_options_cuda={
-        "train_iters": 1 if warp.context.runtime.core.is_debug_enabled() else 3,
+        "train_iters": 1 if warp.context.runtime.core.wp_is_debug_enabled() else 3,
     },
     test_options_cpu={"train_iters": 1},
 )

warp/tests/test_fem.py

@@ -46,6 +46,11 @@ def linear_form(s: Sample, u: Field):
     return u(s)
 
 
+@integrand
+def bilinear_form(s: Sample, u: Field, v: Field):
+    return u(s) * v(s)
+
+
 @integrand
 def scaled_linear_form(s: Sample, u: Field, scale: wp.array(dtype=float)):
     return u(s) * scale[0]
@@ -1868,8 +1873,6 @@ def test_vector_spaces(test, device):
             fields={"field": div_field.trace()},
         )
 
-    return
-
     with wp.ScopedDevice(device):
         positions, tri_vidx = _gen_trimesh(3, 5)
 
@@ -2055,6 +2058,45 @@ def test_array_axpy(test, device):
     assert_np_equal(y.grad.numpy(), beta * np.ones(N))
 
 
+def test_integrate_high_order(test_field, device):
+    with wp.ScopedDevice(device):
+        geo = fem.Grid3D(res=(1, 1, 1))
+        space = fem.make_polynomial_space(geo, degree=4)
+        test_field = fem.make_test(space)
+        trial_field = fem.make_trial(space)
+
+        # compare consistency of tile-based "dispatch" assembly and generic
+        v0 = fem.integrate(
+            linear_form, fields={"u": test_field}, assembly="dispatch", kernel_options={"enable_backward": False}
+        )
+        v1 = fem.integrate(
+            linear_form, fields={"u": test_field}, assembly="generic", kernel_options={"enable_backward": False}
+        )
+
+        assert_np_equal(v0.numpy(), v1.numpy(), tol=1.0e-6)
+
+        h0 = fem.integrate(
+            bilinear_form,
+            fields={"v": test_field, "u": trial_field},
+            assembly="dispatch",
+            kernel_options={"enable_backward": False},
+        )
+        h1 = fem.integrate(
+            bilinear_form,
+            fields={"v": test_field, "u": trial_field},
+            assembly="generic",
+            kernel_options={"enable_backward": False},
+        )
+
+        h0_nnz = h0.nnz_sync()
+        h1_nnz = h1.nnz_sync()
+        assert h0.shape == h1.shape
+        assert h0_nnz == h1_nnz
+        assert_array_equal(h0.offsets[: h0.nrow + 1], h1.offsets[: h1.nrow + 1])
+        assert_array_equal(h0.columns[:h0_nnz], h1.columns[:h1_nnz])
+        assert_np_equal(h0.values[:h0_nnz].numpy(), h1.values[:h1_nnz].numpy(), tol=1.0e-6)
+
+
 devices = get_test_devices()
 cuda_devices = get_selected_cuda_test_devices()
 
@@ -2088,6 +2130,7 @@ add_function_test(TestFem, "test_point_basis", test_point_basis)
 add_function_test(TestFem, "test_particle_quadratures", test_particle_quadratures)
 add_function_test(TestFem, "test_nodal_quadrature", test_nodal_quadrature)
 add_function_test(TestFem, "test_implicit_fields", test_implicit_fields)
+add_function_test(TestFem, "test_integrate_high_order", test_integrate_high_order, devices=cuda_devices)
 
 
 class TestFemUtilities(unittest.TestCase):

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)