kernbench2/src/kernbench/components/builtin/tiling.py

"""Tile plan generators for PE pipeline (ADR-0014 D6).

Generates TilePlan with stage sequences for GEMM and Math operations.
Ported from pe_accel tiling.py with stage-based plan structure.
"""
from __future__ import annotations

from math import ceil, prod

from kernbench.components.builtin.pe_types import (
    PipelinePlan,
    Stage,
    StageType,
    TilePlan,
)


def generate_gemm_plan(
    M: int, K: int, N: int,
    tile_m: int, tile_k: int, tile_n: int,
    bytes_per_element: int,
    A_addr: int, B_addr: int, C_addr: int,
    pe_prefix: str,
    a_pinned: bool = False,
    b_pinned: bool = False,
    epilogue_specs: tuple = (),
    out_space: str = "hbm",
) -> PipelinePlan:
    """Generate GEMM tile plan: M→N→K order.

    Each tile follows stage sequence:
      [DMA_READ(A)] → [DMA_READ(B)] → FETCH → GEMM → [STORE → DMA_WRITE]
      DMA_READ(A) skipped when a_pinned=True (operand pre-staged in TCM).
      DMA_READ(B) skipped when b_pinned=True.
      STORE + DMA_WRITE only emitted on last K-tile per (m,n) — accumulator
      stays in RegFile across K loop.

    Args:
        pe_prefix: e.g. "sip0.cube0.pe0" — used to build component IDs.
        a_pinned: A operand already resident in TCM (via prior tl.load).
        b_pinned: B operand already resident in TCM.
    """
    M_tiles = max(1, ceil(M / tile_m))
    K_tiles = max(1, ceil(K / tile_k))
    N_tiles = max(1, ceil(N / tile_n))
    bpe = bytes_per_element

    dma_id = f"{pe_prefix}.pe_dma"
    fetch_id = f"{pe_prefix}.pe_fetch_store"
    gemm_id = f"{pe_prefix}.pe_gemm"
    math_id = f"{pe_prefix}.pe_math"

    # Split epilogue_specs by scope. Lazy import to avoid circular dep
    # between pe_commands ← pe_types ← tiling.
    from kernbench.common.pe_commands import Scope as _Scope
    k_tile_ops = [o for o in epilogue_specs
                  if getattr(o, "scope", None) == _Scope.K_TILE]
    out_tile_ops = [o for o in epilogue_specs
                    if getattr(o, "scope", None) == _Scope.OUTPUT_TILE]

    tiles: list[TilePlan] = []
    tile_id = 0

    for m in range(M_tiles):
        for n in range(N_tiles):
            c_addr = C_addr + (m * tile_m * N + n * tile_n) * bpe
            for k in range(K_tiles):
                last_k = k == K_tiles - 1
                a_addr = A_addr + (m * tile_m * K + k * tile_k) * bpe
                b_addr = B_addr + (k * tile_k * N + n * tile_n) * bpe

                a_bytes = tile_m * tile_k * bpe
                b_bytes = tile_k * tile_n * bpe
                out_bytes = tile_m * tile_n * bpe

                stages: list[Stage] = []

                # DMA READ: load A and B tiles from HBM → TCM.
                # Skip if the operand is already pre-staged via tl.load.
                if not a_pinned:
                    stages.append(Stage(
                        stage_type=StageType.DMA_READ,
                        component=dma_id,
                        params={
                            "src_addr": a_addr, "nbytes": a_bytes,
                            "operand": "A", "tile_m": tile_m, "tile_k": tile_k,
                        },
                    ))
                if not b_pinned:
                    stages.append(Stage(
                        stage_type=StageType.DMA_READ,
                        component=dma_id,
                        params={
                            "src_addr": b_addr, "nbytes": b_bytes,
                            "operand": "B", "tile_k": tile_k, "tile_n": tile_n,
                        },
                    ))

                # FETCH: TCM → Register File
                stages.append(Stage(
                    stage_type=StageType.FETCH,
                    component=fetch_id,
                    params={
                        "direction": "read",
                        "nbytes": a_bytes + b_bytes,
                    },
                ))

                # GEMM: MAC compute
                stages.append(Stage(
                    stage_type=StageType.GEMM,
                    component=gemm_id,
                    params={
                        "m": tile_m, "k": tile_k, "n": tile_n,
                        "is_last_k": last_k,
                    },
                ))

                # K-tile-scope epilogue MATH stages (e.g. dequant) — run on
                # every K-tile right after GEMM, before the accumulator
                # advances to the next K slice.
                for op in k_tile_ops:
                    stages.append(Stage(
                        stage_type=StageType.MATH,
                        component=math_id,
                        params={
                            "op_kind": op.kind,
                            "num_elements": tile_m * tile_n,
                            "scope": "k_tile",
                        },
                    ))

                # STORE + DMA_WRITE only on last K-tile per (m,n). The C
                # accumulator stays in RegFile across the K loop.
                if last_k:
                    # Output-tile-scope epilogue MATH (bias, relu, ...) runs
                    # ONCE per (m,n) after the final K-tile, before writeback.
                    for op in out_tile_ops:
                        stages.append(Stage(
                            stage_type=StageType.MATH,
                            component=math_id,
                            params={
                                "op_kind": op.kind,
                                "num_elements": tile_m * tile_n,
                                "scope": "output_tile",
                            },
                        ))
                    stages.append(Stage(
                        stage_type=StageType.STORE,
                        component=fetch_id,
                        params={
                            "direction": "write",
                            "nbytes": out_bytes,
                        },
                    ))
                    # DMA_WRITE only when the output is HBM-resident
                    # (ADR-0065 D8). A TCM output stays on-chip (chainable);
                    # writing its high-bit scratch address as an HBM PA would
                    # also fail the DMA decoder.
                    if out_space == "hbm":
                        stages.append(Stage(
                            stage_type=StageType.DMA_WRITE,
                            component=dma_id,
                            params={
                                "dst_addr": c_addr, "nbytes": out_bytes,
                            },
                        ))

                tiles.append(TilePlan(tile_id=tile_id, stages=tuple(stages)))
                tile_id += 1

    return PipelinePlan(
        tiles=tiles, m_tiles=M_tiles, k_tiles=K_tiles, n_tiles=N_tiles,
    )


