warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_array.py

@@ -412,7 +412,7 @@ def test_slicing(test, device):
     assert_array_equal(wp_arr[:5], wp.array(np_arr[:5], dtype=int, device=device))
     assert_array_equal(wp_arr[1:5], wp.array(np_arr[1:5], dtype=int, device=device))
     assert_array_equal(wp_arr[-9:-5:1], wp.array(np_arr[-9:-5:1], dtype=int, device=device))
-    assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device))  # noqa: E231
+    assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device))
 
 
 def test_view(test, device):
@@ -2370,6 +2370,257 @@ def test_array_from_cai(test, device):
     assert_np_equal(arr_warp.numpy(), np.array([[2, 1, 1], [1, 0, 0], [1, 0, 0]]))
 
 
+def test_array_from_data(test, device):
+    with wp.ScopedDevice(device):
+        # =========================================
+        # scalars, reshaping
+
+        data = np.arange(12, dtype=np.float32).reshape((3, 4))
+        src = wp.array(data)
+
+        assert src.device == device
+
+        dtypes = [Any, wp.float32]
+        shapes = [None, (3, 4), (12,), (3, 2, 2)]
+
+        for dtype in dtypes:
+            for shape in shapes:
+                with test.subTest(msg=f"scalar, dtype={dtype}, shape={shape}"):
+                    dst = wp.array(src, dtype=dtype, shape=shape)
+                    assert dst.device == src.device
+                    if dtype is Any:
+                        assert dst.dtype == src.dtype
+                    else:
+                        assert dst.dtype == dtype
+                    if shape is None:
+                        assert dst.shape == src.shape
+                        assert_np_equal(dst.numpy(), data)
+                    else:
+                        assert dst.shape == shape
+                        assert_np_equal(dst.numpy(), data.reshape(shape))
+
+        # =========================================
+        # vectors, reshaping
+
+        with test.subTest(msg="vector, single"):
+            data = np.arange(3, dtype=np.float32)
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.vec3)
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (1,)
+            assert_np_equal(dst.numpy(), data.reshape((1, 3)))
+
+        with test.subTest(msg="vector, multiple in 1d"):
+            data = np.arange(12, dtype=np.float32)
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.vec3)
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (4,)
+            assert_np_equal(dst.numpy(), data.reshape((4, 3)))
+
+        with test.subTest(msg="vector, singles in 2d"):
+            data = np.arange(12, dtype=np.float32).reshape((4, 3))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.vec3)
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (4,)
+            assert_np_equal(dst.numpy(), data.reshape((4, 3)))
+
+        with test.subTest(msg="vector, multiples in 2d"):
+            data = np.arange(24, dtype=np.float32).reshape((4, 6))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.vec3)
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (4, 2)
+            assert_np_equal(dst.numpy(), data.reshape((4, 2, 3)))
+
+        with test.subTest(msg="vector, singles in 2d, reshape"):
+            data = np.arange(12, dtype=np.float32).reshape((4, 3))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.vec3, shape=(2, 2))
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (2, 2)
+            assert_np_equal(dst.numpy(), data.reshape((2, 2, 3)))
+
+        with test.subTest(msg="vector, multiples in 2d, reshape"):
+            data = np.arange(24, dtype=np.float32).reshape((4, 6))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.vec3, shape=(2, 2, 2))
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (2, 2, 2)
+            assert_np_equal(dst.numpy(), data.reshape((2, 2, 2, 3)))
+
+        # =========================================
+        # matrices, reshaping
+
+        with test.subTest(msg="matrix, single in 2d"):
+            # one 2x2 matrix in a 2d array
+            data = np.arange(4, dtype=np.float32).reshape((2, 2))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (1,)
+            assert_np_equal(dst.numpy(), data.reshape((1, 2, 2)))
+
+        with test.subTest(msg="matrix, single in 1d"):
+            # 2x2 matrix in a 1d array
+            data = np.arange(4, dtype=np.float32)
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (1,)
+            assert_np_equal(dst.numpy(), data.reshape((1, 2, 2)))
+
+        with test.subTest(msg="matrix, multiples in 1d"):
+            # 3 2x2 matrices in a 1d array
+            data = np.arange(12, dtype=np.float32)
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (3,)
+            assert_np_equal(dst.numpy(), data.reshape((3, 2, 2)))
+
+        with test.subTest(msg="matrix, multiples in 1d, reshape"):
+            # 4 2x2 matrices in a 1d array
+            data = np.arange(16, dtype=np.float32)
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22, shape=(4,))
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (4,)
+            assert_np_equal(dst.numpy(), data.reshape((4, 2, 2)))
+
+        with test.subTest(msg="matrix, multiples in 2d"):
+            # 3 2x2 matrices in a 2d array
+            data = np.arange(12, dtype=np.float32).reshape((3, 4))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (3,)
+            assert_np_equal(dst.numpy(), data.reshape((3, 2, 2)))
+
+        with test.subTest(msg="matrix, multiples in 2d, reshape"):
+            # 4 2x2 matrices in a 2d array
+            data = np.arange(16, dtype=np.float32).reshape((4, 4))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22, shape=(2, 2))
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (2, 2)
+            assert_np_equal(dst.numpy(), data.reshape((2, 2, 2, 2)))
+
+        with test.subTest(msg="matrix, multiples in 3d"):
+            # 3 2x2 matrices in a 3d array
+            data = np.arange(12, dtype=np.float32).reshape((3, 2, 2))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (3,)
+            assert_np_equal(dst.numpy(), data.reshape((3, 2, 2)))
+
+        with test.subTest(msg="matrix, multiples in 3d, reshape"):
+            # 4 2x2 matrices in a 3d array
+            data = np.arange(16, dtype=np.float32).reshape((4, 2, 2))
+            src = wp.array(data)
+            dst = wp.array(src, dtype=wp.mat22, shape=(2, 2))
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (2, 2)
+            assert_np_equal(dst.numpy(), data.reshape((2, 2, 2, 2)))
+
+        # =========================================
+        # vectors and matrices in strided arrays
+
+        with test.subTest(msg="vector, singles in 2d, strided"):
+            # 4 vec3 in strided 2d array
+            data = np.arange(20, dtype=np.float32).reshape((4, 5))
+            src = wp.array(data)[:, 2:]  # source with strides
+            dst = wp.array(src, dtype=wp.vec3)
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (4,)
+            expected = np.array(
+                [
+                    [2, 3, 4],
+                    [7, 8, 9],
+                    [12, 13, 14],
+                    [17, 18, 19],
+                ],
+                dtype=np.float32,
+            )
+            assert_np_equal(dst.numpy(), expected)
+
+        with test.subTest(msg="vector, multiples in 2d, strided"):
+            # 4 vec3 in strided 2d array
+            data = np.arange(14, dtype=np.float32).reshape((2, 7))
+            src = wp.array(data)[:, 1:]  # source with strides
+            dst = wp.array(src, dtype=wp.vec3)
+            assert dst.dtype == wp.vec3
+            assert dst.shape == (2, 2)
+            expected = np.array(
+                [
+                    [
+                        [1, 2, 3],
+                        [4, 5, 6],
+                    ],
+                    [
+                        [8, 9, 10],
+                        [11, 12, 13],
+                    ],
+                ],
+                dtype=np.float32,
+            )
+            assert_np_equal(dst.numpy(), expected)
+
+        with test.subTest(msg="matrix, multiples in 2d, strided"):
+            # 3 2x2 matrices in a 2d array
+            data = np.arange(15, dtype=np.float32).reshape((3, 5))
+            src = wp.array(data)[:, 1:]  # source with strides
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (3,)
+            expected = np.array(
+                [
+                    [
+                        [1, 2],
+                        [3, 4],
+                    ],
+                    [
+                        [6, 7],
+                        [8, 9],
+                    ],
+                    [
+                        [11, 12],
+                        [13, 14],
+                    ],
+                ],
+                dtype=np.float32,
+            )
+            assert_np_equal(dst.numpy(), expected)
+
+        with test.subTest(msg="matrix, multiples in 3d, strided"):
+            # 3 2x2 matrices in a 3d array
+            data = np.arange(18, dtype=np.float32).reshape((3, 3, 2))
+            src = wp.array(data)[:, 1:]  # source with strides
+            dst = wp.array(src, dtype=wp.mat22)
+            assert dst.dtype == wp.mat22
+            assert dst.shape == (3,)
+            expected = np.array(
+                [
+                    [
+                        [2, 3],
+                        [4, 5],
+                    ],
+                    [
+                        [8, 9],
+                        [10, 11],
+                    ],
+                    [
+                        [14, 15],
+                        [16, 17],
+                    ],
+                ],
+                dtype=np.float32,
+            )
+            assert_np_equal(dst.numpy(), expected)
+
+
 @wp.kernel
 def inplace_add_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
     i = wp.tid()
