warp-lang 1.5.1__py3-none-manylinux2014_x86_64.whl → 1.6.1__py3-none-manylinux2014_x86_64.whl

warp/tests/test_mesh_query_ray.py

@@ -91,114 +91,120 @@ def test_mesh_query_ray_grad(test, device):
     mesh_points = wp.array(np.array(mesh_geom.GetPointsAttr().Get()), dtype=wp.vec3, device=device)
     mesh_indices = wp.array(np.array(tri_indices), dtype=int, device=device)
 
-    p = wp.vec3(50.0, 50.0, 0.0)
-    D = wp.vec3(0.0, -1.0, 0.0)
-
-    # create mesh
-    mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices)
-
-    tape = wp.Tape()
-
-    # analytic gradients
-    with tape:
-        query_points = wp.array(p, dtype=wp.vec3, device=device, requires_grad=True)
-        query_dirs = wp.array(D, dtype=wp.vec3, device=device, requires_grad=True)
-        intersection_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
-        loss = wp.zeros(n=1, dtype=float, device=device, requires_grad=True)
-
-        wp.launch(
-            kernel=mesh_query_ray_loss,
-            dim=1,
-            inputs=[mesh.id, query_points, query_dirs, intersection_points, loss],
-            device=device,
-        )
-
-    tape.backward(loss=loss)
-    q = intersection_points.numpy().flatten()
-    analytic_p = tape.gradients[query_points].numpy().flatten()
-    analytic_D = tape.gradients[query_dirs].numpy().flatten()
-
-    # numeric gradients
-
-    # ray origin
-    eps = 1.0e-3
-    loss_values_p = []
-    numeric_p = np.zeros(3)
-
-    offset_query_points = [
-        wp.vec3(p[0] - eps, p[1], p[2]),
-        wp.vec3(p[0] + eps, p[1], p[2]),
-        wp.vec3(p[0], p[1] - eps, p[2]),
-        wp.vec3(p[0], p[1] + eps, p[2]),
-        wp.vec3(p[0], p[1], p[2] - eps),
-        wp.vec3(p[0], p[1], p[2] + eps),
-    ]
-
-    for i in range(6):
-        q = offset_query_points[i]
-
-        query_points = wp.array(q, dtype=wp.vec3, device=device)
-        query_dirs = wp.array(D, dtype=wp.vec3, device=device)
-        intersection_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
-        loss = wp.zeros(n=1, dtype=float, device=device)
-
-        wp.launch(
-            kernel=mesh_query_ray_loss,
-            dim=1,
-            inputs=[mesh.id, query_points, query_dirs, intersection_points, loss],
-            device=device,
-        )
-
-        loss_values_p.append(loss.numpy()[0])
-
-    for i in range(3):
-        l_0 = loss_values_p[i * 2]
-        l_1 = loss_values_p[i * 2 + 1]
-        gradient = (l_1 - l_0) / (2.0 * eps)
-        numeric_p[i] = gradient
-
-    # ray dir
-    loss_values_D = []
-    numeric_D = np.zeros(3)
-
-    offset_query_dirs = [
-        wp.vec3(D[0] - eps, D[1], D[2]),
-        wp.vec3(D[0] + eps, D[1], D[2]),
-        wp.vec3(D[0], D[1] - eps, D[2]),
-        wp.vec3(D[0], D[1] + eps, D[2]),
-        wp.vec3(D[0], D[1], D[2] - eps),
-        wp.vec3(D[0], D[1], D[2] + eps),
-    ]
-
-    for i in range(6):
-        q = offset_query_dirs[i]
-
-        query_points = wp.array(p, dtype=wp.vec3, device=device)
-        query_dirs = wp.array(q, dtype=wp.vec3, device=device)
-        intersection_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
-        loss = wp.zeros(n=1, dtype=float, device=device)
-
-        wp.launch(
-            kernel=mesh_query_ray_loss,
-            dim=1,
-            inputs=[mesh.id, query_points, query_dirs, intersection_points, loss],
-            device=device,
-        )
-
-        loss_values_D.append(loss.numpy()[0])
-
-    for i in range(3):
-        l_0 = loss_values_D[i * 2]
-        l_1 = loss_values_D[i * 2 + 1]
-        gradient = (l_1 - l_0) / (2.0 * eps)
-        numeric_D[i] = gradient
-
-    error_p = ((analytic_p - numeric_p) * (analytic_p - numeric_p)).sum(axis=0)
-    error_D = ((analytic_D - numeric_D) * (analytic_D - numeric_D)).sum(axis=0)
-
-    tolerance = 1.0e-3
-    test.assertTrue(error_p < tolerance, f"error is {error_p} which is >= {tolerance}")
-    test.assertTrue(error_D < tolerance, f"error is {error_D} which is >= {tolerance}")
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    for constructor in constructors:
+        p = wp.vec3(50.0, 50.0, 0.0)
+        D = wp.vec3(0.0, -1.0, 0.0)
+
+        # create mesh
+        mesh = wp.Mesh(points=mesh_points, velocities=None, indices=mesh_indices, bvh_constructor=constructor)
+
+        tape = wp.Tape()
+
+        # analytic gradients
+        with tape:
+            query_points = wp.array(p, dtype=wp.vec3, device=device, requires_grad=True)
+            query_dirs = wp.array(D, dtype=wp.vec3, device=device, requires_grad=True)
+            intersection_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
+            loss = wp.zeros(n=1, dtype=float, device=device, requires_grad=True)
+
+            wp.launch(
+                kernel=mesh_query_ray_loss,
+                dim=1,
+                inputs=[mesh.id, query_points, query_dirs, intersection_points, loss],
+                device=device,
+            )
+
+        tape.backward(loss=loss)
+        q = intersection_points.numpy().flatten()
+        analytic_p = tape.gradients[query_points].numpy().flatten()
+        analytic_D = tape.gradients[query_dirs].numpy().flatten()
+
+        # numeric gradients
+
+        # ray origin
+        eps = 1.0e-3
+        loss_values_p = []
+        numeric_p = np.zeros(3)
+
+        offset_query_points = [
+            wp.vec3(p[0] - eps, p[1], p[2]),
+            wp.vec3(p[0] + eps, p[1], p[2]),
+            wp.vec3(p[0], p[1] - eps, p[2]),
+            wp.vec3(p[0], p[1] + eps, p[2]),
+            wp.vec3(p[0], p[1], p[2] - eps),
+            wp.vec3(p[0], p[1], p[2] + eps),
+        ]
+
+        for i in range(6):
+            q = offset_query_points[i]
+
+            query_points = wp.array(q, dtype=wp.vec3, device=device)
+            query_dirs = wp.array(D, dtype=wp.vec3, device=device)
+            intersection_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
+            loss = wp.zeros(n=1, dtype=float, device=device)
+
+            wp.launch(
+                kernel=mesh_query_ray_loss,
+                dim=1,
+                inputs=[mesh.id, query_points, query_dirs, intersection_points, loss],
+                device=device,
+            )
+
+            loss_values_p.append(loss.numpy()[0])
+
+        for i in range(3):
+            l_0 = loss_values_p[i * 2]
+            l_1 = loss_values_p[i * 2 + 1]
+            gradient = (l_1 - l_0) / (2.0 * eps)
+            numeric_p[i] = gradient
+
+        # ray dir
+        loss_values_D = []
+        numeric_D = np.zeros(3)
+
+        offset_query_dirs = [
+            wp.vec3(D[0] - eps, D[1], D[2]),
+            wp.vec3(D[0] + eps, D[1], D[2]),
+            wp.vec3(D[0], D[1] - eps, D[2]),
+            wp.vec3(D[0], D[1] + eps, D[2]),
+            wp.vec3(D[0], D[1], D[2] - eps),
+            wp.vec3(D[0], D[1], D[2] + eps),
+        ]
+
+        for i in range(6):
+            q = offset_query_dirs[i]
+
+            query_points = wp.array(p, dtype=wp.vec3, device=device)
+            query_dirs = wp.array(q, dtype=wp.vec3, device=device)
+            intersection_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
+            loss = wp.zeros(n=1, dtype=float, device=device)
+
+            wp.launch(
+                kernel=mesh_query_ray_loss,
+                dim=1,
+                inputs=[mesh.id, query_points, query_dirs, intersection_points, loss],
+                device=device,
+            )
+
+            loss_values_D.append(loss.numpy()[0])
+
+        for i in range(3):
+            l_0 = loss_values_D[i * 2]
+            l_1 = loss_values_D[i * 2 + 1]
+            gradient = (l_1 - l_0) / (2.0 * eps)
+            numeric_D[i] = gradient
+
+        error_p = ((analytic_p - numeric_p) * (analytic_p - numeric_p)).sum(axis=0)
+        error_D = ((analytic_D - numeric_D) * (analytic_D - numeric_D)).sum(axis=0)
+
+        tolerance = 1.0e-3
+        test.assertTrue(error_p < tolerance, f"error is {error_p} which is >= {tolerance}")
+        test.assertTrue(error_D < tolerance, f"error is {error_D} which is >= {tolerance}")
 
 
 @wp.kernel
