Merge branch 'GH-688' into 'main'

Fix 2D tile_load fast path conditions

See merge request omniverse/warp!1276

Author: mmacklin

CHANGELOG.md

@@ -72,6 +72,7 @@
 - Fix `OpenGLRenderer.update_shape_instance()` not having color buffers created for the shape instances.
 - Allow recovering from out-of-memory errors during Volume allocation ([GH-611](https://github.com/NVIDIA/warp/issues/611)).
 - Address `wp.tile_atomic_add()` compiler errors.
+- Fix 2D tile load when source array and tile have incompatible strides
 
 ## [1.7.0] - 2025-03-30
 

warp/native/tile.h

@@ -1160,14 +1160,15 @@ struct tile_shared_t
             const bool contiguous_dest = dest.data.strides[lastdim] == sizeof(T);
             const int elements = min(Layout::Shape::dim(1), (dest.data.shape[lastdim] - dest.offset[lastdim]));
             const bool aligned_size = (elements*sizeof(T))%sizeof(float4) == 0;
-           
+            const bool aligned_stride = (dest.data.strides[0]/sizeof(T))%Layout::Stride::dim(0) == 0;
+
             float4* dest128 = (float4*)&dest.data.data[dest.index_from_coord(tile_coord(0,0))];
             const bool aligned_dst = (uint64_t)(dest128)%sizeof(float4) == 0;
 
             constexpr int M = Layout::Shape::dim(0);
             constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);
 
-            if (contiguous_dest && contiguous_src && aligned_size && aligned_dst && N)
+            if (contiguous_dest && contiguous_src && aligned_size && aligned_dst && aligned_stride && N)
             {                               
                 // alias of shared tile with 128bit type
                 using SrcLayout = tile_layout_strided_t<tile_shape_t<M, N>>;
@@ -1249,14 +1250,15 @@ struct tile_shared_t
             const bool contiguous_src = src.data.strides[lastdim] == sizeof(T);
             const int elements = min(Layout::Shape::dim(1), (src.data.shape[lastdim] - src.offset[lastdim]));
             const bool aligned_size = (elements*sizeof(T))%sizeof(float4) == 0;
+            const bool aligned_stride = (src.data.strides[0]/sizeof(T))%Layout::Stride::dim(0) == 0;
 
             float4* src128 = (float4*)&src.data.data[src.index_from_coord(tile_coord(0,0))];
             const bool aligned_src = (uint64_t)(src128)%sizeof(float4) == 0;
 
             constexpr int M = Layout::Shape::dim(0);
             constexpr int N = (Layout::Shape::dim(1)*sizeof(T))/sizeof(float4);
 
-            if (contiguous_dest && contiguous_src && aligned_size && aligned_src && N)
+            if (contiguous_dest && contiguous_src && aligned_size && aligned_src && aligned_stride && N)
             {
                 // alias of shared tile with 128bit type
                 using DestLayout = tile_layout_strided_t<tile_shape_t<M, N>>;

warp/tests/tile/test_tile_load.py

@@ -274,6 +274,32 @@ def test_tile_load_aligned_offset_unaligned_size(test, device):
     assert_np_equal(output_array.numpy()[TILE_WIDTH:, :], np.zeros((remaining_height, TILE_M)))
 
 
+@wp.kernel
+def test_tile_load_stride_unaligned_kernel(input: wp.array2d(dtype=wp.float32), output: wp.array2d(dtype=wp.float32)):
+    tile = wp.tile_load(input, shape=(4, 4))
+    wp.tile_store(output, tile)
+
+
+# regression test for float4 aligned tiles that load from a source array with an incommensurate stride
+def test_tile_load_stride_unaligned(test, device):
+    DIM = 5
+    input_np = np.eye(DIM) * 2.0
+    input_array = wp.array(input_np, dtype=wp.float32, device=device)
+    output_array = wp.zeros_like(input_array)
+
+    wp.launch_tiled(
+        test_tile_load_stride_unaligned_kernel,
+        dim=(1, 1),
+        inputs=[input_array],
+        outputs=[output_array],
+        block_dim=TILE_DIM,
+        device=device,
+    )
+
+    input_np[DIM - 1, DIM - 1] = 0.0
+    assert_np_equal(output_array.numpy(), input_np)
+
+
 # ----------------------------------------------------------------------------------------
 
 TILE_SIZE = 4
@@ -485,6 +511,7 @@ class TestTileLoad(unittest.TestCase):
     test_tile_load_aligned_offset_unaligned_size,
     devices=devices,
 )
+add_function_test(TestTileLoad, "test_tile_load_stride_unaligned", test_tile_load_stride_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)