"""Single-PE async-loading GEMM (tl.load + tl.dot) for composite-vs-async
dispatch comparison.

Pairs with ``matmul_composite``: both run the same (M,K,N) sweep, but
this kernel uses atomic primitives:

  a = tl.load(a_ptr, ...)   # async DMA into TCM (1 PE_CPU cmd)
  b = tl.load(b_ptr, ...)   # async DMA into TCM (1 PE_CPU cmd), overlaps with a
  out = tl.dot(a, b)        # 1 PE_CPU cmd; GEMM engine starts only after
                            # both loads complete (tl.dot _await_pending)
  tl.store(out_ptr, out)    # 1 PE_CPU cmd

Compared to composite:
  - dispatch count grows to 4 instead of ~2
  - GEMM cannot start until *all* of B is in TCM (no per-tile overlap of
    streaming B-load with GEMM compute)

Env vars: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE (mirror matmul_composite).
"""
import os

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


def _kernel_async(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
    M, K, N = int(M), int(K), int(N)
    # Pre-stage A: load A and block the user kernel until its DMA finishes
    # before dispatching load(B). Without this barrier, load(A) and load(B)
    # run concurrently on the DMA engine and share HBM bandwidth, so the
    # "exclude load(A) from the engine window" methodology does not
    # actually subtract A's time — A's traffic still shows up as halved-BW
    # contention with B. Calling _await_pending makes the load(A) the
    # explicit pre-staging step, matching the composite load_ref semantic
    # (where the composite waits for A before its first tile fetch).
    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)


@bench(
    name="matmul-async",
    description="Single-PE async-loading GEMM (tl.load + tl.dot) for "
                "composite-vs-async dispatch comparison.",
)
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", _kernel_async, a, b, out, M, K, N)
