kernbench2/src/kernbench/common/pe_commands.py

"""PE-internal command types and handles (ADR-0014).

Generated by triton_emu (TLContext) and consumed by PE component
implementations (PE_CPU, PE_SCHEDULER, PE_DMA, PE_GEMM, PE_MATH).

Command lifecycle:
  Triton kernel → TLContext → [PeCommand list] → PE_CPU → PE_SCHEDULER → engines
"""
from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import simpy


class Scope(Enum):
    K_TILE = "k_tile"
    OUTPUT_TILE = "output_tile"
    KERNEL = "kernel"


def _extra_bytes(v: Any) -> int:
    """Type-aware HW-logical byte count for an OpSpec.extra value (ADR-0064 D2).

    bool→1, int/float→4 (scalar), tuple/list→1 + 4·len (length marker + 4
    per element, e.g. shape/axes), str→1 (opcode-like tag). Default→4.
    ``bool`` is checked first because it subclasses ``int``.
    """
    if isinstance(v, bool):
        return 1
    if isinstance(v, (int, float)):
        return 4
    if isinstance(v, (tuple, list)):
        return 1 + 4 * len(v)
    if isinstance(v, str):
        return 1
    return 4


@dataclass(frozen=True)
class OpSpec:
    """One operation in a multi-op composite (head + epilogue, ADR-0014 D3.3).

    The head op (first in CompositeCmd.ops) defines tile geometry; subsequent
    ops are epilogue stages whose ``scope`` controls how often they fire
    (per-K-tile, per-output-tile, or once per kernel).
    """

    kind: str                                      # "gemm" | "bias" | "relu" | ...
    scope: "Scope" = Scope.OUTPUT_TILE
    operands: dict[str, Any] = field(default_factory=dict)  # name → TensorHandle
    scalar: float | None = None
    extra: dict[str, Any] = field(default_factory=dict)
    out: "TensorHandle | None" = None              # explicit write-back handle

    @property
    def logical_bytes(self) -> int:
        """HW-logical byte size (ADR-0064 D2). ``scalar`` is a transitional
        field not part of the flat-ops model (ADR-0065 D2) — excluded,
        pending its removal."""
        return (
            1 + 1                                       # opcode + scope enum
            + 1 + 8 * len(self.operands)                # len marker + handles
            + (8 if self.out is not None else 0)        # out handle
            + 1 + sum(_extra_bytes(v) for v in self.extra.values())
        )


# Epilogue op contracts: kind → (required field names, default scope).
# Used by tl.composite to validate user-provided epilogue dicts at
# command-emit time so typos fail before reaching the scheduler.
EPILOGUE_OPS: dict[str, tuple[tuple[str, ...], "Scope"]] = {
    "bias":    (("bias",),          Scope.OUTPUT_TILE),
    "relu":    ((),                 Scope.OUTPUT_TILE),
    "gelu":    ((),                 Scope.OUTPUT_TILE),
    "sigmoid": ((),                 Scope.OUTPUT_TILE),
    "scale":   (("factor",),        Scope.OUTPUT_TILE),
    "clamp":   (("lo", "hi"),       Scope.OUTPUT_TILE),
    "dequant": (("scale",),         Scope.K_TILE),
    "add":     (("other",),         Scope.OUTPUT_TILE),
}


# ── Handles ───────────────────────────────────────────────────────


@dataclass(frozen=True)
class TensorHandle:
    """Opaque reference to a tensor residing in PE_TCM.

    Returned by tl.load, tl.dot, tl.exp, etc.
    Carries metadata for command generation; data field is reserved
    for future validate mode (numpy array).
    """

    id: str
    addr: int                        # address (VA when MMU enabled, PA otherwise)
    shape: tuple[int, ...]
    dtype: str
    nbytes: int                      # total byte size
    data: object = None              # reserved for validate mode
    space: str = "tcm"               # MemoryStore space ("tcm" | "hbm" | "sram")
    pinned: bool = False             # operand already DMA-staged in TCM (via tl.load)
    # ADR-0062 §D2: lazy tl.load attaches a LoadFuture here. None for
    # handles that have no in-flight DMA (constants, math outputs, etc.).
    # Consumer ops call _await_pending() to yield on the future before
    # emitting their own command. Excluded from eq/hash/repr so handle
    # identity is unaffected by pending state.
    pending: object = field(default=None, compare=False, hash=False, repr=False)


@dataclass(frozen=True)
class CompletionHandle:
    """Opaque handle for a non-blocking composite command.

    Returned by tl.composite, consumed by tl.wait.
    """

    id: str


# ── PE Commands ───────────────────────────────────────────────────