def generate_math_plan(
    M: int, N: int,
    tile_m: int, tile_n: int,
    bytes_per_element: int,
    math_op: str,
    src_addr: int, dst_addr: int,
    pe_prefix: str,
) -> PipelinePlan:
    """Generate element-wise math tile plan.

    Each tile: DMA_READ → FETCH → MATH → STORE → DMA_WRITE
    """
    M_tiles = max(1, ceil(M / tile_m))
    N_tiles = max(1, ceil(N / tile_n))
    bpe = bytes_per_element

    dma_id = f"{pe_prefix}.pe_dma"
    fetch_id = f"{pe_prefix}.pe_fetch_store"
    math_id = f"{pe_prefix}.pe_math"

    tiles: list[TilePlan] = []
    tile_id = 0

    for m in range(M_tiles):
        for n in range(N_tiles):
            offset = (m * tile_m * N + n * tile_n) * bpe
            tile_bytes = tile_m * tile_n * bpe

            stages = [
                Stage(StageType.DMA_READ, dma_id, {
                    "src_addr": src_addr + offset, "nbytes": tile_bytes,
                }),
                Stage(StageType.FETCH, fetch_id, {
                    "direction": "read", "nbytes": tile_bytes,
                }),
                Stage(StageType.MATH, math_id, {
                    "op": math_op, "num_elements": tile_m * tile_n,
                }),
                Stage(StageType.STORE, fetch_id, {
                    "direction": "write", "nbytes": tile_bytes,
                }),
                Stage(StageType.DMA_WRITE, dma_id, {
                    "dst_addr": dst_addr + offset, "nbytes": tile_bytes,
                }),
            ]

            tiles.append(TilePlan(tile_id=tile_id, stages=tuple(stages)))
            tile_id += 1

    return PipelinePlan(tiles=tiles, m_tiles=M_tiles, n_tiles=N_tiles)


