warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__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)

warp/tests/test_struct.py

@@ -5,11 +5,15 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import unittest
+from typing import Any
+
 import numpy as np
+
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
-import unittest
+from warp.fem import Sample as StructFromAnotherModule
 
 wp.init()
 
@@ -47,6 +51,8 @@ def kernel_step_with_copy(state_in: State, state_out: State, model: Model):
 
 
 def test_step(test, device):
+    rng = np.random.default_rng(123)
+
     dim = 5
 
     dt = 0.01
@@ -61,9 +67,8 @@ def test_step(test, device):
     model.dt = dt
     model.gravity = wp.vec3(0, 0, -9.81)
 
-    np.random.seed(0)
-    x = np.random.normal(size=(dim, 3))
-    v = np.random.normal(size=(dim, 3))
+    x = rng.normal(size=(dim, 3))
+    v = rng.normal(size=(dim, 3))
 
     x_expected = x + (v + gravity / m[:, None] * dt) * dt
 
@@ -92,13 +97,14 @@ def kernel_loss(x: wp.array(dtype=wp.vec3), loss: wp.array(dtype=float)):
 
 
 def test_step_grad(test, device):
+    rng = np.random.default_rng(123)
+
     dim = 5
 
     dt = 0.01
     gravity = np.array([0, 0, -9.81])
 
-    np.random.seed(0)
-    m = np.random.rand(dim) + 0.1
+    m = rng.random(size=dim) + 0.1
 
     m_model = wp.array(m, dtype=float, device=device, requires_grad=True)
 
@@ -107,8 +113,8 @@ def test_step_grad(test, device):
     model.dt = dt
     model.gravity = wp.vec3(0, 0, -9.81)
 
-    x = np.random.normal(size=(dim, 3))
-    v = np.random.normal(size=(dim, 3))
+    x = rng.normal(size=(dim, 3))
+    v = rng.normal(size=(dim, 3))
 
     x_in = wp.array(x, dtype=wp.vec3, device=device, requires_grad=True)
     v_in = wp.array(v, dtype=wp.vec3, device=device, requires_grad=True)
@@ -218,6 +224,20 @@ def test_nested_struct(test, device):
     )
 
 
+def test_struct_attribute_error(test, device):
+    @wp.kernel
+    def kernel(foo: Foo):
+        _ = foo.nonexisting
+
+    with test.assertRaisesRegex(AttributeError, r"`nonexisting` is not an attribute of 'foo' \([\w.]+\.Foo\)$"):
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[Foo()],
+            device=device,
+        )
+
+
 @wp.kernel
 def test_struct_instantiate(data: wp.array(dtype=int)):
     baz = Baz(data, wp.vec3(0.0, 0.0, 26.0))
@@ -272,7 +292,7 @@ def test_struct_math_conversions(test, device):
     s.m5 = [10, 20, 30, 40]
     s.m6 = np.array([100, 200, 300, 400])
 
-    wp.launch(check_math_conversions, dim=1, inputs=[s])
+    wp.launch(check_math_conversions, dim=1, inputs=[s], device=device)
 
 
 @wp.struct
@@ -397,74 +417,259 @@ def test_nested_array_struct(test, device):
     var2.i = 2
 
     struct = ArrayStruct()
-    struct.array = wp.array([var1, var2], dtype=InnerStruct)
+    struct.array = wp.array([var1, var2], dtype=InnerStruct, device=device)
 
-    wp.launch(struct2_reader, dim=2, inputs=[struct])
+    wp.launch(struct2_reader, dim=2, inputs=[struct], device=device)
 
 
-def register(parent):
-    devices = get_test_devices()
+@wp.struct
+class EmptyNest1:
+    a: Empty
+    z: int
 
-    class TestStruct(parent):
-        pass
 
