kernbench2/docs/adr/ADR-0036-dev-io-cpu-component-model.md

# ADR-0036: IO_CPU Component Model

## Status

Accepted

## Context

IO_CPU is the IO chiplet's host-facing endpoint inside the simulation
graph. PCIE_EP receives host messages from the runtime API and routes
them via the io_noc; for command-bearing requests (KernelLaunch,
MmuMap/Unmap) the io_noc forwards to IO_CPU, which:

- Fans out the request to per-cube M_CPUs.
- Aggregates per-cube responses into a single host-visible completion.
- For kernel launches, stamps a global `target_start_ns` barrier so
  every PE across every targeted cube begins kernel body execution at
  the same simulated time (ADR-0009 D5).

Memory R/W traffic bypasses IO_CPU per ADR-0015 D4 / ADR-0016 D3;
this component therefore handles only command-plane traffic in normal
operation.

This ADR documents the IO_CPU component implementation that realizes
those responsibilities.

## Decision

### D1. Role

IO_CPU is the host-facing endpoint of the IO chiplet. It has two
primary responsibilities:

1. **Multi-cube fan-out** — distribute KernelLaunchMsg / MmuMapMsg /
   MmuUnmapMsg to per-cube M_CPUs.
2. **Response aggregation** — collect per-cube ResponseMsg, signal
   parent `txn.done` when all targeted cubes have responded.

A third, narrower responsibility applies only to KernelLaunchMsg:
**`target_start_ns` global barrier stamping** (D3).

The component does **not**:

- Decide routing — paths are pre-computed by the router (ADR-0002).
- Decode tensor or kernel internals — those concerns belong to
  M_CPU / PE_CPU / engines.
- Handle PE-level fan-out — M_CPU fans out within a cube (ADR-0009 D3).
- Handle Memory R/W data path — those bypass IO_CPU per ADR-0015 D4
  and ADR-0016 D3 (Memory R/W resolution code in
  `_resolve_cube_targets` exists as a defensive fallback only).

Per invocation (`run()`): applies the configured `overhead_ns` once
per incoming Transaction (D8).

### D2. Forward path — multi-cube fan-out

When a non-response Transaction arrives, the worker:

1. Pays `overhead_ns` via `run()`.
2. Calls `_resolve_cube_targets` to derive the list of `(sip, cube)`
   targets from the request (D5).
3. For each target:
   - Resolves M_CPU node id via `ctx.resolver.find_m_cpu(sip, cube)`.
   - Resolves the path via `ctx.router.find_node_path(io_cpu, m_cpu)`.
   - Creates a per-cube sub-Transaction with `path` populated and
     forwards it to `path[1]` (the first hop on the io_noc).
4. Registers aggregation state: `_pending[request_id] = (expected,
   received=0, parent_done)`.

### D3. KernelLaunch `target_start_ns` global barrier (ADR-0009 D5)

IO_CPU is the canonical stamper for `target_start_ns`. When the
request is a `KernelLaunchMsg`, IO_CPU computes a single global
barrier covering every targeted PE across every targeted cube:

```text
for (sip, cube) in cube_targets:
    leg1 = compute_path_latency_ns(io_cpu → m_cpu(sip, cube), nbytes=0)
    for pe_id in target_pe_ids:
        leg2 = compute_path_latency_ns(m_cpu → pe_cpu(sip, cube, pe_id),
                                       nbytes=0)
        latency = leg1 + leg2 - io_overhead_ns - m_overhead_ns
        global_max = max(global_max, latency)

target_start_ns = env.now + global_max
```

The request is then replaced (via `dataclasses.replace`) so the
stamped value propagates through the fan-out.

Two overhead corrections:

- `io_overhead_ns` is subtracted because IO_CPU has already paid it
  in `run()` before this method runs.
- `m_overhead_ns` is subtracted once because it appears as the
  endpoint of leg1 *and* the start of leg2 in path latency, but
  M_CPU pays it only once at run time.

Every downstream PE_CPU yields until `target_start_ns` before
beginning kernel body execution; all PEs therefore start at the same
simulated time regardless of how long their individual dispatch path
took.

### D4. KernelLaunch sub-Transactions carry `nbytes=0`

Per-cube sub-Transactions for KernelLaunchMsg force `nbytes=0`,
overriding the parent `txn.nbytes`:

- Kernel launch is a control message; payload size is irrelevant at
  the data-fabric level.
- If `nbytes > 0`, every per-cube sub-txn occupies fabric BW on the
  io_noc's shared first hop. With 16 cubes this serializes fan-out,
  pushing far M_CPUs past `target_start_ns` and breaking the D3
  invariant.

