ywkang
a796c1d2f7
Establish English as the canonical ADR language with Korean translations held in a parallel docs/adr-ko/ tree as derived artifacts (1:1 mirror). Promotion from adr-proposed/ to adr/ now writes English to adr/ and the Korean to adr-ko/; bidirectional sync rule documented in CLAUDE.md. - Migrate 30 ADRs in docs/adr/: 28 Korean-only translated to English, 2 bilingual pairs (ADR-0020, ADR-0023) consolidated (.en.md suffix dropped). ADR-0023 EN regenerated against KO source which had newer HW Realization Notes (D16-D23) section. - docs/adr-history/ left frozen by design (transitional state). - CLAUDE.md (Part 2): update ADR Lifecycle for 4-folder layout, mark docs/adr-ko/ as a Derived Artifact, add ADR Translation Discipline section covering bidirectional sync, conflict resolution (EN wins), and proposed-language freedom. - tools/verify_adr_lang_pairs.py: new verification tool checking pair completeness, filename mirroring, ADR-ID match, Status byte-equality. Pre-commit hook intentionally not added; run on demand or in CI. - tests/test_verify_adr_lang_pairs.py: 11 cases including CRLF/LF normalization, em-dash title separator, underscore-slug edge case. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
91 lines
3.2 KiB
Markdown
91 lines
3.2 KiB
Markdown
# 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.
|