gemm_streamk_benchmark.py

"""
Stream K GEMM with Tensor Descriptors
=====================================
Stream K is a approach that aims to resovle quantization inefficiency in GPU work load for GEMM problems.
This script implements a Stream K GEMM with tensor descriptors to achieve better hardware utilization.
"""

import torch
import triton
import triton.language as tl

import triton_kernels_benchmark as benchmark_suite
from triton_kernels_benchmark import sycl_tla_kernel


# pylint: disable=unused-argument
@triton.jit
def swizzle_tile(tile_id,
                 # Matrix dimensions
                 M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                 # Meta-parameters
                 BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
                 GROUP_SIZE_M: tl.constexpr):
    grid_m = tl.cdiv(M, BLOCK_SIZE_M)
    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
    width = GROUP_SIZE_M * grid_n
    group_id = tile_id // width
    group_size = tl.minimum(GROUP_SIZE_M, grid_m - group_id * GROUP_SIZE_M)
    pid_m = group_id * GROUP_SIZE_M + (tile_id % group_size)
    pid_n = (tile_id % width) // group_size
    return pid_m, pid_n


# pylint: disable=unused-argument
@triton.jit
def linear_tile(tile_id,
                # Matrix dimensions
                M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                # Meta-parameters
                BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
                GROUP_SIZE_M: tl.constexpr):
    pid_m = tile_id // tl.cdiv(N, BLOCK_SIZE_N)
    pid_n = tile_id % tl.cdiv(N, BLOCK_SIZE_N)
    return pid_m, pid_n


# Multiply-accumulate loop in GEMM Stream K tiles
@triton.jit
def mac_loop(
        # Pointers to matrices
        a_ptr, b_ptr, c_ptr,
        # Matrix dimensions
        M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,  #
        stride_am: tl.constexpr, stride_ak: tl.constexpr,  #
        stride_bk: tl.constexpr, stride_bn: tl.constexpr,  #
        stride_cm: tl.constexpr, stride_cn: tl.constexpr,
        # Stream-K parameters
        iters_per_tile, start_iter, end_iter,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr):

    tile_id = start_iter // iters_per_tile
    remain_iters = start_iter % iters_per_tile
    if GROUP_SIZE_M > 0:
        # pid swizzle to get better L2 cache performance
        pid_m, pid_n = swizzle_tile(tile_id, M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M)
    else:
        pid_m, pid_n = linear_tile(tile_id, M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M)

    # Create tensor descriptors
    a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(stride_am, stride_ak),
                                       block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
    b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(stride_bk, stride_bn),
                                       block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))

    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    off_k = remain_iters * BLOCK_SIZE_K
    for _ in range(start_iter, end_iter):
        a = a_desc.load([pid_m * BLOCK_SIZE_M, off_k])
        b = b_desc.load([off_k, pid_n * BLOCK_SIZE_N])
        acc = tl.dot(a, b, acc)
        off_k += BLOCK_SIZE_K

    if remain_iters == 0 and end_iter % iters_per_tile == 0:
        c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(stride_cm, stride_cn),
                                           block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
        c_desc.store([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], acc)
    else:
        rm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        rn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        c_ptr_ = c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
        mask = (rm < M)[:, None] & (rn < N)[None, :]
        tl.atomic_add(c_ptr_, acc, mask=mask, sem='relaxed')


@triton.autotune(
    configs=[
        triton.Config(
            {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode': '256'},
            num_stages=2, num_warps=32),
    ],
    key=['M', 'N', 'K'],
    restore_value=['c_ptr'],
)
@triton.jit
def first_wave(
        # Pointers to matrices
        a_ptr, b_ptr, c_ptr,
        # Matrix dimensions
        M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,  #
        stride_am: tl.constexpr, stride_ak: tl.constexpr,  #
        stride_bk: tl.constexpr, stride_bn: tl.constexpr,  #
        stride_cm: tl.constexpr, stride_cn: tl.constexpr,
        # Stream-K parameters
        full_tiles,  # pylint: disable=redefined-outer-name
        partial_tiles, iters_per_tile,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr):

    pid = tl.program_id(axis=0)
    start_iter = pid * full_tiles + tl.minimum(pid, partial_tiles)
    last_iter = (pid + 1) * full_tiles + tl.minimum(pid + 1, partial_tiles)

    while start_iter < last_iter:
        end_iter = start_iter + (iters_per_tile - start_iter % iters_per_tile)
        end_iter = tl.minimum(end_iter, last_iter)
        mac_loop(a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                 iters_per_tile, start_iter, end_iter, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M)

        start_iter = end_iter


@triton.autotune(
    configs=[
        triton.Config(
            {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode': '256'},
            num_stages=2, num_warps=32),
    ],
    key=['M', 'N', 'K'],
    restore_value=['c_ptr'],
)
@triton.jit
def full_tiles(
        # Pointers to matrices
        a_ptr, b_ptr, c_ptr,
        # Matrix dimensions
        M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
        # The stride variables represent how much to increase the ptr by when moving by 1
        # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
        # by to get the element one row down (A has M rows).
        stride_am: tl.constexpr, stride_ak: tl.constexpr,  #
        stride_bk: tl.constexpr, stride_bn: tl.constexpr,  #
        stride_cm: tl.constexpr, stride_cn: tl.constexpr,
        # Stream-K parameters
        streamk_tiles,
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr):

    tile_id = tl.program_id(axis=0) + streamk_tiles
    if GROUP_SIZE_M > 0:
        # pid swizzle to get better L2 cache performance
        pid_m, pid_n = swizzle_tile(tile_id, M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M)
    else:
        pid_m, pid_n = linear_tile(tile_id, M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M)

    a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(stride_am, stride_ak),
                                       block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
    b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(stride_bk, stride_bn),
                                       block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))

    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    off_k = 0
    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = a_desc.load([pid_m * BLOCK_SIZE_M, off_k])
        b = b_desc.load([off_k, pid_n * BLOCK_SIZE_N])
        acc = tl.dot(a, b, acc)
        off_k += BLOCK_SIZE_K

    c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(stride_cm, stride_cn),
                                       block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
    c_desc.store([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], acc)


