warp-lang 1.8.0__py3-none-win_amd64.whl → 1.9.0__py3-none-win_amd64.whl

warp/tests/sim/test_coloring.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
+
 import numpy as np
 
 import warp as wp
@@ -48,6 +50,63 @@ def color_lattice_grid(num_x, num_y):
     return color_groups
 
 
+def create_lattice_grid(N):
+    size = 10
+    position = (0, 0)
+
+    X = np.linspace(-0.5 * size + position[0], 0.5 * size + position[0], N)
+    Y = np.linspace(-0.5 * size + position[1], 0.5 * size + position[1], N)
+
+    X, Y = np.meshgrid(X, Y)
+
+    Z = []
+    for _i in range(N):
+        Z.append(np.linspace(0, size, N))
+
+    Z = np.array(Z)
+
+    vs = []
+    for i, j in itertools.product(range(N), range(N)):
+        vs.append(wp.vec3((X[i, j], Y[i, j], Z[i, j])))
+
+    fs = []
+    for i, j in itertools.product(range(0, N - 1), range(0, N - 1)):
+        vId = j + i * N
+
+        if (j + i) % 2:
+            fs.extend(
+                [
+                    vId,
+                    vId + N + 1,
+                    vId + 1,
+                ]
+            )
+            fs.extend(
+                [
+                    vId,
+                    vId + N,
+                    vId + N + 1,
+                ]
+            )
+        else:
+            fs.extend(
+                [
+                    vId,
+                    vId + N,
+                    vId + 1,
+                ]
+            )
+            fs.extend(
+                [
+                    vId + N,
+                    vId + N + 1,
+                    vId + 1,
+                ]
+            )
+
+    return vs, fs
+
+
 def test_coloring_corner_case(test, device):
     builder_1 = wp.sim.ModelBuilder()
     builder_1.color()
@@ -91,7 +150,7 @@ def test_coloring_trimesh(test, device):
         edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
 
         # coloring without bending
-        num_colors_greedy = wp.context.runtime.core.graph_coloring(
+        num_colors_greedy = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu.__ctype__(),
             ColoringAlgorithm.GREEDY.value,
@@ -104,7 +163,7 @@ def test_coloring_trimesh(test, device):
             device="cpu",
         )
 
-        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+        num_colors_mcs = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu.__ctype__(),
             ColoringAlgorithm.MCS.value,
@@ -119,13 +178,13 @@ def test_coloring_trimesh(test, device):
 
         # coloring with bending
         edge_indices_cpu_with_bending = construct_trimesh_graph_edges(model.edge_indices, True)
-        num_colors_greedy = wp.context.runtime.core.graph_coloring(
+        num_colors_greedy = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             ColoringAlgorithm.GREEDY.value,
             particle_colors.__ctype__(),
         )
-        wp.context.runtime.core.balance_coloring(
+        wp.context.runtime.core.wp_balance_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             num_colors_greedy,
@@ -139,13 +198,13 @@ def test_coloring_trimesh(test, device):
             device="cpu",
         )
 
-        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+        num_colors_mcs = wp.context.runtime.core.wp_graph_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             ColoringAlgorithm.MCS.value,
             particle_colors.__ctype__(),
         )