@@ -229,6 +235,11 @@ def raycast_kernel(
 
 
 def test_mesh_query_ray_edge(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
     # Create raycast starts and directions
     xx, yy = np.meshgrid(np.arange(0.1, 0.4, 0.01), np.arange(0.1, 0.4, 0.01))
     xx = xx.flatten().reshape(-1, 1)
@@ -239,27 +250,29 @@ def test_mesh_query_ray_edge(test, device):
     ray_dirs = np.zeros_like(ray_starts)
     ray_dirs[:, 2] = -1.0
 
+    n = len(ray_starts)
+
+    ray_starts = wp.array(ray_starts, shape=(n,), dtype=wp.vec3, device=device)
+    ray_dirs = wp.array(ray_dirs, shape=(n,), dtype=wp.vec3, device=device)
+
     # Create simple square mesh
     vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.5, 0.0], [0.5, 0.0, 0.0], [0.5, 0.5, 0.0]], dtype=np.float32)
 
     triangles = np.array([[1, 0, 2], [1, 2, 3]], dtype=np.int32)
 
-    mesh = wp.Mesh(
-        points=wp.array(vertices, dtype=wp.vec3, device=device),
-        indices=wp.array(triangles.flatten(), dtype=int, device=device),
-    )
-
-    counts = wp.zeros(1, dtype=int, device=device)
-
-    n = len(ray_starts)
+    for constructor in constructors:
+        mesh = wp.Mesh(
+            points=wp.array(vertices, dtype=wp.vec3, device=device),
+            indices=wp.array(triangles.flatten(), dtype=int, device=device),
+            bvh_constructor=constructor,
+        )
 