@dataclass(frozen=True)
class DmaReadCmd:
    """DMA READ: HBM → PE_TCM. src_addr is VA (translated to PA by PE_DMA)."""

    handle: TensorHandle
    src_addr: int
    nbytes: int
    data_op: bool = True

    @property
    def logical_bytes(self) -> int:
        return 4 + 8 + 4 + 4          # framing + handle + src_addr + nbytes


@dataclass(frozen=True)
class DmaWriteCmd:
    """DMA WRITE: PE_TCM → HBM. dst_addr is VA (translated to PA by PE_DMA)."""

    handle: TensorHandle
    dst_addr: int
    nbytes: int
    data_op: bool = True

    @property
    def logical_bytes(self) -> int:
        return 4 + 8 + 4 + 4          # framing + handle + dst_addr + nbytes


@dataclass(frozen=True)
class GemmCmd:
    """GEMM engine command: matrix multiply on TCM data.

    out = a @ b, all operands in TCM.
    """

    a: TensorHandle
    b: TensorHandle
    out: TensorHandle
    m: int
    k: int
    n: int
    data_op: bool = True

    @property
    def logical_bytes(self) -> int:
        return 4 + 8 * 3 + 4 * 3      # framing + 3 handles + m/k/n scalars


@dataclass(frozen=True)
class MathCmd:
    """MATH engine command: unary/binary/reduction on TCM data.

    op: "exp", "log", "sqrt", "abs", "sigmoid", "cos", "sin",
        "add", "sub", "mul", "div", "where",
        "sum", "max", "min"
    """

    op: str
    inputs: tuple[TensorHandle, ...]
    out: TensorHandle
    axis: int | None = None          # for reductions
    data_op: bool = True

    @property
    def logical_bytes(self) -> int:
        return (
            4 + 1 + 1 + 8 * len(self.inputs) + 8   # framing+opcode+len+inputs+out
            + (4 if self.axis is not None else 0)
        )


@dataclass(frozen=True)
class CopyCmd:
    """TCM-to-TCM byte copy (ADR-0063 §D3.1).

    Emitted by ``tl.copy_to`` to persist a scoped result's bytes to an
    outside-``scratch_scope`` (persistent) address — the two-arena
    pattern for tiled flash attention. Runs on the vector engine;
    op_log classifies as ``op_kind="math"``, ``op_name="copy"``.
    """

    src: TensorHandle
    dst: TensorHandle
    nbytes: int
    data_op: bool = True

    @property
    def logical_bytes(self) -> int:
        return 4 + 8 * 2 + 4         # framing + src/dst handles + nbytes


@dataclass(frozen=True)
class CompositeCmd:
    """Composite command: tiled pipeline of DMA_READ + COMPUTE + DMA_WRITE.

    Non-blocking — submitted to PE_SCHEDULER which manages tile splitting
    and pipeline overlaps (ADR-0014 D3.2).

    Flat-ops shape (ADR-0065 D1): ``ops`` is an ordered list of OpSpecs.
    The GEMM op (if any, ≤1) drives the tile loop; preceding/following
    OpSpecs are placed by position + scope. ``rw_handles`` carries
    cross-composite hazard metadata (ADR-0065 D6.3); unused in P1.
    """

    completion: CompletionHandle
    ops: tuple[OpSpec, ...] = ()
    rw_handles: tuple["TensorHandle", ...] = ()
    data_op: bool = True

    @property
    def logical_bytes(self) -> int:
        """HW-logical byte size (ADR-0064 D2). Per-op summation, no dedup."""
        return (
            4                                       # framing
            + 1 + sum(op.logical_bytes for op in self.ops)
            + 1 + 8 * len(self.rw_handles)
        )


@dataclass(frozen=True)
class WaitCmd:
    """Wait for a specific composite or all pending composites."""

    handle: CompletionHandle | None = None  # None = wait all


@dataclass(frozen=True)
class PeCpuOverheadCmd:
    """PE_CPU scalar execution overhead (cycles)."""

    cycles: int


# Union type for all PE commands
PeCommand = (
    DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd | CopyCmd
    | CompositeCmd | WaitCmd | PeCpuOverheadCmd
)


@dataclass
class PeInternalTxn:
    """PE-internal message flowing PE_CPU → PE_SCHEDULER → engines.

    Carries a single PeCommand and a completion event.  PE_CPU creates one
    PeInternalTxn per command during the replay phase and sends it to
    PE_SCHEDULER, which routes it to the appropriate engine (PE_DMA,
    PE_GEMM, PE_MATH).  The engine signals ``done`` on completion.
    """

    command: PeCommand
    done: simpy.Event  # succeeded when the engine completes this command
    pe_prefix: str = ""  # e.g. "sip0.cube0.pe0" — needed by PE_DMA for path resolution
    result_data: dict[str, Any] = field(default_factory=dict)