warp-lang 1.8.1__py3-none-macosx_10_13_universal2.whl → 1.9.1__py3-none-macosx_10_13_universal2.whl

warp/tests/aux_test_module_aot.py

@@ -0,0 +1,7 @@
+import warp as wp
+
+
+@wp.kernel
+def add_kernel(a: wp.array(dtype=wp.int32), b: wp.array(dtype=wp.int32), res: wp.array(dtype=wp.int32)):
+    i = wp.tid()
+    res[i] = a[i] + b[i]

warp/tests/cuda/test_async.py

@@ -426,13 +426,13 @@ def copy_template(test, src_ctor, dst_ctor, src_device, dst_device, n, params: C
         if expected_error_type is not None:
             # disable error output from Warp if we expect an exception
             try:
-                saved_error_output_enabled = wp.context.runtime.core.is_error_output_enabled()
-                wp.context.runtime.core.set_error_output_enabled(False)
+                saved_error_output_enabled = wp.context.runtime.core.wp_is_error_output_enabled()
+                wp.context.runtime.core.wp_set_error_output_enabled(False)
                 with test.assertRaisesRegex(expected_error_type, expected_error_regex):
                     with Capturable(use_graph=params.use_graph, stream=stream):
                         wp.copy(dst, src, stream=stream_arg)
             finally:
-                wp.context.runtime.core.set_error_output_enabled(saved_error_output_enabled)
+                wp.context.runtime.core.wp_set_error_output_enabled(saved_error_output_enabled)
                 wp.synchronize()
 
                 # print(f"SUCCESSFUL ERROR PREDICTION: {expected_error_regex}")

warp/tests/cuda/test_conditional_captures.py

@@ -979,6 +979,94 @@ def test_graph_debug_dot_print(test, device):
     os.remove(dot_file)
 
 
+# ================================================================================================================
+# test exceptions
+# ================================================================================================================
+
+
+# body graphs with allocations are not supported
+def body_with_alloc():
+    wp.zeros(10)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_if(test, device):
+    with wp.ScopedDevice(device):
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Conditional body graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_true=body_with_alloc)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_else(test, device):
+    with wp.ScopedDevice(device):
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Conditional body graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_false=body_with_alloc)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_while(test, device):
+    with wp.ScopedDevice(device):
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Conditional body graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_while(condition=cond, while_body=body_with_alloc)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_if_subgraph(test, device):
+    with wp.ScopedDevice(device):
+        # capture body subgraph
+        with wp.ScopedCapture() as body_capture:
+            body_with_alloc()
+
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Child graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_true=body_capture.graph)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_else_subgraph(test, device):
+    with wp.ScopedDevice(device):
+        # capture body subgraph
+        with wp.ScopedCapture() as body_capture:
+            body_with_alloc()
+
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Child graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_if(condition=cond, on_false=body_capture.graph)
+
+
+@unittest.skipUnless(check_conditional_graph_support(), "Conditional graph nodes not supported")
+def test_error_alloc_while_subgraph(test, device):
+    with wp.ScopedDevice(device):
+        # capture body subgraph
+        with wp.ScopedCapture() as body_capture:
+            body_with_alloc()
+
+        cond = wp.ones(1, dtype=wp.int32)
+        with test.assertRaisesRegex(
+            RuntimeError, r"Child graph contains an unsupported operation \(memory allocation\)"
+        ):
+            with wp.ScopedCapture():
+                wp.capture_while(condition=cond, while_body=body_capture.graph)
+
+
 devices = get_test_devices()
 cuda_devices = get_cuda_test_devices()
 
@@ -1040,6 +1128,19 @@ add_function_test(
     TestConditionalCaptures, "test_graph_debug_dot_print", test_graph_debug_dot_print, devices=cuda_devices
 )
 