-    ray_starts = wp.array(ray_starts, shape=(n,), dtype=wp.vec3, device=device)
-    ray_dirs = wp.array(ray_dirs, shape=(n,), dtype=wp.vec3, device=device)
+        counts = wp.zeros(1, dtype=int, device=device)
 
-    wp.launch(kernel=raycast_kernel, dim=n, inputs=[mesh.id, ray_starts, ray_dirs, counts], device=device)
-    wp.synchronize()
+        wp.launch(kernel=raycast_kernel, dim=n, inputs=[mesh.id, ray_starts, ray_dirs, counts], device=device)
+        wp.synchronize()
 
-    test.assertEqual(counts.numpy()[0], n)
+        test.assertEqual(counts.numpy()[0], n)
 
 
 devices = get_test_devices()

warp/tests/test_mlp.py

@@ -265,8 +265,8 @@ class TestMLP(unittest.TestCase):
     pass
 
 
-add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices)
-add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices)
+add_function_test(TestMLP, "test_mlp", test_mlp, devices=devices, check_output=False)
+add_function_test(TestMLP, "test_mlp_grad", test_mlp_grad, devices=devices, check_output=False)
 
 
 if __name__ == "__main__":

warp/tests/test_operators.py

@@ -224,6 +224,48 @@ def test_operators_mat44():
     expect_eq(r0[3], wp.vec4(39.0, 42.0, 45.0, 48.0))
 
 
