warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/geometry/test_bvh.py

@@ -0,0 +1,335 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def bvh_query_aabb(bvh_id: wp.uint64, lower: wp.vec3, upper: wp.vec3, bounds_intersected: wp.array(dtype=int)):
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
+    bounds_nr = int(0)
+
+    while wp.bvh_query_next(query, bounds_nr):
+        bounds_intersected[bounds_nr] = 1
+
+
+@wp.kernel
+def bvh_query_ray(bvh_id: wp.uint64, start: wp.vec3, dir: wp.vec3, bounds_intersected: wp.array(dtype=int)):
+    query = wp.bvh_query_ray(bvh_id, start, dir)
+    bounds_nr = int(0)
+
+    while wp.bvh_query_next(query, bounds_nr):
+        bounds_intersected[bounds_nr] = 1
+
+
+def aabb_overlap(a_lower, a_upper, b_lower, b_upper):
+    if (
+        a_lower[0] > b_upper[0]
+        or a_lower[1] > b_upper[1]
+        or a_lower[2] > b_upper[2]
+        or a_upper[0] < b_lower[0]
+        or a_upper[1] < b_lower[1]
+        or a_upper[2] < b_lower[2]
+    ):
+        return 0
+    else:
+        return 1
+
+
+def intersect_ray_aabb(start, rcp_dir, lower, upper):
+    l1 = (lower[0] - start[0]) * rcp_dir[0]
+    l2 = (upper[0] - start[0]) * rcp_dir[0]
+    lmin = min(l1, l2)
+    lmax = max(l1, l2)
+
+    l1 = (lower[1] - start[1]) * rcp_dir[1]
+    l2 = (upper[1] - start[1]) * rcp_dir[1]
+    lmin = max(min(l1, l2), lmin)
+    lmax = min(max(l1, l2), lmax)
+
+    l1 = (lower[2] - start[2]) * rcp_dir[2]
+    l2 = (upper[2] - start[2]) * rcp_dir[2]
+    lmin = max(min(l1, l2), lmin)
+    lmax = min(max(l1, l2), lmax)
+
+    if lmax >= 0.0 and lmax >= lmin:
+        return 1
+    else:
+        return 0
+
+
+def test_bvh(test, type, device, leaf_size):
+    rng = np.random.default_rng(123)
+
+    num_bounds = 100
+    lowers = rng.random(size=(num_bounds, 3)) * 5.0
+    uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
+
+    device_lowers = wp.array(lowers, dtype=wp.vec3, device=device)
+    device_uppers = wp.array(uppers, dtype=wp.vec3, device=device)
+
+    bvh = wp.Bvh(device_lowers, device_uppers, leaf_size=leaf_size)
+
+    bounds_intersected = wp.zeros(shape=(num_bounds), dtype=int, device=device)
+
+    query_lower = wp.vec3(2.0, 2.0, 2.0)
+    query_upper = wp.vec3(8.0, 8.0, 8.0)
+
+    query_start = wp.vec3(0.0, 0.0, 0.0)
+    query_dir = wp.normalize(wp.vec3(1.0, 1.0, 1.0))
+
+    for test_case in range(3):
+        if type == "AABB":
+            wp.launch(
+                bvh_query_aabb,
+                dim=1,
+                inputs=[bvh.id, query_lower, query_upper, bounds_intersected],
+                device=device,
+            )
+        else:
+            wp.launch(bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device)
+
+        device_intersected = bounds_intersected.numpy()
+
+        for i in range(num_bounds):
+            lower = lowers[i]
+            upper = uppers[i]
+            if type == "AABB":
+                host_intersected = aabb_overlap(lower, upper, query_lower, query_upper)
+            else:
+                host_intersected = intersect_ray_aabb(query_start, 1.0 / query_dir, lower, upper)
+
+            test.assertEqual(host_intersected, device_intersected[i])
+
+        if test_case == 0 or test_case == 1:
+            lowers = rng.random(size=(num_bounds, 3)) * 5.0
+            uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
+            wp.copy(device_lowers, wp.array(lowers, dtype=wp.vec3, device=device))
+            wp.copy(device_uppers, wp.array(uppers, dtype=wp.vec3, device=device))
+            bounds_intersected.zero_()
+
+            if test_case == 0:
+                bvh.refit()
+            else:
+                bvh.rebuild()
+
+
+def test_bvh_query_aabb(test, device):
+    for leaf_size in [1, 2, 4]:
+        test_bvh(test, "AABB", device, leaf_size)
+
+
+def test_bvh_query_ray(test, device):
+    for leaf_size in [1, 2, 4]:
+        test_bvh(test, "ray", device, leaf_size)
+
+
+def test_bvh_ray_query_inside_and_outside_bounds(test, device):
+    """Regression test for issue #288: BVH ray queries should detect intersections
+    regardless of whether the ray origin is inside or outside the bounding volumes.
+
+    Previously, rays starting outside the bounds would fail to detect intersections.
+    """
+    # Create a single AABB spanning x=[0.5, 1.0], extending across y and z axes
+    lowers = ((0.5, -1.0, -1.0),)
+    uppers = ((1.0, 1.0, 1.0),)
+
+    device_lowers = wp.array(lowers, dtype=wp.vec3f, device=device)
+    device_uppers = wp.array(uppers, dtype=wp.vec3f, device=device)
+
+    bvh = wp.Bvh(device_lowers, device_uppers)
+
+    bounds_intersected = wp.zeros(shape=1, dtype=int, device=device)
+
+    # Test both ray origins: outside (x=0.0) and inside (x=0.75) the AABB
+    for x in (0.0, 0.75):
+        query_start = wp.vec3(x, 0.0, 0.0)
+        query_dir = wp.vec3(1.0, 0.0, 0.0)  # Ray pointing in +x direction
+
+        wp.launch(bvh_query_ray, dim=1, inputs=[bvh.id, query_start, query_dir, bounds_intersected], device=device)
+
+        device_intersected = bounds_intersected.numpy()
+        # Both cases should detect the single intersection
+        test.assertEqual(device_intersected.sum(), 1)
+
+
+def get_random_aabbs(n, center, relative_shift, relative_size, rng):
+    centers = rng.uniform(-0.5, 0.5, size=n * 3).reshape(n, 3) * relative_shift + center
+    diffs = 0.5 * rng.random(n * 3).reshape(n, 3) * relative_size
+
+    lowers = centers - diffs
+    uppers = centers + diffs
+
+    return lowers, uppers
+
+
+@wp.kernel
+def compute_num_contact_with_checksums(
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+    bvh_id: wp.uint64,
+    counts: wp.array(dtype=int),
+    check_sums: wp.array(dtype=int),
+):
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
+    count = int(0)
+
+    check_sum = int(0)
+    index = int(0)
+    while wp.bvh_query_next(query, index):
+        check_sum = check_sum ^ index
+        count += 1
+
+    counts[tid] = count
+    check_sums[tid] = check_sum
+
+
+def test_capture_bvh_rebuild(test, device):
+    with wp.ScopedDevice(device):
+        rng = np.random.default_rng(123)
+
+        num_item_bounds = 100000
+        item_bound_size = 0.01
+
+        relative_shift = 2
+
+        num_test_bounds = 10000
+        test_bound_relative_size = 0.05
+
+        center = np.array([0.0, 0.0, 0.0])
+
+        item_lowers_np, item_uppers_np = get_random_aabbs(num_item_bounds, center, relative_shift, item_bound_size, rng)
+        item_lowers = wp.array(item_lowers_np, dtype=wp.vec3)
+        item_uppers = wp.array(item_uppers_np, dtype=wp.vec3)
+        bvh_1 = wp.Bvh(item_lowers, item_uppers)
+        item_lowers_2 = wp.zeros_like(item_lowers)
+        item_uppers_2 = wp.zeros_like(item_lowers)
+
+        test_lowers_np, test_uppers_np = get_random_aabbs(
+            num_test_bounds, center, relative_shift, test_bound_relative_size, rng
+        )
+        test_lowers = wp.array(test_lowers_np, dtype=wp.vec3)
+        test_uppers = wp.array(test_uppers_np, dtype=wp.vec3)
+
+        item_lowers_2_np, item_uppers_2_np = get_random_aabbs(
+            num_item_bounds,
+            center,
+            relative_shift,
+            item_bound_size,
+            rng,
+        )
+        item_lowers_2.assign(item_lowers_2_np)
+        item_uppers_2.assign(item_uppers_2_np)
+
+        counts_1 = wp.empty(n=num_test_bounds, dtype=int)
+        checksums_1 = wp.empty(n=num_test_bounds, dtype=int)
+        counts_2 = wp.empty(n=num_test_bounds, dtype=int)
+        checksums_2 = wp.empty(n=num_test_bounds, dtype=int)
+
+        wp.load_module(device=device)
+        with wp.ScopedCapture(force_module_load=False) as capture:
+            wp.copy(item_lowers, item_lowers_2)
+            wp.copy(item_uppers, item_uppers_2)
+            bvh_1.rebuild()
+            wp.launch(
+                compute_num_contact_with_checksums,
+                dim=num_test_bounds,
+                inputs=[test_lowers, test_uppers, bvh_1.id],
+                outputs=[counts_1, checksums_1],
+            )
+
+        cuda_graph = capture.graph
+
+        for _ in range(10):
+            item_lowers_2_np, item_uppers_2_np = get_random_aabbs(
+                num_item_bounds,
+                center,
+                relative_shift,
+                item_bound_size,
+                rng,
+            )
+            item_lowers_2.assign(item_lowers_2_np)
+            item_uppers_2.assign(item_uppers_2_np)
+
+            wp.capture_launch(cuda_graph)
+
+            bvh_2 = wp.Bvh(item_lowers_2, item_uppers_2)
+            wp.launch(
+                compute_num_contact_with_checksums,
+                dim=num_test_bounds,
+                inputs=[test_lowers, test_uppers, bvh_2.id],
+                outputs=[counts_2, checksums_2],
+                device=device,
+            )
+
+            assert_array_equal(counts_1, counts_2)
+            assert_array_equal(checksums_1, checksums_2)
+
+
+devices = get_test_devices()
+cuda_devices = get_cuda_test_devices()
+
+
+class TestBvh(unittest.TestCase):
+    def test_bvh_codegen_adjoints_with_select(self):
+        def kernel_fn(bvh: wp.uint64):
+            v = wp.vec3(0.0, 0.0, 0.0)
+            bounds_nr = int(0)
+
+            if True:
+                query_1 = wp.bvh_query_aabb(bvh, v, v)
+                query_2 = wp.bvh_query_ray(bvh, v, v)
+
+                wp.bvh_query_next(query_1, bounds_nr)
+                wp.bvh_query_next(query_2, bounds_nr)
+            else:
+                query_1 = wp.bvh_query_aabb(bvh, v, v)
+                query_2 = wp.bvh_query_ray(bvh, v, v)
+
+                wp.bvh_query_next(query_1, bounds_nr)
+                wp.bvh_query_next(query_2, bounds_nr)
+
+        wp.Kernel(func=kernel_fn)
+
+    def test_bvh_new_del(self):
+        # test the scenario in which a bvh is created but not initialized before gc
+        instance = wp.Bvh.__new__(wp.Bvh)
+        instance.__del__()
+
+
+add_function_test(TestBvh, "test_bvh_aabb", test_bvh_query_aabb, devices=devices)
+add_function_test(TestBvh, "test_bvh_ray", test_bvh_query_ray, devices=devices)
+add_function_test(
+    TestBvh,
+    "test_bvh_ray_query_inside_and_outside_bounds",
+    test_bvh_ray_query_inside_and_outside_bounds,
+    devices=devices,
+)
+
+add_function_test(TestBvh, "test_capture_bvh_rebuild", test_capture_bvh_rebuild, devices=cuda_devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/geometry/test_hash_grid.py

@@ -0,0 +1,259 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+num_points = 4096
+dim_x = 128
+dim_y = 128
+dim_z = 128
+
+scale = 150.0
+
+cell_radius = 8.0
+query_radius = 8.0
+
+num_runs = 4
+
+print_enabled = False
+
+
+@wp.kernel
+def count_neighbors(grid: wp.uint64, radius: float, points: wp.array(dtype=wp.vec3), counts: wp.array(dtype=int)):
+    tid = wp.tid()
+
+    # order threads by cell
+    i = wp.hash_grid_point_id(grid, tid)
+
+    # query point
+    p = points[i]
+    count = int(0)
+
+    # construct query around point p
+    for index in wp.hash_grid_query(grid, p, radius):
+        # compute distance to point
+        d = wp.length(p - points[index])
+
+        if d <= radius:
+            count += 1
+
+    counts[i] = count
+
+
+@wp.kernel
+def count_neighbors_reference(
+    radius: float, points: wp.array(dtype=wp.vec3), counts: wp.array(dtype=int), num_points: int
+):
+    tid = wp.tid()
+
+    i = tid % num_points
+    j = tid // num_points
+
+    # query point
+    p = points[i]
+    q = points[j]
+
+    # compute distance to point
+    d = wp.length(p - q)
+
+    if d <= radius:
+        wp.atomic_add(counts, i, 1)
+
+
+def particle_grid(dim_x, dim_y, dim_z, lower, radius, jitter):
+    rng = np.random.default_rng(123)
+    points = np.meshgrid(np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z))
+    points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
+    points_t = points_t + rng.random(size=points_t.shape) * radius * jitter
+
+    return points_t.reshape((-1, 3))
+
+
+def test_hashgrid_query(test, device):
+    wp.load_module(device=device)
+
+    grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+
+    for i in range(num_runs):
+        if print_enabled:
+            print(f"Run: {i + 1}")
+            print("---------")
+
+        points = particle_grid(16, 32, 16, (0.0, 0.3, 0.0), cell_radius * 0.25, 0.1)
+
+        points_arr = wp.array(points, dtype=wp.vec3, device=device)
+        counts_arr = wp.zeros(len(points), dtype=int, device=device)
+        counts_arr_ref = wp.zeros(len(points), dtype=int, device=device)
+
+        profiler = {}
+
+        with wp.ScopedTimer("grid operations", print=print_enabled, dict=profiler, synchronize=True):
+            with wp.ScopedTimer("brute", print=print_enabled, dict=profiler, synchronize=True):
+                wp.launch(
+                    kernel=count_neighbors_reference,
+                    dim=len(points) * len(points),
+                    inputs=[query_radius, points_arr, counts_arr_ref, len(points)],
+                    device=device,
+                )
+                wp.synchronize_device(device)
+
+            with wp.ScopedTimer("grid build", print=print_enabled, dict=profiler, synchronize=True):
+                grid.build(points_arr, cell_radius)
+
+            with wp.ScopedTimer("grid query", print=print_enabled, dict=profiler, synchronize=True):
+                wp.launch(
+                    kernel=count_neighbors,
+                    dim=len(points),
+                    inputs=[grid.id, query_radius, points_arr, counts_arr],
+                    device=device,
+                )
+
+        counts = counts_arr.numpy()
+        counts_ref = counts_arr_ref.numpy()
+
+        if print_enabled:
+            print(f"Grid min: {np.min(counts)} max: {np.max(counts)} avg: {np.mean(counts)}")
+            print(f"Ref min: {np.min(counts_ref)} max: {np.max(counts_ref)} avg: {np.mean(counts_ref)}")
+
+            print(f"Passed: {np.array_equal(counts, counts_ref)}")
+
+        assert_np_equal(counts, counts_ref)
+
+
+def test_hashgrid_inputs(test, 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, device=device)
+    counts_ref = wp.zeros(len(points), dtype=int, device=device)
+
+    grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+    grid.build(points_ref, cell_radius)
+
+    # get reference counts
+    wp.launch(
+        kernel=count_neighbors, dim=len(points), inputs=[grid.id, query_radius, points_ref, counts_ref], device=device
+    )
+
+    # test with strided 1d input arrays
+    for stride in [2, 3]:
+        with test.subTest(msg=f"stride_{stride}"):
+            points_buffer = wp.zeros(len(points) * stride, dtype=wp.vec3, device=device)
+            points_strided = points_buffer[::stride]
+            wp.copy(points_strided, points_ref)
+            counts_strided = wp.zeros(len(points), dtype=int, device=device)
+
+            grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+            grid.build(points_strided, cell_radius)
+
+            wp.launch(
+                kernel=count_neighbors,
+                dim=len(points),
+                inputs=[grid.id, query_radius, points_ref, counts_strided],
+                device=device,
+            )
+
+            assert_array_equal(counts_strided, counts_ref)
+
+    # test with multidimensional input arrays
+    for ndim in [2, 3, 4]:
+        with test.subTest(msg=f"ndim_{ndim}"):
+            shape = (len(points) // (2 ** (ndim - 1)), *((ndim - 1) * (2,)))
+            points_ndim = wp.zeros(shape, dtype=wp.vec3, device=device)
+            wp.copy(points_ndim, points_ref)
+            counts_ndim = wp.zeros(len(points), dtype=int, device=device)
+
+            grid = wp.HashGrid(dim_x, dim_y, dim_z, device)
+            grid.build(points_ndim, cell_radius)
+
+            wp.launch(
+                kernel=count_neighbors,
+                dim=len(points),
+                inputs=[grid.id, query_radius, points_ref, counts_ndim],
+                device=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):
+    def test_hashgrid_codegen_adjoints_with_select(self):
+        def kernel_fn(grid: wp.uint64):
+            v = wp.vec3(0.0, 0.0, 0.0)
+
+            if True:
+                query = wp.hash_grid_query(grid, v, 0.0)
+            else:
+                query = wp.hash_grid_query(grid, v, 0.0)
+
+        wp.Kernel(func=kernel_fn)
+
+    def test_hashgrid_new_del(self):
+        # test the scenario in which a hashgrid is created but not initialized before gc
+        instance = wp.HashGrid.__new__(wp.HashGrid)
+        instance.__del__()
+
+
+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__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)