warp-lang 1.5.1__py3-none-manylinux2014_x86_64.whl → 1.6.1__py3-none-manylinux2014_x86_64.whl

warp/tests/test_tile_load.py

@@ -0,0 +1,399 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+TILE_DIM = 64
+
+TILE_M = wp.constant(16)
+TILE_N = wp.constant(8)
+TILE_O = wp.constant(8)
+TILE_P = wp.constant(6)
+
+TILE_OFFSET = 5
+
+
+@wp.kernel
+def tile_load_1d_kernel(
+    input: wp.array1d(dtype=float),
+    out_full: wp.array1d(dtype=float),
+    out_padded: wp.array1d(dtype=float),
+    out_offset: wp.array1d(dtype=float),
+):
+    full0 = wp.tile_load(input, TILE_M)
+    full1 = wp.tile_load(input, shape=TILE_M)
+    full2 = wp.tile_load(input, shape=(TILE_M,))
+
+    padded0 = wp.tile_load(input, TILE_M, TILE_OFFSET)
+    padded1 = wp.tile_load(input, shape=TILE_M, offset=TILE_OFFSET)
+    padded2 = wp.tile_load(input, shape=(TILE_M,), offset=(TILE_OFFSET,))
+
+    wp.tile_store(out_full, full0)
+    wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_offset, full0, offset=(TILE_OFFSET,))
+
+
+@wp.kernel
+def tile_load_2d_kernel(
+    input: wp.array2d(dtype=float),
+    out_full: wp.array2d(dtype=float),
+    out_padded: wp.array2d(dtype=float),
+    out_offset: wp.array2d(dtype=float),
+):
+    full0 = wp.tile_load(input, shape=(TILE_M, TILE_N))
+    padded0 = wp.tile_load(input, shape=(TILE_M, TILE_N), offset=(TILE_OFFSET, TILE_OFFSET))
+
+    wp.tile_store(out_full, full0)
+    wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET))
+
+
+@wp.kernel
+def tile_load_3d_kernel(
+    input: wp.array3d(dtype=float),
+    out_full: wp.array3d(dtype=float),
+    out_padded: wp.array3d(dtype=float),
+    out_offset: wp.array3d(dtype=float),
+):
+    full0 = wp.tile_load(input, shape=(TILE_M, TILE_N, TILE_O))
+    padded0 = wp.tile_load(input, shape=(TILE_M, TILE_N, TILE_O), offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
+
+    wp.tile_store(out_full, full0)
+    wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
+
+
+@wp.kernel
+def tile_load_4d_kernel(
+    input: wp.array4d(dtype=float),
+    out_full: wp.array4d(dtype=float),
+    out_padded: wp.array4d(dtype=float),
+    out_offset: wp.array4d(dtype=float),
+):
+    full0 = wp.tile_load(input, shape=(TILE_M, TILE_N, TILE_O, TILE_P))
+    padded0 = wp.tile_load(
+        input, shape=(TILE_M, TILE_N, TILE_O, TILE_P), offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET, TILE_OFFSET)
+    )
+
+    wp.tile_store(out_full, full0)
+    wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
+
+
+def test_tile_load(kernel, ndim):
+    def test(test, device):
+        rng = np.random.default_rng(42)
+
+        shape = [TILE_M, TILE_N, TILE_O, TILE_P]
+        shape = shape[0:ndim]
+
+        input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
+        output_full = wp.zeros(shape, dtype=float, device=device)
+        output_padded = wp.zeros(shape, dtype=float, device=device)
+        output_offset = wp.zeros(shape, dtype=float, device=device)
+
+        with wp.Tape() as tape:
+            wp.launch_tiled(
+                kernel,
+                dim=[1],
+                inputs=[input, output_full, output_padded, output_offset],
+                block_dim=TILE_DIM,
+                device=device,
+            )
+
+        # construct a slice for the offset portion of the source/dest arrays
+        src_slice = tuple(slice(TILE_OFFSET, dim) for dim in shape)
+        dest_slice = tuple(slice(None, dim - TILE_OFFSET) for dim in shape)
+
+        ref_full = input.numpy()
+        ref_padded = np.zeros_like(ref_full)
+        ref_padded[dest_slice] = ref_full[src_slice]
+
+        ref_offset = np.zeros_like(ref_full)
+        ref_offset[src_slice] = ref_full[dest_slice]
+
+        assert_np_equal(output_full.numpy(), ref_full)
+        assert_np_equal(output_padded.numpy(), ref_padded)
+        assert_np_equal(output_offset.numpy(), ref_offset)
+
+        output_full.grad = wp.ones_like(output_full)
+        tape.backward()
+
+        assert_np_equal(input.grad.numpy(), np.ones_like(input.grad.numpy()))
+
+    return test
+
+
+@wp.kernel
+def tile_load_unaligned_kernel(
+    input: wp.array2d(dtype=float),
+    output: wp.array2d(dtype=float),
+):
+    t = wp.tile_load(input, shape=(TILE_M, TILE_N), offset=(1, 1), storage="shared")
+    wp.tile_store(output, t, offset=(1, 1))
+
+
+def test_tile_load_unaligned(test, device):
+    rng = np.random.default_rng(42)
+
+    shape = [TILE_M + 1, TILE_N + 1]
+
+    input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
+    output = wp.zeros(shape, dtype=float, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_load_unaligned_kernel,
+            dim=[1],
+            inputs=[input, output],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    # first row and column should be zero
+    assert_np_equal(output.numpy()[0, :], np.zeros(TILE_N + 1))
+    assert_np_equal(output.numpy()[:, 0], np.zeros(TILE_M + 1))
+
+    # check output elements
+    assert_np_equal(output.numpy()[1:, 1:], input.numpy()[1:, 1:])
+
+    output.grad = wp.ones_like(output)
+    tape.backward()
+
+    expected_grad = np.ones_like(input.grad.numpy())
+    expected_grad[0, :] = 0.0
+    expected_grad[:, 0] = 0.0
+
+    assert_np_equal(input.grad.numpy(), expected_grad)
+
+
+# ----------------------------------------------------------------------------------------
+
+TILE_SIZE = 4
+
+
+@wp.kernel
+def tile_extract_1d_kernel(input: wp.array1d(dtype=float), output: wp.array1d(dtype=float)):
+    i = wp.tid()
+
+    t = wp.tile_load(input, shape=TILE_SIZE)
+
+    output[i] = t[i]
+
+
+@wp.kernel
+def tile_extract_2d_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+
+    t = wp.tile_load(input, shape=(TILE_SIZE, TILE_SIZE))
+
+    output[i, j] = t[i, j]
+
+
+@wp.kernel
+def tile_extract_3d_kernel(input: wp.array3d(dtype=float), output: wp.array3d(dtype=float)):
+    i, j, k = wp.tid()
+
+    t = wp.tile_load(input, shape=(TILE_SIZE, TILE_SIZE, TILE_SIZE))
+
+    output[i, j, k] = t[i, j, k]
+
+
+@wp.kernel
+def tile_extract_4d_kernel(input: wp.array4d(dtype=float), output: wp.array4d(dtype=float)):
+    i, j, k, l = wp.tid()
+
+    t = wp.tile_load(input, shape=(TILE_SIZE, TILE_SIZE, TILE_SIZE, TILE_SIZE))
+
+    output[i, j, k, l] = t[i, j, k, l]
+
+
+def test_tile_extract(kernel, ndim):
+    shape = (TILE_SIZE,) * ndim
+
+    def test_run(test, device):
+        rng = np.random.default_rng(42)
+
+        input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
+        output = wp.zeros_like(input)
+
+        with wp.Tape() as tape:
+            wp.launch(
+                kernel,
+                dim=shape,
+                inputs=[input, output],
+                block_dim=1024,
+                device=device,
+            )
+
+        assert_np_equal(output.numpy(), input.numpy())
+
+        output.grad = wp.ones_like(output)
+        tape.backward()
+
+        assert_np_equal(input.grad.numpy(), np.ones_like(input.numpy()))
+
+    return test_run
+
+
+# ----------------------------------------------------------------------------------------
+
+TILE_SIZE = 4
+
+
+@wp.kernel
+def tile_assign_1d_kernel(input: wp.array1d(dtype=float), output: wp.array1d(dtype=float)):
+    i = wp.tid()
+
+    t = wp.tile_zeros(shape=(TILE_SIZE,), dtype=float)
+
+    # assign to tile
+    t[i] = input[i] * 2.0
+
+    output[i] = t[i]
+
+
+@wp.kernel
+def tile_assign_2d_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+
+    t = wp.tile_zeros(shape=(TILE_SIZE, TILE_SIZE), dtype=float)
+
+    # assign to tile
+    t[i, j] = input[i, j] * 2.0
+
+    output[i, j] = t[i, j]
+
+
+@wp.kernel
+def tile_assign_3d_kernel(input: wp.array3d(dtype=float), output: wp.array3d(dtype=float)):
+    i, j, k = wp.tid()
+
+    t = wp.tile_zeros(shape=(TILE_SIZE, TILE_SIZE, TILE_SIZE), dtype=float)
+
+    # assign to tile
+    t[i, j, k] = input[i, j, k] * 2.0
+
+    output[i, j, k] = t[i, j, k]
+
+
+@wp.kernel
+def tile_assign_4d_kernel(input: wp.array4d(dtype=float), output: wp.array4d(dtype=float)):
+    i, j, k, l = wp.tid()
+
+    t = wp.tile_zeros(shape=(TILE_SIZE, TILE_SIZE, TILE_SIZE, TILE_SIZE), dtype=float)
+
+    # assign to tile
+    t[i, j, k, l] = input[i, j, k, l] * 2.0
+
+    output[i, j, k, l] = t[i, j, k, l]
+
+
+def test_tile_assign(kernel, ndim):
+    shape = (TILE_SIZE,) * ndim
+
+    def test_run(test, device):
+        rng = np.random.default_rng(42)
+
+        input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
+        output = wp.zeros_like(input)
+
+        with wp.Tape() as tape:
+            wp.launch(
+                kernel,
+                dim=shape,
+                inputs=[input, output],
+                block_dim=1024,
+                device=device,
+            )
+
+        assert_np_equal(output.numpy(), input.numpy() * 2.0)
+
+    return test_run
+
+
+# ----------------------------------------------------------------------------------------
+
+
+@wp.kernel
+def tile_load_fortran_kernel(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float)):
+    # tile index
+    i, j = wp.tid()
+
+    a = wp.tile_load(A, shape=(TILE_M, TILE_N), offset=(i * TILE_M, j * TILE_N))
+    wp.tile_store(B, t=a, offset=(i * TILE_M, j * TILE_N))
+
+
+def test_tile_load_fortran(test, device):
+    rng = np.random.default_rng(42)
+
+    M = TILE_M * 7
+    N = TILE_N * 5
+
+    A = rng.random((M, N), dtype=np.float32)
+    B = rng.random((M, N), dtype=np.float32)
+
+    # convert to column major layout
+    A = np.asfortranarray(A)
+    B = np.asfortranarray(B)
+
+    A_wp = wp.array(A, requires_grad=True, device=device)
+    B_wp = wp.array(B, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_load_fortran_kernel,
+            dim=[int(M / TILE_M), int(N / TILE_N)],
+            inputs=[A_wp, B_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    # verify forward pass
+    assert_array_equal(B_wp, A_wp)
+
+    # verify backward pass
+    B_wp.grad = wp.ones_like(B_wp, device=device)
+    tape.backward()
+
+    assert_array_equal(B_wp.grad, A_wp.grad)
+
+
+devices = get_cuda_test_devices()
+
+
+class TestTileLoad(unittest.TestCase):
+    pass
+
+
+add_function_test(TestTileLoad, "test_tile_load_1d", test_tile_load(tile_load_1d_kernel, 1), devices=devices)
+add_function_test(TestTileLoad, "test_tile_load_2d", test_tile_load(tile_load_2d_kernel, 2), devices=devices)
+add_function_test(TestTileLoad, "test_tile_load_3d", test_tile_load(tile_load_3d_kernel, 3), devices=devices)
+add_function_test(TestTileLoad, "test_tile_load_4d", test_tile_load(tile_load_4d_kernel, 4), devices=devices)
+add_function_test(TestTileLoad, "test_tile_load_unaligned", test_tile_load_unaligned, devices=devices)
+
+add_function_test(TestTileLoad, "test_tile_extract_1d", test_tile_extract(tile_extract_1d_kernel, 1), devices=devices)
+add_function_test(TestTileLoad, "test_tile_extract_2d", test_tile_extract(tile_extract_2d_kernel, 2), devices=devices)
+add_function_test(TestTileLoad, "test_tile_extract_3d", test_tile_extract(tile_extract_3d_kernel, 3), devices=devices)
+add_function_test(TestTileLoad, "test_tile_extract_4d", test_tile_extract(tile_extract_4d_kernel, 4), devices=devices)
+
+add_function_test(TestTileLoad, "test_tile_assign_1d", test_tile_assign(tile_assign_1d_kernel, 1), devices=devices)
+add_function_test(TestTileLoad, "test_tile_assign_2d", test_tile_assign(tile_assign_2d_kernel, 2), devices=devices)
+add_function_test(TestTileLoad, "test_tile_assign_3d", test_tile_assign(tile_assign_3d_kernel, 3), devices=devices)
+add_function_test(TestTileLoad, "test_tile_assign_4d", test_tile_assign(tile_assign_4d_kernel, 4), devices=devices)
+
+add_function_test(TestTileLoad, "test_tile_load_fortran", test_tile_load_fortran, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_tile_mathdx.py

@@ -30,11 +30,11 @@ def tile_math_matmul_kernel(
     ga: wp.array2d(dtype=wp.float16), gb: wp.array2d(dtype=wp.float32), gc: wp.array2d(dtype=wp.float64)
 ):
     i, j = wp.tid()
-    a = wp.tile_load(ga, i, j, m=TILE_M, n=TILE_K)
-    b = wp.tile_load(gb, i, j, m=TILE_K, n=TILE_N)
-    c = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=wp.float64)
+    a = wp.tile_load(ga, shape=(TILE_M, TILE_K), offset=(i * TILE_M, j * TILE_K))
+    b = wp.tile_load(gb, shape=(TILE_K, TILE_N), offset=(i * TILE_K, j * TILE_N))
+    c = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float64)
     wp.tile_matmul(a, b, c)
-    wp.tile_store(gc, i, j, c)
+    wp.tile_store(gc, c, offset=(i * TILE_M, j * TILE_N))
 
 
 def test_tile_math_matmul(test, device):
@@ -71,17 +71,17 @@ def test_tile_math_matmul(test, device):
 @wp.kernel()
 def tile_math_fft_kernel_vec2f(gx: wp.array2d(dtype=wp.vec2f), gy: wp.array2d(dtype=wp.vec2f)):
     i, j = wp.tid()
-    xy = wp.tile_load(gx, i, j, m=FFT_SIZE_FP32, n=FFT_SIZE_FP32)
+    xy = wp.tile_load(gx, shape=(FFT_SIZE_FP32, FFT_SIZE_FP32))
     wp.tile_fft(xy)
-    wp.tile_store(gy, i, j, xy)
+    wp.tile_store(gy, xy)
 
 
 @wp.kernel()
 def tile_math_fft_kernel_vec2d(gx: wp.array2d(dtype=wp.vec2d), gy: wp.array2d(dtype=wp.vec2d)):
     i, j = wp.tid()
-    xy = wp.tile_load(gx, i, j, m=FFT_SIZE_FP64, n=FFT_SIZE_FP64)
+    xy = wp.tile_load(gx, shape=(FFT_SIZE_FP64, FFT_SIZE_FP64))
     wp.tile_fft(xy)
-    wp.tile_store(gy, i, j, xy)
+    wp.tile_store(gy, xy)
 
 
 def test_tile_math_fft(test, device, wp_dtype):
@@ -114,6 +114,56 @@ def test_tile_math_fft(test, device, wp_dtype):
     # TODO: implement and test backward pass
 
 
+@wp.kernel()
+def tile_math_cholesky(
+    gA: wp.array2d(dtype=wp.float64),
+    gD: wp.array1d(dtype=wp.float64),
+    gL: wp.array2d(dtype=wp.float64),
+    gx: wp.array1d(dtype=wp.float64),
+    gy: wp.array1d(dtype=wp.float64),
+):
+    i, j = wp.tid()
+    # Load A, D & x
+    a = wp.tile_load(gA, shape=(TILE_M, TILE_M), storage="shared")
+    d = wp.tile_load(gD, shape=TILE_M, storage="shared")
+    x = wp.tile_load(gx, shape=TILE_M, storage="shared")
+    # Compute L st LL^T = A + diag(D)
+    b = wp.tile_diag_add(a, d)
+    l = wp.tile_cholesky(b)
+    # Solve for y in LL^T y = x
+    y = wp.tile_cholesky_solve(l, x)
+    # Store L & y
+    wp.tile_store(gL, l)
+    wp.tile_store(gy, y)
+
+
+def test_tile_math_cholesky(test, device):
+    A_h = np.ones((TILE_M, TILE_M), dtype=np.float64)
+    D_h = 8.0 * np.ones(TILE_M, dtype=np.float64)
+    L_h = np.zeros_like(A_h)
+    X_h = np.arange(TILE_M, dtype=np.float64)
+    Y_h = np.zeros_like(X_h)
+
+    A_np = A_h + np.diag(D_h)
+    L_np = np.linalg.cholesky(A_np)
+    Y_np = np.linalg.solve(A_np, X_h)
+
+    A_wp = wp.array2d(A_h, requires_grad=True, dtype=wp.float64, device=device)
+    D_wp = wp.array2d(D_h, requires_grad=True, dtype=wp.float64, device=device)
+    L_wp = wp.array2d(L_h, requires_grad=True, dtype=wp.float64, device=device)
+    X_wp = wp.array2d(X_h, requires_grad=True, dtype=wp.float64, device=device)
+    Y_wp = wp.array2d(Y_h, requires_grad=True, dtype=wp.float64, device=device)
+
+    wp.launch_tiled(
+        tile_math_cholesky, dim=[1, 1], inputs=[A_wp, D_wp, L_wp, X_wp, Y_wp], block_dim=TILE_DIM, device=device
+    )
+    wp.synchronize_device()
+
+    assert np.allclose(Y_wp.numpy(), Y_np) and np.allclose(L_wp.numpy(), L_np)
+
+    # TODO: implement and test backward pass
+
+
 devices = get_cuda_test_devices()
 
 
@@ -124,6 +174,9 @@ class TestTileMathDx(unittest.TestCase):
 
 # check_output=False so we can enable libmathdx's logging without failing the tests
 add_function_test(TestTileMathDx, "test_tile_math_matmul", test_tile_math_matmul, devices=devices, check_output=False)
+add_function_test(
+    TestTileMathDx, "test_tile_math_cholesky", test_tile_math_cholesky, devices=devices, check_output=False
+)
 add_function_test(
     TestTileMathDx,
     "test_tile_math_fft_vec2f",

warp/tests/test_tile_mlp.py

@@ -114,23 +114,23 @@ def test_multi_layer_nn(test, device):
         f = wp.tile(local)
 
         # input layer
-        w0 = wp.tile_load(weights_0, 0, 0, m=DIM_HID, n=DIM_IN)
-        b0 = wp.tile_load(bias_0, 0, 0, m=DIM_HID, n=1)
-        z = wp.tile_map(relu, wp.tile_matmul(w0, f) + wp.tile_broadcast(b0, m=DIM_HID, n=NUM_THREADS))
+        w0 = wp.tile_load(weights_0, shape=(DIM_HID, DIM_IN))
+        b0 = wp.tile_load(bias_0, shape=(DIM_HID, 1))
+        z = wp.tile_map(relu, wp.tile_matmul(w0, f) + wp.tile_broadcast(b0, shape=(DIM_HID, NUM_THREADS)))
 
         # hidden layer
-        w1 = wp.tile_load(weights_1, 0, 0, m=DIM_HID, n=DIM_HID)
-        b1 = wp.tile_load(bias_1, 0, 0, m=DIM_HID, n=1)
-        z = wp.tile_map(relu, wp.tile_matmul(w1, z) + wp.tile_broadcast(b1, m=DIM_HID, n=NUM_THREADS))
+        w1 = wp.tile_load(weights_1, shape=(DIM_HID, DIM_HID))
+        b1 = wp.tile_load(bias_1, shape=(DIM_HID, 1))
+        z = wp.tile_map(relu, wp.tile_matmul(w1, z) + wp.tile_broadcast(b1, shape=(DIM_HID, NUM_THREADS)))
 
-        w2 = wp.tile_load(weights_2, 0, 0, m=DIM_HID, n=DIM_HID)
-        b2 = wp.tile_load(bias_2, 0, 0, m=DIM_HID, n=1)
-        z = wp.tile_map(relu, wp.tile_matmul(w2, z) + wp.tile_broadcast(b2, m=DIM_HID, n=NUM_THREADS))
+        w2 = wp.tile_load(weights_2, shape=(DIM_HID, DIM_HID))
+        b2 = wp.tile_load(bias_2, shape=(DIM_HID, 1))
+        z = wp.tile_map(relu, wp.tile_matmul(w2, z) + wp.tile_broadcast(b2, shape=(DIM_HID, NUM_THREADS)))
 
         # output layer
-        w3 = wp.tile_load(weights_3, 0, 0, m=DIM_OUT, n=DIM_HID)
-        b3 = wp.tile_load(bias_3, 0, 0, m=DIM_OUT, n=1)
-        o = wp.tile_map(relu, wp.tile_matmul(w3, z) + wp.tile_broadcast(b3, m=DIM_OUT, n=NUM_THREADS))
+        w3 = wp.tile_load(weights_3, shape=(DIM_OUT, DIM_HID))
+        b3 = wp.tile_load(bias_3, shape=(DIM_OUT, 1))
+        o = wp.tile_map(relu, wp.tile_matmul(w3, z) + wp.tile_broadcast(b3, shape=(DIM_OUT, NUM_THREADS)))
 
         # untile back to SIMT
         output = wp.untile(o)
@@ -292,14 +292,14 @@ def test_single_layer_nn(test, device):
     ):
         i = wp.tid()
 
-        f = wp.tile_load(input, 0, i, m=DIM_IN, n=NUM_THREADS)
+        f = wp.tile_load(input, shape=(DIM_IN, NUM_THREADS), offset=(0, i * NUM_THREADS))
 
-        w = wp.tile_load(weights, 0, 0, DIM_OUT, DIM_IN)
-        b = wp.tile_load(bias, 0, 0, m=DIM_OUT, n=1)
+        w = wp.tile_load(weights, shape=(DIM_OUT, DIM_IN))
+        b = wp.tile_load(bias, shape=(DIM_OUT, 1))
 
-        o = wp.tile_map(relu, wp.tile_matmul(w, f) + wp.tile_broadcast(b, m=DIM_OUT, n=NUM_THREADS))
+        o = wp.tile_map(relu, wp.tile_matmul(w, f) + wp.tile_broadcast(b, shape=(DIM_OUT, NUM_THREADS)))
 
-        wp.tile_store(out, 0, i, o)
+        wp.tile_store(out, o, offset=(0, i * NUM_THREADS))
 
     with wp.ScopedDevice(device):
         rng = np.random.default_rng(45)

warp/tests/test_tile_reduce.py

@@ -28,13 +28,13 @@ def tile_sum_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float
     n = input.shape[1]
     count = int(n / TILE_DIM)
 
-    s = wp.tile_zeros(m=1, n=1, dtype=float)
+    s = wp.tile_zeros(shape=1, dtype=float)
 
     for j in range(count):
-        a = wp.tile_load(input, i, j, m=1, n=TILE_DIM)
+        a = wp.tile_load(input[i], shape=TILE_DIM, offset=j * TILE_DIM)
         s += wp.tile_sum(a) * 0.5
 
-    wp.tile_store(output, i, s)
+    wp.tile_store(output, s, offset=i)
 
 
 def test_tile_reduce_sum(test, device):
@@ -70,10 +70,10 @@ def tile_min_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float
     # output tile index
     i = wp.tid()
 
-    a = wp.tile_load(input, i, 0, m=1, n=TILE_DIM)
+    a = wp.tile_load(input[i], shape=TILE_DIM)
     m = wp.tile_min(a)
 
-    wp.tile_store(output, i, m)
+    wp.tile_store(output, m, offset=i)
 
 
 def test_tile_reduce_min(test, device):
@@ -103,10 +103,10 @@ def tile_max_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float
     # output tile index
     i = wp.tid()
 
-    a = wp.tile_load(input, i, 0, m=1, n=TILE_DIM)
+    a = wp.tile_load(input[i], shape=TILE_DIM)
     m = wp.tile_max(a)
 
-    wp.tile_store(output, i, m)
+    wp.tile_store(output, m, offset=i)
 
 
 def test_tile_reduce_max(test, device):
@@ -136,10 +136,10 @@ def tile_reduce_custom_kernel(input: wp.array2d(dtype=float), output: wp.array(d
     # output tile index
     i = wp.tid()
 
-    a = wp.tile_load(input, i, 0, m=1, n=TILE_DIM)
+    a = wp.tile_load(input[i], shape=TILE_DIM)
     m = wp.tile_reduce(wp.mul, a)
 
-    wp.tile_store(output, i, m)
+    wp.tile_store(output, m, offset=i)
 
 
 def test_tile_reduce_custom(test, device):
@@ -173,10 +173,10 @@ def tile_grouped_sum_kernel(input: wp.array3d(dtype=float), output: wp.array(dty
     # output tile index
     i = wp.tid()
 
-    a = wp.tile_load(input[i], 0, 0, m=TILE_M, n=TILE_N)
+    a = wp.tile_load(input[i], shape=(TILE_M, TILE_N))
     s = wp.tile_sum(a) * 0.5
 
-    wp.tile_store(output, i, s)
+    wp.tile_store(output, s, offset=i)
 
 
 def test_tile_reduce_grouped_sum(test, device):
@@ -217,7 +217,7 @@ def tile_reduce_simt_kernel(output: wp.array(dtype=int)):
     s = wp.tile_sum(t)  # sum over block
 
     # update global sum
-    wp.tile_atomic_add(output, 0, 0, s)
+    wp.tile_atomic_add(output, s)
 
 
 def test_tile_reduce_simt(test, device):
@@ -310,10 +310,10 @@ def test_tile_untile_vector(test, device):
 def tile_ones_kernel(out: wp.array(dtype=float)):
     i = wp.tid()
 
-    t = wp.tile_ones(dtype=float, m=16, n=16)
+    t = wp.tile_ones(dtype=float, shape=(16, 16))
     s = wp.tile_sum(t)
 
-    wp.tile_store(out, 0, s)
+    wp.tile_store(out, s)
 
 
 def test_tile_ones(test, device):
@@ -332,16 +332,20 @@ def tile_arange_kernel(out: wp.array2d(dtype=int)):
     a = wp.tile_arange(17, dtype=int)
     b = wp.tile_arange(5, 23, dtype=int)
     c = wp.tile_arange(0, 34, 2, dtype=int)
+    d = wp.tile_arange(-1, 16, dtype=int)
+    e = wp.tile_arange(17, 0, -1, dtype=int)
 
-    wp.tile_store(out, 0, 0, a)
-    wp.tile_store(out, 1, 0, b)
-    wp.tile_store(out, 2, 0, c)
+    wp.tile_store(out[0], a)
+    wp.tile_store(out[1], b)
+    wp.tile_store(out[2], c)
+    wp.tile_store(out[3], d)
+    wp.tile_store(out[4], e)
 
 
 def test_tile_arange(test, device):
     N = 17
 
-    output = wp.zeros(shape=(3, N), dtype=int, device=device)
+    output = wp.zeros(shape=(5, N), dtype=int, device=device)
 
     with wp.Tape() as tape:
         wp.launch_tiled(tile_arange_kernel, dim=[1], inputs=[output], block_dim=TILE_DIM, device=device)
@@ -349,6 +353,8 @@ def test_tile_arange(test, device):
     assert_np_equal(output.numpy()[0], np.arange(17))
     assert_np_equal(output.numpy()[1], np.arange(5, 22))
     assert_np_equal(output.numpy()[2], np.arange(0, 34, 2))
+    assert_np_equal(output.numpy()[3], np.arange(-1, 16))
+    assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
 
 
 devices = get_cuda_test_devices()