[GLUON][TEST] Finish subslice test and remove all layout helpers from Triton's test_core.py (triton-lang#8049)

Author: Jokeren

python/test/gluon/test_lowerings.py

@@ -1209,3 +1209,48 @@ def test_gather_layouts(axis, src_layout, index_layout, src_shape, idx_shape, de
 
     torch.testing.assert_close(out, ref, rtol=0, atol=0)
     assert ("nvvm.shfl.sync.idx" in obj.asm["llir"]) or ("llvm.amdgcn.ds.bpermute" in obj.asm["llir"])
+
+
+@pytest.mark.parametrize("M, N, M_tile_size, N_tile_size",
+                         [[128, 128, 64, 64], [128, 128, 64, 32], [128, 64, 64, 32], [256, 128, 64, 64]])
+def test_memdesc_subslice(M, N, M_tile_size, N_tile_size, device):
+    if M % M_tile_size != 0 or N % N_tile_size != 0:
+        pytest.skip(f"Shape size ({M}, {N}) must be divisible by tile size ({M_tile_size}, {N_tile_size})")
+
+    num_rows_per_warp = THREADS_PER_WARP // 4
+    blocked_layout = ttgl.BlockedLayout(size_per_thread=[1, 8], threads_per_warp=[num_rows_per_warp, 4],
+                                        warps_per_cta=[4, 1], order=[1, 0])
+    shared_layout = ttgl.SwizzledSharedLayout(vec=8, per_phase=1, max_phase=8, order=[1, 0])
+
+    @gluon.jit
+    def kernel(
+        out,
+        M: ttgl.constexpr,
+        N: ttgl.constexpr,
+        BLOCK_SIZE_M: ttgl.constexpr,
+        BLOCK_SIZE_N: ttgl.constexpr,
+        blocked_layout: ttgl.constexpr,
+        shared_layout: ttgl.constexpr,
+    ):
+        offs_m = ttgl.arange(0, M, layout=ttgl.SliceLayout(1, blocked_layout))[:, None]
+        offs_n = ttgl.arange(0, N, layout=ttgl.SliceLayout(0, blocked_layout))[None, :]
+        vals = ttgl.load(out + offs_m * N + offs_n)
+
+        smem: ttgl.shared_memory_descriptor = ttgl.allocate_shared_memory(vals.dtype, (M, N), shared_layout, value=vals)
+        for i in ttgl.static_range(M // BLOCK_SIZE_M):
+            for j in ttgl.static_range(N // BLOCK_SIZE_N):
+                tile = smem.slice(i * BLOCK_SIZE_M, BLOCK_SIZE_M, dim=0).slice(j * BLOCK_SIZE_N, BLOCK_SIZE_N, dim=1)
+                tile_vals = tile.load(blocked_layout)
+                tile_offs_m = ttgl.arange(0, BLOCK_SIZE_M, layout=ttgl.SliceLayout(1, blocked_layout))[:, None]
+                tile_offs_n = ttgl.arange(0, BLOCK_SIZE_N, layout=ttgl.SliceLayout(0, blocked_layout))[None, :]
+                linear_idx = tile_offs_m * N + tile_offs_n + i * BLOCK_SIZE_M * N + j * BLOCK_SIZE_N
+                tile.store(linear_idx + tile_vals)
+
+        vals = smem.load(blocked_layout)
+        ttgl.store(out + offs_m * N + offs_n, vals)
+
+    out = torch.zeros((M, N), device=device, dtype=torch.float16)
+    kernel[(1, )](out, M, N, M_tile_size, N_tile_size, blocked_layout, shared_layout)
+
+    out_ref = torch.arange(0, M * N, device=device).reshape((M, N)).to(torch.float16)
+    torch.testing.assert_close(out, out_ref, rtol=0, atol=0)

python/test/unit/language/test_core.py

@@ -5,7 +5,6 @@
 from typing import Optional
 import math
 import textwrap
-import pathlib
 
 import numpy as np
 import pytest
@@ -29,7 +28,6 @@
     is_cuda,
     is_interpreter,
     is_hopper,
-    is_hopper_or_newer,
     is_hip,
     is_hip_cdna,
     is_hip_cdna2,
@@ -144,199 +142,6 @@ def get_src_element_ty_size(dtype_str):
     raise ValueError(f"Unknown dtype {dtype_str}")
 
 
-class MfmaLayout:
-
-    def __init__(self, version, warps_per_cta, tiles_per_warp, instr_shape, is_transposed):
-        self.version = version
-        self.warps_per_cta = warps_per_cta
-        self.tiles_per_warp = tiles_per_warp
-        self.instr_shape = instr_shape
-        self.is_transposed = is_transposed
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.amd_mfma<{{versionMajor={self.version[0]}, versionMinor={self.version[1]}, warpsPerCTA = {self.warps_per_cta}, tilesPerWarp = {self.tiles_per_warp}, instrShape={self.instr_shape}, isTransposed = {str(self.is_transposed).lower()}}}>"
-
-
-class WmmaLayout:
-
-    def __init__(self, version, warps_per_cta):
-        self.version = version
-        self.warps_per_cta = warps_per_cta
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.amd_wmma<{{version = {self.version}, warpsPerCTA = {self.warps_per_cta}}}>"
-
-
-class MmaLayout:
-
-    def __init__(self, version, warps_per_cta, ctas_per_cga, cta_split_num, cta_order, instr_shape):
-        self.version = version
-        self.warps_per_cta = warps_per_cta
-        self.ctas_per_cga = ctas_per_cga
-        self.cta_split_num = cta_split_num
-        self.cta_order = cta_order
-        self.instr_shape = instr_shape
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.nvidia_mma<{{versionMajor={self.version[0]}, versionMinor={self.version[1]}, warpsPerCTA={self.warps_per_cta}, CTAsPerCGA={self.ctas_per_cga}, CTASplitNum={self.cta_split_num}, CTAOrder={self.cta_order}, instrShape={self.instr_shape}}}>"
-
-
-class DotOperandLayout:
-
-    def __init__(self, parent, op_idx, k_width):
-        self.parent = parent
-        self.op_idx = op_idx
-        self.k_width = k_width
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.dot_op<{{parent={self.parent}, opIdx={self.op_idx}, kWidth={self.k_width}}}>"
-
-
-class SliceLayout:
-
-    def __init__(self, dim, parent):
-        self.dim = dim
-        self.parent = parent
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.slice<{{dim = {self.dim}, parent = {self.parent}}}>"
-
-
-class BlockedLayout:
-
-    def __init__(self, size_per_thread, threads_per_warp, warps_per_cta, order, ctas_per_cga=[1, 1],
-                 cta_split_num=[1, 1], cta_order=[0, 1]):
-        self.sz_per_thread = size_per_thread
-        self.threads_per_warp = threads_per_warp
-        self.warps_per_cta = warps_per_cta
-        self.order = order
-        self.ctas_per_cga = ctas_per_cga
-        self.cta_split_num = cta_split_num
-        self.cta_order = cta_order
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.blocked<{{sizePerThread={self.sz_per_thread}, threadsPerWarp={self.threads_per_warp}, warpsPerCTA={self.warps_per_cta}, order={self.order}, CTAsPerCGA={self.ctas_per_cga}, CTASplitNum={self.cta_split_num}, CTAOrder={self.cta_order}}}>"
-
-
-class SwizzledSharedLayout:
-
-    def __init__(self, vec, per_phase, max_phase, order, ctas_per_cga, cta_split_num, cta_order):
-        self.vec = vec
-        self.per_phase = per_phase
-        self.max_phase = max_phase
-        self.order = order
-        self.ctas_per_cga = ctas_per_cga
-        self.cta_split_num = cta_split_num
-        self.cta_order = cta_order
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.swizzled_shared<{{vec={self.vec}, perPhase={self.per_phase}, maxPhase={self.max_phase}, order={self.order}, CTAsPerCGA={self.ctas_per_cga}, CTASplitNum={self.cta_split_num}, CTAOrder={self.cta_order}}}>"
-
-
-class PaddedSharedLayout:
-
-    def __init__(self, interval_padding_pairs, linear_layout_offset_bases, linear_layout_block_bases):
-        self.interval_padding_pairs = "[" + ", ".join(f"{v[0]}:{v[1]:+d}" for v in interval_padding_pairs) + "]"
-        self.offset_bases = linear_layout_offset_bases
-        self.block_bases = linear_layout_block_bases
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.padded_shared<{self.interval_padding_pairs} {{offset={self.offset_bases}, block={self.block_bases}}}>"
-
-
-class NVMMASharedLayout:
-
-    def __init__(self, swizzle, transpose, element_bit_width, ctas_per_cga, cta_split_num, cta_order):
-        self.swizzle = swizzle
-        self.transpose = transpose
-        self.element_bit_width = element_bit_width
-        self.ctas_per_cga = ctas_per_cga
-        self.cta_split_num = cta_split_num
-        self.cta_order = cta_order
-
-    def __str__(self):
-        transpose_str = "true" if self.transpose else "false"
-        return f"#{GPU_DIALECT}.nvmma_shared<{{swizzlingByteWidth={self.swizzle}, transposed={transpose_str}, elementBitWidth={self.element_bit_width}, CTAsPerCGA={self.ctas_per_cga}, CTASplitNum={self.cta_split_num}, CTAOrder={self.cta_order}}}>"
-
-
-class LinearLayout:
-
-    def __init__(self, register, lane, warp, block):
-        self.register = register
-        self.lane = lane
-        self.warp = warp
-        self.block = block
-
-    def __str__(self):
-        return f"#{GPU_DIALECT}.linear<{{register={self.register}, lane={self.lane}, warp={self.warp}, block={self.block}}}>"
-
-
-# Python impl of LinearEncodingAttr::basesPerDim
-def bases_per_dim(layout, dim, rank, skip_broadcast=True):
-    assert isinstance(layout, LinearLayout)
-    bases = getattr(layout, dim)
-    result = [1] * rank
-
-    if not bases:
-        return result
-
-    non_zero_idx = None
-
-    for basis in bases:
-        # Find the first non-zero index in the current basis
-        idx = next((i for i, v in enumerate(basis) if v != 0), None)
-        if idx is not None:
-            non_zero_idx = idx
-            result[idx] *= 2
-        elif not skip_broadcast:
-            # If no non-zero found and we're not skipping broadcasts, use the last found non-zero index
-            assert non_zero_idx is not None
-            result[non_zero_idx] *= 2
-
-    return result
-
-
-def warps_per_cta(layout, shape):
-    if isinstance(layout, LinearLayout):
-        return bases_per_dim(layout, 'warp', len(shape))
-    elif isinstance(layout, (SliceLayout, DotOperandLayout)):
-        return warps_per_cta(layout.parent, shape)
-    else:
-        return layout.warps_per_cta
-
-
-def is_layout_applicable(layout) -> bool:
-    if isinstance(layout, (BlockedLayout, SwizzledSharedLayout, LinearLayout)):
-        return True
-    elif isinstance(layout, SliceLayout):
-        return is_layout_applicable(layout.parent)
-    elif is_cuda():
-        mma_layout = layout.parent if isinstance(layout, DotOperandLayout) else layout
-        if not isinstance(mma_layout, MmaLayout):
-            return False
-        if mma_layout.version[0] >= 3 and not is_hopper_or_newer():
-            return False
-        return True
-    elif is_hip():
-        target_arch = triton.runtime.driver.active.get_current_target().arch
-        if isinstance(layout, PaddedSharedLayout):
-            return True
-        elif any(arch for arch in ["gfx11", "gfx12"] if arch in target_arch):
-            # RDNA 3, 4
-            return isinstance(layout, WmmaLayout)
-        elif any(arch for arch in ["gfx8", "gfx9"] if arch in target_arch):
-            # CDNA 1, 2, 3, 4
-            return isinstance(layout, MfmaLayout)
-        else:
-            return False
-    else:
-        return True
-
-
-def filter_layouts(layouts):
-    return [l for l in layouts if is_layout_applicable(l)]
-
-
 @pytest.mark.interpreter
 def test_scalar_overflow(device):
 
@@ -5722,91 +5527,6 @@ def kernel(Out):
     assert h.asm["ptx"].count("%smid") == 1
 
 
-# -----------------------
-# test layout conversions
-# -----------------------
-# TODO: backend should be tested separately
-
-
-@pytest.mark.parametrize("M, N, M_tile_size, N_tile_size",
-                         [[128, 128, 64, 64], [128, 128, 64, 32], [128, 64, 64, 32], [256, 128, 64, 64]])
-def test_split_subview(M, N, M_tile_size, N_tile_size, device, tmp_path: pathlib.Path):
-    num_rows_per_warp = THREADS_PER_WARP // 4
-    num_repeats_M = triton.cdiv(M, M_tile_size)
-    num_repeats_N = triton.cdiv(N, N_tile_size)
-
-    ir = f"""
-    #blocked = #ttg.blocked<{{sizePerThread=[1, 8], threadsPerWarp=[{num_rows_per_warp}, 4], warpsPerCTA=[4, 1], order=[1, 0], CTAsPerCGA=[1, 1], CTASplitNum=[1, 1], CTAOrder=[0, 1]}}>
-    #shared = #ttg.swizzled_shared<{{vec = 8, perPhase = 1, maxPhase = 8, order = [1, 0]}}>
-    #smem = #ttg.shared_memory
-
-    module attributes {{"ttg.num-ctas" = 1, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = {THREADS_PER_WARP} : i32}} {{
-    tt.func public @kernel(%arg0: !tt.ptr<f16> {{tt.divisibility = 16 : i32}}) {{
-        %cst = arith.constant dense<{N}> : tensor<{M}x1xi32, #blocked>
-        %cst_n = arith.constant dense<{N_tile_size}> : tensor<{M_tile_size}x1xi32, #blocked>
-        %0 = tt.make_range {{end = {M} : i32, start = 0 : i32}} : tensor<{M}xi32, #ttg.slice<{{dim = 1, parent = #blocked}}>>
-        %1 = tt.make_range {{end = {N} : i32, start = 0 : i32}} : tensor<{N}xi32, #ttg.slice<{{dim = 0, parent = #blocked}}>>
-        %2 = tt.splat %arg0 : !tt.ptr<f16> -> tensor<{M}x{N}x!tt.ptr<f16>, #blocked>
-        %4 = tt.expand_dims %0 {{axis = 1 : i32}} : tensor<{M}xi32, #ttg.slice<{{dim = 1, parent = #blocked}}>> -> tensor<{M}x1xi32, #blocked>
-        %5 = arith.muli %4, %cst : tensor<{M}x1xi32, #blocked>
-        %6 = tt.expand_dims %1 {{axis = 0 : i32}} : tensor<{N}xi32, #ttg.slice<{{dim = 0, parent = #blocked}}>> -> tensor<1x{N}xi32, #blocked>
-        %7 = tt.broadcast %6 : tensor<1x{N}xi32, #blocked> -> tensor<{M}x{N}xi32, #blocked>
-        %8 = tt.broadcast %5 : tensor<{M}x1xi32, #blocked> -> tensor<{M}x{N}xi32, #blocked>
-        %9 = arith.addi %8, %7 : tensor<{M}x{N}xi32, #blocked>
-        %ptrs = tt.addptr %2, %9 : tensor<{M}x{N}x!tt.ptr<f16>, #blocked>, tensor<{M}x{N}xi32, #blocked>
-        %11 = tt.load %ptrs {{cache = 1 : i32, evict = 1 : i32, isVolatile = false}} : tensor<{M}x{N}x!tt.ptr<f16>, #blocked>
-
-        %c0_i32 = arith.constant 0 : i32
-
-        %12 = ttg.local_alloc : () -> !ttg.memdesc<1x{M}x{N}xf16, #shared, #smem, mutable>
-        %13 = ttg.memdesc_index %12[%c0_i32] : !ttg.memdesc<1x{M}x{N}xf16, #shared, #smem, mutable> -> !ttg.memdesc<{M}x{N}xf16, #shared, #smem, mutable>
-        ttg.local_store %11, %13 : tensor<{M}x{N}xf16, #blocked> -> !ttg.memdesc<{M}x{N}xf16, #shared, #smem, mutable>
-
-    """
-
-    for m in range(num_repeats_M):
-        for n in range(num_repeats_N):
-            linear_idx = n + m * num_repeats_N
-            m_offset = m * M_tile_size
-            n_offset = n * N_tile_size
-            ir += f"""
-        %view{linear_idx} = ttg.memdesc_subslice %13[{m_offset}, {n_offset}] : !ttg.memdesc<{M}x{N}xf16, #shared, #smem, mutable> -> !ttg.memdesc<{M_tile_size}x{N_tile_size}xf16, #shared, #smem, mutable, {M}x{N}>
-        %data{linear_idx} = ttg.local_load %view{linear_idx} : !ttg.memdesc<{M_tile_size}x{N_tile_size}xf16, #shared, #smem, mutable, {M}x{N}> -> tensor<{M_tile_size}x{N_tile_size}xf16, #blocked>
-        %inc{linear_idx} = arith.constant dense<{linear_idx}.0> : tensor<{M_tile_size}x{N_tile_size}xf16, #blocked>
-
-        %res{linear_idx} = arith.addf %data{linear_idx}, %inc{linear_idx} : tensor<{M_tile_size}x{N_tile_size}xf16, #blocked>
-        ttg.local_store %res{linear_idx}, %view{linear_idx} : tensor<{M_tile_size}x{N_tile_size}xf16, #blocked> -> !ttg.memdesc<{M_tile_size}x{N_tile_size}xf16, #shared, #smem, mutable, {M}x{N}>
-        """
-
-    ir += f"""
-        %res = ttg.local_load %13 : !ttg.memdesc<{M}x{N}xf16, #shared, #smem, mutable> -> tensor<{M}x{N}xf16, #blocked>
-        tt.store %ptrs, %res : tensor<{M}x{N}x!tt.ptr<f16>, #blocked>
-        tt.return
-    }}
-    }}
-    """
-
-    temp_file = tmp_path / "test_split_subview.ttgir"
-    temp_file.write_text(ir)
-    kernel = triton.compile(str(temp_file))
-
-    triton_result = torch.zeros((M, N), device=device, dtype=torch.float16)
-    kernel[(1, 1, 1)](triton_result.data_ptr())
-
-    rows = []
-    for m in range(num_repeats_M):
-        columns = []
-        for n in range(num_repeats_N):
-            linear_idx = n + m * num_repeats_N
-            tile = float(linear_idx) * torch.ones((M_tile_size, N_tile_size), device=device, dtype=torch.float16)
-            columns.append(tile)
-        rows.append(torch.cat(columns, dim=1))
-    expected_result = torch.cat(rows, dim=0)
-
-    test_result = torch.equal(triton_result, expected_result)
-    assert test_result
-
-
 @pytest.mark.interpreter
 def test_load_scalar_with_mask(device):