warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_streams.py

@@ -5,11 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
+
 import numpy as np
-import warp as wp
-from warp.tests.test_base import *
 
-import unittest
+import warp as wp
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -45,6 +46,11 @@ def test_stream_arg_implicit_sync(test, device):
 
     new_stream = wp.Stream(device)
 
+    # Exercise code path
+    wp.set_stream(new_stream, device)
+
+    test.assertTrue(wp.get_device(device).has_stream)
+
     # launch work on new stream
     wp.launch(inc, dim=a.size, inputs=[a], stream=new_stream)
     wp.copy(b, a, stream=new_stream)
@@ -278,119 +284,138 @@ def test_stream_scope_wait_stream(test, device):
         assert_np_equal(d.numpy(), np.full(N, fill_value=4.0))
 
 
-def test_stream_arg_graph_mgpu(test, device):
-    # resources on GPU 0
-    stream0 = wp.get_stream("cuda:0")
-    a0 = wp.zeros(N, dtype=float, device="cuda:0")
-    b0 = wp.empty(N, dtype=float, device="cuda:0")
-    c0 = wp.empty(N, dtype=float, device="cuda:0")
+devices = get_unique_cuda_test_devices()
 
-    # resources on GPU 1
-    stream1 = wp.get_stream("cuda:1")
-    a1 = wp.zeros(N, dtype=float, device="cuda:1")
 
-    # start recording on stream0
-    wp.capture_begin(stream=stream0)
+class TestStreams(unittest.TestCase):
+    def test_stream_exceptions(self):
+        cpu_device = wp.get_device("cpu")
 
-    # branch into stream1
-    stream1.wait_stream(stream0)
+        # Can't set the stream on a CPU device
+        with self.assertRaises(RuntimeError):
+            stream0 = wp.Stream()
+            cpu_device.stream = stream0
 
-    # launch concurrent kernels on each stream
-    wp.launch(inc, dim=N, inputs=[a0], stream=stream0)
-    wp.launch(inc, dim=N, inputs=[a1], stream=stream1)
+        # Can't create a stream on the CPU
+        with self.assertRaises(RuntimeError):
+            wp.Stream(device="cpu")
 
-    # wait for stream1 to finish
-    stream0.wait_stream(stream1)
+        # Can't create an event with CPU device
+        with self.assertRaises(RuntimeError):
+            wp.Event(device=cpu_device)
 
-    # copy values from stream1
-    wp.copy(b0, a1, stream=stream0)
+        # Can't get the stream on a CPU device
+        with self.assertRaises(RuntimeError):
+            cpu_stream = cpu_device.stream  # noqa: F841
 
-    # compute sum
-    wp.launch(sum, dim=N, inputs=[a0, b0, c0], stream=stream0)
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_stream_arg_graph_mgpu(self):
+        wp.load_module(device="cuda:0")
+        wp.load_module(device="cuda:1")
 
-    # finish recording on stream0
-    g = wp.capture_end(stream=stream0)
+        # resources on GPU 0
+        stream0 = wp.get_stream("cuda:0")
+        a0 = wp.zeros(N, dtype=float, device="cuda:0")
+        b0 = wp.empty(N, dtype=float, device="cuda:0")
+        c0 = wp.empty(N, dtype=float, device="cuda:0")
 
-    # replay
-    num_iters = 10
-    for _ in range(num_iters):
-        wp.capture_launch(g, stream=stream0)
+        # resources on GPU 1
+        stream1 = wp.get_stream("cuda:1")
+        a1 = wp.zeros(N, dtype=float, device="cuda:1")
 
-    # check results
-    assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
+        # start recording on stream0
+        wp.capture_begin(stream=stream0, force_module_load=False)
+        try:
+            # branch into stream1
+            stream1.wait_stream(stream0)
 
+            # launch concurrent kernels on each stream
+            wp.launch(inc, dim=N, inputs=[a0], stream=stream0)
+            wp.launch(inc, dim=N, inputs=[a1], stream=stream1)
 
-def test_stream_scope_graph_mgpu(test, device):
-    # resources on GPU 0
-    with wp.ScopedDevice("cuda:0"):
-        stream0 = wp.get_stream()
-        a0 = wp.zeros(N, dtype=float)
-        b0 = wp.empty(N, dtype=float)
-        c0 = wp.empty(N, dtype=float)
+            # wait for stream1 to finish
+            stream0.wait_stream(stream1)
 
-    # resources on GPU 1
-    with wp.ScopedDevice("cuda:1"):
-        stream1 = wp.get_stream()
-        a1 = wp.zeros(N, dtype=float)
+            # copy values from stream1
+            wp.copy(b0, a1, stream=stream0)
 
-    # capture graph
-    with wp.ScopedDevice("cuda:0"):
-        # start recording
-        wp.capture_begin()
+            # compute sum
+            wp.launch(sum, dim=N, inputs=[a0, b0, c0], stream=stream0)
+        finally:
+            # finish recording on stream0
+            g = wp.capture_end(stream=stream0)
 
