[TLE][NVIDIA] lower same-warp layout conversions with shuffles

[Kafka-Hatsune]

What For

for Layout Convertion

tensor<16x1xptr, #src> -> tensor<16x1xptr, #dst>
Src：
sizePerThread = [1,1]
threadsPerWarp = [32,1]
warpsPerCTA = [4,1]
order = [1,0]
Dst：
sizePerThread = [1,8]
threadsPerWarp = [4,8]
warpsPerCTA = [4,1]
order = [1,0]

The LayoutConvert Analysis gives a LinearLayout matrix(Src) like：

             lane bits       warp bits
          l0 l1 l2 l3 l4     w0 w1
        ┌──────────────────────────┐
row0    │ 1  0  0  0  0      0  0 │
row1    │ 0  1  0  0  0      0  0 │
row2    │ 0  0  1  0  0      0  0 │
row3    │ 0  0  0  1  0      0  0 │
        └──────────────────────────┘

Giving a logical row = 0d10 = 0b1010，physical position is solved to be that (l0,l1,l2,l3)=(0,1,0,1)，and l4, w0, w1 are free variables. l4 can be 0 or 1, and warp can be 0, 1, 2, or 3. Thus, there are 2 lanes × 4 warps = 8 solutions. When there are multiple answers, all the unnecessary numbered ones should be set to 0. Therefore, it selects warp 0, lane 10. This is represented as warp 2 attempting to read data from warp 0. However, shuffle can only exchange registers within the same warp and cannot read from warp 0 from warp 2. Thus, the compiler wrongly concludes: shuffle is not possible and must go through shared memory.

improvement

Add constraints for Src&Dst LinearLayout matrix in the reverse solving: Among these positions, search for solutions that are in the same block and the same warp.

the linearLayout matrix of Src after adding constraints like：

             lane bits       warp bits
         l0 l1 l2 l3 l4     w0 w1
       ┌──────────────────────────┐
row0    │ 1  0  0  0  0      0  0 │
row1    │ 0  1  0  0  0      0  0 │
row2    │ 0  0  1  0  0      0  0 │
row3    │ 0  0  0  1  0      0  0 │
warp0   │ 0  0  0  0  0      1  0 │
warp1   │ 0  0  0  0  0      0  1 │
       └──────────────────────────┘

So the result P in the S x P = D layout convertion computation is limited that warp ids that are the same in the src&dst. 8 solutions is reduced into 2 solutions. By default we use the first solution in the same warp, it is the time to use warp shuffle instead of smem load&store to do the layout convertion.

Triton code examples

