[TLE] Feature/extrac tile strides new

.github/workflows/hopper-build-and-test.yml

@@ -122,8 +122,12 @@ jobs:
           # python tutorials
           python3 python/tutorials/tle/01-fft.py
           python3 python/tutorials/tle/02-moe_align_block_size.py
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/03-topk.py
-          CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py
+          python3 python/tutorials/tle/03-topk.py
+          python3 python/tutorials/tle/04-cluster-gemm.py
+          python3 python/tutorials/tle/05-glu.py --only_unit_test
+          python3 python/tutorials/tle/06-2D_Depthwise_Conv.py --only_unit_test
+          python3 python/tutorials/tle/07-causal-conv1d.py
+          python3 python/tutorials/tle/08-rope.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py --mode test
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
@@ -182,6 +186,10 @@ jobs:
           python3 python/tutorials/tle/02-moe_align_block_size.py
           python3 python/tutorials/tle/03-topk.py
           python3 python/tutorials/tle/04-cluster-gemm.py
+          python3 python/tutorials/tle/05-glu.py --only_unit_test
+          python3 python/tutorials/tle/06-2D_Depthwise_Conv.py --only_unit_test
+          python3 python/tutorials/tle/07-causal-conv1d.py
+          python3 python/tutorials/tle/08-rope.py
           python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py
           python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py
           # python unit test

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -882,12 +882,13 @@ class TleExtractTileOpPattern : public OpConversionPattern<tle::ExtractTileOp> {
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
+    auto stridesAttr =
+        mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(op->getAttr("strides"));
     auto newOp = rewriter.replaceOpWithNewOp<tle::ExtractTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getIndex());
+        op, retType, adaptor.getSrc(), adaptor.getIndex(), stridesAttr);
 
     if (auto tileShapeAttr = op->getAttr("tile_shape"))
       newOp->setAttr("tile_shape", tileShapeAttr);
-
     addNamedAttrs(newOp, adaptor.getAttributes());
 
     return success();
@@ -925,8 +926,12 @@ class TleInsertTileOpPattern : public OpConversionPattern<tle::InsertTileOp> {
 
     Type retType = op.getType().cloneWithEncoding(srcEnc);
 
+    auto stridesAttr =
+        mlir::dyn_cast_or_null<mlir::DenseI64ArrayAttr>(op->getAttr("strides"));
+
     auto newOp = rewriter.replaceOpWithNewOp<tle::InsertTileOp>(
-        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex());
+        op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex(),
+        stridesAttr);
 
     addNamedAttrs(newOp, adaptor.getAttributes());
 

python/test/tle/unit/test_extract_tile_dynamic_index.py