-    add_function_test(TestStruct, "test_step", test_step, devices=devices)
-    add_function_test(TestStruct, "test_step_grad", test_step_grad, devices=devices)
-    add_kernel_test(TestStruct, kernel=test_empty, name="test_empty", dim=1, inputs=[Empty()], devices=devices)
-    add_kernel_test(
-        TestStruct,
-        kernel=test_uninitialized,
-        name="test_uninitialized",
-        dim=1,
-        inputs=[Uninitialized()],
-        devices=devices,
-    )
-    add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
-    add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
-    add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, devices=devices)
-    add_function_test(TestStruct, "test_struct_math_conversions", test_struct_math_conversions, devices=devices)
-    add_function_test(
-        TestStruct, "test_struct_default_attributes_python", test_struct_default_attributes_python, devices=devices
-    )
+@wp.struct
+class EmptyNest2:
+    a: Empty
+    b: Empty
+    z: int
+
+
+@wp.struct
+class EmptyNest3:
+    a: Empty
+    b: Empty
+    c: Empty
+    z: int
+
+
+@wp.struct
+class EmptyNest4:
+    a: Empty
+    b: Empty
+    c: Empty
+    d: Empty
+    z: int
+
+
+@wp.struct
+class EmptyNest5:
+    a: Empty
+    b: Empty
+    c: Empty
+    d: Empty
+    e: Empty
+    z: int
+
+
+@wp.struct
+class EmptyNest6:
+    a: Empty
+    b: Empty
+    c: Empty
+    d: Empty
+    e: Empty
+    f: Empty
+    z: int
+
+
+@wp.struct
+class EmptyNest7:
+    a: Empty
+    b: Empty
+    c: Empty
+    d: Empty
+    e: Empty
+    f: Empty
+    g: Empty
+    z: int
+
+
+@wp.struct
+class EmptyNest8:
+    a: Empty
+    b: Empty
+    c: Empty
+    d: Empty
+    e: Empty
+    f: Empty
+    g: Empty
+    h: Empty
+    z: int
+
+
+@wp.kernel
+def empty_nest_kernel(s: Any):
+    wp.expect_eq(s.z, 42)
+
+
+wp.overload(empty_nest_kernel, [EmptyNest1])
+wp.overload(empty_nest_kernel, [EmptyNest2])
+wp.overload(empty_nest_kernel, [EmptyNest3])
+wp.overload(empty_nest_kernel, [EmptyNest4])
+wp.overload(empty_nest_kernel, [EmptyNest5])
+wp.overload(empty_nest_kernel, [EmptyNest6])
+wp.overload(empty_nest_kernel, [EmptyNest7])
+wp.overload(empty_nest_kernel, [EmptyNest8])
+
+
+def test_nested_empty_struct(test, device):
+    with wp.ScopedDevice(device):
+        e1 = EmptyNest1()
+        e1.z = 42
+        e2 = EmptyNest2()
+        e2.z = 42
+        e3 = EmptyNest3()
+        e3.z = 42
+        e4 = EmptyNest4()
+        e4.z = 42
+        e5 = EmptyNest5()
+        e5.z = 42
+        e6 = EmptyNest6()
+        e6.z = 42
+        e7 = EmptyNest7()
+        e7.z = 42
+        e8 = EmptyNest8()
+        e8.z = 42
+
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e1])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e2])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e3])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e4])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e5])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e6])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e7])
+        wp.launch(empty_nest_kernel, dim=1, inputs=[e8])
+
+        wp.synchronize_device()
+
+
+@wp.struct
+class DependentModuleImport_A:
+    s: StructFromAnotherModule
+
+
+@wp.struct
+class DependentModuleImport_B:
+    s: StructFromAnotherModule
+
+
+@wp.struct
+class DependentModuleImport_C:
+    a: DependentModuleImport_A
+    b: DependentModuleImport_B
+
+
+@wp.kernel
+def test_dependent_module_import(c: DependentModuleImport_C):
+    wp.tid()  # nop, we're just testing codegen
+
+
+devices = get_test_devices()
+
+
+class TestStruct(unittest.TestCase):
+    pass
+
+
+add_function_test(TestStruct, "test_step", test_step, devices=devices)
+add_function_test(TestStruct, "test_step_grad", test_step_grad, devices=devices)
+add_kernel_test(TestStruct, kernel=test_empty, name="test_empty", dim=1, inputs=[Empty()], devices=devices)
+add_kernel_test(
+    TestStruct,
+    kernel=test_uninitialized,
+    name="test_uninitialized",
+    dim=1,
+    inputs=[Uninitialized()],
+    devices=devices,
+)
+add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
+add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_empty_struct", test_nested_empty_struct, devices=devices)
+add_function_test(TestStruct, "test_struct_math_conversions", test_struct_math_conversions, devices=devices)
+add_function_test(
+    TestStruct, "test_struct_default_attributes_python", test_struct_default_attributes_python, devices=devices
+)
+add_kernel_test(
+    TestStruct,
+    name="test_struct_default_attributes",
+    kernel=test_struct_default_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+
+add_kernel_test(
+    TestStruct,
+    name="test_struct_mutate_attributes",
+    kernel=test_struct_mutate_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+add_kernel_test(
+    TestStruct,
+    kernel=test_uninitialized,
+    name="test_uninitialized",
+    dim=1,
+    inputs=[Uninitialized()],
+    devices=devices,
+)
+add_kernel_test(TestStruct, kernel=test_return, name="test_return", dim=1, inputs=[], devices=devices)
+add_function_test(TestStruct, "test_nested_struct", test_nested_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_array_struct", test_nested_array_struct, devices=devices)
+add_function_test(TestStruct, "test_nested_empty_struct", test_nested_empty_struct, devices=devices)
+add_function_test(TestStruct, "test_struct_math_conversions", test_struct_math_conversions, devices=devices)
+add_function_test(
+    TestStruct, "test_struct_default_attributes_python", test_struct_default_attributes_python, devices=devices
+)
+add_kernel_test(
+    TestStruct,
+    name="test_struct_default_attributes",
+    kernel=test_struct_default_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+
+add_kernel_test(
+    TestStruct,
+    name="test_struct_mutate_attributes",
+    kernel=test_struct_mutate_attributes_kernel,
+    dim=1,
+    inputs=[],
+    devices=devices,
+)
+
+for device in devices:
     add_kernel_test(
         TestStruct,
-        name="test_struct_default_attributes",
-        kernel=test_struct_default_attributes_kernel,
+        kernel=test_struct_instantiate,
+        name="test_struct_instantiate",
         dim=1,
-        inputs=[],
-        devices=devices,
+        inputs=[wp.array([1], dtype=int, device=device)],
+        devices=[device],
     )
-
     add_kernel_test(
         TestStruct,
-        name="test_struct_mutate_attributes",
-        kernel=test_struct_mutate_attributes_kernel,
+        kernel=test_return_struct,
+        name="test_return_struct",
         dim=1,
-        inputs=[],
-        devices=devices,
+        inputs=[wp.zeros(10, dtype=int, device=device)],
+        devices=[device],
     )
 
-    for device in devices:
-        add_kernel_test(
-            TestStruct,
-            kernel=test_struct_instantiate,
-            name="test_struct_instantiate",
-            dim=1,
-            inputs=[wp.array([1], dtype=int, device=device)],
-            devices=[device],
-        )
-        add_kernel_test(
-            TestStruct,
-            kernel=test_return_struct,
-            name="test_return_struct",
-            dim=1,
-            inputs=[wp.zeros(10, dtype=int, device=device)],
-            devices=[device],
-        )
-
-    return TestStruct
+add_kernel_test(
+    TestStruct,
+    kernel=test_dependent_module_import,
+    name="test_dependent_module_import",
+    dim=1,
+    inputs=[DependentModuleImport_C()],
+    devices=devices,
+)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)