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

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)

warp/tests/test_lerp.py

@@ -31,66 +31,14 @@ class TestData:
 
 TEST_DATA = {
     wp.float32: (
-        TestData(
-            a=1.0,
-            b=5.0,
-            t=0.75,
-            expected=4.0,
-            expected_adj_a=0.25,
-            expected_adj_b=0.75,
-            expected_adj_t=4.0,
-        ),
-        TestData(
-            a=-2.0,
-            b=5.0,
-            t=0.25,
-            expected=-0.25,
-            expected_adj_a=0.75,
-            expected_adj_b=0.25,
-            expected_adj_t=7.0,
-        ),
-        TestData(
-            a=1.23,
-            b=2.34,
-            t=0.5,
-            expected=1.785,
-            expected_adj_a=0.5,
-            expected_adj_b=0.5,
-            expected_adj_t=1.11,
-        ),
-    ),
-    wp.vec2: (
-        TestData(
-            a=[1, 2],
-            b=[3, 4],
-            t=0.5,
-            expected=[2, 3],
-        ),
-    ),
-    wp.vec3: (
-        TestData(
-            a=[1, 2, 3],
-            b=[3, 4, 5],
-            t=0.5,
-            expected=[2, 3, 4],
-        ),
-    ),
-    wp.vec4: (
-        TestData(
-            a=[1, 2, 3, 4],
-            b=[3, 4, 5, 6],
-            t=0.5,
-            expected=[2, 3, 4, 5],
-        ),
-    ),
-    wp.mat22: (
-        TestData(
-            a=[[1, 2], [2, 1]],
-            b=[[3, 4], [4, 3]],
-            t=0.5,
-            expected=[[2, 3], [3, 2]],
-        ),
+        TestData(a=1.0, b=5.0, t=0.75, expected=4.0, expected_adj_a=0.25, expected_adj_b=0.75, expected_adj_t=4.0),
+        TestData(a=-2.0, b=5.0, t=0.25, expected=-0.25, expected_adj_a=0.75, expected_adj_b=0.25, expected_adj_t=7.0),
+        TestData(a=1.23, b=2.34, t=0.5, expected=1.785, expected_adj_a=0.5, expected_adj_b=0.5, expected_adj_t=1.11),
     ),
+    wp.vec2: (TestData(a=[1, 2], b=[3, 4], t=0.5, expected=[2, 3]),),
+    wp.vec3: (TestData(a=[1, 2, 3], b=[3, 4, 5], t=0.5, expected=[2, 3, 4]),),
+    wp.vec4: (TestData(a=[1, 2, 3, 4], b=[3, 4, 5, 6], t=0.5, expected=[2, 3, 4, 5]),),
+    wp.mat22: (TestData(a=[[1, 2], [2, 1]], b=[[3, 4], [4, 3]], t=0.5, expected=[[2, 3], [3, 2]]),),
     wp.mat33: (
         TestData(
             a=[[1, 2, 3], [3, 1, 2], [2, 3, 1]],
@@ -107,30 +55,9 @@ TEST_DATA = {
             expected=[[2, 3, 4, 5], [5, 2, 3, 4], [4, 5, 2, 3], [3, 4, 5, 2]],
         ),
     ),
-    wp.quat: (
-        TestData(
-            a=[1, 2, 3, 4],
-            b=[3, 4, 5, 6],
-            t=0.5,
-            expected=[2, 3, 4, 5],
-        ),
-    ),
-    wp.transform: (
-        TestData(
-            a=[1, 2, 3, 4, 5, 6, 7],
-            b=[3, 4, 5, 6, 7, 8, 9],
-            t=0.5,
-            expected=[2, 3, 4, 5, 6, 7, 8],
-        ),
-    ),
-    wp.spatial_vector: (
-        TestData(
-            a=[1, 2, 3, 4, 5, 6],
-            b=[3, 4, 5, 6, 7, 8],
-            t=0.5,
-            expected=[2, 3, 4, 5, 6, 7],
-        ),
-    ),
+    wp.quat: (TestData(a=[1, 2, 3, 4], b=[3, 4, 5, 6], t=0.5, expected=[2, 3, 4, 5]),),
+    wp.transform: (TestData(a=[1, 2, 3, 4, 5, 6, 7], b=[3, 4, 5, 6, 7, 8, 9], t=0.5, expected=[2, 3, 4, 5, 6, 7, 8]),),
+    wp.spatial_vector: (TestData(a=[1, 2, 3, 4, 5, 6], b=[3, 4, 5, 6, 7, 8], t=0.5, expected=[2, 3, 4, 5, 6, 7]),),
     wp.spatial_matrix: (
         TestData(
             a=[
@@ -175,12 +102,12 @@ def test_lerp(test, device):
 
         return fn
 
-    for data_type in TEST_DATA:
+    for data_type, test_data_set in TEST_DATA.items():
         kernel_fn = make_kernel_fn(data_type)
         kernel = wp.Kernel(func=kernel_fn, key=f"test_lerp_{data_type.__name__}_kernel")
 
         with test.subTest(data_type=data_type):
-            for test_data in TEST_DATA[data_type]:
+            for test_data in test_data_set:
                 a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
                 b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
                 t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
@@ -188,8 +115,7 @@ def test_lerp(test, device):
                     [0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True
                 )
 
-                tape = wp.Tape()
-                with tape:
+                with wp.Tape() as tape:
                     wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)
 
                 assert_np_equal(out.numpy(), np.array([test_data.expected]), tol=1e-6)