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:

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)