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

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()
 
@@ -35,7 +39,20 @@ def kernel_step(state_in: State, state_out: State, model: Model):
     state_out.x[i] = state_in.x[i] + state_out.v[i] * model.dt
 
 
+@wp.kernel
+def kernel_step_with_copy(state_in: State, state_out: State, model: Model):
+    i = wp.tid()
+
+    model_rescaled = Model(1.0, model.gravity / model.m[i] * model.dt, model.m)
+
+    state_out_copy = State(state_out.x, state_out.v)
+    state_out_copy.v[i] = state_in.v[i] + model_rescaled.gravity
+    state_out_copy.x[i] = state_in.x[i] + state_out_copy.v[i] * model.dt
+
+
 def test_step(test, device):
+    rng = np.random.default_rng(123)
+
     dim = 5
 
     dt = 0.01
@@ -50,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
 
@@ -67,10 +83,11 @@ def test_step(test, device):
     state_out.x = wp.empty_like(x_in)
     state_out.v = wp.empty_like(v_in)
 
-    with CheckOutput(test):
-        wp.launch(kernel_step, dim=dim, inputs=[state_in, state_out, model], device=device)
+    for step_kernel in [kernel_step, kernel_step_with_copy]:
+        with CheckOutput(test):
+            wp.launch(step_kernel, dim=dim, inputs=[state_in, state_out, model], device=device)
 
-    assert_np_equal(state_out.x.numpy(), x_expected, tol=1e-6)
+        assert_np_equal(state_out.x.numpy(), x_expected, tol=1e-6)
 
 
 @wp.kernel
@@ -80,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)
 
@@ -95,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)
@@ -111,33 +129,34 @@ def test_step_grad(test, device):
 
     loss = wp.empty(1, dtype=float, device=device, requires_grad=True)
 
-    tape = wp.Tape()
+    for step_kernel in [kernel_step, kernel_step_with_copy]:
+        tape = wp.Tape()
 
-    with tape:
-        wp.launch(kernel_step, dim=dim, inputs=[state_in, state_out, model], device=device)
-        wp.launch(kernel_loss, dim=dim, inputs=[state_out.x, loss], device=device)
+        with tape:
+            wp.launch(step_kernel, dim=dim, inputs=[state_in, state_out, model], device=device)
+            wp.launch(kernel_loss, dim=dim, inputs=[state_out.x, loss], device=device)
 
-    tape.backward(loss)
+        tape.backward(loss)
 
-    dl_dx = 2 * state_out.x.numpy()
-    dl_dv = dl_dx * dt
+        dl_dx = 2 * state_out.x.numpy()
+        dl_dv = dl_dx * dt
 
-    dv_dm = -gravity * dt / m[:, None] ** 2
-    dl_dm = (dl_dv * dv_dm).sum(-1)
+        dv_dm = -gravity * dt / m[:, None] ** 2
+        dl_dm = (dl_dv * dv_dm).sum(-1)
 
-    assert_np_equal(state_out.x.grad.numpy(), dl_dx, tol=1e-6)
-    assert_np_equal(state_in.x.grad.numpy(), dl_dx, tol=1e-6)
-    assert_np_equal(state_out.v.grad.numpy(), dl_dv, tol=1e-6)
-    assert_np_equal(state_in.v.grad.numpy(), dl_dv, tol=1e-6)
-    assert_np_equal(model.m.grad.numpy(), dl_dm, tol=1e-6)
+        assert_np_equal(state_out.x.grad.numpy(), dl_dx, tol=1e-6)
+        assert_np_equal(state_in.x.grad.numpy(), dl_dx, tol=1e-6)
+        assert_np_equal(state_out.v.grad.numpy(), dl_dv, tol=1e-6)
+        assert_np_equal(state_in.v.grad.numpy(), dl_dv, tol=1e-6)
+        assert_np_equal(model.m.grad.numpy(), dl_dm, tol=1e-6)
 
-    tape.zero()
+        tape.zero()
 
-    assert state_out.x.grad.numpy().sum() == 0.0
-    assert state_in.x.grad.numpy().sum() == 0.0
-    assert state_out.v.grad.numpy().sum() == 0.0
-    assert state_in.v.grad.numpy().sum() == 0.0
-    assert model.m.grad.numpy().sum() == 0.0
+        assert state_out.x.grad.numpy().sum() == 0.0
+        assert state_in.x.grad.numpy().sum() == 0.0
+        assert state_out.v.grad.numpy().sum() == 0.0
+        assert state_in.v.grad.numpy().sum() == 0.0
+        assert model.m.grad.numpy().sum() == 0.0
 
 
 @wp.struct
@@ -205,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))
@@ -259,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
@@ -346,14 +379,12 @@ def test_struct_default_attributes_kernel():
 
 @wp.struct
 class MutableStruct:
