"""Single-PE chunked-prefetching async GEMM with depth=2 double-buffer.

A TCM-bounded variant of ``matmul_async_chunked``: at most two B-chunk
loads are in flight at any time, so peak TCM occupancy for B is
``2 × TILE_K × N × bpe`` rather than the full ``K × N × bpe``.

This is the architecturally required form for any real LLM workload
where ``K = context_length`` (4 K–32 K tokens): the
``matmul_async_chunked`` variant queues *all* B-chunk loads up front
and therefore needs the full B in TCM; that exceeds the 1 MiB kernel
scratch cap (topology.yaml ``pe_tcm.kernel_scratch_mb``) at
sub-attention scale. Double-buffer prefetch is the closest
user-orchestrated kernel pattern to composite's per-tile streaming.

Structure (after pre-staging A):
  - Issue B_chunk[0] + B_chunk[1] async (2 outstanding)
  - For each chunk i:
    - tl.dot(A_chunk[i], B_chunk[i])   -- blocks on b[i], then issues GEMM
    - tl.add into running accumulator
    - Issue B_chunk[i+2] async (if available) -- keeps in-flight count = 2

Compared to the full-prefetch variant, depth-2 changes *when* DMA
descriptors hit the queue, not their total number; the structural
dispatch cost amortization story therefore stays the same. The
purpose of the variant is to honour the TCM cap, so the comparison
to composite reflects a kernel that could actually run on real
hardware at LLM-scale K.

Env: MATMUL_M, MATMUL_K, MATMUL_N, MATMUL_DTYPE.
"""
import os

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

TILE_K = 64
PREFETCH_DEPTH = 2


def _kernel_async_chunked_db(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:
        # Single-chunk fallback (under-tile or one-tile case): degenerates
        # to naive async, no prefetch opportunity.
        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

    # Pre-stage all A chunks; this is the "activation pre-staged" half of
    # load_ref. A is loaded once at the start and reused across every
    # output dot.
    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)

    # Prime the prefetch pipeline: issue the first PREFETCH_DEPTH B-chunk
    # loads up front. Anything beyond depth-2 stays in HBM until we get
    # to its dot.
    b_chunks: list = [
        tl.load(int(b_ptr) + i * b_chunk_bytes,
                shape=(TILE_K, N), dtype=DTYPE)
        for i in range(min(PREFETCH_DEPTH, n_chunks))
    ]

    # First dot — blocks on b_chunks[0]; while it computes, b_chunks[1]
    # is already in flight.
    out = tl.dot(a_chunks[0], b_chunks[0])

    for i in range(1, n_chunks):
        # Issue the next prefetch BEFORE the dot that consumes b_chunks[i],
        # so peak outstanding B loads stays at PREFETCH_DEPTH.
        if i + PREFETCH_DEPTH - 1 < n_chunks:
            b_chunks.append(
                tl.load(int(b_ptr) + (i + PREFETCH_DEPTH - 1) * b_chunk_bytes,
                        shape=(TILE_K, N), dtype=DTYPE)
            )
        out = out + tl.dot(a_chunks[i], b_chunks[i])

    tl.store(int(out_ptr), out)


@bench(
    name="matmul-async-chunked-db",
    description="Single-PE chunked-prefetching async GEMM with depth=2 "
                "double-buffer (TCM-bounded streaming analog of composite "
                "load_ref).",
)
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_db", _kernel_async_chunked_db,
                 a, b, out, M, K, N)
