#!/usr/bin/env python3
"""TileLang GEMM teaching kernel based on the v0.1.10 quick-start structure."""

import tilelang
import tilelang.language as T


@tilelang.jit
def matmul(a, b, block_m: int = 64, block_n: int = 64, block_k: int = 64):
    m, n, k = T.const("M, N, K")
    a: T.Tensor[[m, k], T.float16]
    b: T.Tensor[[k, n], T.float16]
    output = T.empty([m, n], T.float16)
    with T.Kernel(T.ceildiv(n, block_n), T.ceildiv(m, block_m), threads=128) as (bx, by):
        a_shared = T.alloc_shared((block_m, block_k), T.float16)
        b_shared = T.alloc_shared((block_k, block_n), T.float16)
        accumulator = T.alloc_fragment((block_m, block_n), T.float32)
        T.clear(accumulator)
        for ko in T.Pipelined(T.ceildiv(k, block_k), num_stages=3):
            T.copy(a[by * block_m, ko * block_k], a_shared)
            T.copy(b[ko * block_k, bx * block_n], b_shared)
            T.gemm(a_shared, b_shared, accumulator)
        T.copy(accumulator, output[by * block_m, bx * block_n])
    return output


if __name__ == "__main__":
    raise SystemExit("parser_only: pin a target, shapes, tensors, and reference before running")