kernbench2/docs/adr/ADR-0042-prog-tile-plan-generators.md

# ADR-0042: Tile Plan Generators — GEMM/Math Pipeline Plan Builders

## Status

Accepted (2026-05-20).

This ADR pins down `tiling.py` as a **plan-generator
module**, not a SimPy component.

ADR-0014 (PE Pipeline Execution Model) D6 (tile plan / self-routing) does not
specify the tile-plan generation algorithm itself; this ADR fills that gap.

## First action

When `generate_gemm_plan(M, K, N, tile_m, tile_k, tile_n, ..., pe_prefix,
a_pinned, b_pinned, epilogue_specs)` is called, the very first action is
**computing tile counts and constructing the PE-component ID strings**:

```
M_tiles = max(1, ceil(M / tile_m))
K_tiles = max(1, ceil(K / tile_k))
N_tiles = max(1, ceil(N / tile_n))
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"
```

In short, **the plan generator's first act is "compute ceiling tile counts
and assemble the four sub-component IDs for this PE once"**. No SimPy event
or environment is touched — this module is a pure function.

`generate_math_plan(M, N, tile_m, tile_n, ..., math_op, src_addr, dst_addr,
pe_prefix)` likewise begins by computing `M_tiles`, `N_tiles` and assembling
three component IDs (`dma_id`, `fetch_id`, `math_id`).

## Context

ADR-0014 D6 agreed that "PE_SCHEDULER, on receiving a CompositeCmd, generates
a TilePlan and feeds self-routing tile tokens". But the **concrete plan
generation algorithm** lives in `src/kernbench/components/builtin/tiling.py`,
which:

- Defines no component — it is a pair of **pure functions**
  (`generate_gemm_plan`, `generate_math_plan`).
- Does not depend on the SimPy environment, queues, op_log, or hooks.
- Returns a `PipelinePlan` (dataclass).

The original G4 analysis incorrectly described `tiling.py` as a component;
it is in fact a plan-builder helper consumed by PE_SCHEDULER. Pinning this
down in its own ADR (paired with ADR-0014 D6) prevents:

- Ambiguity over whether plan generation belongs to PE_SCHEDULER or a
  separate module.
- Inconsistent rationale for stage sequences (e.g., FETCH/STORE position)
  between GEMM and Math plans.
- Undocumented branching rationale for `a_pinned` / `b_pinned` /
  `epilogue_specs`.

## Decision

### D1. `tiling` is a pure plan-generator module, not a component

`components/builtin/tiling.py` defines no `ComponentBase` subclass. It exports
two module-level functions:

- `generate_gemm_plan(...) -> PipelinePlan`
- `generate_math_plan(...) -> PipelinePlan`

There is no `tiling` node in the topology graph. It lives in `builtin/`
because it is a direct helper for PE_SCHEDULER (ADR-0014 D6) and is
conceptually a PE_SCHEDULER internal utility.

### D2. GEMM plan stage sequence — `M → N → K` order

For each `(m, n, k)` tile (default — no operand pinning, no epilogue):

```
[DMA_READ(A)] → [DMA_READ(B)] → FETCH → GEMM
                                ↑
                                ↓
(last k tile only)              [MATH(output_tile)]* → STORE → DMA_WRITE
```

`k_tile` epilogue inserts a MATH stage immediately after GEMM on every
K-tile; `output_tile` epilogue inserts MATH stages once per `(m, n)` after
the final K-tile but before STORE/DMA_WRITE. The K-loop accumulator stays
in the register file across K-tiles — STORE/DMA_WRITE happens only when
`last_k`.

### D3. Operand pinning — `a_pinned` / `b_pinned`

If a caller passes `a_pinned=True`, **the A DMA_READ is omitted from every
(m, n, k) tile**. Semantically: the caller (e.g., `tl.composite`) has already
staged all of A in TCM via a prior `tl.load`, and signals so to the plan
generator.