+add_function_test(TestConditionalCaptures, "test_error_alloc_if", test_error_alloc_if, devices=cuda_devices)
+add_function_test(TestConditionalCaptures, "test_error_alloc_else", test_error_alloc_else, devices=cuda_devices)
+add_function_test(TestConditionalCaptures, "test_error_alloc_while", test_error_alloc_while, devices=cuda_devices)
+add_function_test(
+    TestConditionalCaptures, "test_error_alloc_if_subgraph", test_error_alloc_if_subgraph, devices=cuda_devices
+)
+add_function_test(
+    TestConditionalCaptures, "test_error_alloc_else_subgraph", test_error_alloc_else_subgraph, devices=cuda_devices
+)
+add_function_test(
+    TestConditionalCaptures, "test_error_alloc_while_subgraph", test_error_alloc_while_subgraph, devices=cuda_devices
+)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/geometry/test_hash_grid.py

@@ -191,7 +191,44 @@ def test_hashgrid_inputs(test, device):
             assert_array_equal(counts_ndim, counts_ref)
 
 
+def test_hashgrid_multiple_streams(test, device):
+    with wp.ScopedDevice(device):
+        points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)
+        points_ref = wp.array(points, dtype=wp.vec3)
+        counts_ref = wp.zeros(len(points), dtype=int)
+
+        grid_dim = 64
+        grid_ref = wp.HashGrid(grid_dim, grid_dim, grid_dim)
+        grid_ref.build(points_ref, cell_radius)
+
+        # get reference counts
+        wp.launch(kernel=count_neighbors, dim=len(points), inputs=[grid_ref.id, query_radius, points_ref, counts_ref])
+
+        # create multiple streams
+        num_streams = 10
+        streams = [wp.Stream(device=device) for _ in range(num_streams)]
+        counts_per_stream = [wp.zeros(len(points), dtype=int) for _ in range(num_streams)]
+
+        # test whether HashGrid and radix sort work with multiple streams without race conditions
+        for i in range(num_streams):
+            with wp.ScopedStream(streams[i]):
+                grid = wp.HashGrid(grid_dim, grid_dim, grid_dim)
+                grid.build(points_ref, cell_radius)
+
+                # get counts for this stream
+                wp.launch(
+                    kernel=count_neighbors,
+                    dim=len(points),
+                    inputs=[grid.id, query_radius, points_ref, counts_per_stream[i]],
+                )
+
+        # run this loop after all streams are scheduled to ensure asynchronous behaviour above
+        for i in range(num_streams):
+            assert_array_equal(counts_per_stream[i], counts_ref)
+
+
 devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
 
 
 class TestHashGrid(unittest.TestCase):
@@ -214,6 +251,7 @@ class TestHashGrid(unittest.TestCase):
 
 add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
 add_function_test(TestHashGrid, "test_hashgrid_inputs", test_hashgrid_inputs, devices=devices)
+add_function_test(TestHashGrid, "test_hashgrid_multiple_streams", test_hashgrid_multiple_streams, devices=cuda_devices)
 
 
 if __name__ == "__main__":

warp/tests/geometry/test_marching_cubes.py

@@ -22,7 +22,9 @@ from warp.tests.unittest_utils import *
 
 
 @wp.kernel
