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

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)

warp/tests/tile/test_tile_cholesky.py

@@ -20,7 +20,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.wp_is_mathdx_enabled()
+wp.init()  # For wp._src.context.runtime.core.wp_is_mathdx_enabled()
 
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
@@ -490,7 +490,7 @@ def test_tile_upper_solve(L: wp.array2d(dtype=float), y: wp.array(dtype=float),
 
 
 def test_tile_cholesky_singular_matrices(test, device):
-    if not wp.context.runtime.core.wp_is_mathdx_enabled():
+    if not wp._src.context.runtime.core.wp_is_mathdx_enabled():
         test.skipTest("MathDx is not enabled")
 
     rng = np.random.default_rng(42)
@@ -527,8 +527,11 @@ cuda_devices = get_cuda_test_devices()
 
 
 @unittest.skipUnless(
-    not wp.context.runtime.core.wp_is_mathdx_enabled()
-    or (wp.context.runtime.core.wp_is_mathdx_enabled() and wp.context.runtime.core.wp_cuda_toolkit_version() >= 12060),
+    not wp._src.context.runtime.core.wp_is_mathdx_enabled()
+    or (
+        wp._src.context.runtime.core.wp_is_mathdx_enabled()
+        and wp._src.context.runtime.core.wp_cuda_toolkit_version() >= 12060
+    ),
     "MathDx is not enabled or is enabled but CUDA toolkit version is less than 12.6",
 )
 class TestTileCholesky(unittest.TestCase):

warp/tests/tile/test_tile_load.py

@@ -40,6 +40,7 @@ def tile_load_1d_kernel(
     input: wp.array1d(dtype=float),
     out_full: wp.array1d(dtype=float),
     out_padded: wp.array1d(dtype=float),
+    out_sliced: wp.array1d(dtype=float),
     out_offset: wp.array1d(dtype=float),
 ):
     full0 = wp.tile_load(input, TILE_M)
@@ -50,8 +51,13 @@ def tile_load_1d_kernel(
     padded1 = wp.tile_load(input, shape=TILE_M, offset=TILE_OFFSET)
     padded2 = wp.tile_load(input, shape=(TILE_M,), offset=(TILE_OFFSET,))
 
+    sliced0 = wp.tile_load(input[::2], TILE_M)
+    sliced1 = wp.tile_load(input[::2], shape=TILE_M)
+    sliced2 = wp.tile_load(input[::2], shape=(TILE_M,))
+
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET,))
 
 
@@ -60,13 +66,16 @@ def tile_load_2d_kernel(
     input: wp.array2d(dtype=float),
     out_full: wp.array2d(dtype=float),
     out_padded: wp.array2d(dtype=float),
+    out_sliced: 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))
+    sliced0 = wp.tile_load(input[::2, ::2], shape=(TILE_M, TILE_N))
 
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET))
 
 
@@ -75,13 +84,16 @@ def tile_load_3d_kernel(
     input: wp.array3d(dtype=float),
     out_full: wp.array3d(dtype=float),
     out_padded: wp.array3d(dtype=float),
+    out_sliced: 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))
+    sliced0 = wp.tile_load(input[::2, ::2, ::2], shape=(TILE_M, TILE_N, TILE_O))
 
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
 
 
@@ -90,15 +102,18 @@ def tile_load_4d_kernel(
     input: wp.array4d(dtype=float),
     out_full: wp.array4d(dtype=float),
     out_padded: wp.array4d(dtype=float),
+    out_sliced: 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)
     )
+    sliced0 = wp.tile_load(input[::2, ::2, ::2, ::2], shape=(TILE_M, TILE_N, TILE_O, TILE_P))
 
     wp.tile_store(out_full, full0)
     wp.tile_store(out_padded, padded0)
+    wp.tile_store(out_sliced, sliced0)
     wp.tile_store(out_offset, full0, offset=(TILE_OFFSET, TILE_OFFSET, TILE_OFFSET, TILE_OFFSET))
 
 
@@ -112,13 +127,14 @@ def test_tile_load(kernel, 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_sliced = 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],
+                inputs=[input, output_full, output_padded, output_sliced, output_offset],
                 block_dim=TILE_DIM,
                 device=device,
             )
@@ -134,8 +150,16 @@ def test_tile_load(kernel, ndim):
         ref_offset = np.zeros_like(ref_full)
         ref_offset[src_slice] = ref_full[dest_slice]
 