-
     param1: int
     param2: float
 
 
 @wp.kernel
 def test_struct_mutate_attributes_kernel():
-
     t = MutableStruct()
     t.param1 = 1
     t.param2 = 1.1
@@ -362,68 +393,283 @@ def test_struct_mutate_attributes_kernel():
     wp.expect_eq(t.param2, 1.1)
 
 
-def register(parent):
-    devices = get_test_devices()
+@wp.struct
+class InnerStruct:
+    i: 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_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 ArrayStruct:
+    array: wp.array(dtype=InnerStruct)
+
+
+@wp.kernel
+def struct2_reader(test: ArrayStruct):
+    k = wp.tid()
+    wp.expect_eq(k + 1, test.array[k].i)
+
+
+def test_nested_array_struct(test, device):
+    var1 = InnerStruct()
+    var1.i = 1
+
+    var2 = InnerStruct()
+    var2.i = 2
+
+    struct = ArrayStruct()
+    struct.array = wp.array([var1, var2], dtype=InnerStruct, device=device)
+
+    wp.launch(struct2_reader, dim=2, inputs=[struct], device=device)
+
+
+@wp.struct
+class EmptyNest1:
+    a: Empty
+    z: int
+
+
+@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)

warp/tests/test_tape.py

@@ -5,9 +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 *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -19,11 +22,17 @@ def mul_constant(x: wp.array(dtype=float), y: wp.array(dtype=float)):
     y[tid] = x[tid] * 2.0
 
 
+@wp.struct
+class Multiplicands:
+    x: wp.array(dtype=float)
+    y: wp.array(dtype=float)
+
+
 @wp.kernel
-def mul_variable(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+def mul_variable(mutiplicands: Multiplicands, z: wp.array(dtype=float)):
     tid = wp.tid()
 
-    z[tid] = x[tid] * y[tid]
+    z[tid] = mutiplicands.x[tid] * mutiplicands.y[tid]
 
 
 @wp.kernel
@@ -65,12 +74,13 @@ def test_tape_mul_variable(test, device):
 
     # record onto tape
     with tape:
-        # input data
-        x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
-        y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
-        z = wp.zeros_like(x)
+        # input data (Note: We're intentionally testing structs in tapes here)
+        multiplicands = Multiplicands()
+        multiplicands.x = wp.array(np.ones(dim) * 16.0, dtype=wp.float32, device=device, requires_grad=True)
+        multiplicands.y = wp.array(np.ones(dim) * 32.0, dtype=wp.float32, device=device, requires_grad=True)
+        z = wp.zeros_like(multiplicands.x)
 
-        wp.launch(kernel=mul_variable, dim=dim, inputs=[x, y], outputs=[z], device=device)
+        wp.launch(kernel=mul_variable, dim=dim, inputs=[multiplicands], outputs=[z], device=device)
 
     # loss = wp.sum(x)
     z.grad = wp.array(np.ones(dim), device=device, dtype=wp.float32)
@@ -79,16 +89,21 @@ def test_tape_mul_variable(test, device):
     tape.backward()
 
     # grad_x=y, grad_y=x
-    assert_np_equal(tape.gradients[x].numpy(), y.numpy())
-    assert_np_equal(tape.gradients[y].numpy(), x.numpy())
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), multiplicands.y.numpy())
+    assert_np_equal(tape.gradients[multiplicands].y.numpy(), multiplicands.x.numpy())
 
     # run backward again with different incoming gradient
     # should accumulate the same gradients again onto output
     # so gradients = 2.0*prev
     tape.backward()
 
-    assert_np_equal(tape.gradients[x].numpy(), y.numpy() * 2.0)
-    assert_np_equal(tape.gradients[y].numpy(), x.numpy() * 2.0)
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), multiplicands.y.numpy() * 2.0)
+    assert_np_equal(tape.gradients[multiplicands].y.numpy(), multiplicands.x.numpy() * 2.0)
+
+    # Clear launches and zero out the gradients
+    tape.reset()
+    assert_np_equal(tape.gradients[multiplicands].x.numpy(), np.zeros_like(tape.gradients[multiplicands].x.numpy()))
+    test.assertFalse(tape.launches)
 
 
 def test_tape_dot_product(test, device):
@@ -112,19 +127,22 @@ def test_tape_dot_product(test, device):
     assert_np_equal(tape.gradients[y].numpy(), x.numpy())
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestTape(parent):
-        pass
+class TestTape(unittest.TestCase):
+    def test_tape_no_nested_tapes(self):
+        with self.assertRaises(RuntimeError):
+            with wp.Tape():
+                with wp.Tape():
+                    pass
 
-    add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, devices=devices)
-    add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
-    add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
 
-    return TestTape
+add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, devices=devices)
+add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
+add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)