Non-KernelLaunch sub-Transactions preserve `txn.nbytes` (only relevant
for the defensive Memory R/W fallback path, which carries actual
payload sizes).

### D5. Per-request-type cube target resolution

`_resolve_cube_targets` dispatches by request type:

| Request type | Source of `(sip, cube)` | `target_cubes="all"` semantics |
| --- | --- | --- |
| `MemoryWriteMsg` | `dst_sip`, `dst_cube` (or `PhysAddr.decode(dst_pa).die_id` fallback) | single cube derived from PA decode |
| `MemoryReadMsg` | `src_sip`, `src_cube` (or `PhysAddr.decode(src_pa).die_id` fallback) | single cube derived from PA decode |
| `KernelLaunchMsg` | tensor shards filtered by `shard.sip == my_sip` | every cube that owns a shard on this SIP |
| `MmuMapMsg` / `MmuUnmapMsg` | `target_cubes` list, filtered to this SIP | `range(cubes_per_sip)` from spec |

Each IO_CPU instance fans out only within its own SIP — `_my_sip()`
parses the SIP id from the node id (e.g., `sip0.io0.io_cpu` → 0).

The Memory R/W rows exist for defensive completeness; the engine's
normal path routes Memory R/W via `_process_memory_direct()` /
`find_memory_path()`, bypassing IO_CPU entirely (ADR-0015 D4 /
ADR-0016 D3).

### D6. Response aggregation

`_pending: dict[request_id → (expected, received, parent_done)]`:

- On dispatch: register `(len(cube_targets), 0, txn.done)`.
- `_worker` recognises responses by `is_response=True` and routes
  them to `_collect_response`.
- `_collect_response` increments `received`; when `received >=
  expected`, `parent_done.succeed()` is invoked and the entry is
  removed from `_pending`.

This is a simple per-request counter. There is no per-cube identity
tracking and no partial-failure handling — a missing response
indefinitely stalls the parent done. Production-style failure paths
are out of scope for the current simulator model.

### D7. `target_pe` resolution helper

`_resolve_pe_ids(target_pe)`:

- `int` → `[target_pe]`.
- `tuple[int, ...]` → `list(target_pe)`.
- `"all"` → `range(n_slices)`, where `n_slices` comes from cube
  `memory_map.hbm_slices_per_cube` (default 8).

Used in D3's barrier computation to enumerate every PE target per
cube.

### D8. Configurable `overhead_ns`

A single attribute drives per-instance latency:

| Site | impl name | overhead_ns |
| --- | --- | --- |
| IO chiplet `io_cpu` | `builtin.io_cpu` | 10.0 |

Applied once in `run()` per Transaction. Models command
interpretation + dispatch-decision time at IO_CPU.

## Consequences

### Positive

- Cross-cube and cross-SIP kernel launches share a single global
  barrier (D3 + D4) — no per-cube divergence in start time.
- nbytes=0 invariant keeps fan-out off the shared first-hop fabric
  BW, preserving the barrier's accuracy at scale (16 cubes).
- Response aggregation via a single counter → minimal state,
  deterministic ordering of completion.
- Per-SIP scoping (`_my_sip()`) keeps IO_CPUs in different SIPs
  cleanly independent.

### Negative

- No partial-failure semantics — a missing per-cube response
  indefinitely stalls the parent. Adequate for simulation but not
  suitable as a production-style endpoint.
- `_pending` is a regular dict; in-flight requests accumulate state.
  Acceptable for current benchmark workloads (few concurrent
  outstanding launches); unbounded in principle.
- The Memory R/W resolution branches in `_resolve_cube_targets` are
  dead code in the normal engine path. Kept defensively but invite
  drift if the bypass path ever changes.

## Links

- ADR-0002 (Routing distance — path computation)
- ADR-0009 D1 (Kernel launch is an endpoint request to IO_CPU)
- ADR-0009 D3 (M_CPU fans out within a cube; IO_CPU fans out across
  cubes)
- ADR-0009 D5 (target_start_ns canonical stamping at IO_CPU)
- ADR-0011 D-VA3 (MmuMapMsg routes through IO_CPU for cube fan-out)
- ADR-0012 (Host ↔ IO_CPU message schema)
- ADR-0015 D4 (Memory R/W bypasses IO_CPU; Kernel Launch via IO_CPU)
- ADR-0016 D1 (IO chiplet io_noc — IO_CPU attaches here)
- ADR-0016 D3 (Memory R/W path bypasses IO_CPU)
- ADR-0016 D4 (Kernel Launch path through IO_CPU for command
  interpretation)