Transposing 128x128 blocks is suboptimal

[danthe3rd]

Describe the issue

Triton is not really efficient at transposing arrays of 128x128 elements. NCU reports many shared-memory bank conflicts, and memory IO peaks at ~1700GB/s (instead of 2800GB/s+ for other shapes)

Python repro

Python code

import torch
import triton
import triton.language as tl
from typing import List, Optional, Union
import os
def benchmark_fn(
    fn,
    name: Optional[str] = None,
    *,
    cudagraph: bool = False,
    torch_profiler: bool = False,
    warmup_repeats: int = 0,
    repeats: int = 10,
    verbose: bool = False,
    IO: Union[List[torch.Tensor], int, None] = None,
    flops: int = 0,
) -> float:
    if name is None:
        name = fn.__name__
    # warmup
    fn()
    if "PROFILING" in os.environ:
        return 0.0
    if cudagraph:
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            fn()
        def fn_graphed():
            g.replay()
        fn_graphed()
        fn = fn_graphed
    # loop
    ev_start = torch.cuda.Event(enable_timing=True)
    ev_end = torch.cuda.Event(enable_timing=True)
    torch.cuda.synchronize()
    for _ in range(warmup_repeats):
        fn()
    ev_start.record()
    for _ in range(repeats):
        fn()
    ev_end.record()
    torch.cuda.synchronize()
    # (elapsed_time returns in `ms`)
    dt = ev_start.elapsed_time(ev_end) / (1000 * repeats)
    if verbose:
        if dt < 0.005:
            time_str = f"{int(dt * 10000000) / 10}us"
        elif dt < 5.0:
            time_str = f"{int(dt * 10000) / 10}ms"
        else:
            time_str = f"{int(dt * 10) / 10}s"
        IO_bytes = 0
        if isinstance(IO, list):
            IO_bytes = sum(
                x.numel() * x.dtype.itemsize for x in IO if isinstance(x, torch.Tensor)
            )
        elif isinstance(IO, (int, float)):
            IO_bytes = int(IO)
        if IO_bytes <= 0:
            bytes_per_s_str = ""
        else:
            bytes_per_s_str = f" [{int(IO_bytes / dt / (1024 ** 3))} GB/s]"
        if flops <= 0:
            flops_str = ""
        else:
            flops_str = f" [{int(flops / dt / (1000 ** 3))} GFlops]"
        print(f"{name}: {time_str}{flops_str}{bytes_per_s_str}")
    # get a trace
    if torch_profiler:
        profiler = torch.profiler.profile(
            on_trace_ready=lambda p: p.export_chrome_trace(f"benchmark_{name}.json.gz"),
            profile_memory=True,
            record_shapes=True,
            with_stack=True,
            with_flops=True,
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
        )
        torch.cuda.synchronize()
        with profiler:
            fn()
            torch.cuda.synchronize()
    return dt
def _kernel_repr(proxy):
    constants = proxy.constants
    return constants["KERNEL_NAME"]
@triton.jit(repr=_kernel_repr)
def transpose_k(in_ptr, out_ptr, BLOCK_SZ: tl.constexpr, DIM0: tl.constexpr, DIM1: tl.constexpr, TRANSPOSE_GRID: tl.constexpr, KERNEL_NAME: tl.constexpr):
    if TRANSPOSE_GRID:
        b0, b1 = tl.program_id(axis=0) * BLOCK_SZ, tl.program_id(axis=1) * BLOCK_SZ
    else:
        b1, b0 = tl.program_id(axis=0) * BLOCK_SZ, tl.program_id(axis=1) * BLOCK_SZ
    x = tl.load(in_ptr + (b0 + tl.arange(0, BLOCK_SZ)[:, None]) * DIM1 + b1 + tl.arange(0, BLOCK_SZ)[None, :])
    x = tl.trans(x)
    tl.store(out_ptr + (b0 + tl.arange(0, BLOCK_SZ)[:, None]) * DIM1 + b1 + tl.arange(0, BLOCK_SZ)[None, :], x)