-        max_min_ratio = wp.context.runtime.core.balance_coloring(
+        max_min_ratio = wp.context.runtime.core.wp_balance_coloring(
             model.particle_count,
             edge_indices_cpu_with_bending.__ctype__(),
             num_colors_mcs,
@@ -164,6 +223,22 @@ def test_coloring_trimesh(test, device):
         color_sizes = np.array([c.shape[0] for c in color_categories_balanced], dtype=np.float32)
         test.assertTrue(np.max(color_sizes) / np.min(color_sizes) <= max_min_ratio)
 
+        # test if the color balance can quit from equilibrium
+        builder = wp.sim.ModelBuilder()
+
+        vs, fs = create_lattice_grid(100)
+        builder.add_cloth_mesh(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            scale=1.0,
+            vertices=vs,
+            indices=fs,
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            density=0.02,
+        )
+
+        builder.color(include_bending=True)
+
 
 @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
 def test_combine_coloring(test, device):
@@ -263,7 +338,7 @@ class TestColoring(unittest.TestCase):
     pass
 
 
-add_function_test(TestColoring, "test_coloring_trimesh", test_coloring_trimesh, devices=devices)
+add_function_test(TestColoring, "test_coloring_trimesh", test_coloring_trimesh, devices=devices, check_output=False)
 add_function_test(TestColoring, "test_combine_coloring", test_combine_coloring, devices=devices)
 add_function_test(TestColoring, "test_coloring_corner_case", test_coloring_corner_case, devices=devices)
 

warp/tests/test_array.py

@@ -2902,10 +2902,8 @@ def test_direct_from_numpy(test, device):
 
 
 @wp.kernel
-def kernel_array_from_ptr(
-    ptr: wp.uint64,
-):
-    arr = wp.array(ptr=ptr, shape=(2, 3), dtype=wp.float32)
+def kernel_array_from_ptr(arr_orig: wp.array2d(dtype=wp.float32)):
+    arr = wp.array(ptr=arr_orig.ptr, shape=(2, 3), dtype=wp.float32)
     arr[0, 0] = 1.0
     arr[0, 1] = 2.0
     arr[0, 2] = 3.0
@@ -2913,7 +2911,56 @@ def kernel_array_from_ptr(
 
 def test_kernel_array_from_ptr(test, device):
     arr = wp.zeros(shape=(2, 3), dtype=wp.float32, device=device)
-    wp.launch(kernel_array_from_ptr, dim=(1,), inputs=(arr.ptr,), device=device)
+    wp.launch(kernel_array_from_ptr, dim=(1,), inputs=(arr,), device=device)
+    assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 3.0), (0.0, 0.0, 0.0))))
+
+
+@wp.struct
+class MyStruct:
+    a: wp.float32
+    b: wp.float32
+    c: wp.float32
+
+
+@wp.kernel
+def kernel_array_from_ptr_struct(arr_orig: wp.array(dtype=MyStruct)):
+    arr = wp.array(ptr=arr_orig.ptr, shape=(2,), dtype=MyStruct)
+    arr[0].a = 1.0
+    arr[0].b = 2.0
+    arr[0].c = 3.0
+    arr[1].a = 4.0
+    arr[1].b = 5.0
+    arr[1].c = 6.0
+
+
+def test_kernel_array_from_ptr_struct(test, device):
+    arr = wp.zeros(shape=(2,), dtype=MyStruct, device=device)
+    wp.launch(kernel_array_from_ptr_struct, dim=(1,), inputs=(arr,), device=device)
+    arr_np = arr.numpy()
+    expected = np.zeros_like(arr_np)
+    expected[0] = (1.0, 2.0, 3.0)
+    expected[1] = (4.0, 5.0, 6.0)
+    assert_np_equal(arr_np, expected)
+
+
+@wp.kernel
+def kernel_array_from_ptr_variable_shape(
+    ptr: wp.uint64,
+    shape_x: int,
+    shape_y: int,
+):
+    arr = wp.array(ptr=ptr, shape=(shape_x, shape_y), dtype=wp.float32)
+    arr[0, 0] = 1.0
+    arr[0, 1] = 2.0
+    if shape_y > 2:
+        arr[0, 2] = 3.0
+
+
+def test_kernel_array_from_ptr_variable_shape(test, device):
+    arr = wp.zeros(shape=(2, 3), dtype=wp.float32, device=device)
+    wp.launch(kernel_array_from_ptr_variable_shape, dim=(1,), inputs=(arr.ptr, 2, 2), device=device)
+    assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 0.0), (0.0, 0.0, 0.0))))
+    wp.launch(kernel_array_from_ptr_variable_shape, dim=(1,), inputs=(arr.ptr, 2, 3), device=device)
     assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 3.0), (0.0, 0.0, 0.0))))
 
 
