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>

src/kernbench/components/builtin/pe_cpu.py

@@ -42,6 +42,9 @@ class PeCpuComponent(ComponentBase):
             self._cube_idx = int(parts[1].replace("cube", ""))
         except (IndexError, ValueError):
             self._cube_idx = 0
+        # num_cubes from spec (for tl.program_id(axis=1))
+        spec = ctx.spec if ctx else {}
+        self._num_cubes = spec.get("system", {}).get("sips", {}).get("cubes_per_sip", 1)
 
     def _find_shard(self, shards: tuple) -> Any:
         """Find shard matching this PE's (sip, cube, pe). Fallback to positional index."""
@@ -139,6 +142,7 @@ class PeCpuComponent(ComponentBase):
             pe_idx=self._pe_idx,
             sip_idx=self._sip_idx,
             cube_idx=self._cube_idx,
+            num_cubes=self._num_cubes,
             scheduler_id=scheduler_id,
             out_ports=self.out_ports,
             store=store,
@@ -155,7 +159,11 @@ class PeCpuComponent(ComponentBase):
         )
         from kernbench.triton_emu.tl_context import TLContext, run_kernel
 
-        tl = TLContext(pe_id=self._pe_idx, num_programs=num_programs, dispatch_cycles=0)
+        tl = TLContext(
+            pe_id=self._pe_idx, num_programs=num_programs,
+            cube_id=self._cube_idx, num_cubes=self._num_cubes,
+            dispatch_cycles=0,
+        )
         run_kernel(kernel_fn, tl, *kernel_args)
         commands = tl.commands