def transpose(x: torch.Tensor, num_warps: int, BLOCK_SZ: int, TRANSPOSE_GRID: bool, KERNEL_NAME: str):
    y = torch.empty_like(x)
    grid = (x.shape[0] // BLOCK_SZ, x.shape[1] // BLOCK_SZ)
    transpose_k[grid](
        x,
        y,
        BLOCK_SZ=BLOCK_SZ,
        TRANSPOSE_GRID=TRANSPOSE_GRID,
        DIM0=x.shape[0],
        DIM1=x.shape[1],
        num_warps=num_warps,
        KERNEL_NAME=KERNEL_NAME,
    )
    return y
x = torch.randn([8192 * 2, 8192 * 2], device="cuda", dtype=torch.bfloat16)
for dtype_name, dtype, dtype_bytes in [
    # ("fp32", torch.float32, 4),
    ("bf16", torch.bfloat16, 2),
    ("int8", torch.int8, 1),
]:
    print(f"\n  {dtype_name}")
    x = x.to(dtype)
    IO = x.numel() * dtype_bytes * 2
    for BLOCK_SZ in [32, 64, 128]:
        for num_warps in [1, 2, 4, 8, 16]:
            for TRANSPOSE_GRID in [True, False]:
                name = f"transpose_{dtype_name}_{num_warps}_{BLOCK_SZ}_{'T' if TRANSPOSE_GRID else 'N'}"
                latency = benchmark_fn(
                    lambda: transpose(x, num_warps=num_warps, BLOCK_SZ=BLOCK_SZ, TRANSPOSE_GRID=TRANSPOSE_GRID, KERNEL_NAME=name),
                    name=name,
                    cudagraph=True,
                    IO=IO,
                    verbose=True)

Output on H100

Output

  bf16
transpose_bf16_1_32_T: 480.5us [2080 GB/s]
transpose_bf16_1_32_N: 365.0us [2739 GB/s]
transpose_bf16_2_32_T: 479.2us [2086 GB/s]
transpose_bf16_2_32_N: 366.9us [2725 GB/s]
transpose_bf16_4_32_T: 471.8us [2119 GB/s]
transpose_bf16_4_32_N: 372.8us [2681 GB/s]
transpose_bf16_8_32_T: 541.4us [1846 GB/s]
transpose_bf16_8_32_N: 421.1us [2374 GB/s]
transpose_bf16_16_32_T: 691.8us [1445 GB/s]
transpose_bf16_16_32_N: 599.4us [1668 GB/s]
transpose_bf16_1_64_T: 409.3us [2442 GB/s]
transpose_bf16_1_64_N: 374.1us [2673 GB/s]
transpose_bf16_2_64_T: 407.2us [2455 GB/s]
transpose_bf16_2_64_N: 374.4us [2670 GB/s]
transpose_bf16_4_64_T: 405.3us [2466 GB/s]
transpose_bf16_4_64_N: 363.5us [2750 GB/s]
transpose_bf16_8_64_T: 400.4us [2497 GB/s]
transpose_bf16_8_64_N: 360.0us [2777 GB/s]
transpose_bf16_16_64_T: 420.2us [2379 GB/s]
transpose_bf16_16_64_N: 400.1us [2498 GB/s]
transpose_bf16_1_128_T: 607.2us [1646 GB/s]
transpose_bf16_1_128_N: 605.9us [1650 GB/s]
transpose_bf16_2_128_T: 591.6us [1690 GB/s]
transpose_bf16_2_128_N: 592.0us [1689 GB/s]
transpose_bf16_4_128_T: 588.3us [1699 GB/s]
transpose_bf16_4_128_N: 588.4us [1699 GB/s]
transpose_bf16_8_128_T: 590.4us [1693 GB/s]
transpose_bf16_8_128_N: 590.5us [1693 GB/s]
transpose_bf16_16_128_T: 599.6us [1667 GB/s]
transpose_bf16_16_128_N: 599.6us [1667 GB/s]
  int8
transpose_int8_1_32_T: 247.6us [2019 GB/s]
transpose_int8_1_32_N: 205.0us [2438 GB/s]
transpose_int8_2_32_T: 248.0us [2015 GB/s]
transpose_int8_2_32_N: 208.2us [2401 GB/s]
transpose_int8_4_32_T: 261.4us [1912 GB/s]
transpose_int8_4_32_N: 213.1us [2345 GB/s]
transpose_int8_8_32_T: 321.6us [1554 GB/s]
transpose_int8_8_32_N: 305.5us [1636 GB/s]
transpose_int8_16_32_T: 440.7us [1134 GB/s]
transpose_int8_16_32_N: 498.8us [1002 GB/s]
transpose_int8_1_64_T: 242.4us [2062 GB/s]
transpose_int8_1_64_N: 209.3us [2388 GB/s]
transpose_int8_2_64_T: 237.3us [2106 GB/s]
transpose_int8_2_64_N: 199.1us [2510 GB/s]
transpose_int8_4_64_T: 229.9us [2174 GB/s]
transpose_int8_4_64_N: 210.1us [2378 GB/s]
transpose_int8_8_64_T: 234.7us [2129 GB/s]
transpose_int8_8_64_N: 207.8us [2405 GB/s]
transpose_int8_16_64_T: 283.8us [1761 GB/s]
transpose_int8_16_64_N: 248.2us [2013 GB/s]
transpose_int8_1_128_T: 326.0us [1533 GB/s]
transpose_int8_1_128_N: 326.8us [1529 GB/s]
transpose_int8_2_128_T: 298.7us [1673 GB/s]
transpose_int8_2_128_N: 299.4us [1669 GB/s]
transpose_int8_4_128_T: 294.7us [1696 GB/s]
transpose_int8_4_128_N: 294.6us [1696 GB/s]
transpose_int8_8_128_T: 293.6us [1702 GB/s]
transpose_int8_8_128_N: 294.0us [1700 GB/s]
transpose_int8_16_128_T: 295.6us [1691 GB/s]
transpose_int8_16_128_N: 295.8us [1690 GB/s]

Environment details

Triton: 3.4.0
GPU: H100

[CRobeck]

It looks like the bank conflicts are coming from storing the input tensor in SMEM before it is stored out after the transpose. The swizzling layout selection seems to be choosing bad heuristics for this case. There were some improvements to the swizzling layout selection in #7565 but that actually seems to make this test case regress since, I believe, the heuristics now aim to benefit the loads over stores. @lezcano do you have any thoughts here?

[lezcano]

can you paste the linear layouts associated to srcLayout, dstLayout and smem for this particular use case?

[CRobeck]

For the 16x128 case (transpose_bf16_16_128_T in the reproducer):

srcLayout
register=1 -> (1, 0)
register=2 -> (2, 0)
register=4 -> (4, 0)
register=8 -> (0, 32)
register=16 -> (0, 64)
lane=1 -> (8, 0)
lane=2 -> (16, 0)
lane=4 -> (32, 0)
lane=8 -> (64, 0)
lane=16 -> (0, 1)
warp=1 -> (0, 2)
warp=2 -> (0, 4)
warp=4 -> (0, 8)
warp=8 -> (0, 16)
block is a size 1 dimension
where out dims are: [dim0 (size 128), dim1 (size 128)]

dstLayout
register=1 -> (0, 1)
register=2 -> (0, 2)
register=4 -> (0, 4)
register=8 -> (32, 0)
register=16 -> (64, 0)
lane=1 -> (0, 8)
lane=2 -> (0, 16)
lane=4 -> (0, 32)
lane=8 -> (0, 64)
lane=16 -> (1, 0)
warp=1 -> (2, 0)
warp=2 -> (4, 0)
warp=4 -> (8, 0)
warp=8 -> (16, 0)
block is a size 1 dimension
where out dims are: [dim0 (size 128), dim1 (size 128)]

smem
vector=1 -> (2, 0)
vector=2 -> (4, 0)
bank=1 -> (0, 32)
bank=2 -> (0, 64)
bank=4 -> (0, 8)
bank=8 -> (0, 16)
segment=1 -> (0, 1)
segment=2 -> (0, 2)
segment=4 -> (0, 4)
segment=8 -> (16, 8)
segment=16 -> (32, 16)
segment=32 -> (64, 32)
segment=64 -> (8, 0)
segment=128 -> (1, 0)
reps is a size 1 dimension
where out dims are: [dim0 (size 128), dim1 (size 128)]

[lezcano]

mind testing #7833 see if this fixes your issue?
Also, thank you for the report :) If you find any other pairs of layouts that cause bank conflicts please send them over see if we can fix them!