@@ -20,10 +20,7 @@ def simple_extract_kernel(x_ptr, out_ptr, stride_xb, stride_xm, stride_xn, strid
     # 3. Perform extraction based on the layer index
     # Layer 0 extracts from top-left [0, 0], Layer 1 extracts from bottom-right [1, 1]
     # We use a conditional check because 'index' usually requires constant values for TLE hardware
-    if pid_z % 2 == 0:
-        extracted_tile = tle.extract_tile(bg_tile, index=[0, 0], tile_shape=[TILE_M, TILE_N])
-    else:
-        extracted_tile = tle.extract_tile(bg_tile, index=[1, 1], tile_shape=[TILE_M, TILE_N])
+    extracted_tile = tle.extract_tile(bg_tile, index=[pid_z, pid_z], tile_shape=[TILE_M, TILE_N])
 
     # 4. Store the extracted small tile into the corresponding Z layer of the output tensor
     offs_tm = tl.arange(0, TILE_M)

python/test/tle/unit/test_insert_tile_dynamic_index.py

@@ -6,8 +6,22 @@
 
 
 @triton.jit
-def simple_insert_kernel(x_ptr, y_ptr, stride_xb, stride_xm, stride_xn, stride_ym, stride_yn, BLOCK_M: tl.constexpr,
-                         BLOCK_N: tl.constexpr, TILE_M: tl.constexpr, TILE_N: tl.constexpr):
+def simple_insert_dynamic_index_kernel(
+    x_ptr,
+    y_ptr,
+    index_ptr,
+    stride_xb,
+    stride_xm,
+    stride_xn,
+    stride_ym,
+    stride_yn,
+    stride_ib,
+    stride_ic,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+):
     # 1. Get 3D coordinates: z (layer/batch), m (row block), n (col block)
     pid_z = tl.program_id(0)
     pid_m = tl.program_id(1)
@@ -25,13 +39,51 @@ def simple_insert_kernel(x_ptr, y_ptr, stride_xb, stride_xm, stride_xn, stride_y
     y_ptrs = y_ptr + offs_tm[:, None] * stride_ym + offs_tn[None, :] * stride_yn
     small_tile = tl.load(y_ptrs)
 
-    # 4. Determine insertion position:
-    # Layer 0 inserts at top-left [0, 0], Layer 1 inserts at bottom-right [1, 1]
-    # Note: tle.insert_tile 'index' usually must be a constant or determined by static logic
-    if pid_z % 2 == 0:
-        res_tile = tle.insert_tile(bg_tile, small_tile, index=[0, 0])
-    else:
-        res_tile = tle.insert_tile(bg_tile, small_tile, index=[1, 1])
+    # 4. Runtime index per layer: index[pid_z] = [idx_m, idx_n]
+    idx_m = tl.load(index_ptr + pid_z * stride_ib + 0 * stride_ic)
+    idx_n = tl.load(index_ptr + pid_z * stride_ib + 1 * stride_ic)
+    res_tile = tle.insert_tile(bg_tile, small_tile, index=[idx_m, idx_n])
+
+    # 5. Store the resulting tile back to memory
+    tl.store(x_ptrs, res_tile)
+
+
+@triton.jit
+def simple_insert_dynamic_index_stride_kernel(
+    x_ptr,
+    y_ptr,
+    index_ptr,
+    stride_xb,
+    stride_xm,
+    stride_xn,
+    stride_ym,
+    stride_yn,
+    stride_ib,
+    stride_ic,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    STRIDE_M: tl.constexpr,
+    STRIDE_N: tl.constexpr,
+):
+    pid_z = tl.program_id(0)
+    pid_m = tl.program_id(1)
+    pid_n = tl.program_id(2)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    x_ptrs = x_ptr + pid_z * stride_xb + offs_m[:, None] * stride_xm + offs_n[None, :] * stride_xn
+    bg_tile = tl.load(x_ptrs)
+
+    offs_tm = tl.arange(0, TILE_M)
+    offs_tn = tl.arange(0, TILE_N)
+    y_ptrs = y_ptr + offs_tm[:, None] * stride_ym + offs_tn[None, :] * stride_yn
+    small_tile = tl.load(y_ptrs)
+
+    idx_m = tl.load(index_ptr + pid_z * stride_ib + 0 * stride_ic)
+    idx_n = tl.load(index_ptr + pid_z * stride_ib + 1 * stride_ic)
+    res_tile = tle.insert_tile(bg_tile, small_tile, index=[idx_m, idx_n], strides=(STRIDE_M, STRIDE_N))
 
     # 5. Store the resulting tile back to memory
     tl.store(x_ptrs, res_tile)
@@ -43,7 +95,7 @@ def simple_insert_kernel(x_ptr, y_ptr, stride_xb, stride_xm, stride_xn, stride_y
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-def test_simple_insert_kernel_inserts_tiles_correctly():
+def test_simple_insert_kernel_with_dynamic_index():
     B = 2  # 2 layers (Z dimension)
     M, N = 32, 32  # 32x32 size per layer
     TM, TN = 16, 16  # The inserted small tile is 16x16
@@ -53,33 +105,84 @@ def test_simple_insert_kernel_inserts_tiles_correctly():
     # y is an all-99.0 2D small tile
     y = torch.ones((TM, TN), device="cuda", dtype=torch.float32) * 99.0
 
+    # Dynamic insertion indices per layer (no stride argument):
+    # Layer 0 -> [0, 0] => start at [0, 0]
+    # Layer 1 -> [1, 1] => start at [16, 16] (tile size is 16x16)
+    index = torch.tensor([[0, 0], [1, 1]], device="cuda", dtype=torch.int32)
+
     # Launch Kernel: B layers, each needs exactly 1x1 block (since M=32 and BLOCK_M=32)
     grid = (B, 1, 1)
 
-    simple_insert_kernel[grid](
+    simple_insert_dynamic_index_kernel[grid](
         x,
         y,
+        index,
         x.stride(0),
         x.stride(1),
         x.stride(2),
         y.stride(0),
         y.stride(1),
+        index.stride(0),
+        index.stride(1),
         BLOCK_M=32,
         BLOCK_N=32,
         TILE_M=16,
         TILE_N=16,
     )
 
     # --- Verification ---
-    # Layer 0 (Z=0): tile inserted at top-left [0:16, 0:16]
-    layer0_tl_mean = x[0, 0:16, 0:16].mean().item()
-    layer0_br_mean = x[0, 16:32, 16:32].mean().item()
-
-    # Layer 1 (Z=1): tile inserted at bottom-right [16:32, 16:32]
-    layer1_tl_mean = x[1, 0:16, 0:16].mean().item()
-    layer1_br_mean = x[1, 16:32, 16:32].mean().item()
-
-    assert layer0_tl_mean == pytest.approx(99.0)
-    assert layer0_br_mean == pytest.approx(0.0)
-    assert layer1_tl_mean == pytest.approx(0.0)
-    assert layer1_br_mean == pytest.approx(99.0)
+    expected = torch.zeros_like(x)
+    expected[0, 0:16, 0:16] = 99.0  # [idx_m, idx_n] = [0, 0]
+    expected[1, 16:32, 16:32] = 99.0  # [idx_m, idx_n] = [1, 1]
+
+    assert torch.allclose(x, expected)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@pytest.mark.parametrize(
+    "SM,SN,idx_m1,idx_n1",
+    [
+        (4, 4, 3, 2),
+        (8, 8, 2, 1),
+        (16, 8, 1, 1),
+    ],
+)
+def test_simple_insert_kernel_with_dynamic_index_and_stride(SM, SN, idx_m1, idx_n1):
+    B = 2  # 2 layers (Z dimension)
+    M, N = 32, 32  # 32x32 size per layer
+    TM, TN = 16, 16  # The inserted small tile is 16x16
+
+    x = torch.zeros((B, M, N), device="cuda", dtype=torch.float32)
+    y = torch.ones((TM, TN), device="cuda", dtype=torch.float32) * 99.0
+
+    # Layer 0 is fixed at top-left; Layer 1 uses parameterized dynamic index.
+    index = torch.tensor([[0, 0], [idx_m1, idx_n1]], device="cuda", dtype=torch.int32)
+
+    grid = (B, 1, 1)
+
+    simple_insert_dynamic_index_stride_kernel[grid](
+        x,
+        y,
+        index,
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        y.stride(0),
+        y.stride(1),
+        index.stride(0),
+        index.stride(1),
+        BLOCK_M=32,
+        BLOCK_N=32,
+        TILE_M=16,
+        TILE_N=16,
+        STRIDE_M=SM,
+        STRIDE_N=SN,
+    )
+
+    expected = torch.zeros_like(x)
+    expected[0, 0:16, 0:16] = 99.0
+    start_m = idx_m1 * SM
+    start_n = idx_n1 * SN
+    expected[1, start_m:start_m + TM, start_n:start_n + TN] = 99.0
+
+    assert torch.allclose(x, expected)

python/test/tle/unit/test_insert_tile_static_index.py

@@ -5,6 +5,33 @@
 import pytest
 
 
+@triton.jit
+def insert_tile_stride_kernel(
+    x_ptr,
+    y_ptr,
+    out_ptr,
+    M: tl.constexpr,
+    N: tl.constexpr,
+    TM: tl.constexpr,
+    TN: tl.constexpr,
+    SM: tl.constexpr,
+    SN: tl.constexpr,
+    idx_m: tl.constexpr,
+    idx_n: tl.constexpr,
+):
+    offs_m = tl.arange(0, M)
+    offs_n = tl.arange(0, N)
+    x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :])
+
+    tile_m = tl.arange(0, TM)
+    tile_n = tl.arange(0, TN)
+    y = tl.load(y_ptr + tile_m[:, None] * TN + tile_n[None, :])
+
+    z = tle.insert_tile(x, y, index=[idx_m, idx_n], strides=(SM, SN))
+
+    tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], z)
+
+
 @triton.jit
 def insert_tile_kernel(
     x_ptr,
@@ -61,3 +88,32 @@ def test_insert_tile_static_index():
     print(out[128:132, 128:132].cpu().int())
 
     assert torch.allclose(out, expected)
+
+
+# Test insert_tile where stride != tile_shape
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for this test")
+def test_insert_tile_with_stride():
+    M, N = 256, 256
+    TM, TN = 64, 64
+    SM, SN = 32, 32  # stride < tile_shape
+    idx_m, idx_n = 2, 3
+
+    x = torch.arange(M * N, device="cuda", dtype=torch.float32).reshape(M, N)
+    y = (10000 + torch.arange(TM * TN, device="cuda", dtype=torch.float32)).reshape(TM, TN)
+    out = torch.empty_like(x)
+
+    print(
+        f"Running insert_tile kernel with stride: x={M}x{N}, tile={TM}x{TN}, stride={SM}x{SN}, index=[{idx_m},{idx_n}]..."
+    )
+    insert_tile_stride_kernel[(1, )](x, y, out, M, N, TM, TN, SM, SN, idx_m, idx_n)
+    print("Kernel executed.\n")
+
+    expected = x.clone()
+    start_m = idx_m * SM
+    start_n = idx_n * SN
+    expected[start_m:start_m + TM, start_n:start_n + TN] = y
+
+    max_abs_diff = (out - expected).abs().max().item()
+    print(f"max_abs_diff = {max_abs_diff}")
+    assert torch.allclose(out, expected)
+    print("Test passed: insert_tile with stride updated the correct region.")