+        # construct a slice for the source/dest sliced arrays
+        src_slice = tuple(slice(0, dim, 2) for dim in shape)
+        dest_slice = tuple(slice(0, (dim + 1) // 2) for dim in shape)
+
+        ref_sliced = np.zeros_like(ref_full)
+        ref_sliced[dest_slice] = ref_full[src_slice]
+
         assert_np_equal(output_full.numpy(), ref_full)
         assert_np_equal(output_padded.numpy(), ref_padded)
+        assert_np_equal(output_sliced.numpy(), ref_sliced)
         assert_np_equal(output_offset.numpy(), ref_offset)
 
         output_full.grad = wp.ones_like(output_full)
@@ -570,7 +594,7 @@ def test_tile_assign(kernel, ndim):
         input = wp.array(rng.random(shape), dtype=float, requires_grad=True, device=device)
         output = wp.zeros_like(input)
 
-        with wp.Tape() as tape:
+        with wp.Tape():
             wp.launch(
                 kernel,
                 dim=shape,

warp/tests/tile/test_tile_mathdx.py

@@ -21,7 +21,7 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()  # For wp.context.runtime.core.wp_is_mathdx_enabled()
+wp.init()  # For wp._src.context.runtime.core.wp_is_mathdx_enabled()
 
 TILE_M = wp.constant(8)
 TILE_N = wp.constant(4)
@@ -92,7 +92,7 @@ def tile_math_fft_kernel_vec2d(gx: wp.array2d(dtype=wp.vec2d), gy: wp.array2d(dt
     wp.tile_store(gy, xy)
 
 
-@unittest.skipUnless(wp.context.runtime.core.wp_is_mathdx_enabled(), "Warp was not built with MathDx support")
+@unittest.skipUnless(wp._src.context.runtime.core.wp_is_mathdx_enabled(), "Warp was not built with MathDx support")
 def test_tile_math_fft(test, device, wp_dtype):
     np_real_dtype = {wp.vec2f: np.float32, wp.vec2d: np.float64}[wp_dtype]
     np_cplx_dtype = {wp.vec2f: np.complex64, wp.vec2d: np.complex128}[wp_dtype]
@@ -113,7 +113,7 @@ def test_tile_math_fft(test, device, wp_dtype):
     X_c64 = X.view(np_cplx_dtype).reshape(fft_size, fft_size)
     Y_c64 = np.fft.fft(X_c64, axis=-1)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(kernel, dim=[1, 1], inputs=[X_wp, Y_wp], block_dim=TILE_DIM, device=device)
 
     Y_wp_c64 = Y_wp.numpy().view(np_cplx_dtype).reshape(fft_size, fft_size)

warp/tests/tile/test_tile_matmul.py

@@ -60,7 +60,7 @@ def test_tile_grouped_gemm(test, device):
     B_wp = wp.array(B, requires_grad=True, device=device)
     C_wp = wp.zeros((batch_count, TILE_M, TILE_N), requires_grad=True, device=device)
 
-    with wp.Tape() as tape:
+    with wp.Tape():
         wp.launch_tiled(
             tile_grouped_gemm, dim=[batch_count], inputs=[A_wp, B_wp, C_wp], block_dim=TILE_DIM, device=device
         )

warp/tests/tile/test_tile_mlp.py

@@ -43,7 +43,7 @@ def create_array(rng, dim_in, dim_hid, dtype=float):
 def test_multi_layer_nn(test, device):
     import torch as tc
 
-    if device.is_cuda and not wp.context.runtime.core.wp_is_mathdx_enabled():
+    if device.is_cuda and not wp._src.context.runtime.core.wp_is_mathdx_enabled():
         test.skipTest("Skipping test on CUDA device without MathDx (tolerance)")
 
     NUM_FREQ = wp.constant(8)
@@ -63,7 +63,7 @@ def test_multi_layer_nn(test, device):
         NUM_THREADS = 32
 
     dtype = wp.float16
-    npdtype = wp.types.warp_type_to_np_dtype[dtype]
+    npdtype = wp._src.types.warp_type_to_np_dtype[dtype]
 
     @wp.func
     def relu(x: dtype):
@@ -188,7 +188,6 @@ def test_multi_layer_nn(test, device):
         optimizer_inputs = [p.flatten() for p in params]
         optimizer = warp.optim.Adam(optimizer_inputs, lr=0.01)
 
-        num_batches = int((IMG_WIDTH * IMG_HEIGHT) / BATCH_SIZE)
         max_epochs = 30
 
         # create randomized batch indices
@@ -288,7 +287,6 @@ def test_single_layer_nn(test, device):
     import torch as tc
 
     DIM_IN = 8
-    DIM_HID = 32
     DIM_OUT = 16
 
     NUM_BLOCKS = 56