+@wp.struct
+class Complex:
+    real: float
+    imag: float
+
+
+@wp.func
+def add(
+    a: Complex,
+    b: Complex,
+) -> Complex:
+    return Complex(
+        a.real + b.real,
+        a.imag + b.imag,
+    )
+
+
+@wp.func
+def mul(
+    a: Complex,
+    b: Complex,
+) -> Complex:
+    return Complex(
+        a.real * b.real - a.imag * b.imag,
+        a.real * b.imag + a.imag * b.real,
+    )
+
+
+@wp.kernel
+def test_operators_overload():
+    a = Complex(1.0, 2.0)
+    b = Complex(3.0, 4.0)
+
+    c = a + b
+    expect_eq(c.real, 4.0)
+    expect_eq(c.imag, 6.0)
+
+    d = a * b
+    expect_eq(d.real, -5.0)
+    expect_eq(d.imag, 10.0)
+
+
 devices = get_test_devices()
 
 
@@ -241,6 +283,7 @@ add_kernel_test(TestOperators, test_operators_vec4, dim=1, devices=devices)
 add_kernel_test(TestOperators, test_operators_mat22, dim=1, devices=devices)
 add_kernel_test(TestOperators, test_operators_mat33, dim=1, devices=devices)
 add_kernel_test(TestOperators, test_operators_mat44, dim=1, devices=devices)