def _math_stage(op: object, pe_prefix: str) -> Stage:
    """Single-shot MATH stage for a KERNEL-scope flat op (ADR-0065 D3).

    Carries operand + output addresses (ADR-0065 N2) so the DataExecutor can
    compute the recipe's MATH op in data mode. Legacy epilogue MATH stages
    (generated in ``generate_gemm_plan``) do NOT carry these, so their op_log
    records stay op_kind-opaque → byte-equal.
    """
    out = getattr(op, "out", None)
    num_elements = prod(out.shape) if out is not None else 1
    operands = list(op.operands.values())
    params: dict = {
        "op_kind": op.kind, "num_elements": num_elements, "scope": "kernel",
        "op": op.kind,
        "input_addrs": [h.addr for h in operands],
        "input_shapes": [h.shape for h in operands],
        "input_spaces": [getattr(h, "space", "tcm") for h in operands],
        "input_dtypes": [h.dtype for h in operands],
    }
    if out is not None:
        params["dst_addr"] = out.addr
        params["dst_space"] = getattr(out, "space", "tcm")
        params["dtype"] = out.dtype
    axis = op.extra.get("reduce_axis") if hasattr(op, "extra") else None
    if axis is not None:
        params["axis"] = axis
    return Stage(StageType.MATH, f"{pe_prefix}.pe_math", params)


def generate_plan_from_ops(
    ops: tuple,
    tile_m: int, tile_k: int, tile_n: int,
    bytes_per_element: int,
    pe_prefix: str,
) -> PipelinePlan:
    """Generate a PipelinePlan from a flat ops list (ADR-0065 D3).

    Scans for the single GEMM op (≤1). Pre-GEMM KERNEL ops become single-shot
    prologue MATH stages; post-GEMM ops are split by scope — K_TILE/
    OUTPUT_TILE feed the GEMM tile loop's epilogue (as today), KERNEL ops
    become post-loop stages. A composite with no GEMM op is MATH-only.

    DMA insertion keeps the existing ``pinned`` signal (operand already in
    TCM → no DMA_READ). Prologue / post-loop stages are folded into the
    first / last tile so the feeder and completion counting stay unchanged.
    """
    from kernbench.common.pe_commands import Scope as _Scope

    gemm_idx = next((i for i, o in enumerate(ops) if o.kind == "gemm"), None)

    # MATH-only composite (no GEMM): reproduce the legacy math-head plan.
    if gemm_idx is None:
        head = ops[0]
        a = head.operands["a"]
        M = a.shape[-2] if len(a.shape) >= 2 else a.shape[0]
        N = a.shape[-1] if len(a.shape) >= 2 else 1
        return generate_math_plan(
            M=M, N=N, tile_m=tile_m, tile_n=tile_n,
            bytes_per_element=bytes_per_element,
            math_op=head.extra.get("math_op") or "identity",
            src_addr=a.addr, dst_addr=head.out.addr, pe_prefix=pe_prefix,
        )

    head = ops[gemm_idx]
    pre_ops = ops[:gemm_idx]
    post_ops = ops[gemm_idx + 1:]
    a = head.operands["a"]
    b = head.operands["b"]
    M, K = a.shape[-2], a.shape[-1]
    N = b.shape[-1]

    plan = generate_gemm_plan(
        M=M, K=K, N=N,
        tile_m=tile_m, tile_k=tile_k, tile_n=tile_n,
        bytes_per_element=bytes_per_element,
        A_addr=a.addr, B_addr=b.addr, C_addr=head.out.addr,
        pe_prefix=pe_prefix,
        a_pinned=getattr(a, "pinned", False),
        b_pinned=getattr(b, "pinned", False),
        epilogue_specs=tuple(post_ops),
        out_space=getattr(head.out, "space", "hbm"),
    )

    # Prologue (pre-GEMM KERNEL) + post-loop KERNEL ops become single-shot
    # MATH stages. They are fed as standalone 1-stage tiles by the scheduler
    # (PE_SCHEDULER._feed_loop) — NOT folded into the GEMM tiles, because a
    # folded MATH→DMA_READ boundary would require a PE_MATH→PE_DMA token
    # route that the pipeline does not wire. Existing benches have neither →
    # their tiles + feed order are unchanged (byte-equal op_log).
    plan.prologue_stages = tuple(_math_stage(o, pe_prefix) for o in pre_ops)
    plan.epilogue_stages = tuple(_math_stage(o, pe_prefix) for o in post_ops
                                 if o.scope == _Scope.KERNEL)
    return plan