warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/tests/tile/test_tile_shared_memory.py

@@ -0,0 +1,376 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 *
+
+
+# checks that we can configure shared memory to the expected size
+def test_tile_shared_mem_size(test, device):
+    DIM_M = 32
+    DIM_N = 32
+
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(out: wp.array2d(dtype=float)):
+        a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
+        b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
+
+        c = a + b
+        wp.tile_store(out, c)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    # check output
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
+
+    # check required shared memory
+    expected_forward_bytes = DIM_M * DIM_N * 4 * 2
+    expected_backward_bytes = expected_forward_bytes * 2
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    assert hooks.forward_smem_bytes == expected_forward_bytes
+    assert hooks.backward_smem_bytes == expected_backward_bytes
+
+
+# checks that we can configure shared memory > 48kb default
+def test_tile_shared_mem_large(test, device):
+    # set dimensions that require 64kb for the forward kernel
+    DIM_M = 64
+    DIM_N = 128
+
+    BLOCK_DIM = 256
+
+    # we disable backward kernel gen since 128k is not supported on most architectures
+    @wp.kernel(enable_backward=False, module="unique")
+    def compute(out: wp.array2d(dtype=float)):
+        a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
+        b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
+
+        c = a + b
+        wp.tile_store(out, c)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    # check output
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
+
+    # check required shared memory
+    expected_forward_bytes = DIM_M * DIM_N * 4 * 2
+    expected_backward_bytes = 0
+
+    assert expected_forward_bytes == 2**16
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    assert hooks.forward_smem_bytes == expected_forward_bytes
+    assert hooks.backward_smem_bytes == expected_backward_bytes
+
+
+# checks that we can configure dynamic shared memory during graph capture
+def test_tile_shared_mem_graph(test, device):
+    DIM_M = 32
+    DIM_N = 32
+
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(out: wp.array2d(dtype=float)):
+        a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
+        b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
+
+        c = a + b
+        wp.tile_store(out, c)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    # preload the unique module
+    wp.load_module(compute.module, device=device, block_dim=BLOCK_DIM)
+
+    with wp.ScopedCapture(device, force_module_load=False) as capture:
+        wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    wp.capture_launch(capture.graph)
+
+    # check output
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N)) * 3.0)
+
+    # check required shared memory
+    expected_forward_bytes = DIM_M * DIM_N * 4 * 2
+    expected_backward_bytes = expected_forward_bytes * 2
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    assert hooks.forward_smem_bytes == expected_forward_bytes
+    assert hooks.backward_smem_bytes == expected_backward_bytes
+
+
+# checks that stack allocations work for user functions
+def test_tile_shared_mem_func(test, device):
+    DIM_M = 64
+    DIM_N = 64
+
+    SMALL_DIM_M = 64 // 4
+    SMALL_DIM_N = 64 // 4
+
+    BLOCK_DIM = 256
+
+    @wp.func
+    def add_tile_small():
+        a = wp.tile_ones(shape=(SMALL_DIM_M, SMALL_DIM_N), dtype=float, storage="shared")
+        b = wp.tile_ones(shape=(SMALL_DIM_M, SMALL_DIM_N), dtype=float, storage="shared") * 2.0
+
+        return a + b
+
+    @wp.func
+    def add_tile_big():
+        a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
+        b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
+
+        return a + b
+
+    @wp.kernel(module="unique")
+    def compute(out: wp.array2d(dtype=float)):
+        s = add_tile_small()
+        b = add_tile_big()
+
+        wp.tile_store(out, b)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    # ensure that total required dynamic shared is the larger of the two tiles
+    expected_required_shared = 64 * 64 * 4 * 2
+
+    assert hooks.forward_smem_bytes == expected_required_shared
+    assert hooks.backward_smem_bytes == expected_required_shared * 2
+
+
+def round_up(a, b):
+    return b * ((a + b - 1) // b)
+
+
+# checks that using non-16B aligned sizes work
+def test_tile_shared_non_aligned(test, device):
+    # Tile size = 4 (float) * 1 * 3 = 12B % 16 != 0
+    DIM_M = 1
+    DIM_N = 3
+
+    BLOCK_DIM = 256
+
+    @wp.func
+    def foo():
+        a = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 2.0
+        b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared") * 3.0
+        return a + b
+
+    @wp.kernel(module="unique")
+    def compute(out: wp.array2d(dtype=float)):
+        # This test the logic in the stack allocator, which should increment and
+        # decrement the stack pointer each time foo() is called
+        # Failing to do so correct will make b out of bounds and corrupt the results
+        for _ in range(4096):
+            foo()
+        b = wp.tile_ones(shape=(DIM_M, DIM_N), dtype=float, storage="shared")
+        wp.tile_store(out, b)
+
+    out = wp.empty((DIM_M, DIM_N), dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=[1], inputs=[out], block_dim=BLOCK_DIM, device=device)
+
+    assert_np_equal(out.numpy(), np.ones((DIM_M, DIM_N), dtype=float))
+
+    # check shared memory for kernel on the device
+    module_exec = compute.module.load(device, BLOCK_DIM)
+    hooks = module_exec.get_kernel_hooks(compute)
+
+    # ensure that total required dynamic shared is the larger of the two tiles
+    expected_required_shared = 3 * round_up(DIM_M * DIM_N * 4, 16)
+
+    assert hooks.forward_smem_bytes == expected_required_shared
+    assert hooks.backward_smem_bytes == expected_required_shared * 2
+
+
+def test_tile_shared_vec_accumulation(test, device):
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(indices: wp.array(dtype=int), vecs: wp.array(dtype=wp.vec3), output: wp.array2d(dtype=float)):
+        i, j = wp.tid()
+
+        idx_tile = wp.tile_load(indices, shape=BLOCK_DIM, offset=i * BLOCK_DIM)
+        idx = idx_tile[j]
+
+        s = wp.tile_zeros(shape=(1, 3), dtype=float)
+
+        s[0, 0] += vecs[idx].x
+        s[0, 1] += vecs[idx].y
+        s[0, 2] += vecs[idx].z
+
+        wp.tile_store(output, s, offset=(i, 0))
+
+    N = BLOCK_DIM * 3
+
+    basis_vecs = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32)
+    vecs = wp.array(basis_vecs, dtype=wp.vec3, requires_grad=True, device=device)
+
+    rng = np.random.default_rng(42)
+    indices_np = rng.integers(0, 3, size=N)
+
+    indices = wp.array(indices_np, dtype=int, requires_grad=True, device=device)
+
+    output = wp.zeros(shape=(3, 3), dtype=float, requires_grad=True, device=device)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch_tiled(compute, dim=3, inputs=[indices, vecs, output], block_dim=BLOCK_DIM, device=device)
+
+    output.grad = wp.ones_like(output)
+
+    tape.backward()
+
+    n0 = np.count_nonzero(indices_np == 0)
+    n1 = np.count_nonzero(indices_np == 1)
+    n2 = np.count_nonzero(indices_np == 2)
+    true_grads = np.array([[n0, n0, n0], [n1, n1, n1], [n2, n2, n2]])
+
+    indices_np = indices_np.reshape((3, BLOCK_DIM))
+
+    def compute_row(idx):
+        n0 = np.count_nonzero(indices_np[idx, :] == 0)
+        n1 = np.count_nonzero(indices_np[idx, :] == 1)
+        n2 = np.count_nonzero(indices_np[idx, :] == 2)
+        return np.array([1, 0, 0]) * n0 + np.array([0, 1, 0]) * n1 + np.array([0, 0, 1]) * n2
+
+    row_0 = compute_row(0)
+    row_1 = compute_row(1)
+    row_2 = compute_row(2)
+
+    true_vecs = np.stack([row_0, row_1, row_2])
+
+    assert_np_equal(output.numpy(), true_vecs)
+    assert_np_equal(vecs.grad.numpy(), true_grads)
+
+
+def test_tile_shared_simple_reduction_add(test, device):
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+        i, j = wp.tid()
+
+        t = wp.tile_load(x, shape=BLOCK_DIM, offset=BLOCK_DIM * i)
+
+        k = BLOCK_DIM // 2
+        while k > 0:
+            if j < k:
+                t[j] += t[j + k]
+            k //= 2
+
+        wp.tile_store(y, wp.tile_view(t, offset=(0,), shape=(1,)), i)
+
+    N = BLOCK_DIM * 4
+    x_np = np.arange(N, dtype=np.float32)
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(4, dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=4, inputs=[x], outputs=[y], block_dim=BLOCK_DIM, device=device)
+
+    assert_np_equal(np.sum(y.numpy()), np.sum(x_np))
+
+
+def test_tile_shared_simple_reduction_sub(test, device):
+    BLOCK_DIM = 256
+
+    @wp.kernel(module="unique")
+    def compute(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+        i, j = wp.tid()
+
+        t = wp.tile_load(x, shape=BLOCK_DIM, offset=BLOCK_DIM * i)
+
+        k = BLOCK_DIM // 2
+        while k > 0:
+            if j < k:
+                t[j] -= t[j + k]
+            k //= 2
+
+        wp.tile_store(y, wp.tile_view(t, offset=(0,), shape=(1,)), i)
+
+    N = BLOCK_DIM * 4
+    x_np = np.arange(N, dtype=np.float32)
+    x = wp.array(x_np, dtype=float, device=device)
+    y = wp.zeros(4, dtype=float, device=device)
+
+    wp.launch_tiled(compute, dim=4, inputs=[x], outputs=[y], block_dim=BLOCK_DIM, device=device)
+
+    assert_np_equal(np.sum(y.numpy()), 0.0)
+
+
+devices = get_cuda_test_devices()
+
+
+class TestTileSharedMemory(unittest.TestCase):
+    pass
+
+
+add_function_test(
+    TestTileSharedMemory, "test_tile_shared_mem_size", test_tile_shared_mem_size, devices=devices, check_output=False
+)
+add_function_test(
+    TestTileSharedMemory, "test_tile_shared_mem_large", test_tile_shared_mem_large, devices=devices, check_output=False
+)
+add_function_test(TestTileSharedMemory, "test_tile_shared_mem_graph", test_tile_shared_mem_graph, devices=devices)
+add_function_test(TestTileSharedMemory, "test_tile_shared_mem_func", test_tile_shared_mem_func, devices=devices)
+add_function_test(TestTileSharedMemory, "test_tile_shared_non_aligned", test_tile_shared_non_aligned, devices=devices)
+add_function_test(
+    TestTileSharedMemory, "test_tile_shared_vec_accumulation", test_tile_shared_vec_accumulation, devices=devices
+)
+add_function_test(
+    TestTileSharedMemory,
+    "test_tile_shared_simple_reduction_add",
+    test_tile_shared_simple_reduction_add,
+    devices=devices,
+)
+add_function_test(
+    TestTileSharedMemory,
+    "test_tile_shared_simple_reduction_sub",
+    test_tile_shared_simple_reduction_sub,
+    devices=devices,
+)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/tile/test_tile_sort.py

@@ -0,0 +1,121 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 *
+
+
+def create_sort_kernel(KEY_TYPE, MAX_SORT_LENGTH):
+    @wp.kernel
+    def tile_sort_kernel(
+        input_keys: wp.array(dtype=KEY_TYPE),
+        input_values: wp.array(dtype=wp.int32),
+        output_keys: wp.array(dtype=KEY_TYPE),
+        output_values: wp.array(dtype=wp.int32),
+    ):
+        # Load input into shared memory
+        keys = wp.tile_load(input_keys, shape=MAX_SORT_LENGTH, storage="shared")
+        values = wp.tile_load(input_values, shape=MAX_SORT_LENGTH, storage="shared")
+
+        # Perform in-place sorting
+        wp.tile_sort(keys, values)
+
+        # Store sorted shared memory into output arrays
+        wp.tile_store(output_keys, keys)
+        wp.tile_store(output_values, values)
+
+    return tile_sort_kernel
+
+
+def test_tile_sort(test, device):
+    # Forward-declare kernels for more efficient compilation
+    kernels = {}
+    for dtype in [int, float]:
+        for i in range(0, 11):
+            length = 2**i + 1
+            kernels[(dtype, length)] = create_sort_kernel(dtype, length)
+
+    for (dtype, length), kernel in kernels.items():
+        for j in range(5, 10):
+            TILE_DIM = 2**j
+
+            rng = np.random.default_rng(42)  # Create a random generator instance
+
+            if dtype == int:
+                np_keys = rng.choice(1000000000, size=length, replace=False)
+            else:  # dtype == float
+                np_keys = rng.uniform(0, 1000000000, size=length).astype(dtype)
+
+            np_values = np.arange(length)
+
+            # Generate random keys and iota indexer
+            input_keys = wp.array(np_keys, dtype=dtype, device=device)
+            input_values = wp.array(np_values, dtype=int, device=device)
+            output_keys = wp.zeros_like(input_keys, device=device)
+            output_values = wp.zeros_like(input_values, device=device)
+
+            # Execute sorting kernel
+            wp.launch_tiled(
+                kernel,
+                dim=1,
+                inputs=[input_keys, input_values, output_keys, output_values],
+                block_dim=TILE_DIM,
+                device=device,
+            )
+            wp.synchronize()
+
+            # Sort using NumPy for validation
+            sorted_indices = np.argsort(np_keys)
+            np_sorted_keys = np_keys[sorted_indices]
+            np_sorted_values = np_values[sorted_indices]
+
+            if dtype == int:
+                keys_match = np.array_equal(output_keys.numpy(), np_sorted_keys)
+            else:  # dtype == float
+                keys_match = np.allclose(output_keys.numpy(), np_sorted_keys, atol=1e-6)  # Use tolerance for floats
+
+            values_match = np.array_equal(output_values.numpy(), np_sorted_values)
+
+            if not keys_match or not values_match:
+                print(f"Test failed for dtype={dtype}, TILE_DIM={TILE_DIM}, length={length}")
+                print("")
+                print(output_keys.numpy())
+                print(np_sorted_keys)
+                print("")
+                print(output_values.numpy())
+                print(np_sorted_values)
+                print("")
+
+            # Validate results
+            test.assertTrue(keys_match, f"Key sorting mismatch for dtype={dtype}!")
+            test.assertTrue(values_match, f"Value sorting mismatch for dtype={dtype}!")
+
+
+devices = get_test_devices()
+
+
+class TestTileSort(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTileSort, "test_tile_sort", test_tile_sort, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/tile/test_tile_view.py

@@ -0,0 +1,173 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 *
+
+TILE_DIM = 64
+TILE_M = 16
+TILE_N = 32
+TILE_O = 8
+
+
+@wp.kernel
+def test_tile_view_kernel(src: wp.array2d(dtype=float), dst: wp.array2d(dtype=float)):
+    # load whole source into local memory
+    a = wp.tile_load(src, shape=(TILE_M, TILE_N))
+
+    # copy the source array row by row
+    for i in range(TILE_M):
+        # create a view on original array and store
+        row = a[i]
+        wp.tile_store(dst[i], row)
+
+
+def test_tile_view(test, device):
+    rng = np.random.default_rng(42)
+
+    a = wp.array(rng.random((TILE_M, TILE_N), dtype=np.float32), requires_grad=True, device=device)
+    b = wp.array(np.zeros((TILE_M, TILE_N), dtype=np.float32), requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(test_tile_view_kernel, dim=[1], inputs=[a, b], block_dim=32, device=device)
+
+    assert_np_equal(b.numpy(), a.numpy())
+    b.grad = wp.ones_like(b, device=device)
+    tape.backward()
+
+    assert_np_equal(a.grad.numpy(), np.ones_like(a.numpy()))
+
+
+@wp.kernel
+def test_tile_assign_1d_kernel(src: wp.array2d(dtype=float), dst: wp.array2d(dtype=float)):
+    # load whole source into local memory
+    a = wp.tile_load(src, shape=(TILE_M, TILE_N))
+    b = wp.tile_zeros(dtype=float, shape=(TILE_M, TILE_N))
+
+    # copy the source array row by row
+    for i in range(int(TILE_M)):
+        # create views onto source and dest rows
+        row_src = a[i]
+        row_dst = b[i]
+
+        # copy onto dest row
+        wp.tile_assign(row_dst, row_src)
+
+    wp.tile_store(dst, b)
+
+
+def test_tile_assign_1d(test, device):
+    rng = np.random.default_rng(42)
+
+    a = wp.array(rng.random((TILE_M, TILE_N), dtype=np.float32), requires_grad=True, device=device)
+    b = wp.array(np.zeros((TILE_M, TILE_N), dtype=np.float32), requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(test_tile_assign_1d_kernel, dim=[1], inputs=[a, b], block_dim=32, device=device)
+
+    assert_np_equal(b.numpy(), a.numpy())
+    b.grad = wp.ones_like(b, device=device)
+    tape.backward()
+
+    assert_np_equal(a.grad.numpy(), np.ones_like(a.numpy()))
+
+
+@wp.kernel
+def test_tile_assign_2d_kernel(src: wp.array3d(dtype=float), dst: wp.array3d(dtype=float)):
+    # load whole source into local memory
+    a = wp.tile_load(src, shape=(TILE_M, TILE_N, TILE_O))
+    b = wp.tile_zeros(dtype=float, shape=(TILE_M, TILE_N, TILE_O))
+
+    # copy the source array slice by slice
+    for i in range(TILE_M):
+        # create views onto source and dest slice
+        row_src = a[i]
+        row_dst = b[i]
+
+        # copy onto dest slice
+        wp.tile_assign(row_dst, row_src)
+
+    wp.tile_store(dst, b)
+
+
+def test_tile_assign_2d(test, device):
+    rng = np.random.default_rng(42)
+
+    a = wp.array(rng.random((TILE_M, TILE_N, TILE_O), dtype=np.float32), requires_grad=True, device=device)
+    b = wp.array(np.zeros((TILE_M, TILE_N, TILE_O), dtype=np.float32), requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(test_tile_assign_2d_kernel, dim=[1], inputs=[a, b], block_dim=32, device=device)
+
+    assert_np_equal(b.numpy(), a.numpy())
+    b.grad = wp.ones_like(b, device=device)
+    tape.backward()
+
+    assert_np_equal(a.grad.numpy(), np.ones_like(a.numpy()))
+
+
+@wp.kernel
+def test_tile_view_offset_kernel(src: wp.array2d(dtype=float), dst: wp.array2d(dtype=float)):
+    # load whole source into local memory
+    a = wp.tile_load(src, shape=(TILE_M, TILE_N))
+    b = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=float)
+
+    # copy the source array slice by slice
+    for i in range(TILE_M // 4):
+        # create views onto source and dest slice 4 rows at a time
+        v = wp.tile_view(a, offset=(i * 4, 0), shape=(4, TILE_N))
+
+        # copy onto dest slice
+        wp.tile_assign(b, v, offset=(i * 4, 0))
+
+    wp.tile_store(dst, b)
+
+
+def test_tile_view_offset(test, device):
+    rng = np.random.default_rng(42)
+
+    a = wp.array(rng.random((TILE_M, TILE_N), dtype=np.float32), requires_grad=True, device=device)
+    b = wp.array(np.zeros((TILE_M, TILE_N), dtype=np.float32), requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(test_tile_view_offset_kernel, dim=[1], inputs=[a, b], block_dim=32, device=device)
+
+    assert_np_equal(b.numpy(), a.numpy())
+    b.grad = wp.ones_like(b, device=device)
+    tape.backward()
+
+    assert_np_equal(a.grad.numpy(), np.ones_like(a.numpy()))
+
+
+devices = get_test_devices()
+
+
+class TestTileView(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTileView, "test_tile_view", test_tile_view, devices=devices)
+add_function_test(TestTileView, "test_tile_view_offset", test_tile_view_offset, devices=devices)
+add_function_test(TestTileView, "test_tile_assign_1d", test_tile_assign_1d, devices=devices)
+add_function_test(TestTileView, "test_tile_assign_2d", test_tile_assign_2d, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)