The branch is made at plan time (not at runtime). Therefore the stage record
count in op_log changes deterministically with pinning, and sweep analyses
(e.g., gemm_sweep's stage record count) see this decision directly.

### D4. Epilogue scope — `k_tile` vs `output_tile`

`epilogue_specs` is an iterable of op-spec objects. Each op object is
expected to have:

- `op.kind: str` — math op name (e.g., `"dequant"`, `"bias"`, `"relu"`,
  `"scale"`). Placed into the stage's `params["op_kind"]`.
- `op.scope: Scope` — `Scope.K_TILE` or `Scope.OUTPUT_TILE` (`Scope` enum
  in `kernbench.common.pe_commands`).
- Op-specific extras (e.g., `bias`, `scale`, `factor`) — currently not used
  by the plan generator; consumed at runtime by PE_MATH.

The plan generator partitions by `getattr(o, "scope", None)`:

- `scope == Scope.K_TILE`: adds a MATH stage right after GEMM on every K-tile.
- `scope == Scope.OUTPUT_TILE`: adds MATH stages just before STORE on the
  last K-tile per `(m, n)`.

Ops with neither `scope` value (e.g., missing attribute) are **dropped
silently** — `getattr(..., None) == Scope.X` is False for both. Picking a
default (`output_tile`) is the **caller's responsibility** (e.g.,
`tl.composite`), not the plan generator's. This aligns with ADR-0014's
composite epilogue contract.

`Scope` is imported lazily inside the function to avoid the circular path
`pe_commands ← pe_types ← tiling`. This is intentional and not a refactor
target — keeping `tiling` free of compile-time `pe_commands` dependencies
preserves the module boundary (D1).

### D5. Math plan stage sequence — `M → N` order

For each `(m, n)` tile:

```
DMA_READ → FETCH → MATH → STORE → DMA_WRITE
```

There is no K dimension, so concepts like epilogue or accumulator residency
do not apply. PE_FETCH_STORE's register-file accounting follows the same
pattern as the GEMM plan.

### D6. Plans are data — no SimPy dependency

`PipelinePlan` is a dataclass in `pe_types.py` holding `tiles:
list[TilePlan]`. Each `TilePlan` holds `stages: tuple[Stage, ...]`. The plan
itself is near-immutable (only `Stage.params: dict` is mutable) and holds no
SimPy objects.

At runtime, PE_SCHEDULER consumes the plan's first stage, builds a `TileToken`,
and feeds it into the pipeline. The TileToken carries `plan: TilePlan`,
`stage_idx: int`, and a cached `params: dict`. Self-routing proceeds by
`TileToken.advance()` caching the next stage's `params` (ADR-0014 D6).

### D7. Plan generator contract — pure, deterministic, idempotent

Two calls with identical inputs return identical `PipelinePlan` instances
(including `TilePlan.stages` order). This contract aligns with ADR-0014 D6's
"deterministic tile dispatch order".

No side effects (no SimPy events, no file I/O, no global state) — tests can
call the generators directly without an environment object (some cases in
`tests/test_pe_pipeline.py` rely on this).

## Alternatives Considered

### A1. Make tiling a component (e.g., PE_PLANNER)

Rejected. Plan generation consumes no SimPy time — it is a pure decision
algorithm. Making it a component would (a) add unnecessary infrastructure
(inbox, resources), and (b) split PE_SCHEDULER's flow into "receive plan"
plus "feed tiles", inserting a meaningless hop.

### A2. Move plan generation into PE_SCHEDULER as methods

Rejected (currently). Module separation provides (1) testability and
(2) extensibility for additional plan algorithms (e.g., DTensor-aware) —
add a new function. If plan kinds proliferate enough to require explicit
dispatch, a future ADR can introduce a plan factory on PE_SCHEDULER.

### A3. Make plans fully immutable (frozen dataclass + tuple)

Partially adopted. `Stage` and `TilePlan` are dataclasses but not frozen,
because `Stage.params: dict` is populated at plan-generation time and read
at runtime (cached by TileToken on advance). Moving dict → frozendict pays
migration cost without enough benefit. Convention: do not mutate after
generation.

## Consequences

- `tiling.py` is documented as a plan-generator module, not a component —
  preempting future G4-style "this component lacks an ADR" analyses.
- The GEMM plan's stage sequence (D2) and pinning / epilogue branching
  (D3 / D4) are pinned, providing a clear interpretation basis for sweep
  analyses (e.g., `scripts/gemm_sweep.py`'s stage record counts).
- The plan generator's pure contract (D7) enables environment-free testing
  in line with ADR-0013 (verification strategy).
- Future plan kinds (DTensor-aware, K-major, ...) follow D1 / D6 / D7 as a
  baseline — just add a new function.