-        with wp.ScopedDevice("cuda:1"):
-            # branch into stream1
-            wp.wait_stream(stream0)
+        # replay
+        num_iters = 10
+        for _ in range(num_iters):
+            wp.capture_launch(g, stream=stream0)
 
-            wp.launch(inc, dim=N, inputs=[a1])
+        # check results
+        assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
-        wp.launch(inc, dim=N, inputs=[a0])
+    @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    def test_stream_scope_graph_mgpu(self):
+        wp.load_module(device="cuda:0")
+        wp.load_module(device="cuda:1")
 
-        # wait for stream1 to finish
-        wp.wait_stream(stream1)
+        # resources on GPU 0
+        with wp.ScopedDevice("cuda:0"):
+            stream0 = wp.get_stream()
+            a0 = wp.zeros(N, dtype=float)
+            b0 = wp.empty(N, dtype=float)
+            c0 = wp.empty(N, dtype=float)
 
-        # copy values from stream1
-        wp.copy(b0, a1)
+        # resources on GPU 1
+        with wp.ScopedDevice("cuda:1"):
+            stream1 = wp.get_stream()
+            a1 = wp.zeros(N, dtype=float)
 
-        # compute sum
-        wp.launch(sum, dim=N, inputs=[a0, b0, c0])
+        # capture graph
+        with wp.ScopedDevice("cuda:0"):
+            # start recording
+            wp.capture_begin(force_module_load=False)
+            try:
+                with wp.ScopedDevice("cuda:1"):
+                    # branch into stream1
+                    wp.wait_stream(stream0)
 
-        # finish recording
-        g = wp.capture_end()
+                    wp.launch(inc, dim=N, inputs=[a1])
 
-    # replay
-    with wp.ScopedDevice("cuda:0"):
-        num_iters = 10
-        for _ in range(num_iters):
-            wp.capture_launch(g)
+                wp.launch(inc, dim=N, inputs=[a0])
 
-    # check results
-    assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
+                # wait for stream1 to finish
+                wp.wait_stream(stream1)
 
+                # copy values from stream1
+                wp.copy(b0, a1)
 
-def register(parent):
-    devices = wp.get_cuda_devices()
+                # compute sum
+                wp.launch(sum, dim=N, inputs=[a0, b0, c0])
+            finally:
+                # finish recording
+                g = wp.capture_end()
 
-    class TestStreams(parent):
-        pass
+        # replay
+        with wp.ScopedDevice("cuda:0"):
+            num_iters = 10
+            for _ in range(num_iters):
+                wp.capture_launch(g)
 
-    add_function_test(TestStreams, "test_stream_arg_implicit_sync", test_stream_arg_implicit_sync, devices=devices)
-    add_function_test(TestStreams, "test_stream_scope_implicit_sync", test_stream_scope_implicit_sync, devices=devices)
+        # check results
+        assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
-    add_function_test(TestStreams, "test_stream_arg_synchronize", test_stream_arg_synchronize, devices=devices)
-    add_function_test(TestStreams, "test_stream_arg_wait_event", test_stream_arg_wait_event, devices=devices)
-    add_function_test(TestStreams, "test_stream_arg_wait_stream", test_stream_arg_wait_stream, devices=devices)
-    add_function_test(TestStreams, "test_stream_scope_synchronize", test_stream_scope_synchronize, devices=devices)
-    add_function_test(TestStreams, "test_stream_scope_wait_event", test_stream_scope_wait_event, devices=devices)
-    add_function_test(TestStreams, "test_stream_scope_wait_stream", test_stream_scope_wait_stream, devices=devices)
 
-    if len(devices) > 1:
-        add_function_test(TestStreams, "test_stream_arg_graph_mgpu", test_stream_arg_graph_mgpu)
-        add_function_test(TestStreams, "test_stream_scope_graph_mgpu", test_stream_scope_graph_mgpu)
+add_function_test(TestStreams, "test_stream_arg_implicit_sync", test_stream_arg_implicit_sync, devices=devices)
+add_function_test(TestStreams, "test_stream_scope_implicit_sync", test_stream_scope_implicit_sync, devices=devices)
 
-    return TestStreams
+add_function_test(TestStreams, "test_stream_arg_synchronize", test_stream_arg_synchronize, devices=devices)
+add_function_test(TestStreams, "test_stream_arg_wait_event", test_stream_arg_wait_event, devices=devices)
+add_function_test(TestStreams, "test_stream_arg_wait_stream", test_stream_arg_wait_stream, devices=devices)
+add_function_test(TestStreams, "test_stream_scope_synchronize", test_stream_scope_synchronize, devices=devices)
+add_function_test(TestStreams, "test_stream_scope_wait_event", test_stream_scope_wait_event, devices=devices)
+add_function_test(TestStreams, "test_stream_scope_wait_stream", test_stream_scope_wait_stream, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)