#!/root/miniconda3/envs/flagtree/bin/python
"""Runnable A/B test for the same-warp pointer-layout conversion.

The Triton frontend is run once for each shape. The resulting TTGIR is then
lowered by the parent and current triton-opt binaries, so the comparison only
changes the lowering of ttg.convert_layout.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import hashlib
import json
import os
from pathlib import Path
import re
import subprocess
import sys
from typing import Any


SCRIPT_DIR = Path(__file__).resolve().parent
WORKSPACE = Path(__file__).resolve().parents[2]
DEFAULT_REPO = Path(
    os.environ.get(
        "FLAGTREE_REPO",
        WORKSPACE / "FlagTree_fork-reorg-20260820",
    )
)
sys.path.insert(0, str(DEFAULT_REPO / "python"))
os.environ.setdefault("TRITON_PTXAS_PATH", "/usr/local/cuda-13.3/bin/ptxas")

import cupy as cp  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import triton  # noqa: E402
import triton.language as tl  # noqa: E402
import triton.experimental.tle.language as tle  # noqa: E402


@triton.jit
def paged_copy_shuffle_kernel(
    page_table,
    src,
    out,
    row_stride,
    page_stride_rows,
    max_index,
    BLOCK_SIZE: tl.constexpr,
    ROWS: tl.constexpr,
):
    """Copy a paged [ROWS, 64] tile through SMEM and expose the result."""
    PAGES: tl.constexpr = ROWS // BLOCK_SIZE
    pid = tl.program_id(0)
    rows = tl.arange(0, ROWS)
    cols = tl.arange(0, 64)

    # This is the page-table expansion used by the attention paged-gather
    # path. Each page-table entry is broadcast over its 16 token rows.
    page_offsets = tl.arange(0, PAGES)
    first_idx = page_offsets * BLOCK_SIZE
    page_blocks = tl.load(
        page_table + pid * PAGES + page_offsets,
        mask=first_idx < max_index,
        other=0,
    ).to(tl.int32)
    page_blocks = tl.reshape(
        tl.broadcast_to(page_blocks[:, None], (PAGES, BLOCK_SIZE)),
        (ROWS,),
        can_reorder=True,
    )

    cache_idx = (
        page_blocks * page_stride_rows.to(tl.int64)
        + rows % BLOCK_SIZE
    )
    src_ptrs = (
        src
        + cache_idx[:, None] * row_stride.to(tl.int64)
        + cols[None, :]
    )
    mask = (rows[:, None] < max_index) & (cols[None, :] < 64)

    smem = tle.gpu.alloc(
        [ROWS, 64],
        dtype=tl.float16,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    tle.gpu.copy(
        src_ptrs,
        smem,
        [ROWS, 64],
        mask=mask,
        other=0.0,
    )

    values = tl.load(tle.gpu.local_ptr(smem))
    out_base = out + pid * ROWS * 64
    out_ptrs = out_base + rows[:, None] * 64 + cols[None, :]
    tl.store(out_ptrs, values)


def run_checked(
    command: list[str | Path], input_text: str | None = None
) -> tuple[str, str]:
    argv = [str(value) for value in command]
    process = subprocess.run(
        argv,
        input=input_text,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if process.returncode != 0:
        quoted = " ".join(argv)
        raise RuntimeError(
            f"command failed ({process.returncode}): {quoted}\n"
            f"--- stdout ---\n{process.stdout[-8000:]}\n"
            f"--- stderr ---\n{process.stderr[-8000:]}"
        )
    return process.stdout, process.stderr


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def git_head(path: Path) -> str:
    stdout, _ = run_checked(["git", "-C", path, "rev-parse", "HEAD"])
    return stdout.strip()


def count_lines(text: str, pattern: str) -> int:
    return len(re.findall(pattern, text, flags=re.MULTILINE))


def collect_stats(lowered_mlir: str, llvm_ir: str, ptx: str) -> dict[str, int]:
    shared_match = re.search(
        r"\bttg\.shared\s*=\s*(\d+)\s*:\s*i32\b", lowered_mlir
    )
    if shared_match is None:
        raise RuntimeError("lowered MLIR does not contain ttg.shared")

    optional_predicate = r"(?:@!?%\w+\s+)?"
    return {
        "shared_bytes": int(shared_match.group(1)),
        "mlir_shuffle": count_lines(
            lowered_mlir,
            r"^\s*%\S+\s*=\s*nvvm\.shfl\.sync\s+idx\b",
        ),
        "mlir_shared_store": count_lines(
            lowered_mlir,
            r"^\s*llvm\.store\b.*!llvm\.ptr<3>\s*$",
        ),
        "mlir_shared_load": count_lines(
            lowered_mlir,
            r"^\s*%\S+\s*=\s*llvm\.load\b.*!llvm\.ptr<3>.*$",
        ),
        "mlir_barrier": count_lines(
            lowered_mlir,
            r"^\s*nvvm\.barrier0\b",
        ),
        "llvm_shuffle": count_lines(
            llvm_ir,
            r"^\s*%\S+\s*=\s*(?:tail\s+)?call\s+i32\s+"
            r"@llvm\.nvvm\.shfl\.sync\.idx\.i32\(",
        ),
        "ptx_shuffle": count_lines(
            ptx,
            rf"^\s*{optional_predicate}shfl\.sync\.[A-Za-z0-9_.]+\s",
        ),
        "ptx_shared_store": count_lines(
            ptx,
            rf"^\s*{optional_predicate}st\.shared(?:\.[A-Za-z0-9_]+)*\s",
        ),
        "ptx_shared_load": count_lines(
            ptx,
            rf"^\s*{optional_predicate}ld\.shared(?:\.[A-Za-z0-9_]+)*\s",
        ),
        "ptx_barrier": count_lines(
            ptx,
            rf"^\s*{optional_predicate}bar\.sync\b",
        ),
        "ptx_cp_async_copy": count_lines(
            ptx,
            rf"^\s*{optional_predicate}cp\.async\.(?:ca|cg)\.shared\.global\b",
        ),
    }


def extract_ptx_abi(ptx: str) -> tuple[str, tuple[tuple[str, int], ...], int]:
    entry = re.search(
        r"\.visible\s+\.entry\s+([^\s(]+)\s*\((.*?)\)\s*"
        r"(?:\.reqntid|\{)",
        ptx,
        flags=re.DOTALL,
    )
    if entry is None:
        raise RuntimeError("could not find .visible .entry in PTX")

    specs: list[tuple[str, int]] = []
    for parameter in re.findall(r"\.param\s+([^,\n]+)", entry.group(2)):
        width = re.search(r"\.(?:u|s|b)(\d+)\b", parameter)
        if width is None:
            raise RuntimeError(f"unsupported PTX parameter: {parameter.strip()}")
        specs.append(
            ("ptr" if ".ptr" in parameter else "scalar", int(width.group(1)))
        )

    reqntid = re.search(r"\.reqntid\s+(\d+)", ptx[entry.end() - 16 :])
    threads = int(reqntid.group(1)) if reqntid is not None else 0
    return entry.group(1), tuple(specs), threads


def compile_frontend(rows: int, artifact_dir: Path) -> str:
    pages = rows // 16
    page_table = torch.arange(pages, device="cuda", dtype=torch.int32)
    src = torch.randn(rows, 64, device="cuda", dtype=torch.float16)
    out = torch.empty_like(src)
    compiled = paged_copy_shuffle_kernel.warmup(
        page_table,
        src,
        out,
        64,
        16,
        rows,
        BLOCK_SIZE=16,
        ROWS=rows,
        grid=(1,),
        num_warps=4,
    )
    ttgir = compiled.asm["ttgir"]
    (artifact_dir / f"rows{rows}.frontend.ttir").write_text(
        compiled.asm["ttir"]
    )
    (artifact_dir / f"rows{rows}.frontend.ttgir").write_text(ttgir)
    return ttgir


def lower_variant(
    *,
    label: str,
    rows: int,
    ttgir: str,
    triton_opt: Path,
    mlir_translate: Path,
    llvm_opt: Path,
    llc: Path,
    ptxas: Path,
    ptx_version: int,
    artifact_dir: Path,
) -> dict[str, Any]:
    prefix = artifact_dir / f"rows{rows}.{label}"
    lower_flags = [
        str(triton_opt),
        "-",
        "--allocate-shared-memory-nv="
        f"compute-capability=90 ptx-version={ptx_version}",
        "--convert-triton-gpu-to-llvm="
        f"compute-capability=90 ptx-version={ptx_version}",
    ]
    lowered, lower_stderr = run_checked(
        lower_flags + ["-reconcile-unrealized-casts"], ttgir
    )
    full_mlir, full_stderr = run_checked(
        lower_flags
        + ["--convert-nv-gpu-to-llvm", "-reconcile-unrealized-casts"],
        ttgir,
    )
    raw_llvm, translate_stderr = run_checked(
        [mlir_translate, "--mlir-to-llvmir"], full_mlir
    )
    llvm_ir, opt_stderr = run_checked(
        [llvm_opt, "-O3", "-S", "-"], raw_llvm
    )
    ptx, llc_stderr = run_checked(
        [
            llc,
            "-O3",
            "-mtriple=nvptx64-nvidia-cuda",
            "-mcpu=sm_90a",
            f"-mattr=+ptx{ptx_version}",
            "-o",
            "-",
            "-",
        ],
        llvm_ir,
    )

    lowered_path = Path(f"{prefix}.lowered.mlir")
    full_path = Path(f"{prefix}.full.mlir")
    llvm_path = Path(f"{prefix}.ll")
    ptx_path = Path(f"{prefix}.ptx")
    cubin_path = Path(f"{prefix}.cubin")
    lowered_path.write_text(lowered)
    full_path.write_text(full_mlir)
    llvm_path.write_text(llvm_ir)
    ptx_path.write_text(ptx)

    _, ptxas_stderr = run_checked(
        [ptxas, "-arch=sm_90a", "-v", ptx_path, "-o", cubin_path]
    )
    Path(f"{prefix}.compiler.log").write_text(
        lower_stderr
        + full_stderr
        + translate_stderr
        + opt_stderr
        + llc_stderr
        + ptxas_stderr
    )

    symbol, abi, threads = extract_ptx_abi(ptx)
    return {
        "label": label,
        "triton_opt": str(triton_opt),
        "triton_opt_sha256": sha256_file(triton_opt),
        "cubin": str(cubin_path),
        "symbol": symbol,
        "abi": abi,
        "threads": threads,
        "stats": collect_stats(lowered, llvm_ir, ptx),
    }


def load_kernel(variant: dict[str, Any]) -> None:
    module = cp.RawModule(path=variant["cubin"])
    variant["module"] = module
    variant["kernel"] = module.get_function(variant["symbol"])


def kernel_args(
    variant: dict[str, Any],
    page_table: cp.ndarray,
    src: cp.ndarray,
    out: cp.ndarray,
    dummy: cp.ndarray,
    rows: int,
) -> tuple[Any, ...]:
    expected = (
        ("ptr", 64),
        ("ptr", 64),
        ("ptr", 64),
        ("scalar", 32),
        ("scalar", 32),
        ("scalar", 32),
        ("ptr", 64),
        ("ptr", 64),
    )
    if variant["abi"] != expected:
        raise RuntimeError(
            f"unexpected PTX ABI for {variant['label']}: {variant['abi']}"
        )
    if variant["threads"] not in (0, 128):
        raise RuntimeError(
            f"unexpected reqntid for {variant['label']}: {variant['threads']}"
        )
    return (
        page_table,
        src,
        out,
        np.int32(64),
        np.int32(16),
        np.int32(rows),
        dummy,
        dummy,
    )


def launch(
    variant: dict[str, Any],
    args: tuple[Any, ...],
    grid: int,
    stream: cp.cuda.Stream,
) -> None:
    variant["kernel"](
        (grid,),
        (128,),
        args,
        shared_mem=variant["stats"]["shared_bytes"],
        stream=stream,
    )


def capture_graph(
    variant: dict[str, Any],
    args: tuple[Any, ...],
    grid: int,
    nodes: int,
    graph_warmup: int,
    stream: cp.cuda.Stream,
):
    stream.begin_capture()
    for _ in range(nodes):
        launch(variant, args, grid, stream)
    graph = stream.end_capture()
    graph.upload(stream)
    for _ in range(graph_warmup):
        graph.launch(stream)
    stream.synchronize()
    return graph


def summarize_samples(samples: list[float]) -> dict[str, float]:
    values = np.asarray(samples, dtype=np.float64)
    return {
        "median_us": float(np.median(values)),
        "mean_us": float(np.mean(values)),
        "p20_us": float(np.percentile(values, 20)),
        "p80_us": float(np.percentile(values, 80)),
        "cv": float(np.std(values) / np.mean(values)),
    }


def benchmark_rows(
    rows: int,
    variants: dict[str, dict[str, Any]],
    grids: list[int],
    nodes: int,
    samples: int,
    graph_warmup: int,
) -> list[dict[str, Any]]:
    max_grid = max(grids)
    pages = rows // 16
    elements = max_grid * rows * 64
    page_table = cp.arange(max_grid * pages, dtype=cp.int32)
    src = (cp.arange(elements, dtype=cp.int32) % 251).astype(cp.float16)
    out = cp.empty(elements, dtype=cp.float16)
    dummy = cp.empty(1, dtype=cp.uint8)
    stream = cp.cuda.Stream(non_blocking=True)

    args = {
        label: kernel_args(variant, page_table, src, out, dummy, rows)
        for label, variant in variants.items()
    }
    for label in ("before", "after"):
        with stream:
            out.fill(np.float16(np.nan))
        launch(variants[label], args[label], 1, stream)
        stream.synchronize()
        cp.testing.assert_array_equal(out[: rows * 64], src[: rows * 64])

    results: list[dict[str, Any]] = []
    for grid in grids:
        graphs = {
            label: capture_graph(
                variants[label],
                args[label],
                grid,
                nodes,
                graph_warmup,
                stream,
            )
            for label in ("before", "after")
        }
        timings: dict[str, list[float]] = {"before": [], "after": []}
        for sample in range(samples):
            order = ("before", "after")
            if sample % 2:
                order = ("after", "before")
            for label in order:
                start = cp.cuda.Event()
                end = cp.cuda.Event()
                start.record(stream)
                graphs[label].launch(stream)
                end.record(stream)
                end.synchronize()
                elapsed_us = (
                    cp.cuda.get_elapsed_time(start, end) * 1000.0 / nodes
                )
                timings[label].append(elapsed_us)

        before = summarize_samples(timings["before"])
        after = summarize_samples(timings["after"])
        results.append(
            {
                "rows": rows,
                "grid": grid,
                "before": before,
                "after": after,
                "delta_us": before["median_us"] - after["median_us"],
                "speedup": before["median_us"] / after["median_us"],
            }
        )
    return results


def parse_csv_ints(value: str) -> list[int]:
    result = [int(item.strip()) for item in value.split(",") if item.strip()]
    if not result or any(item <= 0 for item in result):
        raise argparse.ArgumentTypeError("expected comma-separated positive integers")
    return result


def device_name() -> str:
    value = cp.cuda.runtime.getDeviceProperties(0)["name"]
    return value.decode() if isinstance(value, bytes) else str(value)


def print_ir_table(rows: int, ttgir: str, variants: dict[str, dict[str, Any]]) -> None:
    print(f"\nN={rows}: TTGIR convert_layout = {ttgir.count('ttg.convert_layout')}")
    print(
        "variant  smem_B  MLIR(shfl/st3/ld3/bar)  "
        "PTX(shfl/st.shared/ld.shared/bar/cp.async)"
    )
    for label in ("before", "after"):
        stats = variants[label]["stats"]
        print(
            f"{label:7s}  {stats['shared_bytes']:6d}  "
            f"{stats['mlir_shuffle']:4d}/{stats['mlir_shared_store']:3d}/"
            f"{stats['mlir_shared_load']:3d}/{stats['mlir_barrier']:3d}"
            "                 "
            f"{stats['ptx_shuffle']:4d}/{stats['ptx_shared_store']:3d}/"
            f"{stats['ptx_shared_load']:3d}/{stats['ptx_barrier']:3d}/"
            f"{stats['ptx_cp_async_copy']:3d}"
        )


def print_perf_table(results: list[dict[str, Any]]) -> None:
    print("\nMedian time per kernel (CUDA Graph; lower is better)")
    print("N     grid     before_us   after_us    delta_us   speedup")
    for row in results:
        print(
            f"{row['rows']:<5d} {row['grid']:<8d} "
            f"{row['before']['median_us']:10.5f} "
            f"{row['after']['median_us']:10.5f} "
            f"{row['delta_us']:10.5f} "
            f"{row['speedup']:8.3f}x"
        )


def main() -> None:
    llvm_root = WORKSPACE / "llvm" / "llvm-f6ded0be-ubuntu-x64" / "bin"
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--rows", type=parse_csv_ints, default=parse_csv_ints("16,32"))
    parser.add_argument("--grids", type=parse_csv_ints, default=parse_csv_ints("1,114,2112"))
    parser.add_argument("--nodes", type=int, default=200)
    parser.add_argument("--samples", type=int, default=100)
    parser.add_argument("--graph-warmup", type=int, default=10)
    parser.add_argument("--ptx-version", type=int, default=83)
    parser.add_argument("--compile-only", action="store_true")
    parser.add_argument(
        "--artifact-dir",
        type=Path,
        default=SCRIPT_DIR / "runnable_out",
    )
    parser.add_argument(
        "--before-opt", type=Path, default=SCRIPT_DIR / "triton-opt-before"
    )
    parser.add_argument(
        "--before-source",
        type=Path,
        default=WORKSPACE / "exp/flagtree-before-99baa",
    )
    parser.add_argument(
        "--after-opt",
        type=Path,
        default=DEFAULT_REPO / "build/cmake.flagtree-tle/bin/triton-opt",
    )
    parser.add_argument("--after-source", type=Path, default=DEFAULT_REPO)
    parser.add_argument(
        "--mlir-translate", type=Path, default=llvm_root / "mlir-translate"
    )
    parser.add_argument("--llvm-opt", type=Path, default=llvm_root / "opt")
    parser.add_argument("--llc", type=Path, default=llvm_root / "llc")
    parser.add_argument(
        "--ptxas", type=Path, default=Path(os.environ["TRITON_PTXAS_PATH"])
    )
    args = parser.parse_args()

    unsupported = [rows for rows in args.rows if rows not in (16, 32)]
    if unsupported:
        parser.error(f"this checked test currently supports --rows 16,32; got {unsupported}")
    if args.nodes <= 0 or args.samples <= 0 or args.graph_warmup < 0:
        parser.error("nodes/samples must be positive and graph-warmup non-negative")

    tools = [
        args.before_opt,
        args.after_opt,
        args.mlir_translate,
        args.llvm_opt,
        args.llc,
        args.ptxas,
    ]
    missing = [str(path) for path in tools if not path.is_file()]
    if missing:
        parser.error("missing tools: " + ", ".join(missing))
    missing_sources = [
        str(path)
        for path in (args.before_source, args.after_source)
        if not path.is_dir()
    ]
    if missing_sources:
        parser.error("missing source worktrees: " + ", ".join(missing_sources))

    args.artifact_dir.mkdir(parents=True, exist_ok=True)
    commits = {
        "before": git_head(args.before_source),
        "after": git_head(args.after_source),
    }
    print(f"device: {device_name()}")
    print(
        f"before commit: {commits['before']}\n"
        "before triton-opt: "
        f"{args.before_opt} ({sha256_file(args.before_opt)[:12]})"
    )
    print(
        f"after  commit: {commits['after']}\n"
        "after  triton-opt: "
        f"{args.after_opt} ({sha256_file(args.after_opt)[:12]})"
    )
    print("frontend rule: one TTGIR is shared byte-for-byte by both lowerings")

    all_variants: dict[int, dict[str, dict[str, Any]]] = {}
    ttgirs: dict[int, str] = {}
    for rows in args.rows:
        ttgir = compile_frontend(rows, args.artifact_dir)
        ttgirs[rows] = ttgir
        jobs: dict[str, concurrent.futures.Future[dict[str, Any]]] = {}
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            for label, binary in (
                ("before", args.before_opt),
                ("after", args.after_opt),
            ):
                jobs[label] = executor.submit(
                    lower_variant,
                    label=label,
                    rows=rows,
                    ttgir=ttgir,
                    triton_opt=binary,
                    mlir_translate=args.mlir_translate,
                    llvm_opt=args.llvm_opt,
                    llc=args.llc,
                    ptxas=args.ptxas,
                    ptx_version=args.ptx_version,
                    artifact_dir=args.artifact_dir,
                )
            variants = {label: job.result() for label, job in jobs.items()}
        if variants["before"]["abi"] != variants["after"]["abi"]:
            raise RuntimeError("before/after PTX ABIs differ")
        all_variants[rows] = variants
        print_ir_table(rows, ttgir, variants)

    performance: list[dict[str, Any]] = []
    if not args.compile_only:
        for rows in args.rows:
            variants = all_variants[rows]
            for variant in variants.values():
                load_kernel(variant)
            performance.extend(
                benchmark_rows(
                    rows,
                    variants,
                    args.grids,
                    args.nodes,
                    args.samples,
                    args.graph_warmup,
                )
            )
            print(f"N={rows}: correctness PASS (before and after, bit exact)")
        print_perf_table(performance)

    serializable_variants = {
        str(rows): {
            label: {
                key: value
                for key, value in variant.items()
                if key not in ("module", "kernel")
            }
            for label, variant in variants.items()
        }
        for rows, variants in all_variants.items()
    }
    report = {
        "device": device_name(),
        "commits": commits,
        "rows": args.rows,
        "grids": args.grids,
        "nodes": args.nodes,
        "samples": args.samples,
        "ttgir_sha256": {
            str(rows): hashlib.sha256(ttgir.encode()).hexdigest()
            for rows, ttgir in ttgirs.items()
        },
        "variants": serializable_variants,
        "performance": performance,
    }
    report_path = args.artifact_dir / "report.json"
    report_path.write_text(json.dumps(report, indent=2, default=list) + "\n")
    print(f"\nartifacts: {args.artifact_dir}")
    print(f"machine-readable report: {report_path}")


if __name__ == "__main__":
    main()

Res:

device: NVIDIA H100 PCIe
before commit: 99baa99306c78f9f8cccfbd0a98623461e9e39ea
before triton-opt: /workspace/exp/aaaee0_ir_compare/triton-opt-before (f69bcafb8a3f)
after  commit: aaaee0dc9bf7b8099f2b1d58f06b54b4c67fb168
after  triton-opt: /workspace/FlagTree_fork-reorg-20260820/build/cmake.flagtree-tle/bin/triton-opt (2d28bd8813ea)
frontend rule: one TTGIR is shared byte-for-byte by both lowerings

N=16: TTGIR convert_layout = 1
variant  smem_B  MLIR(shfl/st3/ld3/bar)  PTX(shfl/st.shared/ld.shared/bar/cp.async)
before     2048     0/  1/  2/  3                    0/  1/  2/  3/  1
after      2048     2/  0/  1/  1                    2/  0/  1/  1/  1

N=32: TTGIR convert_layout = 1
variant  smem_B  MLIR(shfl/st3/ld3/bar)  PTX(shfl/st.shared/ld.shared/bar/cp.async)
before     4096     0/  1/  4/  3                    0/  1/  4/  3/  2
after      4096     4/  0/  2/  1                    4/  0/  2/  1/  2
N=16: correctness PASS (before and after, bit exact)
N=32: correctness PASS (before and after, bit exact)

Median time per kernel (CUDA Graph; lower is better)
N     grid     before_us   after_us    delta_us   speedup
16    1           1.28496    1.24560    0.03936    1.032x
16    114         1.51120    1.47264    0.03856    1.026x
16    2112        3.02448    2.93224    0.09224    1.031x
32    1           1.52016    1.49152    0.02864    1.019x
32    114         1.81808    1.79096    0.02712    1.015x
32    2112        4.69880    4.68664    0.01216    1.003x

[sunnycase]

Thanks for working on this! Since SameWarpShufflePlan and planSameWarpShuffleConversion are specific to the TritonGPU-to-LLVM conversion and also encode NVIDIA lowering/profitability assumptions, would it be more appropriate to place them in triton/Conversion/TritonGPUToLLVM/Utility.h (with the implementation in the corresponding Utility.cpp) rather than in the generic Analysis utilities? This would keep the analysis layer backend-agnostic and place the logic closer to its consumers. What do you think?

[Kafka-Hatsune]

Thanks for working on this! Since SameWarpShufflePlan and planSameWarpShuffleConversion are specific to the TritonGPU-to-LLVM conversion and also encode NVIDIA lowering/profitability assumptions, would it be more appropriate to place them in triton/Conversion/TritonGPUToLLVM/Utility.h (with the implementation in the corresponding Utility.cpp) rather than in the generic Analysis utilities? This would keep the analysis layer backend-agnostic and place the logic closer to its consumers. What do you think?

Thank you for the suggestions. Based on those suggestions, the following code fixes have been made.

• The global triton/Conversion module adds the "layout transformation planner" and data structures.Code in third_party/nvidia is responsible for actually invoking the planner, generating NVIDIA shuffle instructions, and synchronously canceling the shared-memory scratch.Due to the excessive workload, other platforms have not been modified and still follow the original layout conversion path. Therefore, the function is not affected, but this optimization is not available for the time being.

• The new content is enclosed by __FLAGTREE_SAME_WARP_LAYOUT_SHUFFLE__.