@@ -2604,7 +2855,7 @@ def test_array_inplace_non_diff_ops(test, device):
     wp.launch(inplace_div_1d, N, inputs=[x1, y1], device=device)
     assert_np_equal(x1.numpy(), np.full(N, fill_value=2.0, dtype=float))
 
-    for dtype in wp.types.non_atomic_types + (wp.vec2b, wp.vec2ub, wp.vec2s, wp.vec2us, uint16vec3):
+    for dtype in (*wp.types.non_atomic_types, wp.vec2b, wp.vec2ub, wp.vec2s, wp.vec2us, uint16vec3):
         x = wp.full(N, value=0, dtype=dtype, device=device)
         y = wp.full(N, value=1, dtype=dtype, device=device)
 
@@ -2943,6 +3194,7 @@ add_function_test(TestArray, "test_alloc_strides", test_alloc_strides, devices=d
 add_function_test(TestArray, "test_casting", test_casting, devices=devices)
 add_function_test(TestArray, "test_array_len", test_array_len, devices=devices)
 add_function_test(TestArray, "test_cuda_interface_conversion", test_cuda_interface_conversion, devices=devices)
+add_function_test(TestArray, "test_array_from_data", test_array_from_data, devices=devices)
 
 try:
     import torch

warp/tests/test_array_reduce.py

@@ -28,7 +28,7 @@ def make_test_array_sum(dtype):
     def test_array_sum(test, device):
         rng = np.random.default_rng(123)
 
-        cols = wp.types.type_length(dtype)
+        cols = wp.types.type_size(dtype)
 
         values_np = rng.random(size=(N, cols))
         values = wp.array(values_np, device=device, dtype=dtype)
@@ -77,7 +77,7 @@ def make_test_array_inner(dtype):
     def test_array_inner(test, device):
         rng = np.random.default_rng(123)
 
-        cols = wp.types.type_length(dtype)
+        cols = wp.types.type_size(dtype)
 
         a_np = rng.random(size=(N, cols))
         b_np = rng.random(size=(N, cols))

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

@@ -0,0 +1,312 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 *
+
+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)
+
+    @wp.func
+    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, warp_type(0), warp_type(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release_1d(lock: wp.array(dtype=warp_type)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, warp_type(0))
+
+    @wp.func
+    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
+
+    def test_spinlock_counter_1d(counter: wp.array(dtype=warp_type), lock: wp.array(dtype=warp_type)):
+        # Try to acquire the 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_1d(counter, 0)
+        counter[0] = value + warp_type(1)
+
+        # Release the lock
+        spinlock_release_1d(lock)
+
+    kernel = getkernel(test_spinlock_counter_1d, 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=dtype)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+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_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, warp_type(0), warp_type(1)) == 1:
+            pass
+
+    @wp.func
+    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, warp_type(0))
+
+    @wp.func
+    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
+
+    def test_spinlock_counter_2d(counter: wp.array(dtype=warp_type), lock: wp.array2d(dtype=warp_type)):
+        # Try to acquire the 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_2d(counter, 0)
+        counter[0] = value + warp_type(1)
+
+        # Release the lock
+        spinlock_release_2d(lock)
+
+    kernel = getkernel(test_spinlock_counter_2d, 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=dtype)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+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_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, warp_type(0), warp_type(1)) == 1:
+            pass
+
+    @wp.func
+    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, warp_type(0))
+
+    @wp.func
+    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
+
+    def test_spinlock_counter_3d(counter: wp.array(dtype=warp_type), lock: wp.array3d(dtype=warp_type)):
+        # Try to acquire the 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_3d(counter, 0)
+        counter[0] = value + warp_type(1)
+
+        # Release the lock
+        spinlock_release_3d(lock)
+
+    kernel = getkernel(test_spinlock_counter_3d, 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=dtype)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+def create_spinlock_test_4d(dtype):
+    @wp.func
+    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=4)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, 0, 0, 0, dtype(0), dtype(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release(lock: wp.array(dtype=dtype, ndim=4)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, 0, 0, 0, dtype(0))
+
+    @wp.func
+    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
+        value = wp.atomic_exch(ptr, index, dtype(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=4)):
+        # Try to acquire the lock
+        spinlock_acquire(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)
+
+        # Release the lock
+        spinlock_release(lock)
+
+    return test_spinlock_counter
+
+
+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)
+
+    @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=dtype)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+devices = get_test_devices()
+
+
+class TestAtomicCAS(unittest.TestCase):
+    pass
+
+
+# Test all supported types
+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_register_kernel(
+        TestAtomicCAS, f"test_cas_2d_{type_name}", test_atomic_cas_2d, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_3d_{type_name}", test_atomic_cas_3d, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestAtomicCAS, f"test_cas_4d_{type_name}", test_atomic_cas_4d, devices=devices, dtype=dtype
+    )
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)