+add_kernel_test(TestOperators, test_operators_overload, dim=1, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_overwrite.py

@@ -1,6 +1,7 @@
 import contextlib
 import io
 import unittest
+from typing import Any
 
 import numpy as np
 
@@ -164,18 +165,18 @@ def test_kernel_writeread_kernel_write(test, device):
 
 
 @wp.func
-def read_func(a: wp.array(dtype=float), idx: int):
+def read_func(a: wp.array(dtype=Any), idx: int):
     x = a[idx]
     return x
 
 
 @wp.func
-def read_return_func(b: wp.array(dtype=float), idx: int):
+def read_return_func(b: wp.array(dtype=Any), idx: int):
     return 1.0, b[idx]
 
 
 @wp.func
-def write_func(c: wp.array(dtype=float), idx: int):
+def write_func(c: wp.array(dtype=Any), idx: int):
     c[idx] = 1.0
 
 
@@ -577,8 +578,8 @@ add_function_test(TestOverwrite, "test_views", test_views, devices=devices)
 add_function_test(TestOverwrite, "test_reset", test_reset, devices=devices)
 
 add_function_test(TestOverwrite, "test_copy", test_copy, devices=devices)
-add_function_test(TestOverwrite, "test_matmul", test_matmul, devices=devices)
-add_function_test(TestOverwrite, "test_batched_matmul", test_batched_matmul, devices=devices)
+add_function_test(TestOverwrite, "test_matmul", test_matmul, devices=devices, check_output=False)
+add_function_test(TestOverwrite, "test_batched_matmul", test_batched_matmul, devices=devices, check_output=False)
 add_function_test(TestOverwrite, "test_atomic_operations", test_atomic_operations, devices=devices)
 
 # Some warning are only issued during codegen, and codegen only runs on cuda_0 in the MGPU case.

warp/tests/test_quat.py

@@ -2095,6 +2095,81 @@ def test_py_arithmetic_ops(test, device, dtype):
     test.assertSequenceEqual(wptype(24) / v, make_quat(12, 6, 4, 3))
 
 
+@wp.kernel
+def quat_len_kernel(
+    q: wp.quat,
+    out: wp.array(dtype=int),
+):
+    length = wp.static(len(q))
+    wp.expect_eq(wp.static(len(q)), 4)
+    out[0] = wp.static(len(q))
+
+    foo = wp.quat()
+    length = len(foo)
+    wp.expect_eq(len(foo), 4)
+    out[1] = len(foo)
+
+
+def test_quat_len(test, device):
+    q = wp.quat()
+    out = wp.empty(2, dtype=int, device=device)
+    wp.launch(quat_len_kernel, dim=(1,), inputs=(q,), outputs=(out,), device=device)
+
+    test.assertEqual(out.numpy()[0], 4)
+    test.assertEqual(out.numpy()[1], 4)
+
+
+@wp.kernel
+def vector_augassign_kernel(
+    a: wp.array(dtype=wp.quat), b: wp.array(dtype=wp.quat), c: wp.array(dtype=wp.quat), d: wp.array(dtype=wp.quat)
+):
+    i = wp.tid()
+
+    q1 = wp.quat()
+    q2 = b[i]
+
+    q1[0] += q2[0]
+    q1[1] += q2[1]
+    q1[2] += q2[2]
+    q1[3] += q2[3]
+
+    a[i] = q1
+
+    q3 = wp.quat()
+    q4 = d[i]
+
+    q3[0] += q4[0]
+    q3[1] += q4[1]
+    q3[2] += q4[2]
+    q3[3] += q4[3]
+
+    c[i] = q1
+
+
+def test_vector_augassign(test, device):
+    N = 3
+
+    a = wp.zeros(N, dtype=wp.quat, requires_grad=True)
+    b = wp.ones(N, dtype=wp.quat, requires_grad=True)
+
+    c = wp.zeros(N, dtype=wp.quat, requires_grad=True)
+    d = wp.ones(N, dtype=wp.quat, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(vector_augassign_kernel, N, inputs=[a, b, c, d])
+
+    tape.backward(grads={a: wp.ones_like(a), c: wp.ones_like(c)})
+
+    assert_np_equal(a.numpy(), wp.ones_like(a).numpy())
+    assert_np_equal(a.grad.numpy(), wp.ones_like(a).numpy())
+    assert_np_equal(b.grad.numpy(), wp.ones_like(a).numpy())
+
+    assert_np_equal(c.numpy(), -wp.ones_like(c).numpy())
+    assert_np_equal(c.grad.numpy(), wp.ones_like(c).numpy())
+    assert_np_equal(d.grad.numpy(), -wp.ones_like(d).numpy())
+
+
 devices = get_test_devices()
 
 
@@ -2203,6 +2278,8 @@ for dtype in np_float_types:
         TestQuat, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
     )
 
+add_function_test(TestQuat, "test_quat_len", test_quat_len, devices=devices)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_reload.py

@@ -241,6 +241,32 @@ def test_graph_launch_after_module_reload(test, device):
         test.assertEqual(a.numpy()[0], 42)
 
 
+def test_module_unload_during_graph_capture(test, device):
+    @wp.kernel
+    def foo(a: wp.array(dtype=int)):
+        a[0] = 42
+
+    # preload module before graph capture
+    wp.load_module(device=device)
+
+    # load another module to test unloading during graph capture
+    other_module = wp.get_module("warp.tests.aux_test_module_unload")
+    other_module.load(device)
+
+    with wp.ScopedDevice(device):
+        a = wp.zeros(1, dtype=int)
+
+        with wp.ScopedCapture(force_module_load=False) as capture:
+            wp.launch(foo, dim=1, inputs=[a])
+
+            # unloading a module during graph capture should be fine (deferred until capture completes)
+            other_module.unload()
+
+        wp.capture_launch(capture.graph)
+
+        test.assertEqual(a.numpy()[0], 42)
+
+
 devices = get_test_devices()
 cuda_devices = get_cuda_test_devices()
 
@@ -258,6 +284,9 @@ add_function_test(TestReload, "test_reload_references", test_reload_references,
 add_function_test(
     TestReload, "test_graph_launch_after_module_reload", test_graph_launch_after_module_reload, devices=cuda_devices
 )
+add_function_test(
+    TestReload, "test_module_unload_during_graph_capture", test_module_unload_during_graph_capture, devices=cuda_devices
+)
 
 
 if __name__ == "__main__":