warp-lang 1.5.1__py3-none-manylinux2014_aarch64.whl → 1.6.0__py3-none-manylinux2014_aarch64.whl

warp/tests/test_assert.py

@@ -0,0 +1,242 @@
+# 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.
+
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def expect_ones(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] == 1
+
+
+@wp.kernel
+def expect_ones_with_msg(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] == 1, "Array element must be 1"
+
+
+@wp.kernel
+def expect_ones_compound(a: wp.array(dtype=int)):
+    i = wp.tid()
+
+    assert a[i] > 0 and a[i] < 2
+
+
+@wp.func
+def expect_ones_function(value: int):
+    assert value == 1, "Array element must be 1"
+
+
+@wp.kernel
+def expect_ones_call_function(a: wp.array(dtype=int)):
+    i = wp.tid()
+    expect_ones_function(a[i])
+
+
+class TestAssertRelease(unittest.TestCase):
+    """Assert test cases that are to be run with Warp in release mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        wp.config.mode = "release"
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.set_module_options({"mode": cls._saved_mode})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_basic_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_with_msg(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_with_msg, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_compound_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_false_condition_function(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_call_function, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+
+# NOTE: Failed assertions on CUDA devices leaves the CUDA context in an unrecoverable state,
+# so we currently do not test them.
+class TestAssertDebug(unittest.TestCase):
+    """Assert test cases that are to be run with Warp in debug mode."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        wp.set_module_options({"mode": "debug"})
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.set_module_options({"mode": cls._saved_mode})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_basic_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1")
+
+    def test_basic_assert_true_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.ones(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_basic_assert_with_msg(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.zeros(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_with_msg, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1.*Array element must be 1")
+
+    def test_compound_assert_true_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.ones(1, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr, got {output}")
+
+    def test_compound_assert_false_condition(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_compound, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] > 0 and a\[i\] < 2")
+
+    def test_basic_assert_false_condition_function(self):
+        with wp.ScopedDevice("cpu"):
+            wp.load_module(device=wp.get_device())
+
+            input_array = wp.full(1, value=3, dtype=int)
+
+            capture = StdErrCapture()
+            capture.begin()
+
+            wp.launch(expect_ones_call_function, input_array.shape, inputs=[input_array])
+
+            output = capture.end()
+
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert value == 1.*Array element must be 1")
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_codegen.py

@@ -396,48 +396,29 @@ def test_unresolved_symbol(test, device):
 
 
 def test_error_global_var(test, device):
-    arr = wp.array(
-        (1.0, 2.0, 3.0),
-        dtype=float,
-        device=device,
-    )
+    arr = wp.array((1.0, 2.0, 3.0), dtype=float, device=device)
 
-    def kernel_1_fn(
-        out: wp.array(dtype=float),
-    ):
+    def kernel_1_fn(out: wp.array(dtype=float)):
         out[0] = arr[0]
 
-    def kernel_2_fn(
-        out: wp.array(dtype=float),
-    ):
+    def kernel_2_fn(out: wp.array(dtype=float)):
         out[0] = arr
 
-    def kernel_3_fn(
-        out: wp.array(dtype=float),
-    ):
+    def kernel_3_fn(out: wp.array(dtype=float)):
         out[0] = wp.lower_bound(arr, 2.0)
 
     out = wp.empty_like(arr)
 
     kernel = wp.Kernel(func=kernel_1_fn)
-    with test.assertRaisesRegex(
-        TypeError,
-        r"Invalid external reference type: <class 'warp.types.array'>",
-    ):
+    with test.assertRaisesRegex(TypeError, r"Invalid external reference type: <class 'warp.types.array'>"):
         wp.launch(kernel, dim=out.shape, inputs=(), outputs=(out,), device=device)
 
     kernel = wp.Kernel(func=kernel_2_fn)
-    with test.assertRaisesRegex(
-        TypeError,
-        r"Invalid external reference type: <class 'warp.types.array'>",
-    ):
+    with test.assertRaisesRegex(TypeError, r"Invalid external reference type: <class 'warp.types.array'>"):
         wp.launch(kernel, dim=out.shape, inputs=(), outputs=(out,), device=device)
 
     kernel = wp.Kernel(func=kernel_3_fn)
-    with test.assertRaisesRegex(
-        TypeError,
-        r"Invalid external reference type: <class 'warp.types.array'>",
-    ):
+    with test.assertRaisesRegex(TypeError, r"Invalid external reference type: <class 'warp.types.array'>"):
         wp.launch(kernel, dim=out.shape, inputs=(), outputs=(out,), device=device)
 
 
@@ -469,16 +450,12 @@ def test_error_collection_construct(test, device):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_3_fn)
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Construct `ast.Dict` not supported in kernels.",
-    ):
+    with test.assertRaisesRegex(RuntimeError, r"Construct `ast.Dict` not supported in kernels."):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_4_fn)
     with test.assertRaisesRegex(
-        RuntimeError,
-        r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` instead.",
+        RuntimeError, r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` instead."
     ):
         wp.launch(kernel, dim=1, device=device)
 
@@ -491,10 +468,7 @@ def test_error_unmatched_arguments(test, device):
         x = wp.dot(wp.vec2(1.0, 2.0), wp.vec2h(wp.float16(1.0), wp.float16(2.0)))
 
     kernel = wp.Kernel(func=kernel_1_fn)
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Input types must be the same, got \['int32', 'float32'\]",
-    ):
+    with test.assertRaisesRegex(RuntimeError, r"Input types must be the same, got \['int32', 'float32'\]"):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_2_fn)
@@ -704,12 +678,7 @@ add_kernel_test(
     TestCodeGen, name="test_dynamic_for_rename", kernel=test_dynamic_for_rename, inputs=[10], dim=1, devices=devices
 )
 add_kernel_test(
-    TestCodeGen,
-    name="test_dynamic_for_inplace",
-    kernel=test_dynamic_for_inplace,
-    inputs=[10],
-    dim=1,
-    devices=devices,
+    TestCodeGen, name="test_dynamic_for_inplace", kernel=test_dynamic_for_inplace, inputs=[10], dim=1, devices=devices
 )
 add_kernel_test(TestCodeGen, name="test_reassign", kernel=test_reassign, dim=1, devices=devices)
 add_kernel_test(
@@ -754,12 +723,7 @@ add_kernel_test(
 )
 
 add_kernel_test(
-    TestCodeGen,
-    name="test_range_static_sum",
-    kernel=test_range_static_sum,
-    dim=1,
-    expect=[10, 10, 10],
-    devices=devices,
+    TestCodeGen, name="test_range_static_sum", kernel=test_range_static_sum, dim=1, expect=[10, 10, 10], devices=devices
 )
 add_kernel_test(
     TestCodeGen,
@@ -789,20 +753,9 @@ add_kernel_test(
     devices=devices,
 )
 add_kernel_test(
-    TestCodeGen,
-    name="test_range_dynamic_nested",
-    kernel=test_range_dynamic_nested,
-    dim=1,
-    inputs=[4],
-    devices=devices,
-)
-add_kernel_test(
-    TestCodeGen,
-    name="test_range_expression",
-    kernel=test_range_expression,
-    dim=1,
-    devices=devices,
+    TestCodeGen, name="test_range_dynamic_nested", kernel=test_range_dynamic_nested, dim=1, inputs=[4], devices=devices
 )
+add_kernel_test(TestCodeGen, name="test_range_expression", kernel=test_range_expression, dim=1, devices=devices)
 
 add_kernel_test(TestCodeGen, name="test_while_zero", kernel=test_while, dim=1, inputs=[0], devices=devices)
 add_kernel_test(TestCodeGen, name="test_while_positive", kernel=test_while, dim=1, inputs=[16], devices=devices)

warp/tests/test_collision.py

@@ -430,8 +430,8 @@ class Example:
     def set_points_fixed(self, model, fixed_particles):
         if len(fixed_particles):
             flags = model.particle_flags.numpy()
-            for fixed_v_id in fixed_particles:
-                flags[fixed_v_id] = wp.uint32(int(flags[fixed_v_id]) & ~int(PARTICLE_FLAG_ACTIVE))
+            for fixed_vertex_id in fixed_particles:
+                flags[fixed_vertex_id] = wp.uint32(int(flags[fixed_vertex_id]) & ~int(PARTICLE_FLAG_ACTIVE))
 
             model.particle_flags = wp.array(flags, device=model.device)
 

warp/tests/test_examples.py

@@ -304,6 +304,15 @@ add_example_test(
     },
     test_options_cpu={"train_iters": 1, "num_frames": 30},
 )
+add_example_test(
+    TestOptimExamples,
+    name="optim.example_softbody_properties",
+    devices=test_devices,
+    test_options_cuda={
+        "train_iters": 1 if warp.context.runtime.core.is_debug_enabled() else 3,
+    },
+    test_options_cpu={"train_iters": 1},
+)
 
 
 class TestSimExamples(unittest.TestCase):

warp/tests/test_grad_debug.py

@@ -8,7 +8,12 @@
 import unittest
 
 import warp as wp
-from warp.autograd import gradcheck, gradcheck_tape, jacobian, jacobian_fd
+from warp.autograd import (
+    gradcheck,
+    gradcheck_tape,
+    jacobian,
+    jacobian_fd,
+)
 from warp.tests.unittest_utils import *
 
 
@@ -43,7 +48,7 @@ def vec_length_kernel(a: wp.array(dtype=wp.vec3), out: wp.array(dtype=float)):
     tid = wp.tid()
     v = a[tid]
     # instead of wp.length(v), we use a trivial implementation that
-    # fails when a division by zero is occurs in the backward pass of sqrt
+    # fails when a division by zero occurs in the backward pass of sqrt
     out[tid] = wp.sqrt(v[0] ** 2.0 + v[1] ** 2.0 + v[2] ** 2.0)
 
 
@@ -63,6 +68,16 @@ def wrong_grad_kernel(a: wp.array(dtype=float), out: wp.array(dtype=float)):
     out[tid] = wrong_grad_func(a[tid])
 
 
+@wp.kernel
+def transform_point_kernel(
+    transforms: wp.array(dtype=wp.transform),
+    points: wp.array(dtype=wp.vec3),
+    out: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    out[tid] = wp.transform_point(transforms[tid], points[tid])
+
+
 def test_gradcheck_3d(test, device):
     a_3d = wp.array([((2.0, 0.0), (1.0, 0.0), (2.0, 0.0))], dtype=float, requires_grad=True, device=device)
     b_3d = wp.array([((3.0, 0.0), (1.0, 0.0), (2.0, 0.0))], dtype=float, requires_grad=True, device=device)
@@ -229,6 +244,76 @@ def test_gradcheck_tape(test, device):
     assert passed
 
 
+def test_gradcheck_function(test, device):
+    def compute_transformed_point_norms(transforms, points):
+        tf_points = wp.empty_like(points)
+        norms = wp.empty(len(points), dtype=float, requires_grad=points.requires_grad, device=points.device)
+
+        wp.launch(
+            transform_point_kernel,
+            dim=len(points),
+            inputs=[transforms, points],
+            outputs=[tf_points],
+            device=device,
+        )
+        wp.launch(
+            vec_length_kernel,
+            dim=len(points),
+            inputs=[tf_points],
+            outputs=[norms],
+            device=device,
+        )
+        return tf_points, norms
+
+    transforms = wp.array(
+        [
+            wp.transform(wp.vec3(1.0, 0.6, -2.0), wp.quat_rpy(-0.5, 0.1, 0.8)),
+            wp.transform(wp.vec3(0.2, 1.4, -0.4), wp.quat_rpy(0.5, 0.65, -0.3)),
+            wp.transform(wp.vec3(0.5, 0.2, 0.0), wp.quat_rpy(-0.5, -0.3, 0.4)),
+        ],
+        dtype=wp.transform,
+        requires_grad=True,
+        device=device,
+    )
+    points = wp.array(
+        [
+            (1.0, -0.5, 2.0),
+            (-0.95, -0.1, 0.0),
+            (9.1, 9.7, 3.8),
+        ],
+        dtype=wp.vec3,
+        requires_grad=True,
+        device=device,
+    )
+
+    jacs_ad = jacobian(
+        kernel_mixed,
+        dim=len(points),
+        inputs=[transforms, points],
+    )
+    jacs_fd = jacobian_fd(
+        kernel_mixed,
+        dim=len(points),
+        inputs=[transforms, points],
+        eps=1e-4,
+    )
+
+    # manual gradcheck
+    for i in range(2):
+        for j in range(2):
+            assert np.allclose(jacs_ad[(i, j)].numpy(), jacs_fd[(i, j)].numpy(), atol=1e-2, rtol=1e-2)
+
+    passed = gradcheck(
+        kernel_mixed,
+        dim=len(points),
+        inputs=[transforms, points],
+        raise_exception=False,
+        show_summary=False,
+    )
+
+    assert passed
+
+
 devices = get_test_devices()
 
 

warp/tests/test_hash_grid.py

@@ -85,7 +85,7 @@ def test_hashgrid_query(test, device):
 
     for i in range(num_runs):
         if print_enabled:
-            print(f"Run: {i+1}")
+            print(f"Run: {i + 1}")
             print("---------")
 
         points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)

warp/tests/test_ipc.py

@@ -0,0 +1,116 @@
+# 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.
+
+import multiprocessing as mp
+import unittest
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+def test_ipc_get_memory_handle(test, device):
+    if device.is_ipc_supported is False:
+        test.skipTest(f"IPC is not supported on {device}")
+
+    with wp.ScopedMempool(device, False):
+        test_array = wp.full(10, value=42.0, dtype=wp.float32, device=device)
+        ipc_handle = test_array.ipc_handle()
+
+    test.assertNotEqual(ipc_handle, bytes(64), "IPC memory handle appears to be invalid")
+
+
+def test_ipc_get_event_handle(test, device):
+    if device.is_ipc_supported is False:
+        test.skipTest(f"IPC is not supported on {device}")
+
+    e1 = wp.Event(device, interprocess=True)
+
+    ipc_handle = e1.ipc_handle()
+
+    test.assertNotEqual(ipc_handle, bytes(64), "IPC event handle appears to be invalid")
+
+
+def test_ipc_event_missing_interprocess_flag(test, device):
+    if device.is_ipc_supported is False:
+        test.skipTest(f"IPC is not supported on {device}")
+
+    e1 = wp.Event(device, interprocess=False)
+
+    try:
+        capture = StdOutCapture()
+        capture.begin()
+        ipc_handle = e1.ipc_handle()
+    finally:
+        output = capture.end()
+
+    # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+    if sys.platform != "win32":
+        test.assertRegex(output, r"Warp UserWarning: IPC event handle appears to be invalid.")
+
+
+@wp.kernel
+def multiply_by_two(a: wp.array(dtype=wp.float32)):
+    i = wp.tid()
+    a[i] = 2.0 * a[i]
+
+
+def child_task(array_handle, dtype, shape, device, event_handle):
+    with wp.ScopedDevice(device):
+        ipc_array = wp.from_ipc_handle(array_handle, dtype, shape, device=device)
+        ipc_event = wp.event_from_ipc_handle(event_handle, device=device)
+        stream = wp.get_stream()
+        wp.launch(multiply_by_two, ipc_array.shape, inputs=[ipc_array])
+        stream.record_event(ipc_event)
+        stream.wait_event(ipc_event)
+        wp.synchronize_device()
+
+
+def test_ipc_multiprocess_write(test, device):
+    if device.is_ipc_supported is False:
+        test.skipTest(f"IPC is not supported on {device}")
+
+    stream = wp.get_stream(device)
+    e1 = wp.Event(device, interprocess=True)
+
+    with wp.ScopedMempool(device, False):
+        test_array = wp.full(1024, value=42.0, dtype=wp.float32, device=device)
+        ipc_handle = test_array.ipc_handle()
+
+    wp.launch(multiply_by_two, test_array.shape, inputs=[test_array], device=device)
+
+    ctx = mp.get_context("spawn")
+
+    process = ctx.Process(
+        target=child_task, args=(ipc_handle, test_array.dtype, test_array.shape, str(device), e1.ipc_handle())
+    )
+
+    process.start()
+    process.join()
+
+    assert_np_equal(test_array.numpy(), np.full(test_array.shape, 168.0, dtype=np.float32))
+
+
+cuda_devices = get_cuda_test_devices()
+
+
+class TestIpc(unittest.TestCase):
+    pass
+
+
+add_function_test(TestIpc, "test_ipc_get_memory_handle", test_ipc_get_memory_handle, devices=cuda_devices)
+add_function_test(TestIpc, "test_ipc_get_event_handle", test_ipc_get_event_handle, devices=cuda_devices)
+add_function_test(
+    TestIpc, "test_ipc_event_missing_interprocess_flag", test_ipc_event_missing_interprocess_flag, devices=cuda_devices
+)
+add_function_test(
+    TestIpc, "test_ipc_multiprocess_write", test_ipc_multiprocess_write, devices=cuda_devices, check_output=False
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)