kernbench2/docs/adr/ADR-0022-prog-program-id-2d-grid.md

# ADR-0022: 2D Grid program_id Semantics

## Status

Accepted

## Context

Triton kernels use `tl.program_id(axis)` to identify their position in a launch grid.
Our hardware has a 2-level hierarchy: **cubes** contain **PEs**.
The previous implementation ignored the `axis` parameter and always returned a flat PE index,
making it impossible for kernels to distinguish their cube-local position from their cube identity.

## Decision

Map `tl.program_id` and `tl.num_programs` to the 2D hardware grid:

| Call | Returns | Description |
|------|---------|-------------|
| `tl.program_id(axis=0)` | `local_pe_id` | PE index within cube |
| `tl.program_id(axis=1)` | `cube_id` | Cube index |
| `tl.num_programs(axis=0)` | `num_pes_per_cube` | PEs per cube |
| `tl.num_programs(axis=1)` | `num_cubes` | Total cubes |

Global PID is derived as:

```python
global_pid = tl.program_id(axis=1) * tl.num_programs(axis=0) + tl.program_id(axis=0)
```

### Axis mapping rationale

- **axis=0 = PE (innermost)**: PEs within a cube share HBM and communicate via local NOC mesh. This is the fast, tightly-coupled dimension — analogous to threads within a block.
- **axis=1 = Cube (outer)**: Cross-cube communication goes through UCIe with higher latency. This is the coarser scheduling dimension — analogous to blocks in a grid.

## Implementation

### TLContext (`triton_emu/tl_context.py`)

Added `cube_id` and `num_cubes` constructor parameters. `program_id()` and `num_programs()` dispatch on `axis`:

```python
def program_id(self, axis: int = 0) -> int:
    if axis == 1:
        return self._cube_id
    return self._pe_id

def num_programs(self, axis: int = 0) -> int:
    if axis == 1:
        return self._num_cubes
    return self._num_programs
```

### PE_CPU (`components/builtin/pe_cpu.py`)

- Extracts `num_cubes` from `ctx.spec["system"]["sips"]["cubes_per_sip"]`
- Passes `cube_id` (already available as `self._cube_idx`) and `num_cubes` to TLContext

### KernelRunner (`triton_emu/kernel_runner.py`)

- Receives `num_cubes` from PE_CPU
- Passes `cube_id` and `num_cubes` to TLContext in greenlet mode

## Backward Compatibility

- Existing code using `tl.program_id(0)` or `tl.program_id()` is unchanged — returns the same PE index as before.
- `cube_id` and `num_cubes` default to `0` and `1`, so callers that don't provide them (e.g. unit tests) continue to work.

## Usage Example

```python
def sharded_gemm_kernel(a_ptr, b_ptr, out_ptr, M, K, N, tl):
    local_pid = tl.program_id(axis=0)      # PE within cube
    cube_id   = tl.program_id(axis=1)      # which cube
    global_pid = cube_id * tl.num_programs(axis=0) + local_pid

    # Column-wise sharding across global PID
    n_per_pid = N // (tl.num_programs(axis=1) * tl.num_programs(axis=0))
    col_start = global_pid * n_per_pid

    a = tl.load(a_ptr, shape=(M, K), dtype="f16")
    b = tl.ref(b_ptr + col_start * K * 2, shape=(K, n_per_pid), dtype="f16")
    h = tl.composite(op="gemm", a=a, b=b, out_ptr=out_ptr + col_start * M * 2)
    tl.wait(h)
```

## Consequences

- Benchmarks can now express cube-aware sharding and addressing without hardcoding topology dimensions.
- Future axis=2 (SIP-level) can be added following the same pattern if needed.