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

warp/tests/test_matmul_lite.py

@@ -0,0 +1,410 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# 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.unittest_utils import *
+
+wp.init()
+
+from warp.context import runtime  # noqa: E402
+
+
+class gemm_test_bed_runner:
+    def __init__(self, dtype, device):
+        self.dtype = dtype
+        self.device = device
+
+    def alloc(self, m, n, k, batch_count):
+        rng = np.random.default_rng(42)
+        low = -4.5
+        high = 3.5
+        if batch_count == 1:
+            A = wp.array2d(
+                np.ceil(rng.uniform(low=low, high=high, size=(m, k))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            B = wp.array2d(
+                np.ceil(rng.uniform(low=low, high=high, size=(k, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            C = wp.array2d(
+                np.ceil(rng.uniform(low=low, high=high, size=(m, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            D = wp.array2d(np.zeros((m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
+        else:
+            A = wp.array3d(
+                np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            B = wp.array3d(
+                np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            C = wp.array3d(
+                np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            D = wp.array3d(np.zeros((batch_count, m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
+        return A, B, C, D
+
+    def run_and_verify(self, m, n, k, batch_count, alpha, beta):
+        A, B, C, D = self.alloc(m, n, k, batch_count)
+        ones = wp.zeros_like(D)
+        ones.fill_(1.0)
+
+        if batch_count == 1:
+            tape = wp.Tape()
+            with tape:
+                wp.matmul(A, B, C, D, alpha, beta, False, self.device)
+            tape.backward(grads={D: ones})
+
+            D_np = alpha * (A.numpy() @ B.numpy()) + beta * C.numpy()
+            assert np.array_equal(D_np, D.numpy())
+
+            adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose())
+            adj_B_np = alpha * (A.numpy().transpose() @ ones.numpy())
+            adj_C_np = beta * ones.numpy()
+
+        else:
+            tape = wp.Tape()
+            with tape:
+                wp.batched_matmul(A, B, C, D, alpha, beta, False, self.device)
+            tape.backward(grads={D: ones})
+
+            D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
+            assert np.array_equal(D_np, D.numpy())
+
+            adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
+            adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
+            adj_C_np = beta * ones.numpy()
+
+        assert np.array_equal(adj_A_np, A.grad.numpy())
+        assert np.array_equal(adj_B_np, B.grad.numpy())
+        assert np.array_equal(adj_C_np, C.grad.numpy())
+
+    def run(self):
+        Ms = [8]
+        Ns = [16]
+        Ks = [32]
+        batch_counts = [1]
+        betas = [1.0]
+        alpha = 1.0
+
+        for batch_count in batch_counts:
+            for m in Ms:
+                for n in Ns:
+                    for k in Ks:
+                        for beta in betas:
+                            self.run_and_verify(m, n, k, batch_count, alpha, beta)
+
+
+class gemm_test_bed_runner_transpose:
+    def __init__(self, dtype, device):
+        self.dtype = dtype
+        self.device = device
+
+    def alloc(self, m, n, k, batch_count):
+        rng = np.random.default_rng(42)
+        low = -4.5
+        high = 3.5
+        if batch_count == 1:
+            A = wp.array2d(
+                np.ceil(rng.uniform(low=low, high=high, size=(m, k))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            B = wp.array2d(
+                np.ceil(rng.uniform(low=low, high=high, size=(k, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            C = wp.array2d(
+                np.ceil(rng.uniform(low=low, high=high, size=(m, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            D = wp.array2d(np.zeros((m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
+            AT = wp.array2d(A.numpy().transpose([1, 0]), dtype=self.dtype, device=self.device, requires_grad=True)
+            BT = wp.array2d(B.numpy().transpose([1, 0]), dtype=self.dtype, device=self.device, requires_grad=True)
+        else:
+            A = wp.array3d(
+                np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            B = wp.array3d(
+                np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            C = wp.array3d(
+                np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))),
+                dtype=self.dtype,
+                device=self.device,
+                requires_grad=True,
+            )
+            D = wp.array3d(np.zeros((batch_count, m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
+            AT = wp.array3d(A.numpy().transpose([0, 2, 1]), dtype=self.dtype, device=self.device, requires_grad=True)
+            BT = wp.array3d(B.numpy().transpose([0, 2, 1]), dtype=self.dtype, device=self.device, requires_grad=True)
+        return A, B, C, D, AT, BT
+
+    def run_and_verify(self, m, n, k, batch_count, alpha, beta):
+        A, B, C1, D1, AT1, BT1 = self.alloc(m, n, k, batch_count)
+        C2 = wp.clone(C1)
+        C3 = wp.clone(C1)
+        D2 = wp.clone(D1)
+        D3 = wp.clone(D1)
+        AT2 = wp.clone(AT1)
+        BT2 = wp.clone(BT1)
+        ones1 = wp.zeros_like(D1)
+        ones1.fill_(1.0)
+        ones2 = wp.zeros_like(D2)
+        ones2.fill_(1.0)
+        ones3 = wp.zeros_like(D3)
+        ones3.fill_(1.0)
+
+        if batch_count == 1:
+            ATT1 = AT1.transpose([1, 0])
+            BTT1 = BT1.transpose([1, 0])
+            ATT2 = AT2.transpose([1, 0])
+            BTT2 = BT2.transpose([1, 0])
+            tape = wp.Tape()
+            with tape:
+                wp.matmul(A, BTT1, C1, D1, alpha, beta, False, self.device)
+                wp.matmul(ATT1, B, C2, D2, alpha, beta, False, self.device)
+                wp.matmul(ATT2, BTT2, C3, D3, alpha, beta, False, self.device)
+            tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
+
+            D_np = alpha * (A.numpy() @ B.numpy()) + beta * C1.numpy()
+            assert np.array_equal(D_np, D1.numpy())
+            assert np.array_equal(D_np, D2.numpy())
+            assert np.array_equal(D_np, D3.numpy())
+
+            adj_A_np = alpha * (ones1.numpy() @ B.numpy().transpose())
+            adj_B_np = alpha * (A.numpy().transpose() @ ones1.numpy())
+            adj_C_np = beta * ones1.numpy()
+
+        else:
+            ATT1 = AT1.transpose([0, 2, 1])
+            BTT1 = BT1.transpose([0, 2, 1])
+            ATT2 = AT2.transpose([0, 2, 1])
+            BTT2 = BT2.transpose([0, 2, 1])
+            tape = wp.Tape()
+            with tape:
+                wp.batched_matmul(A, BTT1, C1, D1, alpha, beta, False, self.device)
+                wp.batched_matmul(ATT1, B, C2, D2, alpha, beta, False, self.device)
+                wp.batched_matmul(ATT2, BTT2, C3, D3, alpha, beta, False, self.device)
+            tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
+
+            D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C1.numpy()
+            assert np.array_equal(D_np, D1.numpy())
+            assert np.array_equal(D_np, D2.numpy())
+            assert np.array_equal(D_np, D3.numpy())
+
+            adj_A_np = alpha * np.matmul(ones1.numpy(), B.numpy().transpose((0, 2, 1)))
+            adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones1.numpy())
+            adj_C_np = beta * ones1.numpy()
+
+        assert np.array_equal(adj_A_np, A.grad.numpy())
+        assert np.array_equal(adj_A_np, ATT1.grad.numpy())
+        assert np.array_equal(adj_A_np, ATT2.grad.numpy())
+        assert np.array_equal(adj_B_np, B.grad.numpy())
+        assert np.array_equal(adj_B_np, BTT1.grad.numpy())
+        assert np.array_equal(adj_B_np, BTT2.grad.numpy())
+        assert np.array_equal(adj_C_np, C1.grad.numpy())
+        assert np.array_equal(adj_C_np, C2.grad.numpy())
+        assert np.array_equal(adj_C_np, C3.grad.numpy())
+
+    def run(self):
+        m = 8
+        n = 16
+        k = 32
+        batch_counts = [1, 4]
+        beta = 1.0
+        alpha = 1.0
+
+        for batch_count in batch_counts:
+            self.run_and_verify(m, n, k, batch_count, alpha, beta)
+
+
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+def test_f32(test, device):
+    gemm_test_bed_runner(wp.float32, device).run()
+    gemm_test_bed_runner_transpose(wp.float32, device).run()
+
+
+@wp.kernel
+def matrix_sum_kernel(arr: wp.array2d(dtype=float), loss: wp.array(dtype=float)):
+    i, j = wp.tid()
+    wp.atomic_add(loss, 0, arr[i, j])
+
+
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+def test_tape(test, device):
+    rng = np.random.default_rng(42)
+    low = -4.5
+    high = 3.5
+    m = 8
+    n = 16
+    k = 32
+    A = wp.array2d(
+        np.ceil(rng.uniform(low=low, high=high, size=(m, k))), dtype=float, device=device, requires_grad=True
+    )
+    B = wp.array2d(
+        np.ceil(rng.uniform(low=low, high=high, size=(k, n))), dtype=float, device=device, requires_grad=True
+    )
+    C = wp.array2d(
+        np.ceil(rng.uniform(low=low, high=high, size=(m, n))), dtype=float, device=device, requires_grad=True
+    )
+    D = wp.array2d(np.zeros((m, n)), dtype=float, device=device, requires_grad=True)
+    loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
+
+    # test tape
+    tape = wp.Tape()
+    with tape:
+        wp.matmul(A, B, C, D, device=device)
+        wp.launch(matrix_sum_kernel, dim=(m, n), inputs=[D, loss], device=device)
+
+    tape.backward(loss=loss)
+    A_grad = A.grad.numpy()
+    tape.reset()
+
+    # test adjoint
+    D.grad = wp.array2d(np.ones((m, n)), dtype=float, device=device)
+    wp.adj_matmul(A, B, C, A.grad, B.grad, C.grad, D.grad, device=device)
+    assert_np_equal(A_grad, A.grad.numpy())
+
+    # test zero
+    tape.zero()
+    assert_array_equal(A.grad, wp.zeros_like(A))
+
+
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+def test_operator(test, device):
+    rng = np.random.default_rng(42)
+    low = -4.5
+    high = 3.5
+    m = 8
+    n = 16
+    k = 32
+    A = wp.array2d(
+        np.ceil(rng.uniform(low=low, high=high, size=(m, k))), dtype=float, device=device, requires_grad=True
+    )
+    B = wp.array2d(
+        np.ceil(rng.uniform(low=low, high=high, size=(k, n))), dtype=float, device=device, requires_grad=True
+    )
+    loss = wp.zeros(1, dtype=float, device=device, requires_grad=True)
+
+    # test tape
+    tape = wp.Tape()
+    with tape:
+        D = A @ B
+        wp.launch(matrix_sum_kernel, dim=(m, n), inputs=[D, loss], device=device)
+
+    tape.backward(loss=loss)
+
+    # test adjoint
+    D.grad = wp.array2d(np.ones((m, n)), dtype=float, device=device)
+    B_transpose = wp.array2d(B.transpose().numpy(), dtype=float, device=device)
+
+    adj_A = D.grad @ B_transpose
+    assert_array_equal(adj_A, A.grad)
+
+    # test zero
+    tape.zero()
+    assert_array_equal(A.grad, wp.zeros_like(A))
+
+
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
+def test_large_batch_count(test, device):
+    rng = np.random.default_rng(42)
+    low = -4.5
+    high = 3.5
+    m = 2
+    n = 3
+    k = 4
+    batch_count = 65535 * 2 + int(65535 / 2)
+    A = wp.array3d(
+        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))),
+        dtype=float,
+        device=device,
+        requires_grad=True,
+    )
+    B = wp.array3d(
+        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))),
+        dtype=float,
+        device=device,
+        requires_grad=True,
+    )
+    C = wp.array3d(
+        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))),
+        dtype=float,
+        device=device,
+        requires_grad=True,
+    )
+    D = wp.array3d(np.zeros((batch_count, m, n)), dtype=float, device=device, requires_grad=True)
+    ones = wp.zeros_like(D)
+    ones.fill_(1.0)
+
+    alpha = 1.0
+    beta = 1.0
+
+    tape = wp.Tape()
+    with tape:
+        wp.batched_matmul(A, B, C, D, alpha=alpha, beta=beta, allow_tf32x3_arith=False, device=device)
+    tape.backward(grads={D: ones})
+
+    D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
+    assert np.array_equal(D_np, D.numpy())
+
+    adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
+    adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
+    adj_C_np = beta * ones.numpy()
+
+    assert np.array_equal(adj_A_np, A.grad.numpy())
+    assert np.array_equal(adj_B_np, B.grad.numpy())
+    assert np.array_equal(adj_C_np, C.grad.numpy())
+
+
+devices = get_test_devices()
+
+
+class TestMatmulLite(unittest.TestCase):
+    pass
+
+
+add_function_test(TestMatmulLite, "test_f32", test_f32, devices=devices)
+add_function_test(TestMatmulLite, "test_tape", test_tape, devices=devices)
+add_function_test(TestMatmulLite, "test_operator", test_operator, devices=devices)
+add_function_test(TestMatmulLite, "test_large_batch_count", test_large_batch_count, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)

warp/tests/test_mesh.py

@@ -10,8 +10,7 @@ import unittest
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
-
+from warp.tests.unittest_utils import *
 
 # fmt: off
 
@@ -223,9 +222,9 @@ def query_ray_kernel(
 
 
 def test_mesh_query_ray(test, device):
-    points = wp.array(POINT_POSITIONS, dtype=wp.vec3)
+    points = wp.array(POINT_POSITIONS, dtype=wp.vec3, device=device)
 
-    indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int)
+    indices = wp.array(RIGHT_HANDED_FACE_VERTEX_INDICES, dtype=int, device=device)
     mesh = wp.Mesh(points=points, indices=indices)
     expected_sign = -1.0
     wp.launch(
@@ -235,9 +234,10 @@ def test_mesh_query_ray(test, device):
             mesh.id,
             expected_sign,
         ],
+        device=device,
     )
 
-    indices = wp.array(LEFT_HANDED_FACE_VERTEX_INDICES, dtype=int)
+    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(
@@ -247,21 +247,78 @@ def test_mesh_query_ray(test, device):
             mesh.id,
             expected_sign,
         ],
+        device=device,
     )
 
 
-def register(parent):
-    devices = get_test_devices()
+def test_mesh_refit_graph(test, device):
+    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)
+
+    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):
+    pass
 
-    class TestMesh(parent):
-        pass
 
-    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)
-    return TestMesh
+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_unique_cuda_test_devices())
+add_function_test(TestMesh, "test_mesh_exceptions", test_mesh_exceptions, devices=get_unique_cuda_test_devices())
 
 
 if __name__ == "__main__":
-    _ = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_mesh_query_aabb.py

@@ -5,10 +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()
 
@@ -96,7 +98,6 @@ def test_compute_bounds(test, device):
 
     lower_view = lowers.numpy()
     upper_view = uppers.numpy()
-    wp.synchronize()
 
     # Confirm the bounds of each triangle are correct.
     test.assertTrue(lower_view[0][0] == 0)
@@ -148,8 +149,6 @@ def test_mesh_query_aabb_count_overlap(test, device):
         device=device,
     )
 
-    wp.synchronize()
-
     view = counts.numpy()
 
     # 2 triangles that share a vertex having overlapping AABBs.
@@ -188,8 +187,6 @@ def test_mesh_query_aabb_count_nonoverlap(test, device):
         device=device,
     )
 
-    wp.synchronize()
-
     view = counts.numpy()
 
     # AABB query only returns one triangle at a time, the triangles are not close enough to overlap.
@@ -197,29 +194,48 @@ def test_mesh_query_aabb_count_nonoverlap(test, device):
         test.assertTrue(c == 1)
 
 
-def register(parent):
-    devices = get_test_devices()
+def test_mesh_query_aabb_codegen_adjoints_with_select(test, device):
+    def kernel_fn(
+        mesh: wp.uint64,
+    ):
+        v = wp.vec3(0.0, 0.0, 0.0)
 
-    class TestMeshQueryAABBMethods(parent):
-        pass
+        if True:
+            query = wp.mesh_query_aabb(mesh, v, v)
+        else:
+            query = wp.mesh_query_aabb(mesh, v, v)
+
+    wp.Kernel(func=kernel_fn)
+
+
+devices = get_test_devices()
+
+
+class TestMeshQueryAABBMethods(unittest.TestCase):
+    pass
 
-    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,
-    )
 
-    return TestMeshQueryAABBMethods
+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_codegen_adjoints_with_select",
+    test_mesh_query_aabb_codegen_adjoints_with_select,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)