-def make_field(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
+def make_field_sphere_sdf(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
+    """Make a sphere SDF for nodes on the integer domain with node coordinates 0,1,2,3,..."""
+
     i, j, k = wp.tid()
 
     p = wp.vec3(float(i), float(j), float(k))
@@ -32,31 +34,242 @@ def make_field(field: wp.array3d(dtype=float), center: wp.vec3, radius: float):
     field[i, j, k] = d
 
 
-def test_marching_cubes(test, device):
-    dim = 64
-    max_verts = 10**6
-    max_tris = 10**6
+@wp.kernel
+def make_field_sphere_sdf_unit_domain(
+    field: wp.array3d(dtype=float), center: wp.vec3, radius: wp.array(dtype=wp.float32)
+):
+    """Makes a sphere SDF for nodes on the unit domain [-1, 1]^3."""
+    i, j, k = wp.tid()
 
-    field = wp.zeros(shape=(dim, dim, dim), dtype=float, device=device)
+    nx, ny, nz = field.shape[0], field.shape[1], field.shape[2]
+
+    p = wp.vec3(
+        2.0 * wp.float32(i) / (wp.float32(nx) - 1.0) - 1.0,
+        2.0 * wp.float32(j) / (wp.float32(ny) - 1.0) - 1.0,
+        2.0 * wp.float32(k) / (wp.float32(nz) - 1.0) - 1.0,
+    )
+
+    d = wp.length(p - center) - radius[0]
+
+    field[i, j, k] = d
+
+
+@wp.kernel
+def compute_surface_area(
+    verts: wp.array(dtype=wp.vec3), faces: wp.array(dtype=wp.int32), out_area: wp.array(dtype=wp.float32)
+):
+    tid = wp.tid()
+    vi = faces[3 * tid + 0]
+    vj = faces[3 * tid + 1]
+    vk = faces[3 * tid + 2]
+
+    p0 = verts[vi]
+    p1 = verts[vj]
+    p2 = verts[vk]
+
+    # Heron's formula for triangle area
+    a = wp.length(p1 - p0)
+    b = wp.length(p2 - p0)
+    c = wp.length(p2 - p1)
+    s = (a + b + c) / 2.0
+    area = wp.sqrt(s * (s - a) * (s - b) * (s - c))
 
-    iso = wp.MarchingCubes(nx=dim, ny=dim, nz=dim, max_verts=max_verts, max_tris=max_tris, device=device)
+    wp.atomic_add(out_area, 0, area)
 
-    radius = dim / 4.0
 
-    wp.launch(make_field, dim=field.shape, inputs=[field, wp.vec3(dim / 2, dim / 2, dim / 2), radius], device=device)
+def validate_marching_cubes_output(test, verts_np, faces_np, check_nonempty=True):
+    # check that the face array seems valid
+    if check_nonempty:
+        test.assertGreater(faces_np.shape[0], 0)  # at least one face
+    test.assertEqual(faces_np.shape[1], 3)  # all faces triangular
+    test.assertTrue((faces_np >= 0).all())  # all face inds nonnegative
+    test.assertTrue((faces_np < verts_np.shape[0]).all())  # all face inds are in-bounds on the vertex array
+    test.assertTrue((faces_np[:, 0] != faces_np[:, 1]).all())  # all faces have unique vertices
+    test.assertTrue((faces_np[:, 0] != faces_np[:, 2]).all())  # all faces have unique vertices
+    test.assertTrue((faces_np[:, 1] != faces_np[:, 2]).all())  # all faces have unique vertices
+    test.assertTrue(
+        (np.unique(faces_np.flatten()) == np.arange(verts_np.shape[0])).all()
+    )  # all vertices are used in at least one face
+
+    # check that the vertex array seems valid
+    if check_nonempty:
+        test.assertGreater(verts_np.shape[0], 0)  # at least one vertex
+    test.assertEqual(verts_np.shape[1], 3)  # all vertices are 3D
+    test.assertTrue(np.isfinite(verts_np).all())  # all vertices are finite
+
+
+def test_marching_cubes(test, device):
+    """Basic test of typical usage."""
+    node_dim = 64
+    cell_dim = node_dim - 1
+    field = wp.zeros(shape=(node_dim, node_dim, node_dim), dtype=float, device=device)
+    bounds_low = (0.0, 0.0, 0.0)
+    bounds_high = (float(cell_dim), float(cell_dim), float(cell_dim))
+
+    iso = wp.MarchingCubes(
+        nx=node_dim,
+        ny=node_dim,
+        nz=node_dim,
+        device=device,
+        domain_bounds_lower_corner=bounds_low,
+        domain_bounds_upper_corner=bounds_high,
+    )
+
+    radius = node_dim / 4.0
+
+    wp.launch(
+        make_field_sphere_sdf,
+        dim=field.shape,
+        inputs=[field, wp.vec3(node_dim / 2, node_dim / 2, node_dim / 2), radius],
+        device=device,
+    )
 
     iso.surface(field=field, threshold=0.0)
+    verts_np = iso.verts.numpy()
+    faces_np = iso.indices.numpy().reshape(-1, 3)
+    test.assertEqual(
+        iso.indices.dtype, wp.int32
+    )  # make sure we are following Warp convention of using a flat array of indices
+    validate_marching_cubes_output(test, verts_np, faces_np)
 
     # check that all returned vertices lie on the surface of the sphere
-    length = np.linalg.norm(iso.verts.numpy() - np.array([dim / 2, dim / 2, dim / 2]), axis=1)
+    length = np.linalg.norm(verts_np - np.array([node_dim / 2, node_dim / 2, node_dim / 2]), axis=1)
     error = np.abs(length - radius)
+    test.assertTrue(np.max(error) < 1.0)
+
+    iso.resize(nx=node_dim * 2, ny=node_dim * 2, nz=node_dim * 2)  # smoke test for deprecated function
+
+
+def test_marching_cubes_functional(test, device):
+    """Ensure the single-function interface works as expected."""
+    node_dim = 64
+    cell_dim = node_dim - 1
+    field = wp.zeros(shape=(node_dim, node_dim, node_dim), dtype=float, device=device)
+    bounds_low = (0.0, 0.0, 0.0)
+    bounds_high = (float(cell_dim), float(cell_dim), float(cell_dim))
+
+    radius = node_dim / 4.0
+    wp.launch(
+        make_field_sphere_sdf,
+        dim=field.shape,
+        inputs=[field, wp.vec3(node_dim / 2, node_dim / 2, node_dim / 2), radius],
+        device=device,
+    )
+
+    # call via the functional interface
+    verts, faces = wp.MarchingCubes.extract_surface_marching_cubes(
+        field, threshold=0.0, domain_bounds_lower_corner=bounds_low, domain_bounds_upper_corner=bounds_high
+    )
+
+    verts_np = verts.numpy()
+    faces_np = faces.numpy().reshape(-1, 3)
+    validate_marching_cubes_output(test, verts_np, faces_np)
 
+    # check that all returned vertices lie on the surface of the sphere
+    length = np.linalg.norm(verts_np - np.array([node_dim / 2, node_dim / 2, node_dim / 2]), axis=1)
+    error = np.abs(length - radius)
     test.assertTrue(np.max(error) < 1.0)
 
-    iso.resize(nx=dim * 2, ny=dim * 2, nz=dim * 2, max_verts=max_verts, max_tris=max_tris)
+
+def test_marching_cubes_nonuniform(test, device):
+    """Test the logic for when the dimensions of the grid are not uniform."""
+
+    dimX = 64
+    dimY = 48
+    dimZ = 72
+    field = wp.zeros(shape=(dimX, dimY, dimZ), dtype=float, device=device)
+
+    bounds_low = wp.vec3(0.0, 0.0, 0.0)
+    bounds_high = wp.vec3(float(dimX), float(dimY), float(dimZ))
+
+    iso = wp.MarchingCubes(
+        nx=dimX,
+        ny=dimY,
+        nz=dimZ,
+        device=device,
+        domain_bounds_lower_corner=bounds_low,
+        domain_bounds_upper_corner=bounds_high,
+    )
+
+    radius = dimX / 4.0
+    wp.launch(
+        make_field_sphere_sdf,
+        dim=field.shape,
+        inputs=[field, wp.vec3(dimX / 2, dimY / 2, dimZ / 2), radius],
+        device=device,
+    )
+
+    iso.surface(field=field, threshold=0.0)
+    verts_np = iso.verts.numpy()
+    faces_np = iso.indices.numpy().reshape(-1, 3)
+    validate_marching_cubes_output(test, verts_np, faces_np)
+
+
+def test_marching_cubes_empty_output(test, device):
+    """Make sure we handle the empty-output case correctly."""
+
+    dim = 64
+    field = wp.zeros(shape=(dim, dim, dim), dtype=float, device=device)
+
+    iso = wp.MarchingCubes(nx=dim, ny=dim, nz=dim, device=device)
+
+    wp.launch(make_field_sphere_sdf, dim=field.shape, inputs=[field, wp.vec3(0.5, 0.5, 0.5), 0.25], device=device)
+
+    iso.surface(field=field, threshold=1000.0)  # set threshold to a large value so that no vertices are generated
+    verts_np = iso.verts.numpy()
+    faces_np = iso.indices.numpy().reshape(-1, 3)
+    validate_marching_cubes_output(test, verts_np, faces_np, check_nonempty=False)
+
+    test.assertEqual(faces_np.shape[0], 0)  # no faces
+    test.assertEqual(verts_np.shape[0], 0)  # no vertices
+
+
+def test_marching_cubes_differentiable(test, device):
+    """Check that marching cubes has reasonable gradients.
+
+    This test constructs an SDF of sphere, extracts a surface, computes its
+    surface area, and then differentiates the surface area with respect to
+    the sphere's radius.
+    """
+    node_dim = 64
+    cell_dim = node_dim - 1
+    bounds_low = wp.vec3(-1.0, -1.0, -1.0)
+    bounds_high = wp.vec3(1.0, 1.0, 1.0)
+
+    radius = 0.5
+    radius_wp = wp.full((1,), value=0.5, dtype=wp.float32, device=device, requires_grad=True)
+
+    with wp.Tape() as tape:
+        field = wp.zeros(shape=(node_dim, node_dim, node_dim), dtype=float, device=device, requires_grad=True)
+        wp.launch(
+            make_field_sphere_sdf_unit_domain,
+            dim=field.shape,
+            inputs=[field, wp.vec3(0.0, 0.0, 0.0), radius_wp],
+            device=device,
+        )
+
+        # call via the functional interface
+        verts, faces = wp.MarchingCubes.extract_surface_marching_cubes(
+            field, threshold=0.0, domain_bounds_lower_corner=bounds_low, domain_bounds_upper_corner=bounds_high
+        )
+
+        # compute surface area
+        area = wp.zeros(shape=(1,), dtype=float, device=device, requires_grad=True)
+        wp.launch(compute_surface_area, dim=faces.shape[0] // 3, inputs=[verts, faces, area], device=device)
+
+        # confirm surface area is correct vs. the analytical ground truth
+        area_np = area.numpy()[0]
+        test.assertTrue(np.abs(area_np - 4.0 * np.pi * radius * radius) < 1e-2)
+
+    # compute the gradient of the surface area with respect to the radius
+    tape.backward(area)
+
+    # confirm the gradient is correct vs. the analytical ground truth
+    grad_np = radius_wp.grad.numpy()[0]
+    test.assertTrue(np.abs(grad_np - 8.0 * np.pi * radius) < 1e-2)
 
 
-devices = get_selected_cuda_test_devices()
+devices = get_test_devices()
 
 
 class TestMarchingCubes(unittest.TestCase):
@@ -67,6 +280,14 @@ class TestMarchingCubes(unittest.TestCase):
 
 
 add_function_test(TestMarchingCubes, "test_marching_cubes", test_marching_cubes, devices=devices)
+add_function_test(TestMarchingCubes, "test_marching_cubes_functional", test_marching_cubes_functional, devices=devices)
+add_function_test(TestMarchingCubes, "test_marching_cubes_nonuniform", test_marching_cubes_nonuniform, devices=devices)
+add_function_test(
+    TestMarchingCubes, "test_marching_cubes_empty_output", test_marching_cubes_empty_output, devices=devices
+)
+add_function_test(
+    TestMarchingCubes, "test_marching_cubes_differentiable", test_marching_cubes_differentiable, devices=devices
+)
 
 
 if __name__ == "__main__":