# ---------------------------------------------------------------------------
# Wrapper
# ---------------------------------------------------------------------------


def matmul(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor):
    num_xe_core = torch.xpu.get_device_capability(0)['gpu_subslice_count']
    streamk_programs = num_xe_core

    # TODO: use autotune config instread of hardcoding
    BLOCK_SIZE_M = 256
    BLOCK_SIZE_N = 256
    BLOCK_SIZE_K = 32

    # Check constraints.
    assert a.shape[1] == b.shape[0], 'Incompatible dimensions'
    assert a.is_contiguous(), 'Matrix A must be contiguous'
    assert b.is_contiguous(), 'Matrix B must be contiguous'
    M, K = a.shape
    K, N = b.shape

    num_block_m = triton.cdiv(M, BLOCK_SIZE_M)
    num_block_n = triton.cdiv(N, BLOCK_SIZE_N)
    iters_per_tile = triton.cdiv(K, BLOCK_SIZE_K)
    total_tiles = num_block_m * num_block_n

    # Two-tile SK + DP
    streamk_tiles = total_tiles % streamk_programs
    if total_tiles - streamk_tiles > streamk_programs:  # (total_tiles // total_programs > 1)
        streamk_tiles += streamk_programs

    blocking_tiles = total_tiles - streamk_tiles
    streamk_iters = streamk_tiles * iters_per_tile

    streamk_full_tiles = streamk_iters // streamk_programs
    streamk_partial_tiles = streamk_iters % streamk_programs

    first_wave[(streamk_programs, )](
        a, b, c,  #
        M, N, K,  #
        a.stride(0), a.stride(1),  #
        b.stride(0), b.stride(1),  #
        c.stride(0), c.stride(1),  #
        streamk_full_tiles, streamk_partial_tiles, iters_per_tile)
    full_tiles[(blocking_tiles, )](
        a, b, c,  #
        M, N, K,  #
        a.stride(0), a.stride(1),  #
        b.stride(0), b.stride(1),  #
        c.stride(0), c.stride(1),  #
        streamk_tiles)

    return c


# Benchmark Performance
@benchmark_suite.perf_report(
    benchmark_suite.Benchmark(
        # argument names to use as an x-axis for the plot
        x_names=['M', 'K', 'N'],
        x_vals=[[3072, 4096, 3072]],
        line_arg='provider',
        # argument name whose value corresponds to a different line in the plot
        # possible values for `line_arg``
        line_vals=['triton', 'onednn', 'sycl-tla'],
        # label name for the lines
        line_names=['Triton', 'OneDNN', 'SYCL-TLA'],
        # line styles
        styles=[('green', '-'), ('green', '--'), ('blue', '-')],
        ylabel=['GB/s', 'TFlops'],  # label name for the y-axis
        plot_name='matmul-streamk-performance',
        # name for the plot. Used also as a file name for saving the plot.
        args={},
    ))
def benchmark(M, N, K, provider):
    # n_warmup set to maximum across providers: onednn=10, triton=1000, sycl-tla=100
    do_bench = benchmark_suite.get_do_bench(n_warmup=1000, n_repeat=10, quantiles=[0.5, 0.0, 1.0])
    torch.manual_seed(0)
    a = torch.rand((M, K), device='xpu', dtype=torch.bfloat16)
    b = torch.rand((K, N), device='xpu', dtype=torch.bfloat16)

    if provider == 'onednn':
        _, min_ms, max_ms, mean_ms, cv = do_bench(lambda: torch.matmul(a, b))
    elif provider == 'triton':
        c = torch.zeros((M, N), device=a.device, dtype=torch.float32)
        triton_fn = lambda: matmul(a, b, c)
        torch_fn = lambda: torch.matmul(a, b).to(torch.float32)
        benchmark_suite.assert_close(triton_fn, torch_fn, atol=1e-4, rtol=1e-2, err_msg='triton to torch')
        _, min_ms, max_ms, mean_ms, cv = do_bench(triton_fn)
    elif provider == 'sycl-tla':
        # NOTE: SYCL-TLA doesn't have a stream-k matmul op
        c = torch.zeros((M, N), device='xpu', dtype=torch.float32)
        func = getattr(sycl_tla_kernel, 'gemm')
        sycl_tla_fn = lambda: (func(a, b, c, M, N, K, 1), c)[1]
        torch_fn = lambda: torch.matmul(a, b).to(torch.float32)
        benchmark_suite.assert_close(sycl_tla_fn, torch_fn, atol=1e-4, rtol=1e-2, err_msg='sycl-tla to torch')
        _, min_ms, max_ms, mean_ms, cv = do_bench(sycl_tla_fn)
    else:
        raise NotImplementedError(f'Unsupported provider {provider}')

    tflops = lambda mean: 2 * M * N * K * (1e-12) / (mean * 1e-3)
    gbps = lambda mean: (2 * (M * K + K * N) + 4.0 * (M * N)) * (1e-9) / (mean * 1e-3)

    return (gbps(mean_ms), gbps(max_ms), gbps(min_ms)), (tflops(mean_ms), tflops(max_ms), tflops(min_ms)), cv


def get_benchmark(providers_filter=None):
    return benchmark


if __name__ == '__main__':
    benchmark.run(show_plots=False, print_data=True)