warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/tests/tile/test_tile.py

@@ -738,6 +738,28 @@ def test_tile_extract_kernel(a: wp.array2d(dtype=float), b: wp.array2d(dtype=flo
     wp.atomic_add(b, i, j, wp.tile_extract(tile, x, y))
 
 
+@wp.kernel
+def test_tile_extract_vec_kernel(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=float)):
+    i = wp.tid()
+
+    tile = wp.tile_load(x, shape=(TILE_M))
+
+    a = tile[i][1]
+
+    y[i] = a
+
+
+@wp.kernel
+def test_tile_extract_mat_kernel(x: wp.array(dtype=wp.mat33), y: wp.array(dtype=float)):
+    i = wp.tid()
+
+    tile = wp.tile_load(x, shape=(TILE_M))
+
+    a = tile[i][1, 1]
+
+    y[i] = a
+
+
 def test_tile_extract(test, device):
     block_dim = 16
 
@@ -763,6 +785,40 @@ def test_tile_extract(test, device):
     expected_grad = np.ones_like(input)
     assert_np_equal(a.grad.numpy(), expected_grad)
 
+    # vector element test
+    x = wp.ones(TILE_M, dtype=wp.vec3, requires_grad=True, device=device)
+    y = wp.zeros(TILE_M, dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch(test_tile_extract_vec_kernel, dim=[TILE_M], inputs=[x, y], block_dim=TILE_DIM, device=device)
+
+    y.grad = wp.ones_like(y)
+
+    tape.backward()
+
+    x_grad_np = np.zeros((TILE_M, 3), dtype=float)
+    x_grad_np[:, 1] = 1.0
+
+    assert_np_equal(x.grad.numpy(), x_grad_np)
+    assert_np_equal(y.numpy(), np.ones(TILE_M, dtype=float))
+
+    # matrix element test
+    x = wp.ones(TILE_M, dtype=wp.mat33, requires_grad=True, device=device)
+    y = wp.zeros(TILE_M, dtype=float, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch(test_tile_extract_mat_kernel, dim=[TILE_M], inputs=[x, y], block_dim=TILE_DIM, device=device)
+
+    y.grad = wp.ones_like(y)
+
+    tape.backward()
+
+    x_grad_np = np.zeros((TILE_M, 3, 3), dtype=float)
+    x_grad_np[:, 1, 1] = 1.0
+
+    assert_np_equal(y.numpy(), np.ones(TILE_M, dtype=float))
+    assert_np_equal(x.grad.numpy(), x_grad_np)
+
 
 @wp.kernel(module="unique")
 def test_tile_extract_repeated_kernel(a: wp.array2d(dtype=float), b: wp.array2d(dtype=float)):
@@ -822,6 +878,28 @@ def test_tile_assign_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
     wp.tile_atomic_add(y, a, offset=(0,))
 
 
+@wp.kernel
+def test_tile_assign_vec_kernel(x: wp.array(dtype=float), y: wp.array(dtype=wp.vec3)):
+    i = wp.tid()
+
+    a = wp.tile_zeros(shape=(TILE_M,), dtype=wp.vec3)
+
+    a[i][1] = x[i]
+
+    wp.tile_atomic_add(y, a, offset=(0,))
+
+
+@wp.kernel
+def test_tile_assign_mat_kernel(x: wp.array(dtype=float), y: wp.array(dtype=wp.mat33)):
+    i = wp.tid()
+
+    a = wp.tile_zeros(shape=(TILE_M,), dtype=wp.mat33)
+
+    a[i][1, 1] = x[i]
+
+    wp.tile_atomic_add(y, a, offset=(0,))
+
+
 def test_tile_assign(test, device):
     x = wp.full(TILE_M, 2.0, dtype=float, device=device, requires_grad=True)
     y = wp.zeros(TILE_M, dtype=float, device=device, requires_grad=True)
@@ -836,6 +914,100 @@ def test_tile_assign(test, device):
     assert_np_equal(y.numpy(), np.full(TILE_M, 2.0, dtype=np.float32))
     assert_np_equal(x.grad.numpy(), np.full(TILE_M, 1.0, dtype=np.float32))
 
+    # vector element test
+    x = wp.full(TILE_M, 2.0, dtype=float, device=device, requires_grad=True)
+    y = wp.zeros(TILE_M, dtype=wp.vec3, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(test_tile_assign_vec_kernel, dim=[TILE_M], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    y_np = np.zeros((TILE_M, 3), dtype=float)
+    y_np[:, 1] = 2.0
+
+    assert_np_equal(y.numpy(), y_np)
+    assert_np_equal(x.grad.numpy(), np.full(TILE_M, 1.0, dtype=np.float32))
+
+    # matrix element test
+    x = wp.full(TILE_M, 2.0, dtype=float, device=device, requires_grad=True)
+    y = wp.zeros(TILE_M, dtype=wp.mat33, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(test_tile_assign_mat_kernel, dim=[TILE_M], inputs=[x], outputs=[y], block_dim=TILE_DIM, device=device)
+
+    y.grad = wp.ones_like(y)
+    tape.backward()
+
+    y_np = np.zeros((TILE_M, 3, 3), dtype=float)
+    y_np[:, 1, 1] = 2.0
+
+    assert_np_equal(y.numpy(), y_np)
+    assert_np_equal(x.grad.numpy(), np.full(TILE_M, 1.0, dtype=np.float32))
+
+
+@wp.kernel
+def test_tile_where_kernel(select: int, x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+    x_reg = wp.tile_load(x, shape=(TILE_M,), storage="register")
+    y_reg = wp.tile_load(y, shape=(TILE_M,), storage="register")
+
+    x_shared = wp.tile_load(x, shape=(TILE_M,), storage="shared")
+    y_shared = wp.tile_load(y, shape=(TILE_M,), storage="shared")
+
+    if select == 0:
+        s = x_reg
+    elif select == 1:
+        s = y_reg
+    elif select == 2:
+        s = x_shared
+    else:
+        s = y_shared
+
+    wp.tile_store(z, s)
+
+
+def test_tile_where(test, device):
+    x = wp.full((TILE_M,), 1.0, dtype=float, device=device, requires_grad=True)
+    y = wp.full((TILE_M,), 2.0, dtype=float, device=device, requires_grad=True)
+    z = wp.zeros((TILE_M), dtype=float, device=device, requires_grad=True)
+
+    z_expected = [
+        np.full(TILE_M, 1.0, dtype=np.float32),
+        np.full(TILE_M, 2.0, dtype=np.float32),
+        np.full(TILE_M, 1.0, dtype=np.float32),
+        np.full(TILE_M, 2.0, dtype=np.float32),
+    ]
+    x_grad_expected = [
+        np.full(TILE_M, 1.0, dtype=np.float32),
+        np.full(TILE_M, 0.0, dtype=np.float32),
+        np.full(TILE_M, 1.0, dtype=np.float32),
+        np.full(TILE_M, 0.0, dtype=np.float32),
+    ]
+    y_grad_expected = [
+        np.full(TILE_M, 0.0, dtype=np.float32),
+        np.full(TILE_M, 1.0, dtype=np.float32),
+        np.full(TILE_M, 0.0, dtype=np.float32),
+        np.full(TILE_M, 1.0, dtype=np.float32),
+    ]
+
+    for i in range(4):
+        tape = wp.Tape()
+        with tape:
+            wp.launch_tiled(test_tile_where_kernel, dim=[1], inputs=[i, x, y], outputs=[z], block_dim=32, device=device)
+
+        z.grad = wp.ones_like(z)
+
+        tape.backward()
+
+        assert_np_equal(z.numpy(), z_expected[i])
+        assert_np_equal(x.grad.numpy(), x_grad_expected[i])
+        assert_np_equal(y.grad.numpy(), y_grad_expected[i])
+
+        tape.zero()
+
 
 @wp.kernel
 def test_tile_transpose_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
@@ -1118,6 +1290,71 @@ def test_tile_len(test, device):
     test.assertEqual(out.numpy()[0], TILE_M)
 
 
+@wp.struct
+class TestStruct:
+    x: wp.float32
+    y: wp.vec3
+
+
+@wp.kernel
+def test_tile_construction_kernel(
+    out_zeros: wp.array(dtype=float),
+    out_ones: wp.array(dtype=float),
+    out_arange: wp.array(dtype=float),
+    out_full_twos: wp.array(dtype=float),
+    out_full_vecs: wp.array(dtype=wp.vec3),
+    out_full_mats: wp.array(dtype=wp.mat33),
+    out_full_structs: wp.array(dtype=TestStruct),
+):
+    zeros = wp.tile_zeros(TILE_M, dtype=float)
+    ones = wp.tile_ones(TILE_M, dtype=float)
+    arange = wp.tile_arange(TILE_M, dtype=float)
+    full_twos = wp.tile_full(TILE_M, value=2.0, dtype=float)
+    full_vecs = wp.tile_full(TILE_M, value=wp.vec3(1.0), dtype=wp.vec3)
+    full_mats = wp.tile_full(TILE_M, value=wp.mat33(1.0), dtype=wp.mat33)
+
+    ts = TestStruct()
+    ts.x = wp.float32(2.0)
+    ts.y = wp.vec3(1.0)
+    full_structs = wp.tile_full(TILE_M, value=ts, dtype=TestStruct)
+
+    wp.tile_store(out_zeros, zeros)
+    wp.tile_store(out_ones, ones)
+    wp.tile_store(out_arange, arange)
+    wp.tile_store(out_full_twos, full_twos)
+    wp.tile_store(out_full_vecs, full_vecs)
+    wp.tile_store(out_full_mats, full_mats)
+    wp.tile_store(out_full_structs, full_structs)
+
+
+def test_tile_construction(test, device):
+    zeros = wp.empty(TILE_M, dtype=float, device=device)
+    ones = wp.empty(TILE_M, dtype=float, device=device)
+    arange = wp.empty(TILE_M, dtype=float, device=device)
+    full_twos = wp.empty(TILE_M, dtype=float, device=device)
+    full_vecs = wp.empty(TILE_M, dtype=wp.vec3, device=device)
+    full_mats = wp.empty(TILE_M, dtype=wp.mat33, device=device)
+    full_structs = wp.empty(TILE_M, dtype=TestStruct, device=device)
+
+    wp.launch_tiled(
+        test_tile_construction_kernel,
+        dim=1,
+        inputs=[],
+        outputs=[zeros, ones, arange, full_twos, full_vecs, full_mats, full_structs],
+        block_dim=TILE_DIM,
+        device=device,
+    )
+
+    assert_np_equal(zeros.numpy(), np.zeros(TILE_M, dtype=float))
+    assert_np_equal(ones.numpy(), np.ones(TILE_M, dtype=float))
+    assert_np_equal(full_twos.numpy(), np.full(TILE_M, 2.0, dtype=float))
+    assert_np_equal(full_vecs.numpy(), np.ones((TILE_M, 3), dtype=float))
+    assert_np_equal(full_mats.numpy(), np.ones((TILE_M, 3, 3), dtype=float))
+    assert_np_equal(full_structs.numpy()["x"], np.full(TILE_M, 2.0, dtype=float))
+    assert_np_equal(full_structs.numpy()["y"], np.ones((TILE_M, 3), dtype=float))
+    assert_np_equal(arange.numpy(), np.arange(TILE_M, dtype=float))
+
+
 @wp.kernel
 def test_tile_print_kernel():
     # shared tile
@@ -1261,6 +1498,7 @@ add_function_test(TestTile, "test_tile_sum_launch", test_tile_sum_launch, device
 add_function_test(TestTile, "test_tile_extract", test_tile_extract, devices=devices)
 add_function_test(TestTile, "test_tile_extract_repeated", test_tile_extract_repeated, devices=devices)
 add_function_test(TestTile, "test_tile_assign", test_tile_assign, devices=devices)
+add_function_test(TestTile, "test_tile_where", test_tile_where, devices=devices)
 add_function_test(TestTile, "test_tile_broadcast_add_1d", test_tile_broadcast_add_1d, devices=devices)
 add_function_test(TestTile, "test_tile_broadcast_add_2d", test_tile_broadcast_add_2d, devices=devices)
 add_function_test(TestTile, "test_tile_broadcast_add_3d", test_tile_broadcast_add_3d, devices=devices)
@@ -1269,6 +1507,7 @@ add_function_test(TestTile, "test_tile_broadcast_grad", test_tile_broadcast_grad
 add_function_test(TestTile, "test_tile_squeeze", test_tile_squeeze, devices=devices)
 add_function_test(TestTile, "test_tile_reshape", test_tile_reshape, devices=devices)
 add_function_test(TestTile, "test_tile_len", test_tile_len, devices=devices)
+add_function_test(TestTile, "test_tile_construction", test_tile_construction, devices=devices)
 # add_function_test(TestTile, "test_tile_print", test_tile_print, devices=devices, check_output=False)
 # add_function_test(TestTile, "test_tile_inplace", test_tile_inplace, devices=devices)
 # add_function_test(TestTile, "test_tile_astype", test_tile_astype, devices=devices)

warp/tests/tile/test_tile_atomic_bitwise.py

@@ -0,0 +1,403 @@
+# 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 *
+
+
+@wp.kernel
+def test_tile_atomic_bitwise_scalar_kernel(
+    a: wp.array(dtype=wp.uint32), b: wp.array(dtype=wp.uint32), c: wp.array(dtype=wp.uint32), op_type: int
+):
+    word_idx, bit_idx = wp.tid()
+    block_dim = wp.block_dim()
+    assert block_dim == 32
+    s = wp.tile_zeros(shape=1, dtype=wp.uint32)
+    # write to tile first, then write only once to the array
+    s[0] = a[word_idx]
+    if op_type < 3:
+        bit_mask = wp.uint32(1) << wp.uint32(bit_idx)
+        if op_type == 0:
+            s[0] &= (b[word_idx] & bit_mask) | ~bit_mask
+        elif op_type == 1:
+            s[0] |= b[word_idx] & bit_mask
+        elif op_type == 2:
+            s[0] ^= b[word_idx] & bit_mask
+    else:
+        # inter-tile operations
+        s_bit_mask = wp.tile_zeros(shape=32, dtype=wp.uint32)
+        s_bit_mask[(bit_idx + 1) % 32] = wp.uint32(1) << wp.uint32((bit_idx + 1) % 32)
+        if op_type == 3:
+            s[0] &= (b[word_idx] & s_bit_mask[bit_idx]) | ~s_bit_mask[bit_idx]
+        elif op_type == 4:
+            s[0] |= b[word_idx] & s_bit_mask[bit_idx]
+        elif op_type == 5:
+            s[0] ^= b[word_idx] & s_bit_mask[bit_idx]
+    c[word_idx] = s[0]
+
+
+@wp.kernel
+def test_tile_atomic_bitwise_scalar_tilewise_kernel(
+    a: wp.array(dtype=wp.uint32), b: wp.array(dtype=wp.uint32), c: wp.array(dtype=wp.uint32), op_type: int
+):
+    batch_idx, _ = wp.tid()
+    block_dim = wp.block_dim()
+    assert block_dim == 32
+    # Each tile is responsible for a batch of 32 elements
+    s1 = wp.tile_load(a, shape=32, offset=batch_idx * 32)
+    s2 = wp.tile_load(b, shape=32, offset=batch_idx * 32)
+    # inter-tile operations (batch-wise)
+    if op_type < 9:
+        if op_type == 6:
+            s1 &= s2
+        elif op_type == 7:
+            s1 |= s2
+        elif op_type == 8:
+            s1 ^= s2
+        wp.tile_store(c, s1, offset=batch_idx * 32)
+    else:
+        if op_type == 9:
+            s3 = s1 & s2
+        elif op_type == 10:
+            s3 = s1 | s2
+        elif op_type == 11:
+            s3 = s1 ^ s2
+        wp.tile_store(c, s3, offset=batch_idx * 32)
+
+
+def test_tile_atomic_bitwise_scalar(test, device):
+    n = 1024
+    rng = np.random.default_rng(42)
+
+    a = rng.integers(0, np.iinfo(np.uint32).max, size=n, dtype=np.uint32)
+    b = rng.integers(0, np.iinfo(np.uint32).max, size=n, dtype=np.uint32)
+
+    expected_and = a & b
+    expected_or = a | b
+    expected_xor = a ^ b
+
+    with wp.ScopedDevice(device):
+        a_wp = wp.array(a, dtype=wp.uint32, device=device)
+        b_wp = wp.array(b, dtype=wp.uint32, device=device)
+        c_wp = wp.zeros(shape=n, dtype=wp.uint32, device=device)
+
+        wp.launch_tiled(test_tile_atomic_bitwise_scalar_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 0], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(test_tile_atomic_bitwise_scalar_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 1], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(test_tile_atomic_bitwise_scalar_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 2], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_xor)
+        wp.launch_tiled(test_tile_atomic_bitwise_scalar_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 3], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(test_tile_atomic_bitwise_scalar_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 4], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(test_tile_atomic_bitwise_scalar_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 5], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_xor)
+
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_scalar_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 6], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_scalar_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 7], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_scalar_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 8], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_xor)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_scalar_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 9], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_scalar_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 10], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_scalar_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 11], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_xor)
+
+
+@wp.kernel
+def test_tile_atomic_bitwise_vector_kernel(
+    a: wp.array(dtype=wp.vec3ui), b: wp.array(dtype=wp.vec3ui), c: wp.array(dtype=wp.vec3ui), op_type: int
+):
+    word_idx, bit_idx = wp.tid()
+    block_dim = wp.block_dim()
+    assert block_dim == 32
+    s = wp.tile_zeros(shape=1, dtype=wp.vec3ui)
+    # write to tile first, then write only once to the array
+    s[0] = a[word_idx]
+    if op_type < 3:
+        bit_mask = wp.vec3ui(wp.uint32(1)) << wp.vec3ui(wp.uint32(bit_idx))
+        if op_type == 0:
+            s[0] &= (b[word_idx] & bit_mask) | ~bit_mask
+        elif op_type == 1:
+            s[0] |= b[word_idx] & bit_mask
+        elif op_type == 2:
+            s[0] ^= b[word_idx] & bit_mask
+    else:
+        # inter-tile operations
+        s_bit_mask = wp.tile_zeros(shape=32, dtype=wp.vec3ui)
+        s_bit_mask[(bit_idx + 1) % 32] = wp.vec3ui(wp.uint32(1) << wp.uint32((bit_idx + 1) % 32))
+        if op_type == 3:
+            s[0] &= (b[word_idx] & s_bit_mask[bit_idx]) | ~s_bit_mask[bit_idx]
+        elif op_type == 4:
+            s[0] |= b[word_idx] & s_bit_mask[bit_idx]
+        elif op_type == 5:
+            s[0] ^= b[word_idx] & s_bit_mask[bit_idx]
+    c[word_idx] = s[0]
+
+
+@wp.kernel
+def test_tile_atomic_bitwise_vector_tilewise_kernel(
+    a: wp.array(dtype=wp.vec3ui), b: wp.array(dtype=wp.vec3ui), c: wp.array(dtype=wp.vec3ui), op_type: int
+):
+    batch_idx, _ = wp.tid()
+    block_dim = wp.block_dim()
+    assert block_dim == 32
+    # Each tile is responsible for a batch of 32 elements
+    s1 = wp.tile_load(a, shape=32, offset=batch_idx * 32)
+    s2 = wp.tile_load(b, shape=32, offset=batch_idx * 32)
+    # inter-tile operations (batch-wise)
+    if op_type < 9:
+        if op_type == 6:
+            s1 &= s2
+        elif op_type == 7:
+            s1 |= s2
+        elif op_type == 8:
+            s1 ^= s2
+        wp.tile_store(c, s1, offset=batch_idx * 32)
+    else:
+        if op_type == 9:
+            s3 = s1 & s2
+        elif op_type == 10:
+            s3 = s1 | s2
+        elif op_type == 11:
+            s3 = s1 ^ s2
+        wp.tile_store(c, s3, offset=batch_idx * 32)
+
+
+def test_tile_atomic_bitwise_vector(test, device):
+    n = 1024
+    rng = np.random.default_rng(42)
+
+    a = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3), dtype=np.uint32)
+    b = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3), dtype=np.uint32)
+
+    expected_and = a & b
+    expected_or = a | b
+    expected_xor = a ^ b
+
+    with wp.ScopedDevice(device):
+        a_wp = wp.array(a, dtype=wp.vec3ui, device=device)
+        b_wp = wp.array(b, dtype=wp.vec3ui, device=device)
+        c_wp = wp.zeros(shape=n, dtype=wp.vec3ui, device=device)
+
+        wp.launch_tiled(test_tile_atomic_bitwise_vector_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 0], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(test_tile_atomic_bitwise_vector_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 1], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(test_tile_atomic_bitwise_vector_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 2], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_xor)
+        wp.launch_tiled(test_tile_atomic_bitwise_vector_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 3], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(test_tile_atomic_bitwise_vector_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 4], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(test_tile_atomic_bitwise_vector_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 5], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_xor)
+
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_vector_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 6], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_vector_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 7], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_vector_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 8], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_xor)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_vector_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 9], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_vector_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 10], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_vector_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 11], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_xor)
+
+
+mat33ui = wp._src.types.matrix(shape=(3, 3), dtype=wp.uint32)
+
+
+@wp.kernel
+def test_tile_atomic_bitwise_matrix_kernel(
+    a: wp.array(dtype=mat33ui), b: wp.array(dtype=mat33ui), c: wp.array(dtype=mat33ui), op_type: int
+):
+    word_idx, bit_idx = wp.tid()
+    block_dim = wp.block_dim()
+    assert block_dim == 32
+    s = wp.tile_zeros(shape=1, dtype=mat33ui)
+    # write to tile first, then write only once to the array
+    s[0] = a[word_idx]
+    if op_type < 3:
+        bit_mask = mat33ui(wp.uint32(1)) << mat33ui(wp.uint32(bit_idx))
+        if op_type == 0:
+            s[0] &= (b[word_idx] & bit_mask) | ~bit_mask
+        elif op_type == 1:
+            s[0] |= b[word_idx] & bit_mask
+        elif op_type == 2:
+            s[0] ^= b[word_idx] & bit_mask
+    else:
+        # inter-tile operations
+        s_bit_mask = wp.tile_zeros(shape=32, dtype=mat33ui)
+        s_bit_mask[(bit_idx + 1) % 32] = mat33ui(wp.uint32(1) << wp.uint32((bit_idx + 1) % 32))
+        if op_type == 3:
+            s[0] &= (b[word_idx] & s_bit_mask[bit_idx]) | ~s_bit_mask[bit_idx]
+        elif op_type == 4:
+            s[0] |= b[word_idx] & s_bit_mask[bit_idx]
+        elif op_type == 5:
+            s[0] ^= b[word_idx] & s_bit_mask[bit_idx]
+    c[word_idx] = s[0]
+
+
+@wp.kernel
+def test_tile_atomic_bitwise_matrix_tilewise_kernel(
+    a: wp.array(dtype=mat33ui), b: wp.array(dtype=mat33ui), c: wp.array(dtype=mat33ui), op_type: int
+):
+    batch_idx, _ = wp.tid()
+    block_dim = wp.block_dim()
+    assert block_dim == 32
+    # Each tile is responsible for a batch of 32 elements
+    s1 = wp.tile_load(a, shape=32, offset=batch_idx * 32)
+    s2 = wp.tile_load(b, shape=32, offset=batch_idx * 32)
+    # inter-tile operations (batch-wise)
+    if op_type < 9:
+        if op_type == 6:
+            s1 &= s2
+        elif op_type == 7:
+            s1 |= s2
+        elif op_type == 8:
+            s1 ^= s2
+        wp.tile_store(c, s1, offset=batch_idx * 32)
+    else:
+        if op_type == 9:
+            s3 = s1 & s2
+        elif op_type == 10:
+            s3 = s1 | s2
+        elif op_type == 11:
+            s3 = s1 ^ s2
+        wp.tile_store(c, s3, offset=batch_idx * 32)
+
+
+def test_tile_atomic_bitwise_matrix(test, device):
+    n = 1024
+    rng = np.random.default_rng(42)
+
+    a = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3, 3), dtype=np.uint32)
+    b = rng.integers(0, np.iinfo(np.uint32).max, size=(n, 3, 3), dtype=np.uint32)
+
+    expected_and = a & b
+    expected_or = a | b
+    expected_xor = a ^ b
+
+    with wp.ScopedDevice(device):
+        a_wp = wp.array(a, dtype=mat33ui, device=device)
+        b_wp = wp.array(b, dtype=mat33ui, device=device)
+        c_wp = wp.zeros(shape=n, dtype=mat33ui, device=device)
+
+        wp.launch_tiled(test_tile_atomic_bitwise_matrix_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 0], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(test_tile_atomic_bitwise_matrix_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 1], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(test_tile_atomic_bitwise_matrix_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 2], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_xor)
+        wp.launch_tiled(test_tile_atomic_bitwise_matrix_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 3], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(test_tile_atomic_bitwise_matrix_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 4], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(test_tile_atomic_bitwise_matrix_kernel, dim=n, inputs=[a_wp, b_wp, c_wp, 5], block_dim=32)
+        assert_np_equal(c_wp.numpy(), expected_xor)
+
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_matrix_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 6], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_matrix_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 7], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_matrix_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 8], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_xor)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_matrix_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 9], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_and)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_matrix_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 10], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_or)
+        wp.launch_tiled(
+            test_tile_atomic_bitwise_matrix_tilewise_kernel, dim=n // 32, inputs=[a_wp, b_wp, c_wp, 11], block_dim=32
+        )
+        assert_np_equal(c_wp.numpy(), expected_xor)
+
+
+devices = get_cuda_test_devices()
+
+
+class TestTileAtomicBitwise(unittest.TestCase):
+    pass
+
+
+add_function_test(
+    TestTileAtomicBitwise,
+    "test_tile_atomic_bitwise_scalar",
+    test_tile_atomic_bitwise_scalar,
+    devices=devices,
+)
+
+add_function_test(
+    TestTileAtomicBitwise,
+    "test_tile_atomic_bitwise_vector",
+    test_tile_atomic_bitwise_vector,
+    devices=devices,
+)
+
+add_function_test(
+    TestTileAtomicBitwise,
+    "test_tile_atomic_bitwise_matrix",
+    test_tile_atomic_bitwise_matrix,
+    devices=devices,
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)