Add 2D grid program_id semantics (ADR-0022)

tl.program_id(axis=0) returns local PE id within cube,
tl.program_id(axis=1) returns cube id. Enables cube-aware
sharding in benchmark kernels.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-09 16:49:56 -07:00
parent dc3fb02aed
commit ff2c677a9c
5 changed files with 138 additions and 3 deletions
+17
View File
@@ -196,6 +196,23 @@ def test_tl_program_id():
assert tl.num_programs(0) == 8
def test_tl_program_id_axis1():
"""axis=1 returns cube_id and num_cubes."""
tl = TLContext(pe_id=3, num_programs=8, cube_id=7, num_cubes=16)
assert tl.program_id(0) == 3
assert tl.program_id(1) == 7
assert tl.num_programs(0) == 8
assert tl.num_programs(1) == 16
def test_tl_program_id_global():
"""global_pid = cube_id * num_pes_per_cube + local_pe_id."""
pe_id, cube_id, num_pes = 5, 3, 8
tl = TLContext(pe_id=pe_id, num_programs=num_pes, cube_id=cube_id, num_cubes=16)
global_pid = tl.program_id(1) * tl.num_programs(0) + tl.program_id(0)
assert global_pid == cube_id * num_pes + pe_id
# ── 12. tl.arange, tl.zeros, tl.full ─────────────────────────────