warp-lang 1.5.1__py3-none-manylinux2014_x86_64.whl → 1.6.1__py3-none-manylinux2014_x86_64.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_launch.py

@@ -46,6 +46,12 @@ def kernel4d(a: wp.array(dtype=int, ndim=4)):
     wp.expect_eq(a[i, j, k, l], i * dim_y * dim_z * dim_w + j * dim_z * dim_w + k * dim_w + l)
 
 
+@wp.kernel
+def square_kernel(input: wp.array(dtype=float), output: wp.array(dtype=float)):
+    i = wp.tid()
+    output[i] = input[i] * input[i]
+
+
 def test1d(test, device):
     a = np.arange(0, dim_x).reshape(dim_x)
 
@@ -98,8 +104,19 @@ def kernel_cmd(params: Params, i: int, f: float, v: wp.vec3, m: wp.mat33, out: w
 
 
 def test_launch_cmd(test, device):
+    """Tests recording and executing a kernel launch command.
+
+    Verifies that:
+    - A kernel can be recorded as a command without immediate execution
+    - The recorded command can be launched later
+    - Parameters are correctly passed to the kernel
+    - Output matches expected results for both immediate and delayed launches
+
+    Args:
+        test: Test context
+        device: Device to run the test on
+    """
     n = 1
-
     ref = np.arange(0, n)
     out = wp.zeros(n, dtype=int, device=device)
 
@@ -274,12 +291,62 @@ def test_launch_cmd_empty(test, device):
     assert_np_equal(out.numpy(), ref)
 
 
+def test_launch_cmd_adjoint(test, device):
+    """Test recording an adjoint launch with record_cmd=True."""
+    input_arr = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=True, device=device)
+    output_arr = wp.empty_like(input_arr)
+
+    output_arr.grad.fill_(1.0)
+
+    cmd = wp.launch(
+        square_kernel,
+        dim=input_arr.size,
+        inputs=[input_arr, output_arr],
+        adj_inputs=[None, None],
+        adjoint=True,
+        device=device,
+        record_cmd=True,
+    )
+
+    cmd.launch()
+
+    assert_np_equal(input_arr.grad.numpy(), np.array([2.0, 4.0, 6.0]))
+
+
+def test_launch_cmd_adjoint_empty(test, device):
+    """Test constructing a Launch object for an adjoint kernel."""
+    input_arr = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=True, device=device)
+    output_arr = wp.empty_like(input_arr)
+    output_arr.grad.fill_(1.0)
+
+    cmd = wp.Launch(square_kernel, device, adjoint=True)
+    cmd.set_param_by_name("input", input_arr)
+    cmd.set_param_by_name("output", output_arr)
+    cmd.set_dim(input_arr.size)
+    cmd.launch()
+
+    assert_np_equal(input_arr.grad.numpy(), np.array([2.0, 4.0, 6.0]))
+
+    # Now update the launch object's parameters with arrays of different sizes and values
+    # and check that the adjoints are correctly computed
+    input_arr_updated = wp.array([4.0, 5.0, 6.0, 7.0], dtype=float, device=device)
+    input_arr_updated_grad = wp.zeros_like(input_arr_updated)
+
+    output_arr_updated = wp.empty_like(input_arr_updated)
+    output_arr_updated_grad = wp.full_like(output_arr_updated, 1.0)
+
+    cmd.set_param_by_name("input", input_arr_updated)
+    cmd.set_param_by_name("output", output_arr_updated)
+    cmd.set_param_by_name("input", input_arr_updated_grad, adjoint=True)
+    cmd.set_param_by_name("output", output_arr_updated_grad, adjoint=True)
+    cmd.set_dim(input_arr_updated.size)
+    cmd.launch()
+
+    assert_np_equal(input_arr_updated_grad.numpy(), np.array([8.0, 10.0, 12.0, 14.0]))
+
+
 @wp.kernel
-def kernel_mul(
-    values: wp.array(dtype=int),
-    coeff: int,
-    out: wp.array(dtype=int),
-):
+def kernel_mul(values: wp.array(dtype=int), coeff: int, out: wp.array(dtype=int)):
     tid = wp.tid()
     out[tid] = values[tid] * coeff
 
@@ -301,28 +368,10 @@ def test_launch_tuple_args(test, device):
     )
     assert_np_equal(out.numpy(), np.array((0, 3, 6, 9)))
 
-    wp.launch(
-        kernel_mul,
-        dim=len(values),
-        inputs=(
-            values,
-            coeff,
-            out,
-        ),
-        device=device,
-    )
+    wp.launch(kernel_mul, dim=len(values), inputs=(values, coeff, out), device=device)
     assert_np_equal(out.numpy(), np.array((0, 3, 6, 9)))
 
-    wp.launch(
-        kernel_mul,
-        dim=len(values),
-        outputs=(
-            values,
-            coeff,
-            out,
-        ),
-        device=device,
-    )
+    wp.launch(kernel_mul, dim=len(values), outputs=(values, coeff, out), device=device)
     assert_np_equal(out.numpy(), np.array((0, 3, 6, 9)))
 
 
@@ -343,6 +392,8 @@ add_function_test(TestLaunch, "test_launch_cmd_set_param", test_launch_cmd_set_p
 add_function_test(TestLaunch, "test_launch_cmd_set_ctype", test_launch_cmd_set_ctype, devices=devices)
 add_function_test(TestLaunch, "test_launch_cmd_set_dim", test_launch_cmd_set_dim, devices=devices)
 add_function_test(TestLaunch, "test_launch_cmd_empty", test_launch_cmd_empty, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_adjoint", test_launch_cmd_adjoint, devices=devices)
+add_function_test(TestLaunch, "test_launch_cmd_adjoint_empty", test_launch_cmd_adjoint_empty, devices=devices)
 
 add_function_test(TestLaunch, "test_launch_tuple_args", test_launch_tuple_args, devices=devices)