@@ -3185,6 +3232,10 @@ add_function_test(TestArray, "test_array_inplace_diff_ops", test_array_inplace_d
 add_function_test(TestArray, "test_array_inplace_non_diff_ops", test_array_inplace_non_diff_ops, devices=devices)
 add_function_test(TestArray, "test_direct_from_numpy", test_direct_from_numpy, devices=["cpu"])
 add_function_test(TestArray, "test_kernel_array_from_ptr", test_kernel_array_from_ptr, devices=devices)
+add_function_test(TestArray, "test_kernel_array_from_ptr_struct", test_kernel_array_from_ptr_struct, devices=devices)
+add_function_test(
+    TestArray, "test_kernel_array_from_ptr_variable_shape", test_kernel_array_from_ptr_variable_shape, devices=devices
+)
 
 add_function_test(TestArray, "test_array_from_int32_domain", test_array_from_int32_domain, devices=devices)
 add_function_test(TestArray, "test_array_from_int64_domain", test_array_from_int64_domain, devices=devices)

warp/tests/test_assert.py

@@ -245,6 +245,59 @@ class TestAssertDebug(unittest.TestCase):
                 self.assertRegex(output, r"Assertion failed: .*assert value == 1.*Array element must be 1")
 
 
+class TestAssertModeSwitch(unittest.TestCase):
+    """Test that switching from release mode to debug mode rebuilds the module with assertions enabled."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls._saved_mode = wp.config.mode
+        cls._saved_mode_module = wp.get_module_options()["mode"]
+        cls._saved_cache_kernels = wp.config.cache_kernels
+
+        # Don't set any mode initially - use whatever the default is
+        wp.config.cache_kernels = False
+
+    @classmethod
+    def tearDownClass(cls):
+        wp.config.mode = cls._saved_mode
+        wp.set_module_options({"mode": cls._saved_mode_module})
+        wp.config.cache_kernels = cls._saved_cache_kernels
+
+    def test_switch_to_debug_mode(self):
+        """Test that switching from release mode to debug mode rebuilds the module with assertions enabled."""
+        with wp.ScopedDevice("cpu"):
+            # Create an array that will trigger an assertion
+            input_array = wp.zeros(1, dtype=int)
+
+            # In default mode, this should not assert
+            capture = StdErrCapture()
+            capture.begin()
+            wp.launch(expect_ones, input_array.shape, inputs=[input_array])
+            output = capture.end()
+
+            # Should not have any assertion output in release mode
+            self.assertEqual(output, "", f"Kernel should not print anything to stderr in release mode, got {output}")
+
+            # Now switch to debug mode and have it compile a new kernel
+            wp.config.mode = "debug"
+
+            @wp.kernel
+            def expect_ones_debug(a: wp.array(dtype=int)):
+                i = wp.tid()
+                assert a[i] == 1
+
+            # In debug mode, this should assert
+            capture = StdErrCapture()
+            capture.begin()
+            wp.launch(expect_ones_debug, input_array.shape, inputs=[input_array])
+            output = capture.end()
+
+            # Should have assertion output in debug mode
+            # Older Windows C runtimes have a bug where stdout sometimes does not get properly flushed.
+            if output != "" or sys.platform != "win32":
+                self.assertRegex(output, r"Assertion failed: .*assert a\[i\] == 1")
+
+
 if __name__ == "__main__":
     wp.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_atomic_cas.py

@@ -19,54 +19,63 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
+kernel_cache = {}
+
+
+def getkernel(func, suffix=""):
+    key = func.__name__ + "_" + suffix
+    if key not in kernel_cache:
+        kernel_cache[key] = wp.Kernel(func=func, key=key)
+    return kernel_cache[key]
+
+
+def test_atomic_cas(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.array([0], dtype=warp_type, device=device)
 
-def create_spinlock_test(dtype):
     @wp.func
-    def spinlock_acquire(lock: wp.array(dtype=dtype)):
+    def spinlock_acquire_1d(lock: wp.array(dtype=warp_type)):
         # Try to acquire the lock by setting it to 1 if it's 0
-        while wp.atomic_cas(lock, 0, dtype(0), dtype(1)) == 1:
+        while wp.atomic_cas(lock, 0, warp_type(0), warp_type(1)) == 1:
             pass
 
     @wp.func
-    def spinlock_release(lock: wp.array(dtype=dtype)):
+    def spinlock_release_1d(lock: wp.array(dtype=warp_type)):
         # Release the lock by setting it back to 0
-        wp.atomic_exch(lock, 0, dtype(0))
+        wp.atomic_exch(lock, 0, warp_type(0))
 
     @wp.func
-    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
-        value = wp.atomic_exch(ptr, index, dtype(0))
+    def volatile_read_1d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
         wp.atomic_exch(ptr, index, value)
         return value
 
-    @wp.kernel
-    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype)):
+    def test_spinlock_counter_1d(counter: wp.array(dtype=warp_type), lock: wp.array(dtype=warp_type)):
         # Try to acquire the lock
-        spinlock_acquire(lock)
+        spinlock_acquire_1d(lock)
 
         # Critical section - increment counter
         # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
 
         # Work around since warp arrays cannot be marked as volatile
-        value = volatile_read(counter, 0)
-        counter[0] = value + dtype(1)
+        value = volatile_read_1d(counter, 0)
+        counter[0] = value + warp_type(1)
 
         # Release the lock
-        spinlock_release(lock)
-
-    return test_spinlock_counter
+        spinlock_release_1d(lock)
 
+    kernel = getkernel(test_spinlock_counter_1d, suffix=dtype.__name__)
 
-def test_atomic_cas(test, device, warp_type, numpy_type):
-    n = 100
-    counter = wp.array([0], dtype=warp_type, device=device)
-    lock = wp.array([0], dtype=warp_type, device=device)
+    if register_kernels:
+        return
 
-    test_spinlock_counter = create_spinlock_test(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -74,53 +83,53 @@ def test_atomic_cas(test, device, warp_type, numpy_type):
     assert_np_equal(counter_np, expected)
 
 
-def create_spinlock_test_2d(dtype):
+def test_atomic_cas_2d(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1), dtype=warp_type, device=device)
+
     @wp.func
-    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=2)):
+    def spinlock_acquire_2d(lock: wp.array2d(dtype=warp_type)):
         # Try to acquire the lock by setting it to 1 if it's 0
-        while wp.atomic_cas(lock, 0, 0, dtype(0), dtype(1)) == 1:
+        while wp.atomic_cas(lock, 0, 0, warp_type(0), warp_type(1)) == 1:
             pass
 
     @wp.func
-    def spinlock_release(lock: wp.array(dtype=dtype, ndim=2)):
+    def spinlock_release_2d(lock: wp.array2d(dtype=warp_type)):
         # Release the lock by setting it back to 0
-        wp.atomic_exch(lock, 0, 0, dtype(0))
+        wp.atomic_exch(lock, 0, 0, warp_type(0))
 
     @wp.func
-    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
-        value = wp.atomic_exch(ptr, index, dtype(0))
+    def volatile_read_2d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
         wp.atomic_exch(ptr, index, value)
         return value
 
-    @wp.kernel
-    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=2)):
+    def test_spinlock_counter_2d(counter: wp.array(dtype=warp_type), lock: wp.array2d(dtype=warp_type)):
         # Try to acquire the lock
-        spinlock_acquire(lock)
+        spinlock_acquire_2d(lock)
 
         # Critical section - increment counter
         # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
 
         # Work around since warp arrays cannot be marked as volatile
-        value = volatile_read(counter, 0)
-        counter[0] = value + dtype(1)
+        value = volatile_read_2d(counter, 0)
+        counter[0] = value + warp_type(1)
 
         # Release the lock
-        spinlock_release(lock)
-
-    return test_spinlock_counter
+        spinlock_release_2d(lock)
 
+    kernel = getkernel(test_spinlock_counter_2d, suffix=dtype.__name__)
 
-def test_atomic_cas_2d(test, device, warp_type, numpy_type):
-    n = 100
-    counter = wp.array([0], dtype=warp_type, device=device)
-    lock = wp.zeros(shape=(1, 1), dtype=warp_type, device=device)
+    if register_kernels:
+        return
 
-    test_spinlock_counter = create_spinlock_test_2d(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -128,53 +137,53 @@ def test_atomic_cas_2d(test, device, warp_type, numpy_type):
     assert_np_equal(counter_np, expected)
 
 
-def create_spinlock_test_3d(dtype):
+def test_atomic_cas_3d(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1, 1), dtype=warp_type, device=device)
+
     @wp.func
-    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=3)):
+    def spinlock_acquire_3d(lock: wp.array3d(dtype=warp_type)):
         # Try to acquire the lock by setting it to 1 if it's 0
-        while wp.atomic_cas(lock, 0, 0, 0, dtype(0), dtype(1)) == 1:
+        while wp.atomic_cas(lock, 0, 0, 0, warp_type(0), warp_type(1)) == 1:
             pass
 
     @wp.func
-    def spinlock_release(lock: wp.array(dtype=dtype, ndim=3)):
+    def spinlock_release_3d(lock: wp.array3d(dtype=warp_type)):
         # Release the lock by setting it back to 0
-        wp.atomic_exch(lock, 0, 0, 0, dtype(0))
+        wp.atomic_exch(lock, 0, 0, 0, warp_type(0))
 
     @wp.func
-    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
-        value = wp.atomic_exch(ptr, index, dtype(0))
+    def volatile_read_3d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
         wp.atomic_exch(ptr, index, value)
         return value
 
-    @wp.kernel
-    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=3)):
+    def test_spinlock_counter_3d(counter: wp.array(dtype=warp_type), lock: wp.array3d(dtype=warp_type)):
         # Try to acquire the lock
-        spinlock_acquire(lock)
+        spinlock_acquire_3d(lock)
 
         # Critical section - increment counter
         # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
 
         # Work around since warp arrays cannot be marked as volatile
-        value = volatile_read(counter, 0)
-        counter[0] = value + dtype(1)
+        value = volatile_read_3d(counter, 0)
+        counter[0] = value + warp_type(1)
 
         # Release the lock
-        spinlock_release(lock)
+        spinlock_release_3d(lock)
 
-    return test_spinlock_counter
+    kernel = getkernel(test_spinlock_counter_3d, suffix=dtype.__name__)
 
+    if register_kernels:
+        return
 
-def test_atomic_cas_3d(test, device, warp_type, numpy_type):
-    n = 100
-    counter = wp.array([0], dtype=warp_type, device=device)
-    lock = wp.zeros(shape=(1, 1, 1), dtype=warp_type, device=device)
-
-    test_spinlock_counter = create_spinlock_test_3d(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -218,17 +227,53 @@ def create_spinlock_test_4d(dtype):
     return test_spinlock_counter
 
 
-def test_atomic_cas_4d(test, device, warp_type, numpy_type):
+def test_atomic_cas_4d(test, device, dtype, register_kernels=False):
+    warp_type = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
     n = 100
     counter = wp.array([0], dtype=warp_type, device=device)
     lock = wp.zeros(shape=(1, 1, 1, 1), dtype=warp_type, device=device)
 
-    test_spinlock_counter = create_spinlock_test_4d(warp_type)
-    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+    @wp.func
+    def spinlock_acquire_4d(lock: wp.array4d(dtype=warp_type)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, 0, 0, 0, warp_type(0), warp_type(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release_4d(lock: wp.array4d(dtype=warp_type)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, 0, 0, 0, warp_type(0))
+
+    @wp.func
+    def volatile_read_4d(ptr: wp.array(dtype=warp_type), index: int):
+        value = wp.atomic_exch(ptr, index, warp_type(0))
+        wp.atomic_exch(ptr, index, value)
+        return value
+
+    def test_spinlock_counter_4d(counter: wp.array(dtype=warp_type), lock: wp.array4d(dtype=warp_type)):
+        # Try to acquire the lock
+        spinlock_acquire_4d(lock)
+
+        # Critical section - increment counter
+        # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+        # Work around since warp arrays cannot be marked as volatile
+        value = volatile_read_4d(counter, 0)
+        counter[0] = value + warp_type(1)
+
+        # Release the lock
+        spinlock_release_4d(lock)
+
+    kernel = getkernel(test_spinlock_counter_4d, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    wp.launch(kernel, dim=n, inputs=[counter, lock], device=device)
 
     # Verify counter reached n
     counter_np = counter.numpy()
-    expected = np.array([n], dtype=numpy_type)
+    expected = np.array([n], dtype=dtype)
 
     if not np.array_equal(counter_np, expected):
         print(f"Counter mismatch: expected {expected}, got {counter_np}")
@@ -244,54 +289,22 @@ class TestAtomicCAS(unittest.TestCase):
 
 
 # Test all supported types
-test_types = [
-    (wp.int32, np.int32),
-    (wp.uint32, np.uint32),
-    (wp.int64, np.int64),
-    (wp.uint64, np.uint64),
-    (wp.float32, np.float32),
-    (wp.float64, np.float64),
-]
-
-for warp_type, numpy_type in test_types:
-    type_name = warp_type.__name__
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_{type_name}",
-        test_atomic_cas,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
-    )
+np_test_types = (np.int32, np.uint32, np.int64, np.uint64, np.float32, np.float64)
 
+for dtype in np_test_types:
+    type_name = dtype.__name__
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_{type_name}", test_atomic_cas, devices=devices, dtype=dtype
+    )
     # Add 2D test for each type
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_2d_{type_name}",
-        test_atomic_cas_2d,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_2d_{type_name}", test_atomic_cas_2d, devices=devices, dtype=dtype
     )
-
-    # Add 3D test for each type
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_3d_{type_name}",
-        test_atomic_cas_3d,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_3d_{type_name}", test_atomic_cas_3d, devices=devices, dtype=dtype
     )
-
-    # Add 4D test for each type
-    add_function_test(
-        TestAtomicCAS,
-        f"test_cas_4d_{type_name}",
-        test_atomic_cas_4d,
-        devices=devices,
-        warp_type=warp_type,
-        numpy_type=numpy_type,
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_4d_{type_name}", test_atomic_cas_4d, devices=devices, dtype=dtype
     )
 
 if __name__ == "__main__":

warp/tests/test_codegen.py

@@ -756,6 +756,7 @@ def test_multiple_return_values(test, device):
         test_multiple_return_values_quat_to_axis_angle_kernel,
         dim=1,
         inputs=(q, expected_axis, expected_angle),
+        device=device,
     )
 
     # fmt: off
@@ -791,9 +792,9 @@ def test_multiple_return_values(test, device):
 
     test.assertAlmostEqual(V[0][0], expected_V[0][0], places=5)
     test.assertAlmostEqual(V[0][1], expected_V[0][1], places=5)
-    test.assertAlmostEqual(V[0][2], expected_V[0][2], places=5)
+    test.assertAlmostEqual(V[0][2], expected_V[0][2], places=4)  # precision issue on ARM64 (GH-905)
     test.assertAlmostEqual(V[1][0], expected_V[1][0], places=5)
-    test.assertAlmostEqual(V[1][1], expected_V[1][1], places=5)
+    test.assertAlmostEqual(V[1][1], expected_V[1][1], places=4)  # precision issue on ARM64 (GH-905)
     test.assertAlmostEqual(V[1][2], expected_V[1][2], places=5)
     test.assertAlmostEqual(V[2][0], expected_V[2][0], places=5)
     test.assertAlmostEqual(V[2][1], expected_V[2][1], places=5)