[CRobeck]

Looks like perf is identical
Before #7833: transpose_bf16_16_128_T: 517.4us [1932 GB/s]
After #7833: transpose_bf16_16_128_T: 516.9us [1934 GB/s]
I'll have to fire up NCU to look at bank conflicts.

[CRobeck]

Looks like things have mostly stayed the same:

Before #7833:
transpose_bf16_1_32_T: 623.1us [1604 GB/s]
transpose_bf16_2_32_T: 649.1us [1540 GB/s]
transpose_bf16_4_32_T: 647.1us [1545 GB/s]
transpose_bf16_8_32_T: 701.1us [1426 GB/s]
transpose_bf16_16_32_T: 870.5us [1148 GB/s]
transpose_bf16_1_64_T: 562.9us [1776 GB/s]
transpose_bf16_2_64_T: 563.0us [1776 GB/s]
transpose_bf16_4_64_T: 562.7us [1777 GB/s]
transpose_bf16_8_64_T: 561.8us [1779 GB/s]
transpose_bf16_16_64_T: 559.3us [1787 GB/s]
transpose_bf16_1_128_T: 601.4us [1662 GB/s]
transpose_bf16_2_128_T: 539.2us [1854 GB/s]
transpose_bf16_4_128_T: 539.0us [1855 GB/s]
transpose_bf16_8_128_T: 537.4us [1860 GB/s]
transpose_bf16_16_128_T: 538.8us [1855 GB/s]

