"""Single-PE chunked-prefetching async-loading GEMM for composite-vs-async
dispatch comparison.

Splits K into hardware-tile-sized chunks (CHUNK_K = TILE_K = 64) so that:
  - A chunks are loaded first (A pre-staged), then awaited
  - B chunks are all issued as async tl.load at once → DMA engine queues
    them FIFO and starts streaming
  - Per-chunk tl.dot is issued; each dot blocks on its specific b_chunk
    only, so dot[i] can start as soon as b[i] lands in TCM, overlapping
    GEMM[i] with DMA of b[i+1..] still in flight

Per-chunk dot outputs are accumulated with tl.add via a running sum
(``out = out + tl.dot(...)``), which serializes on the MATH engine due
to accumulator data dependency. This is the closest user-level analog to
composite's per-tile pipeline that is achievable with the current tl
DSL (which lacks a GEMM accumulator parameter and slicing primitive).

Env: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE (same as matmul_composite).
"""
import os

from kernbench.benches.registry import bench
from kernbench.policy.placement.dp import DPPolicy

# Mirror PeSchedulerComponent.TILE_K (also TILE_K in milestone_1h_gemm.py).
TILE_K = 64


def _kernel_async_chunked(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
    M, K, N = int(M), int(K), int(N)
    bpe = 2  # f16

    if K <= TILE_K:
        # Under-tile in K: chunking would make each dot smaller, even
        # worse off. Fall back to a single dot (naive async behaviour).
        a = tl.load(int(a_ptr), shape=(M, K), dtype=DTYPE)
        tl._await_pending(a)
        b = tl.load(int(b_ptr), shape=(K, N), dtype=DTYPE)
        out = tl.dot(a, b)
        tl.store(int(out_ptr), out)
        return

    n_chunks = K // TILE_K
    a_chunk_bytes = M * TILE_K * bpe
    b_chunk_bytes = TILE_K * N * bpe

    # ── Phase 1: pre-stage all A chunks (issue + wait) ──
    a_chunks = [
        tl.load(int(a_ptr) + i * a_chunk_bytes,
                shape=(M, TILE_K), dtype=DTYPE)
        for i in range(n_chunks)
    ]
    tl._await_pending(*a_chunks)

    # ── Phase 2: queue all B-chunk async loads up front ──
    # The DMA engine processes them FIFO; later chunks are in flight
    # while the early dots run on GEMM engine.
    b_chunks = [
        tl.load(int(b_ptr) + i * b_chunk_bytes,
                shape=(TILE_K, N), dtype=DTYPE)
        for i in range(n_chunks)
    ]

    # ── Phase 3: per-chunk dot with running accumulator ──
    # dot[i] only blocks on b_chunks[i] (tl.dot's _await_pending). While
    # dot[i] runs on GEMM, b_chunks[i+1..] continue draining off DMA.
    out = tl.dot(a_chunks[0], b_chunks[0])
    for i in range(1, n_chunks):
        out = out + tl.dot(a_chunks[i], b_chunks[i])

    tl.store(int(out_ptr), out)


@bench(
    name="matmul-async-chunked",
    description="Single-PE chunked-prefetching async-loading GEMM "
                "(CHUNK_K = TILE_K = 64; A pre-staged, B chunks "
                "prefetched, per-chunk tl.dot + tl.add accumulator).",
)
def run(torch):
    M = int(os.environ.get("MATMUL_M", "256"))
    K = int(os.environ.get("MATMUL_K", "256"))
    N = int(os.environ.get("MATMUL_N", "256"))
    DTYPE = os.environ.get("MATMUL_DTYPE", "f16")
    dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
    a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a")
    b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b")
    out = torch.empty((M, N), dtype=DTYPE, dp=dp, name="out")
    torch.launch("matmul_async_chunked", _kernel_async_chunked,
                 a, b, out, M, K, N)
