warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/tests/geometry/test_mesh.py

@@ -0,0 +1,316 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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 *
+
+# fmt: off
+
+POINT_POSITIONS = (
+    ( 0.5, -0.5,  0.5),
+    (-0.5, -0.5,  0.5),
+    ( 0.5,  0.5,  0.5),
+    (-0.5,  0.5,  0.5),
+    (-0.5, -0.5, -0.5),
+    ( 0.5, -0.5, -0.5),
+    (-0.5,  0.5, -0.5),
+    ( 0.5,  0.5, -0.5),
+)
+
+# Right-hand winding order. This corresponds to USD's (and others).
+RIGHT_HANDED_FACE_VERTEX_INDICES = (
+    0, 3, 1,
+    0, 2, 3,
+    4, 7, 5,
+    4, 6, 7,
+    6, 2, 7,
+    6, 3, 2,
+    5, 1, 4,
+    5, 0, 1,
+    5, 2, 0,
+    5, 7, 2,
+    1, 6, 4,
+    1, 3, 6,
+)
+
+
+# Left-hand winding order. This corresponds to Houdini's (and others).
+LEFT_HANDED_FACE_VERTEX_INDICES = (
+    0, 1, 3,
+    0, 3, 2,
+    4, 5, 7,
+    4, 7, 6,
+    6, 7, 2,
+    6, 2, 3,
+    5, 4, 1,
+    5, 1, 0,
+    5, 0, 2,
+    5, 2, 7,
+    1, 4, 6,
+    1, 6, 3,
+)
+
+# fmt: on
+
+POINT_COUNT = 8
+VERTEX_COUNT = 36
+FACE_COUNT = 12
+
+
+@wp.kernel(enable_backward=False)
+def read_points_kernel(
+    mesh_id: wp.uint64,
+    out_points: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    mesh = wp.mesh_get(mesh_id)
+    out_points[tid] = mesh.points[tid]
+
+
+@wp.kernel(enable_backward=False)
+def read_indices_kernel(
+    mesh_id: wp.uint64,
+    out_indices: wp.array(dtype=int),
+):
+    tid = wp.tid()
+    mesh = wp.mesh_get(mesh_id)
+    out_indices[tid * 3 + 0] = mesh.indices[tid * 3 + 0]
+    out_indices[tid * 3 + 1] = mesh.indices[tid * 3 + 1]
+    out_indices[tid * 3 + 2] = mesh.indices[tid * 3 + 2]
+
+
+def test_mesh_read_properties(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+
+        assert mesh.points.size == POINT_COUNT
+        assert mesh.indices.size == VERTEX_COUNT
+        assert int(mesh.indices.size / 3) == FACE_COUNT
+
+        out_points = wp.empty(POINT_COUNT, dtype=wp.vec3, device=device)
+        wp.launch(read_points_kernel, dim=POINT_COUNT, inputs=[mesh.id], outputs=[out_points], device=device)
+        assert_np_equal(out_points.numpy(), np.array(POINT_POSITIONS))
+
+        out_indices = wp.empty(VERTEX_COUNT, dtype=int, device=device)
+        wp.launch(read_indices_kernel, dim=FACE_COUNT, inputs=[mesh.id], outputs=[out_indices], device=device)
+        assert_np_equal(out_indices.numpy(), np.array(RIGHT_HANDED_FACE_VERTEX_INDICES))
+
+
+@wp.kernel(enable_backward=False)
+def query_point_kernel(
+    mesh_id: wp.uint64,
+    expected_sign: float,
+):
+    point = wp.vec3(0.1, 0.2, 0.3)
+    expected_pos = wp.vec3(0.1, 0.2, 0.5)
+
+    sign = float(0.0)
+    face = int(0)
+    bary_u = float(0.0)
+    bary_v = float(0.0)
+
+    wp.mesh_query_point(mesh_id, point, 1e6, sign, face, bary_u, bary_v)
+    pos = wp.mesh_eval_position(mesh_id, face, bary_u, bary_v)
+
+    wp.expect_eq(wp.sign(sign), expected_sign)
+    wp.expect_eq(face, 1)
+    wp.expect_near(wp.length(pos - expected_pos), 0.0)
+
+
+def test_mesh_query_point(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+        expected_sign = -1.0
+        wp.launch(query_point_kernel, dim=1, inputs=[mesh.id, expected_sign], device=device)
+
+        indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices)
+        expected_sign = 1.0
+        wp.launch(query_point_kernel, dim=1, inputs=[mesh.id, expected_sign], device=device)
+
+
+@wp.kernel(enable_backward=False)
+def query_ray_kernel(
+    mesh_id: wp.uint64,
+    expected_sign: float,
+):
+    start = wp.vec3(0.1, 0.2, 0.3)
+    dir = wp.normalize(wp.vec3(-1.2, 2.3, -3.4))
+    expected_t = 0.557828
+    expected_pos = wp.vec3(-0.0565217, 0.5, -0.143478)
+
+    t = float(0.0)
+    bary_u = float(0.0)
+    bary_v = float(0.0)
+    sign = float(0.0)
+    normal = wp.vec3(0.0, 0.0, 0.0)
+    face = int(0)
+
+    wp.mesh_query_ray(
+        mesh_id,
+        start,
+        dir,
+        1e6,
+        t,
+        bary_u,
+        bary_v,
+        sign,
+        normal,
+        face,
+    )
+    pos = wp.mesh_eval_position(mesh_id, face, bary_u, bary_v)
+
+    wp.expect_near(t, expected_t)
+    wp.expect_near(t, wp.length(pos - start), 1e-6)
+    wp.expect_eq(wp.sign(sign), expected_sign)
+    wp.expect_eq(face, 4)
+    wp.expect_near(wp.length(pos - expected_pos), 0.0, 1e-6)
+
+
+def test_mesh_query_ray(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+        expected_sign = -1.0
+        wp.launch(
+            query_ray_kernel,
+            dim=1,
+            inputs=[
+                mesh.id,
+                expected_sign,
+            ],
+            device=device,
+        )
+
+        indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices)
+        expected_sign = 1.0
+        wp.launch(
+            query_ray_kernel,
+            dim=1,
+            inputs=[
+                mesh.id,
+                expected_sign,
+            ],
+            device=device,
+        )
+
+
+def test_mesh_refit_graph(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    for constructor in constructors:
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        mesh = wp.Mesh(points=points, indices=indices, bvh_constructor=constructor)
+
+        wp.capture_begin(device, force_module_load=False)
+        try:
+            mesh.refit()
+        finally:
+            graph = wp.capture_end(device)
+
+        # replay
+        num_iters = 10
+        for _ in range(num_iters):
+            wp.capture_launch(graph)
+
+        wp.synchronize_device(device)
+
+
+def test_mesh_exceptions(test, device):
+    # points and indices must be on same device
+    with test.assertRaises(RuntimeError):
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device="cpu")
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        wp.Mesh(points=points, indices=indices)
+
+    # points must be vec3
+    with test.assertRaises(RuntimeError):
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3d, device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        wp.Mesh(points=points, indices=indices)
+
+    # velocities must be vec3
+    with test.assertRaises(RuntimeError):
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+        velocities = wp.zeros(points.shape, dtype=wp.vec3d, device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        wp.Mesh(points=points, indices=indices, velocities=velocities)
+
+    # indices must be int32
+    with test.assertRaises(RuntimeError):
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=wp.int64, device=device)
+        wp.Mesh(points=points, indices=indices)
+
+    # indices must be 1d
+    with test.assertRaises(RuntimeError):
+        points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
+        indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
+        indices = indices.reshape((3, -1))
+        wp.Mesh(points=points, indices=indices)
+
+
+devices = get_test_devices()
+
+
+class TestMesh(unittest.TestCase):
+    def test_mesh_new_del(self):
+        # test the scenario in which a mesh is created but not initialized before gc
+        instance = wp.Mesh.__new__(wp.Mesh)
+        instance.__del__()
+
+
+add_function_test(TestMesh, "test_mesh_read_properties", test_mesh_read_properties, devices=devices)
+add_function_test(TestMesh, "test_mesh_query_point", test_mesh_query_point, devices=devices)
+add_function_test(TestMesh, "test_mesh_query_ray", test_mesh_query_ray, devices=devices)
+add_function_test(TestMesh, "test_mesh_refit_graph", test_mesh_refit_graph, devices=get_selected_cuda_test_devices())
+add_function_test(TestMesh, "test_mesh_exceptions", test_mesh_exceptions, devices=get_selected_cuda_test_devices())
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/geometry/test_mesh_query_aabb.py

@@ -0,0 +1,399 @@
+# 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 os
+import unittest
+
+import numpy as np
+
+import warp as wp
+import warp.examples
+from warp.tests.unittest_utils import *
+
+
+@wp.func
+def min_vec3(a: wp.vec3, b: wp.vec3):
+    return wp.vec3(wp.min(a[0], b[0]), wp.min(a[1], b[1]), wp.min(a[2], b[2]))
+
+
+@wp.func
+def max_vec3(a: wp.vec3, b: wp.vec3):
+    return wp.vec3(wp.max(a[0], b[0]), wp.max(a[1], b[1]), wp.max(a[2], b[2]))
+
+
+@wp.kernel
+def compute_bounds(
+    indices: wp.array(dtype=int),
+    positions: wp.array(dtype=wp.vec3),
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+):
+    tid = wp.tid()
+    i = indices[tid * 3 + 0]
+    j = indices[tid * 3 + 1]
+    k = indices[tid * 3 + 2]
+
+    x0 = positions[i]  # point zero
+    x1 = positions[j]  # point one
+    x2 = positions[k]  # point two
+
+    lower = min_vec3(min_vec3(x0, x1), x2)
+    upper = max_vec3(max_vec3(x0, x1), x2)
+
+    lowers[tid] = lower
+    uppers[tid] = upper
+
+
+@wp.kernel
+def compute_num_contacts(
+    lowers: wp.array(dtype=wp.vec3), uppers: wp.array(dtype=wp.vec3), mesh_id: wp.uint64, counts: wp.array(dtype=int)
+):
+    face_index = int(0)
+    face_u = float(0.0)
+    face_v = float(0.0)
+    sign = float(0.0)
+
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    query = wp.mesh_query_aabb(mesh_id, lower, upper)
+    count = int(0)
+
+    # index = int(-1)
+    # while wp.mesh_query_aabb_next(query, index):
+
+    for _index in query:
+        count = count + 1
+
+    counts[tid] = count
+
+
+def test_compute_bounds(test, device):
+    # create two touching triangles.
+    points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [-1, -1, 1]])
+    indices = np.array([0, 1, 2, 1, 2, 3])
+    m = wp.Mesh(
+        points=wp.array(points, dtype=wp.vec3, device=device),
+        indices=wp.array(indices, dtype=int, device=device),
+    )
+
+    num_tris = int(len(indices) / 3)
+
+    # First compute bounds of each of the triangles.
+    lowers = wp.empty(n=num_tris, dtype=wp.vec3, device=device)
+    uppers = wp.empty_like(lowers)
+    wp.launch(
+        kernel=compute_bounds,
+        dim=num_tris,
+        inputs=[m.indices, m.points],
+        outputs=[lowers, uppers],
+        device=device,
+    )
+
+    lower_view = lowers.numpy()
+    upper_view = uppers.numpy()
+
+    # Confirm the bounds of each triangle are correct.
+    test.assertTrue(lower_view[0][0] == 0)
+    test.assertTrue(lower_view[0][1] == 0)
+    test.assertTrue(lower_view[0][2] == 0)
+
+    test.assertTrue(upper_view[0][0] == 1)
+    test.assertTrue(upper_view[0][1] == 1)
+    test.assertTrue(upper_view[0][2] == 0)
+
+    test.assertTrue(lower_view[1][0] == -1)
+    test.assertTrue(lower_view[1][1] == -1)
+    test.assertTrue(lower_view[1][2] == 0)
+
+    test.assertTrue(upper_view[1][0] == 1)
+    test.assertTrue(upper_view[1][1] == 1)
+    test.assertTrue(upper_view[1][2] == 1)
+
+
+def test_mesh_query_aabb_count_overlap(test, device):
+    # create two touching triangles.
+    points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [-1, -1, 1]])
+    indices = np.array([0, 1, 2, 1, 2, 3])
+    m = wp.Mesh(
+        points=wp.array(points, dtype=wp.vec3, device=device),
+        indices=wp.array(indices, dtype=int, device=device),
+    )
+
+    num_tris = int(len(indices) / 3)
+
+    # Compute AABB of each of the triangles.
+    lowers = wp.empty(n=num_tris, dtype=wp.vec3, device=device)
+    uppers = wp.empty_like(lowers)
+    wp.launch(
+        kernel=compute_bounds,
+        dim=num_tris,
+        inputs=[m.indices, m.points],
+        outputs=[lowers, uppers],
+        device=device,
+    )
+
+    counts = wp.empty(n=num_tris, dtype=int, device=device)
+
+    wp.launch(
+        kernel=compute_num_contacts,
+        dim=num_tris,
+        inputs=[lowers, uppers, m.id],
+        outputs=[counts],
+        device=device,
+    )
+
+    view = counts.numpy()
+
+    # 2 triangles that share a vertex having overlapping AABBs.
+    for c in view:
+        test.assertTrue(c == 2)
+
+
+def test_mesh_query_aabb_count_nonoverlap(test, device):
+    # create two separate triangles.
+    points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [10, 0, 0], [10, 1, 0], [10, 0, 1]])
+    indices = np.array([0, 1, 2, 3, 4, 5])
+    m = wp.Mesh(
+        points=wp.array(points, dtype=wp.vec3, device=device),
+        indices=wp.array(indices, dtype=int, device=device),
+    )
+
+    num_tris = int(len(indices) / 3)
+
+    lowers = wp.empty(n=num_tris, dtype=wp.vec3, device=device)
+    uppers = wp.empty_like(lowers)
+    wp.launch(
+        kernel=compute_bounds,
+        dim=num_tris,
+        inputs=[m.indices, m.points],
+        outputs=[lowers, uppers],
+        device=device,
+    )
+
+    counts = wp.empty(n=num_tris, dtype=int, device=device)
+
+    wp.launch(
+        kernel=compute_num_contacts,
+        dim=num_tris,
+        inputs=[lowers, uppers, m.id],
+        outputs=[counts],
+        device=device,
+    )
+
+    view = counts.numpy()
+
+    # AABB query only returns one triangle at a time, the triangles are not close enough to overlap.
+    for c in view:
+        test.assertTrue(c == 1)
+
+
+@wp.kernel
+def compute_num_contact_with_checksums(
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+    mesh_id: wp.uint64,
+    counts: wp.array(dtype=int),
+    check_sums: wp.array(dtype=int),
+):
+    face_index = int(0)
+    face_u = float(0.0)
+    face_v = float(0.0)
+    sign = float(0.0)
+
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    query = wp.mesh_query_aabb(mesh_id, lower, upper)
+    count = int(0)
+
+    check_sum = int(0)
+    for _index in query:
+        check_sum = check_sum ^ _index
+        count = count + 1
+
+    counts[tid] = count
+    check_sums[tid] = check_sum
+
+
+@wp.func
+def intersect_aabb_aabb(a_lower: wp.vec3, a_upper: wp.vec3, b_lower: wp.vec3, b_upper: wp.vec3):
+    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 False
+    else:
+        return True
+
+
+@wp.kernel
+def compute_num_contact_with_checksums_brutal(
+    lowers: wp.array(dtype=wp.vec3),
+    uppers: wp.array(dtype=wp.vec3),
+    mesh_points: wp.array(dtype=wp.vec3),
+    mesh_indices: wp.array(dtype=int),
+    counts: wp.array(dtype=int),
+    check_sums: wp.array(dtype=int),
+):
+    tid = wp.tid()
+
+    upper = uppers[tid]
+    lower = lowers[tid]
+
+    check_sum = int(0)
+    count = int(0)
+    num_faces = mesh_indices.shape[0] / 3
+
+    for face in range(num_faces):
+        i = mesh_indices[face * 3 + 0]
+        j = mesh_indices[face * 3 + 1]
+        k = mesh_indices[face * 3 + 2]
+
+        x0 = mesh_points[i]  # point zero
+        x1 = mesh_points[j]  # point one
+        x2 = mesh_points[k]  # point two
+
+        tri_lower = min_vec3(min_vec3(x0, x1), x2)
+        tri_upper = max_vec3(max_vec3(x0, x1), x2)
+
+        if intersect_aabb_aabb(lower, upper, tri_lower, tri_upper):
+            check_sum = check_sum ^ face
+            count = count + 1
+
+    counts[tid] = count
+    check_sums[tid] = check_sum
+
+
+def load_mesh():
+    from pxr import Usd, UsdGeom
+
+    usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+    usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+    vertices = np.array(usd_geom.GetPointsAttr().Get())
+    faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+    return vertices, faces
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_mesh_query_aabb_count_overlap_with_checksum(test, device):
+    if device.is_cpu:
+        constructors = ["sah", "median"]
+    else:
+        constructors = ["sah", "median", "lbvh"]
+
+    points, indices = load_mesh()
+    points_wp = wp.array(points, dtype=wp.vec3, device=device)
+    indices_wp = wp.array(indices, dtype=int, device=device)
+
+    for constructor in constructors:
+        m = wp.Mesh(points=points_wp, indices=indices_wp, bvh_constructor=constructor)
+
+        num_test_bounds = 10000
+        test_bound_relative_size = 0.01
+
+        world_min = np.min(points, axis=0)
+        world_max = np.max(points, axis=0)
+
+        world_center = 0.5 * (world_min + world_max)
+        world_size = world_max - world_min
+
+        rng = np.random.default_rng(123)
+
+        centers = (
+            rng.uniform(-0.5, 0.5, size=num_test_bounds * 3).reshape(num_test_bounds, 3) * world_size + world_center
+        )
+        diffs = (
+            0.5 * test_bound_relative_size * rng.random(num_test_bounds * 3).reshape(num_test_bounds, 3) * world_size
+        )
+
+        lowers = wp.array(centers - diffs, dtype=wp.vec3, device=device)
+        uppers = wp.array(centers + diffs, dtype=wp.vec3, device=device)
+
+        counts = wp.empty(n=num_test_bounds, dtype=int, device=device)
+        checksums = wp.empty(n=num_test_bounds, dtype=int, device=device)
+
+        wp.launch(
+            kernel=compute_num_contact_with_checksums,
+            dim=num_test_bounds,
+            inputs=[lowers, uppers, m.id],
+            outputs=[counts, checksums],
+            device=device,
+        )
+
+        counts_brutal = wp.empty(n=num_test_bounds, dtype=int, device=device)
+        checksums_brutal = wp.empty(n=num_test_bounds, dtype=int, device=device)
+
+        wp.launch(
+            kernel=compute_num_contact_with_checksums_brutal,
+            dim=num_test_bounds,
+            inputs=[lowers, uppers, points_wp, indices_wp],
+            outputs=[counts_brutal, checksums_brutal],
+            device=device,
+        )
+
+        assert_array_equal(counts, counts_brutal)
+        assert_array_equal(checksums, checksums_brutal)
+
+
+devices = get_test_devices()
+
+
+class TestMeshQueryAABBMethods(unittest.TestCase):
+    def test_mesh_query_aabb_codegen_adjoints_with_select(self):
+        def kernel_fn(
+            mesh: wp.uint64,
+        ):
+            v = wp.vec3(0.0, 0.0, 0.0)
+
+            if True:
+                query = wp.mesh_query_aabb(mesh, v, v)
+            else:
+                query = wp.mesh_query_aabb(mesh, v, v)
+
+        wp.Kernel(func=kernel_fn)
+
+
+add_function_test(TestMeshQueryAABBMethods, "test_compute_bounds", test_compute_bounds, devices=devices)
+add_function_test(
+    TestMeshQueryAABBMethods, "test_mesh_query_aabb_count_overlap", test_mesh_query_aabb_count_overlap, devices=devices
+)
+add_function_test(
+    TestMeshQueryAABBMethods,
+    "test_mesh_query_aabb_count_nonoverlap",
+    test_mesh_query_aabb_count_nonoverlap,
+    devices=devices,
+)
+add_function_test(
+    TestMeshQueryAABBMethods,
+    "test_mesh_query_aabb_count_overlap_with_checksum",
+    test_mesh_query_aabb_count_overlap_with_checksum,
+    devices=devices,
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)