After #7833:
transpose_bf16_1_32_T: 627.4us [1593 GB/s]
transpose_bf16_2_32_T: 646.4us [1546 GB/s]
transpose_bf16_4_32_T: 650.2us [1537 GB/s]
transpose_bf16_8_32_T: 701.8us [1424 GB/s]
transpose_bf16_16_32_T: 871.2us [1147 GB/s]
transpose_bf16_1_64_T: 564.3us [1772 GB/s]
transpose_bf16_2_64_T: 563.1us [1775 GB/s]
transpose_bf16_4_64_T: 563.8us [1773 GB/s]
transpose_bf16_8_64_T: 563.7us [1773 GB/s]
transpose_bf16_16_64_T: 559.6us [1786 GB/s]
transpose_bf16_1_128_T: 610.3us [1638 GB/s]
transpose_bf16_2_128_T: 538.7us [1856 GB/s]
transpose_bf16_4_128_T: 539.1us [1854 GB/s]
transpose_bf16_8_128_T: 538.7us [1856 GB/s]
transpose_bf16_16_128_T: 542.5us [1843 GB/s]

[lezcano]

this should not have bank conflicts, at least in nvidia's hardware. Would be good to check whether there's anything different in AMD hardware, as I don't know what details are different.

The assumptions are:
different warps do not bank conflict with each other, banks of 32 bits, 32 banks... and I think that's it?

[CRobeck]

OK. I reran things and updated above. Things have mostly stayed the same. I'll have to dig further into the NCU outputs/traces to see why clearing the bank conflicts doesn't boost perf. This is all on H100.

[lezcano]

Reopening as the fix did not fully fix the perf issues.