ADR-0026: DPPolicy intra-device only + ShardSpec structural coords

DPPolicy no longer carries a cross-SIP axis. SIP-level placement is
solely controlled by torch.ahbm.set_device(rank) (ADR-0024); DPPolicy
itself describes only the cube × PE layout within one SIP. ShardSpec
switches to structural (sip, cube, pe) coordinates; the flat pe_index
field/property is fully removed — silent drift between global-flat and
SIP-local interpretations was a foot-gun flagged by ADR-0024 D11.

Breaking API (explicit TypeError / AttributeError):
- DPPolicy(sip=...) / DPPolicy(num_sips=...) -> TypeError
- ShardSpec.pe_index -> AttributeError
- ShardSpec(pe_index=...) -> TypeError
- resolve_dp_policy now takes target_sip= (required), no num_sips.

Downstream migration:
- PE allocator dict keyed by (sip, cube, pe) tuples, in both
  _ensure_allocators and _free_tensor. deploy_tensor uses tuple lookup.
- _create_tensor passes target_sip=current_sip; post-hoc pe_index
  shifting removed entirely.
- launch._compute_local_shape drops the dp.sip branch.
- Internal resolvers (column_wise / row_wise / replicate / tiled_*)
  return _LocalPeShard (cube-local identifier) instead of ShardSpec —
  resolve_dp_policy lifts them to full structural coords.

Tests:
- New tests/test_adr0026_dppolicy_intra_device.py (12 tests) pins the
  contract end-to-end.
- test_sip_parallel.py rewritten: SIP composition now modeled as two
  resolve_dp_policy(target_sip=...) calls (ADR-0024 launcher style).
- Call-site migration: test_tensor, test_va_integration, test_va_offset,
  test_runtime_api_tensor, test_tl_recv_async, test_ccl_* and benches
  gemm_single_pe, gpt3_qkv, va_offset_verify, ccl_allreduce (legacy
  branch) all use intra-device DPPolicy and structural ShardSpec.

Result: 523 passed, 1 strict xfail (ring_default_ws — unchanged
ADR-0024 Phase B blocker; architectural fix deferred to ADR-0027).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-14 13:02:19 -07:00
parent 787409ced1
commit 357cab525b
20 changed files with 549 additions and 328 deletions
+10 -13
View File
@@ -32,29 +32,26 @@ def _derive_dp(spec: dict, world_size: int) -> DPPolicy:
"""Legacy DPPolicy for world_size > SIP count (rank = flat PE index). """Legacy DPPolicy for world_size > SIP count (rank = flat PE index).
Used only in the ccl.yaml-override path so the existing matrix tests Used only in the ccl.yaml-override path so the existing matrix tests
with explicit world_size (8, 16, 7 etc.) keep working. The new with explicit world_size (8, 16, 7 etc.) keep working. ADR-0026:
ADR-0024 TP path (rank = SIP) uses a per-rank DPPolicy inside the DPPolicy is intra-device only, so this legacy path now always stays
worker instead. within a single SIP and distributes the override world_size across
that SIP's cubes and PEs.
""" """
sips = int(spec["system"]["sips"]["count"])
cm = spec["sip"]["cube_mesh"]
pl = spec["cube"]["pe_layout"] pl = spec["cube"]["pe_layout"]
pes_per_cube = int(pl["pe_per_corner"]) * len(pl["corners"]) pes_per_cube = int(pl["pe_per_corner"]) * len(pl["corners"])
cm = spec["sip"]["cube_mesh"]
cubes_per_sip = int(cm["w"]) * int(cm["h"]) cubes_per_sip = int(cm["w"]) * int(cm["h"])
total = sips * cubes_per_sip * pes_per_cube
if world_size == total:
return DPPolicy(sip="column_wise", cube="column_wise", pe="column_wise")
if world_size <= pes_per_cube: if world_size <= pes_per_cube:
return DPPolicy( return DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_pes=world_size, num_cubes=1, num_pes=world_size,
) )
if world_size <= cubes_per_sip * pes_per_cube: if world_size <= cubes_per_sip * pes_per_cube:
return DPPolicy( return DPPolicy(
sip="replicate", cube="column_wise", pe="column_wise", cube="column_wise", pe="column_wise",
num_sips=1, num_cubes=world_size // pes_per_cube, num_cubes=world_size // pes_per_cube,
) )
return DPPolicy(sip="column_wise", cube="column_wise", pe="column_wise") return DPPolicy(cube="column_wise", pe="column_wise")
def worker(rank: int, world_size: int, torch) -> None: def worker(rank: int, world_size: int, torch) -> None:
+2 -2
View File
@@ -3,7 +3,7 @@
Full host-to-PE pipeline: Full host-to-PE pipeline:
Host → PCIE_EP → IO_CPU → M_CPU → PE_CPU → SchedulerV2 → PE_DMA → HBM Host → PCIE_EP → IO_CPU → M_CPU → PE_CPU → SchedulerV2 → PE_DMA → HBM
Single PE: num_sips=1, num_cubes=1, num_pes=1 via DPPolicy override. Single PE: num_cubes=1, num_pes=1 via DPPolicy override.
Both operands use tl.ref (HBM-resident); scheduler_v2 tiles and streams Both operands use tl.ref (HBM-resident); scheduler_v2 tiles and streams
per-tile DMA internally. per-tile DMA internally.
@@ -30,7 +30,7 @@ def _gemm_kernel(a_ptr, b_ptr, out_ptr, M, K, N, tl, DTYPE="f16"):
def run(torch): def run(torch):
"""Run the single-PE GEMM benchmark.""" """Run the single-PE GEMM benchmark."""
dp = DPPolicy(cube="replicate", pe="replicate", dp = DPPolicy(cube="replicate", pe="replicate",
num_sips=1, num_cubes=1, num_pes=1) num_cubes=1, num_pes=1)
a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a") a = torch.empty((M, K), dtype=DTYPE, dp=dp, name="a")
b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b") b = torch.empty((K, N), dtype=DTYPE, dp=dp, name="b")
+8 -4
View File
@@ -72,12 +72,16 @@ def run(torch):
K = GPT3_D_MODEL K = GPT3_D_MODEL
N = COLS_PER_PE N = COLS_PER_PE
# X: replicated across all PEs # ADR-0026: DPPolicy is intra-device only. For multi-SIP execution the
# ADR-0024 launcher calls this bench once per SIP (each worker via
# torch.ahbm.set_device(rank)); here the policy describes only the
# cube × PE layout within a single SIP.
# X: replicated across all PEs within the SIP
dp_replicate = DPPolicy(cube="replicate", pe="replicate", dp_replicate = DPPolicy(cube="replicate", pe="replicate",
num_sips=N_SIPS, num_cubes=N_CUBES, num_pes=N_PE_PER_CUBE) num_cubes=N_CUBES, num_pes=N_PE_PER_CUBE)
# W_Q/K/V, out_Q/K/V: column-wise sharded across all PEs # W_Q/K/V, out_Q/K/V: column-wise sharded across all PEs within the SIP
dp_sharded = DPPolicy(cube="column_wise", pe="column_wise", dp_sharded = DPPolicy(cube="column_wise", pe="column_wise",
num_sips=N_SIPS, num_cubes=N_CUBES, num_pes=N_PE_PER_CUBE) num_cubes=N_CUBES, num_pes=N_PE_PER_CUBE)
x = torch.empty((M, K), dtype=DTYPE, dp=dp_replicate, name="x") x = torch.empty((M, K), dtype=DTYPE, dp=dp_replicate, name="x")
wq = torch.empty((K, GPT3_D_MODEL), dtype=DTYPE, dp=dp_sharded, name="wq") wq = torch.empty((K, GPT3_D_MODEL), dtype=DTYPE, dp=dp_sharded, name="wq")
+2 -2
View File
@@ -1,7 +1,7 @@
"""VA offset verification benchmark. """VA offset verification benchmark.
Verifies that Triton-style base_ptr + pid * stride addressing works correctly Verifies that Triton-style base_ptr + pid * stride addressing works correctly
with full TP sharding (sip/cube/pe all column_wise). Each PE loads its own with intra-SIP TP sharding (cube/pe column_wise). Each PE loads its own
block from a sharded tensor and stores it back. block from a sharded tensor and stores it back.
The kernel uses standard Triton patterns: The kernel uses standard Triton patterns:
@@ -28,7 +28,7 @@ def _copy_kernel(src_ptr, dst_ptr, M, K, tl, DTYPE="f16"):
def run(torch): def run(torch):
"""Run the VA offset verification benchmark with full TP sharding.""" """Run the VA offset verification benchmark with full TP sharding."""
dp = DPPolicy(sip="column_wise", cube="column_wise", pe="column_wise") dp = DPPolicy(cube="column_wise", pe="column_wise")
src = torch.zeros((M, K), dtype=DTYPE, dp=dp, name="src") src = torch.zeros((M, K), dtype=DTYPE, dp=dp, name="src")
dst = torch.empty((M, K), dtype=DTYPE, dp=dp, name="dst") dst = torch.empty((M, K), dtype=DTYPE, dp=dp, name="dst")
+1 -1
View File
@@ -2,7 +2,7 @@
## Status ## Status
Proposed (Revision 4문서 일관성 + grep audit 구체화) Accepted (Revision 5Phase 2 landed 2026-04-14, 523 passed + 1 strict xfail)
## Context ## Context
+2 -2
View File
@@ -129,8 +129,8 @@ N_ELEM = 8
def worker(rank: int, world_size: int, torch) -> None: def worker(rank: int, world_size: int, torch) -> None:
"""Per-rank business logic — mirrors a real PyTorch DDP worker.""" """Per-rank business logic — mirrors a real PyTorch DDP worker."""
dp = DPPolicy( dp = DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_pes=world_size, num_cubes=1, num_pes=world_size,
) )
tensor = torch.zeros( tensor = torch.zeros(
(1, world_size * N_ELEM), dtype="f16", dp=dp, name="hello_in", (1, world_size * N_ELEM), dtype="f16", dp=dp, name="hello_in",
+2 -2
View File
@@ -114,8 +114,8 @@ def run(torch):
a = torch.zeros( a = torch.zeros(
(1, WORLD_SIZE * N_ELEM), dtype="f16", (1, WORLD_SIZE * N_ELEM), dtype="f16",
dp=DPPolicy( dp=DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_cubes=1,
), ),
name="hello_in", name="hello_in",
) )
+87 -67
View File
@@ -1,3 +1,14 @@
"""Data-parallel placement policy (ADR-0026: intra-device only).
``DPPolicy`` describes how a tensor is sharded *within a single SIP* across
that SIP's cubes and PEs. Crossing the SIP boundary is not a DPPolicy
concern: ADR-0024's ``torch.ahbm.set_device(rank)`` picks the SIP, and
Megatron-style TP (ADR-0027) expresses multi-SIP tensors when needed.
``ShardSpec`` is expressed in structural ``(sip, cube, pe)`` coordinates.
The former flat ``pe_index`` field/property is fully removed — callers
needing a flat integer key compute it explicitly at the call site.
"""
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
@@ -7,25 +18,58 @@ from typing import Literal
@dataclass(frozen=True) @dataclass(frozen=True)
class DPPolicy: class DPPolicy:
"""Three-level data-parallel policy: sip-level + cube-level + pe-level. """Intra-device (cube × PE) data-parallel policy.
Policies: SIP-level placement is controlled by ``torch.ahbm.set_device(rank)``
(ADR-0024). For tensors that must cross SIP boundaries, use
Megatron-style parallel layers (ADR-0027). DPPolicy itself never
crosses a SIP boundary.
Policies (per axis):
- "replicate": full copy at each unit - "replicate": full copy at each unit
- "column_wise": split K (column) axis across units - "column_wise": split K (column) axis across units
- "row_wise": split M (row) axis across units - "row_wise": split M (row) axis across units
Optional overrides (default None = use topology dimensions): Optional overrides (``None`` = use topology dimensions):
- num_pes: override PEs per cube (e.g., 1 for single-PE test) - num_pes: override PEs per cube
- num_cubes: override cubes per SIP (e.g., 1 for single-cube test) - num_cubes: override cubes per SIP
- num_sips: override SIP count
""" """
sip: Literal["replicate", "column_wise", "row_wise"] = "replicate"
cube: Literal["replicate", "column_wise", "row_wise"] = "replicate" cube: Literal["replicate", "column_wise", "row_wise"] = "replicate"
pe: Literal["replicate", "column_wise", "row_wise"] = "replicate" pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
num_pes: int | None = None num_pes: int | None = None
num_cubes: int | None = None num_cubes: int | None = None
num_sips: int | None = None
@dataclass(frozen=True)
class ShardSpec:
"""Structural shard placement — ``(sip, cube, pe)`` coord (ADR-0026).
Global-flat ``pe_index`` was removed: callers must use structural
coords directly. If a flat integer key is needed in a local context
(e.g. internal dict lookup), compute it explicitly at the call site
and do not expose it in any public API.
"""
sip: int
cube: int
pe: int
offset_bytes: int
nbytes: int
@dataclass(frozen=True)
class _LocalPeShard:
"""Internal — PE resolver's return type (ADR-0026 D3).
Holds a cube-local PE identifier (``local_pe``) plus the shard's
byte payload. Lifted into ``ShardSpec`` with full ``(sip, cube, pe)``
coordinates inside ``resolve_dp_policy``.
"""
local_pe: int
offset_bytes: int
nbytes: int
def _split_shape( def _split_shape(
@@ -52,14 +96,13 @@ def resolve_dp_policy(
itemsize: int, itemsize: int,
num_pe: int, num_pe: int,
num_cubes: int = 1, num_cubes: int = 1,
num_sips: int = 1, target_sip: int,
) -> list[ShardSpec]: ) -> list[ShardSpec]:
"""Resolve a DPPolicy into a list[ShardSpec] with three-level resolution. """Resolve a DPPolicy into a list[ShardSpec] on a single SIP.
SIP-level → cube-level → pe-level. Two-level resolution (cube × PE) within ``target_sip``. Each returned
num_cubes is cubes per SIP (not total). ``ShardSpec`` carries ``sip=target_sip`` and cube/pe local to the SIP.
ShardSpec.pe_index uses flat indexing: No SIP-level split — DPPolicy is intra-device only (ADR-0026).
sip_id * num_cubes * num_pe + cube_id * num_pe + pe_id
""" """
_PE_RESOLVERS = { _PE_RESOLVERS = {
"replicate": replicate, "replicate": replicate,
@@ -70,84 +113,61 @@ def resolve_dp_policy(
if resolver is None: if resolver is None:
raise ValueError(f"Unknown pe-level policy: {policy.pe}") raise ValueError(f"Unknown pe-level policy: {policy.pe}")
cubes_per_sip = num_cubes
all_shards: list[ShardSpec] = [] all_shards: list[ShardSpec] = []
# Level 1: SIP # Level 1: cube within SIP
sip_splits = _split_shape(policy.sip, shape, num_sips, itemsize) cube_splits = _split_shape(policy.cube, shape, num_cubes, itemsize)
for sip_id, (sip_shape, sip_offset) in enumerate(sip_splits):
# Level 2: Cube within SIP
cube_splits = _split_shape(policy.cube, sip_shape, cubes_per_sip, itemsize)
for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits): for cube_id, (cube_shape, cube_offset) in enumerate(cube_splits):
# Level 3: PE within cube # Level 2: PE within cube — resolver returns _LocalPeShard
pe_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe) local_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)
for ps in pe_shards: for ls in local_shards:
flat_idx = (
sip_id * cubes_per_sip * num_pe
+ cube_id * num_pe
+ ps.pe_index
)
all_shards.append(ShardSpec( all_shards.append(ShardSpec(
pe_index=flat_idx, sip=target_sip,
offset_bytes=sip_offset + cube_offset + ps.offset_bytes, cube=cube_id,
nbytes=ps.nbytes, pe=ls.local_pe,
offset_bytes=cube_offset + ls.offset_bytes,
nbytes=ls.nbytes,
)) ))
return all_shards return all_shards
@dataclass(frozen=True)
class ShardSpec:
pe_index: int
offset_bytes: int
nbytes: int
def column_wise( def column_wise(
*, shape: tuple[int, int], itemsize: int, num_pe: int, *, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[ShardSpec]: ) -> list[_LocalPeShard]:
"""Split K axis into num_pe equal parts. Each PE gets (M, K/P).""" """Split K axis into num_pe equal parts. Each PE gets (M, K/P)."""
M, K = shape M, K = shape
chunk_k = K // num_pe chunk_k = K // num_pe
chunk_bytes = M * chunk_k * itemsize chunk_bytes = M * chunk_k * itemsize
shards = [] return [
for i in range(num_pe): _LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
shards.append(ShardSpec( for i in range(num_pe)
pe_index=i, ]
offset_bytes=i * chunk_bytes,
nbytes=chunk_bytes,
))
return shards
def row_wise( def row_wise(
*, shape: tuple[int, int], itemsize: int, num_pe: int, *, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[ShardSpec]: ) -> list[_LocalPeShard]:
"""Split M axis into num_pe equal parts. Each PE gets (M/P, K).""" """Split M axis into num_pe equal parts. Each PE gets (M/P, K)."""
M, K = shape M, K = shape
chunk_m = M // num_pe chunk_m = M // num_pe
chunk_bytes = chunk_m * K * itemsize chunk_bytes = chunk_m * K * itemsize
shards = [] return [
for i in range(num_pe): _LocalPeShard(local_pe=i, offset_bytes=i * chunk_bytes, nbytes=chunk_bytes)
shards.append(ShardSpec( for i in range(num_pe)
pe_index=i, ]
offset_bytes=i * chunk_bytes,
nbytes=chunk_bytes,
))
return shards
def replicate( def replicate(
*, shape: tuple[int, int], itemsize: int, num_pe: int, *, shape: tuple[int, int], itemsize: int, num_pe: int,
) -> list[ShardSpec]: ) -> list[_LocalPeShard]:
"""Full copy per PE. Each PE gets (M, K).""" """Full copy per PE. Each PE gets (M, K)."""
M, K = shape M, K = shape
full_bytes = M * K * itemsize full_bytes = M * K * itemsize
return [ return [
ShardSpec(pe_index=i, offset_bytes=0, nbytes=full_bytes) _LocalPeShard(local_pe=i, offset_bytes=0, nbytes=full_bytes)
for i in range(num_pe) for i in range(num_pe)
] ]
@@ -155,20 +175,20 @@ def replicate(
def tiled_column_major( def tiled_column_major(
*, shape: tuple[int, int], itemsize: int, num_pe: int, *, shape: tuple[int, int], itemsize: int, num_pe: int,
tile_m: int, tile_k: int, tile_m: int, tile_k: int,
) -> list[ShardSpec]: ) -> list[_LocalPeShard]:
"""2D tiling, column-major order (K axis first), round-robin across PEs.""" """2D tiling, column-major order (K axis first), round-robin across PEs."""
M, K = shape M, K = shape
tiles_m = ceil(M / tile_m) tiles_m = ceil(M / tile_m)
tiles_k = ceil(K / tile_k) tiles_k = ceil(K / tile_k)
tile_bytes = tile_m * tile_k * itemsize tile_bytes = tile_m * tile_k * itemsize
row_bytes = K * itemsize row_bytes = K * itemsize
shards = [] shards: list[_LocalPeShard] = []
idx = 0 idx = 0
for mi in range(tiles_m): for mi in range(tiles_m):
for ki in range(tiles_k): for ki in range(tiles_k):
offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize) offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
shards.append(ShardSpec( shards.append(_LocalPeShard(
pe_index=idx % num_pe, local_pe=idx % num_pe,
offset_bytes=offset, offset_bytes=offset,
nbytes=tile_bytes, nbytes=tile_bytes,
)) ))
@@ -179,20 +199,20 @@ def tiled_column_major(
def tiled_row_major( def tiled_row_major(
*, shape: tuple[int, int], itemsize: int, num_pe: int, *, shape: tuple[int, int], itemsize: int, num_pe: int,
tile_m: int, tile_k: int, tile_m: int, tile_k: int,
) -> list[ShardSpec]: ) -> list[_LocalPeShard]:
"""2D tiling, row-major order (M axis first), round-robin across PEs.""" """2D tiling, row-major order (M axis first), round-robin across PEs."""
M, K = shape M, K = shape
tiles_m = ceil(M / tile_m) tiles_m = ceil(M / tile_m)
tiles_k = ceil(K / tile_k) tiles_k = ceil(K / tile_k)
tile_bytes = tile_m * tile_k * itemsize tile_bytes = tile_m * tile_k * itemsize
row_bytes = K * itemsize row_bytes = K * itemsize
shards = [] shards: list[_LocalPeShard] = []
idx = 0 idx = 0
for ki in range(tiles_k): for ki in range(tiles_k):
for mi in range(tiles_m): for mi in range(tiles_m):
offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize) offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
shards.append(ShardSpec( shards.append(_LocalPeShard(
pe_index=idx % num_pe, local_pe=idx % num_pe,
offset_bytes=offset, offset_bytes=offset,
nbytes=tile_bytes, nbytes=tile_bytes,
)) ))
+14 -45
View File
@@ -89,7 +89,7 @@ class RuntimeContext:
_handles: list[RequestHandle] = field(default_factory=list, init=False) _handles: list[RequestHandle] = field(default_factory=list, init=False)
_completed: set[RequestHandle] = field(default_factory=set, init=False) _completed: set[RequestHandle] = field(default_factory=set, init=False)
_allocators: dict[int, Any] = field(default_factory=dict, init=False) _allocators: dict[tuple[int, int, int], Any] = field(default_factory=dict, init=False)
_va_allocator: Any = field(default=None, init=False) _va_allocator: Any = field(default=None, init=False)
_tensor_counter: int = field(default=0, init=False) _tensor_counter: int = field(default=0, init=False)
_traces: list[dict] = field(default_factory=list, init=False) _traces: list[dict] = field(default_factory=list, init=False)
@@ -270,12 +270,7 @@ class RuntimeContext:
# Return PA space # Return PA space
if self._allocators: if self._allocators:
for shard in handle.shards: for shard in handle.shards:
flat_idx = ( alloc = self._allocators.get((shard.sip, shard.cube, shard.pe))
shard.sip * self._num_cubes * self._pes_per_cube
+ shard.cube * self._pes_per_cube
+ shard.pe
)
alloc = self._allocators.get(flat_idx)
if alloc is not None: if alloc is not None:
from kernbench.policy.address.phyaddr import PhysAddr from kernbench.policy.address.phyaddr import PhysAddr
alloc.free_hbm(PhysAddr.decode(shard.pa), shard.nbytes) alloc.free_hbm(PhysAddr.decode(shard.pa), shard.nbytes)
@@ -339,17 +334,15 @@ class RuntimeContext:
tcm_scheduler_reserved_bytes=4 * (1 << 20), tcm_scheduler_reserved_bytes=4 * (1 << 20),
sram_bytes_per_cube=32 * (1 << 20), sram_bytes_per_cube=32 * (1 << 20),
) )
# Create allocators scoped to target SIP(s) only # Create allocators scoped to target SIP(s) only.
# Flat index: sip_id * cubes_per_sip * pes_per_cube + cube_id * pes_per_cube + pe_id # ADR-0026 D5: dict key is the structural (sip, cube, pe) tuple.
self._pes_per_cube = pes_per_cube self._pes_per_cube = pes_per_cube
self._num_cubes = cubes_per_sip self._num_cubes = cubes_per_sip
self._num_sips = sip_count self._num_sips = sip_count
cubes_x_pes = cubes_per_sip * pes_per_cube
for sip_id in sip_range: for sip_id in sip_range:
for cube_id in range(cubes_per_sip): for cube_id in range(cubes_per_sip):
for pe_id in range(pes_per_cube): for pe_id in range(pes_per_cube):
flat_idx = sip_id * cubes_x_pes + cube_id * pes_per_cube + pe_id self._allocators[(sip_id, cube_id, pe_id)] = PEMemAllocator(
self._allocators[flat_idx] = PEMemAllocator(
rack_id=0, sip_id=sip_id, cube_id=cube_id, pe_id=pe_id, cfg=cfg, rack_id=0, sip_id=sip_id, cube_id=cube_id, pe_id=pe_id, cfg=cfg,
) )
@@ -436,44 +429,23 @@ class RuntimeContext:
# DPPolicy overrides take precedence over topology dimensions # DPPolicy overrides take precedence over topology dimensions
eff_num_pe = dp.num_pes if dp.num_pes is not None else self._pes_per_cube eff_num_pe = dp.num_pes if dp.num_pes is not None else self._pes_per_cube
eff_num_cubes = dp.num_cubes if dp.num_cubes is not None else self._num_cubes eff_num_cubes = dp.num_cubes if dp.num_cubes is not None else self._num_cubes
# ADR-0024 D11: if torch.ahbm.set_device(r) is active AND DPPolicy # ADR-0026 D4: resolve structural coords directly at resolve time.
# leaves the SIP dimension at its default (replicate + no num_sips # ``torch.ahbm.set_device(rank)`` (ADR-0024 D10) selects the target
# override), scope the tensor to SIP r only. # SIP; if unset, fall back to SIP 0 for single-driver compatibility.
# NOTE: this path uses post-hoc pe_index shifting as a temporary
# measure; ADR-0026 replaces it with structural (sip, cube, pe)
# coords in ShardSpec.
current_sip = ( current_sip = (
self.ahbm.current_device() if hasattr(self, "ahbm") else None self.ahbm.current_device() if hasattr(self, "ahbm") else None
) )
scope_to_current_sip = ( if current_sip is None:
current_sip is not None current_sip = 0
and dp.sip == "replicate"
and dp.num_sips is None
)
if scope_to_current_sip:
eff_num_sips = 1
else:
eff_num_sips = (
dp.num_sips if dp.num_sips is not None else self._num_sips
)
placement = resolve_dp_policy( placement = resolve_dp_policy(
dp, shape=shape_2d, itemsize=itemsize, dp, shape=shape_2d, itemsize=itemsize,
num_pe=eff_num_pe, num_cubes=eff_num_cubes, num_pe=eff_num_pe, num_cubes=eff_num_cubes,
num_sips=eff_num_sips, target_sip=int(current_sip),
) )
if scope_to_current_sip:
from kernbench.policy.placement.dp import ShardSpec as _SS
sip_stride = self._num_cubes * self._pes_per_cube
offset = int(current_sip) * sip_stride
placement = [
_SS(pe_index=s.pe_index + offset,
offset_bytes=s.offset_bytes, nbytes=s.nbytes)
for s in placement
]
# Infer target_pe from placement using local (within-cube) PE IDs. # Infer target_pe from placement using local (within-cube) PE IDs.
# This ensures M_CPU only fans out to PEs that own shards, not all PEs. # This ensures M_CPU only fans out to PEs that own shards, not all PEs.
local_pe_ids = sorted({s.pe_index % eff_num_pe for s in placement}) local_pe_ids = sorted({s.pe for s in placement})
if len(local_pe_ids) == 1: if len(local_pe_ids) == 1:
target_pe: int | tuple[int, ...] | str = local_pe_ids[0] target_pe: int | tuple[int, ...] | str = local_pe_ids[0]
elif len(local_pe_ids) == eff_num_pe and eff_num_pe == self._pes_per_cube: elif len(local_pe_ids) == eff_num_pe and eff_num_pe == self._pes_per_cube:
@@ -669,11 +641,8 @@ class RuntimeContext:
dp = t._dp_metadata.dp_policy if t._dp_metadata else None dp = t._dp_metadata.dp_policy if t._dp_metadata else None
if dp is None: if dp is None:
return t.shape return t.shape
if dp.sip != "replicate": # ADR-0026: DPPolicy no longer crosses SIP boundaries; cube + PE
if dp.sip == "column_wise": # are the only axes that shrink the local shape.
K = K // self._num_sips
elif dp.sip == "row_wise":
M = M // self._num_sips
if dp.cube != "replicate": if dp.cube != "replicate":
if dp.cube == "column_wise": if dp.cube == "column_wise":
K = K // self._num_cubes K = K // self._num_cubes
+7 -6
View File
@@ -72,7 +72,7 @@ def deploy_tensor(
shape: tuple[int, ...], shape: tuple[int, ...],
dtype: str, dtype: str,
placement: list[ShardSpec], placement: list[ShardSpec],
allocators: dict[int, PEMemAllocator], allocators: dict[tuple[int, int, int], PEMemAllocator],
mem_kind: Literal["hbm", "tcm"] = "hbm", mem_kind: Literal["hbm", "tcm"] = "hbm",
va_allocator=None, va_allocator=None,
) -> TensorHandle: ) -> TensorHandle:
@@ -86,15 +86,15 @@ def deploy_tensor(
shards: list[TensorShard] = [] shards: list[TensorShard] = []
for spec in placement: for spec in placement:
alloc = allocators[spec.pe_index] alloc = allocators[(spec.sip, spec.cube, spec.pe)]
if mem_kind == "hbm": if mem_kind == "hbm":
pa = alloc.alloc_hbm(spec.nbytes) pa = alloc.alloc_hbm(spec.nbytes)
else: else:
pa = alloc.alloc_tcm(spec.nbytes) pa = alloc.alloc_tcm(spec.nbytes)
shards.append(TensorShard( shards.append(TensorShard(
sip=alloc._sip_id, sip=spec.sip,
cube=alloc._cube_id, cube=spec.cube,
pe=alloc._pe_id, pe=spec.pe,
pa=pa.encode(), pa=pa.encode(),
nbytes=spec.nbytes, nbytes=spec.nbytes,
offset_bytes=spec.offset_bytes, offset_bytes=spec.offset_bytes,
@@ -394,7 +394,8 @@ class Tensor:
) -> Tensor: ) -> Tensor:
"""Set DP placement metadata (like torch.Tensor.to()).""" """Set DP placement metadata (like torch.Tensor.to())."""
if placement is None: if placement is None:
placement = [ShardSpec(pe_index=0, offset_bytes=0, nbytes=self.nbytes)] placement = [ShardSpec(sip=0, cube=0, pe=0,
offset_bytes=0, nbytes=self.nbytes)]
self._dp_metadata = DPMetadata( self._dp_metadata = DPMetadata(
placement=placement, dp_policy=dp_policy, placement=placement, dp_policy=dp_policy,
sip=sip, cube=cube, target_pe=target_pe, sip=sip, cube=cube, target_pe=target_pe,
+239
View File
@@ -0,0 +1,239 @@
"""ADR-0026 Phase 1 tests: DPPolicy intra-device only + ShardSpec structural.
These tests encode the contract from ADR-0026:
- DPPolicy no longer accepts ``sip`` or ``num_sips`` kwargs (TypeError).
- ShardSpec carries structural ``(sip, cube, pe)`` coordinates; the old flat
``pe_index`` field/property is fully removed (AttributeError).
- ``resolve_dp_policy(..., target_sip=N)`` stamps every returned ShardSpec
with ``sip=N``; cube and pe fields are local.
- ``RuntimeContext._allocators`` is keyed by ``(sip, cube, pe)`` tuples.
Phase 1: production code is unchanged → these tests SHOULD FAIL until the
Phase 2 diff lands. Phase 2 makes all of them pass.
"""
from __future__ import annotations
import pytest
from kernbench.policy.address.allocator import AddressConfig, PEMemAllocator
from kernbench.policy.placement.dp import DPPolicy, ShardSpec, resolve_dp_policy
from kernbench.runtime_api.tensor import deploy_tensor
# ── D1: DPPolicy no longer accepts sip / num_sips ─────────────────────
def test_dppolicy_rejects_sip_kwarg():
"""DPPolicy(sip=...) must raise TypeError after field removal."""
with pytest.raises(TypeError):
DPPolicy(sip="column_wise", cube="replicate", pe="replicate")
def test_dppolicy_rejects_num_sips_kwarg():
"""DPPolicy(num_sips=...) must raise TypeError after field removal."""
with pytest.raises(TypeError):
DPPolicy(cube="replicate", pe="replicate", num_sips=2)
def test_dppolicy_accepts_only_intra_device_fields():
"""Intra-device fields still work: cube, pe, num_cubes, num_pes."""
dp = DPPolicy(cube="column_wise", pe="column_wise",
num_cubes=2, num_pes=4)
assert dp.cube == "column_wise"
assert dp.pe == "column_wise"
assert dp.num_cubes == 2
assert dp.num_pes == 4
# No sip / num_sips attributes — even reading them must fail.
assert not hasattr(dp, "sip"), "DPPolicy.sip must be removed"
assert not hasattr(dp, "num_sips"), "DPPolicy.num_sips must be removed"
# ── D2: ShardSpec structural coords, no pe_index ──────────────────────
def test_shardspec_has_structural_coords():
"""ShardSpec constructs from (sip, cube, pe, offset_bytes, nbytes)."""
s = ShardSpec(sip=1, cube=2, pe=3, offset_bytes=128, nbytes=64)
assert s.sip == 1
assert s.cube == 2
assert s.pe == 3
assert s.offset_bytes == 128
assert s.nbytes == 64
def test_shardspec_has_no_pe_index_attr():
"""Flat pe_index must be fully removed — no field, no property."""
s = ShardSpec(sip=0, cube=0, pe=0, offset_bytes=0, nbytes=8)
with pytest.raises(AttributeError):
_ = s.pe_index # noqa: F841
def test_shardspec_rejects_pe_index_kwarg():
"""ShardSpec(pe_index=...) must raise TypeError."""
with pytest.raises(TypeError):
ShardSpec(pe_index=0, offset_bytes=0, nbytes=8) # type: ignore[call-arg]
# ── D3: resolve_dp_policy(target_sip=...) structural semantics ────────
def test_resolve_dp_policy_target_sip_stamps_shards():
"""All returned shards must carry sip == target_sip; cube/pe local."""
dp = DPPolicy(cube="column_wise", pe="column_wise")
shards = resolve_dp_policy(
dp, shape=(4, 32), itemsize=2,
num_pe=4, num_cubes=2, target_sip=1,
)
assert len(shards) == 2 * 4
assert all(s.sip == 1 for s in shards)
assert all(0 <= s.cube < 2 for s in shards)
assert all(0 <= s.pe < 4 for s in shards)
def test_resolve_dp_policy_target_sip_differ_only_in_sip():
"""Same policy + dims on two SIPs → shards identical except .sip."""
dp = DPPolicy(cube="replicate", pe="column_wise")
shards_0 = resolve_dp_policy(
dp, shape=(4, 32), itemsize=2,
num_pe=4, num_cubes=1, target_sip=0,
)
shards_1 = resolve_dp_policy(
dp, shape=(4, 32), itemsize=2,
num_pe=4, num_cubes=1, target_sip=1,
)
assert len(shards_0) == len(shards_1)
for a, b in zip(shards_0, shards_1):
assert a.sip == 0 and b.sip == 1
assert a.cube == b.cube
assert a.pe == b.pe
assert a.offset_bytes == b.offset_bytes
assert a.nbytes == b.nbytes
def test_resolve_dp_policy_no_num_sips_param():
"""resolve_dp_policy must not accept num_sips anymore.
Post-Phase-2 signature drops ``num_sips`` (DPPolicy no longer crosses
SIP boundaries) and adds required ``target_sip``. Calling with
``num_sips=...`` must raise TypeError (unexpected keyword argument).
"""
dp = DPPolicy(cube="replicate", pe="replicate")
with pytest.raises(TypeError, match="num_sips"):
resolve_dp_policy(
dp, shape=(4, 8), itemsize=2,
num_pe=1, num_cubes=1, num_sips=2, # type: ignore[call-arg]
)
# ── D5: Allocator dict keyed by (sip, cube, pe) tuples ────────────────
_MB = 1 << 20
_GB = 1 << 30
_CFG = AddressConfig(
sip_count=2,
cubes_per_sip=2,
pes_per_cube=4,
hbm_bytes_per_cube=_GB,
hbm_slices_per_cube=4,
tcm_bytes_per_pe=_MB,
tcm_scheduler_reserved_bytes=0,
sram_bytes_per_cube=_MB,
)
def _make_tuple_allocators(
num_sips: int = 1, num_cubes: int = 1, num_pe: int = 4,
) -> dict[tuple[int, int, int], PEMemAllocator]:
return {
(s, c, p): PEMemAllocator(
rack_id=0, sip_id=s, cube_id=c, pe_id=p, cfg=_CFG,
)
for s in range(num_sips)
for c in range(num_cubes)
for p in range(num_pe)
}
def test_deploy_tensor_uses_tuple_lookup():
"""deploy_tensor(allocators={(sip,cube,pe): alloc, ...}) succeeds."""
dp = DPPolicy(cube="replicate", pe="column_wise")
placement = resolve_dp_policy(
dp, shape=(4, 16), itemsize=2,
num_pe=4, num_cubes=1, target_sip=0,
)
allocators = _make_tuple_allocators(num_sips=1, num_cubes=1, num_pe=4)
handle = deploy_tensor(
name="t", shape=(4, 16), dtype="f16",
placement=placement, allocators=allocators,
)
assert len(handle.shards) == 4
# Each shard's TensorShard carries structural coords; those coords
# must match the shard's ShardSpec (sip, cube, pe).
for spec, shard in zip(placement, handle.shards):
assert shard.sip == spec.sip
assert shard.cube == spec.cube
assert shard.pe == spec.pe
def test_runtime_context_allocator_keys_are_tuples(topology):
"""After ctx tensor op, ctx._allocators keys are (sip, cube, pe) tuples.
Ensures D5 migration landed (allocator population + lookup).
"""
from kernbench.runtime_api.context import RuntimeContext
from kernbench.runtime_api.types import DeviceSelector
from kernbench.sim_engine.engine import GraphEngine
engine = GraphEngine(topology.topology_obj, enable_data=True)
ctx = RuntimeContext(
engine=engine,
target_device=DeviceSelector("sip:0"),
correlation_id="test_adr0026_tuple_keys",
spec=topology.topology_obj.spec,
)
dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
_ = ctx.zeros((1, 16), dtype="f16", dp=dp)
assert ctx._allocators, "allocators dict should be populated"
keys = list(ctx._allocators.keys())
assert all(isinstance(k, tuple) and len(k) == 3 for k in keys), (
f"_allocators keys must be (sip, cube, pe) tuples; got {keys[:5]}"
)
# ── D4 (via regression): no SIP-crossing tensor without set_device ────
def test_create_tensor_on_target_sip_via_set_device(topology):
"""torch.ahbm.set_device(1) + DPPolicy(cube=replicate, pe=replicate)
→ all shards land on SIP 1 structurally (no post-hoc shifting needed)."""
from kernbench.runtime_api.context import RuntimeContext
from kernbench.runtime_api.types import DeviceSelector
from kernbench.sim_engine.engine import GraphEngine
# Skip the test if topology has only 1 SIP (nothing to verify).
n_sips = int(
topology.topology_obj.spec.get("system", {})
.get("sips", {}).get("count", 1)
)
if n_sips < 2:
pytest.skip("topology has <2 SIPs; set_device(1) not meaningful")
engine = GraphEngine(topology.topology_obj, enable_data=True)
ctx = RuntimeContext(
engine=engine,
target_device=DeviceSelector("sip:1"),
correlation_id="test_adr0026_set_device",
spec=topology.topology_obj.spec,
)
ctx.ahbm.set_device(1)
dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
t = ctx.zeros((1, 16), dtype="f16", dp=dp)
assert t._handle is not None
assert all(s.sip == 1 for s in t._handle.shards), (
f"expected all shards on SIP 1; got {[s.sip for s in t._handle.shards]}"
)
+2 -2
View File
@@ -108,8 +108,8 @@ def test_deadlock_detection_recv_without_send():
(1, 8 * 8), (1, 8 * 8),
dtype="f16", dtype="f16",
dp=DPPolicy( dp=DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_cubes=1,
), ),
name="dl_in", name="dl_in",
) )
+2 -2
View File
@@ -51,8 +51,8 @@ def test_hello_send_via_simpy_runner():
a = torch.zeros( a = torch.zeros(
(1, world_size * n_elem), dtype="f16", (1, world_size * n_elem), dtype="f16",
dp=DPPolicy( dp=DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_cubes=1,
), ),
name="hello_in", name="hello_in",
) )
+2 -2
View File
@@ -48,8 +48,8 @@ def test_recv_copy_to_dst_via_simpy_runner():
(1, 8 * 8), (1, 8 * 8),
dtype="f16", dtype="f16",
dp=DPPolicy( dp=DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_cubes=1,
), ),
name="copy_in", name="copy_in",
) )
+16 -16
View File
@@ -48,8 +48,8 @@ def test_from_numpy_creates_host_tensor():
assert h._handle is None assert h._handle is None
# Submit a no-op so run_bench has at least one handle. # Submit a no-op so run_bench has at least one handle.
torch.zeros((1, 8), dtype="f16", torch.zeros((1, 8), dtype="f16",
dp=DPPolicy(sip="replicate", cube="replicate", pe="replicate", dp=DPPolicy(cube="replicate", pe="replicate",
num_sips=1, num_cubes=1, num_pes=1), num_cubes=1, num_pes=1),
name="dummy") name="dummy")
_run_with(body) _run_with(body)
@@ -63,8 +63,8 @@ def test_copy_and_numpy_single_pe():
a single-PE (no real sharding) tensor.""" a single-PE (no real sharding) tensor."""
def body(torch): def body(torch):
dp = DPPolicy(sip="replicate", cube="replicate", pe="replicate", dp = DPPolicy(cube="replicate", pe="replicate",
num_sips=1, num_cubes=1, num_pes=1) num_cubes=1, num_pes=1)
t = torch.zeros((1, 16), dtype="f16", dp=dp, name="t") t = torch.zeros((1, 16), dtype="f16", dp=dp, name="t")
src = np.arange(16, dtype=np.float16).reshape(1, 16) src = np.arange(16, dtype=np.float16).reshape(1, 16)
t.copy_(torch.from_numpy(src)) t.copy_(torch.from_numpy(src))
@@ -83,8 +83,8 @@ def test_copy_and_numpy_multi_pe_column_wise():
def body(torch): def body(torch):
n_pe = 8 n_pe = 8
dp = DPPolicy(sip="replicate", cube="replicate", pe="column_wise", dp = DPPolicy(cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_pes=n_pe) num_cubes=1, num_pes=n_pe)
t = torch.zeros((1, n_pe * 4), dtype="f16", dp=dp, name="t") t = torch.zeros((1, n_pe * 4), dtype="f16", dp=dp, name="t")
src = np.arange(n_pe * 4, dtype=np.float16).reshape(1, n_pe * 4) src = np.arange(n_pe * 4, dtype=np.float16).reshape(1, n_pe * 4)
t.copy_(torch.from_numpy(src)) t.copy_(torch.from_numpy(src))
@@ -107,8 +107,8 @@ def test_copy_and_numpy_multi_cube():
n_pe_per_cube = 8 n_pe_per_cube = 8
n_cubes = 2 n_cubes = 2
total = n_cubes * n_pe_per_cube # 16 total = n_cubes * n_pe_per_cube # 16
dp = DPPolicy(sip="replicate", cube="column_wise", pe="column_wise", dp = DPPolicy(cube="column_wise", pe="column_wise",
num_sips=1, num_cubes=n_cubes) num_cubes=n_cubes)
t = torch.zeros((1, total * 4), dtype="f16", dp=dp, name="t") t = torch.zeros((1, total * 4), dtype="f16", dp=dp, name="t")
src = np.arange(total * 4, dtype=np.float16).reshape(1, total * 4) src = np.arange(total * 4, dtype=np.float16).reshape(1, total * 4)
t.copy_(torch.from_numpy(src)) t.copy_(torch.from_numpy(src))
@@ -126,8 +126,8 @@ def test_copy_shape_mismatch_raises():
"""copy_ with mismatched shapes raises ValueError.""" """copy_ with mismatched shapes raises ValueError."""
def body(torch): def body(torch):
dp = DPPolicy(sip="replicate", cube="replicate", pe="replicate", dp = DPPolicy(cube="replicate", pe="replicate",
num_sips=1, num_cubes=1, num_pes=1) num_cubes=1, num_pes=1)
t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t") t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t")
src = np.zeros((1, 16), dtype=np.float16) src = np.zeros((1, 16), dtype=np.float16)
with pytest.raises(ValueError, match="copy_ shape mismatch"): with pytest.raises(ValueError, match="copy_ shape mismatch"):
@@ -143,8 +143,8 @@ def test_setitem_getitem_single_pe():
"""Scalar and slice assignment on a single-PE tensor round-trips.""" """Scalar and slice assignment on a single-PE tensor round-trips."""
def body(torch): def body(torch):
dp = DPPolicy(sip="replicate", cube="replicate", pe="replicate", dp = DPPolicy(cube="replicate", pe="replicate",
num_sips=1, num_cubes=1, num_pes=1) num_cubes=1, num_pes=1)
t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t") t = torch.zeros((1, 8), dtype="f16", dp=dp, name="t")
# Scalar broadcast # Scalar broadcast
@@ -169,8 +169,8 @@ def test_setitem_getitem_multi_pe_shard_aligned():
def body(torch): def body(torch):
n_pe = 8 n_pe = 8
n_elem = 4 # per shard n_elem = 4 # per shard
dp = DPPolicy(sip="replicate", cube="replicate", pe="column_wise", dp = DPPolicy(cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_pes=n_pe) num_cubes=1, num_pes=n_pe)
t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t") t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t")
# Write each shard with its rank value # Write each shard with its rank value
@@ -197,8 +197,8 @@ def test_setitem_cross_shard_raises():
def body(torch): def body(torch):
n_pe = 4 n_pe = 4
n_elem = 4 n_elem = 4
dp = DPPolicy(sip="replicate", cube="replicate", pe="column_wise", dp = DPPolicy(cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_pes=n_pe) num_cubes=1, num_pes=n_pe)
t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t") t = torch.zeros((1, n_pe * n_elem), dtype="f16", dp=dp, name="t")
with pytest.raises(NotImplementedError, match="spans multiple shards"): with pytest.raises(NotImplementedError, match="spans multiple shards"):
t[0, 2:6] = 1.0 # crosses shard 0 (0:4) and shard 1 (4:8) t[0, 2:6] = 1.0 # crosses shard 0 (0:4) and shard 1 (4:8)
+90 -127
View File
@@ -1,157 +1,120 @@
"""Tests for SIP-level tensor parallelism. """Tests for SIP-level tensor parallelism — ADR-0026 structural model.
Validates: DPPolicy no longer carries a ``sip`` axis (ADR-0026 D1). SIP placement is
SP1. DPPolicy accepts sip field (default "replicate", backward compat) now expressed structurally: each call to ``resolve_dp_policy(target_sip=N)``
SP2. sip="column_wise": tensor K-axis split across SIPs, each SIP gets K//num_sips emits shards pinned to SIP N. Multi-SIP parallelism is composed by calling
SP3. sip="row_wise": tensor M-axis split across SIPs the resolver once per SIP (typically driven by the ADR-0024 launcher, one
SP4. 3-level resolve: sip × cube × pe produces correct flat indices and offsets worker greenlet per rank, each worker using ``torch.ahbm.set_device(rank)``).
SP5. sip="replicate": all SIPs get full copy (existing behavior)
SP6. PE_CPU sets num_programs from shard count per cube Covered here:
SP7. End-to-end: TP kernel with sip="column_wise" completes on multi-SIP topology SP1. ``target_sip`` stamps every shard.
SP2. Two-SIP placement: union of two resolver calls covers the whole
tensor K-axis when the combined bench treats them as column-split.
SP3. Same for row-wise.
SP4. Cube + PE sharding within a SIP remains correct across SIPs.
SP5. PE_CPU num_programs contract (unchanged by ADR-0026).
""" """
import pytest from __future__ import annotations
from pathlib import Path
from kernbench.policy.placement.dp import DPPolicy, ShardSpec, resolve_dp_policy from kernbench.policy.placement.dp import DPPolicy, resolve_dp_policy
# ── SP1. DPPolicy sip field ────────────────────────────────────────── # ── SP1. target_sip stamps shards ────────────────────────────────────
def test_dp_policy_sip_default_replicate(): def test_target_sip_stamps_all_shards():
"""DPPolicy without sip= defaults to 'replicate'."""
dp = DPPolicy(cube="replicate", pe="column_wise") dp = DPPolicy(cube="replicate", pe="column_wise")
assert dp.sip == "replicate"
def test_dp_policy_sip_column_wise():
"""DPPolicy accepts sip='column_wise'."""
dp = DPPolicy(sip="column_wise", cube="replicate", pe="column_wise")
assert dp.sip == "column_wise"
# ── SP2. sip="column_wise" ──────────────────────────────────────────────
def test_sip_column_wise_splits_across_sips():
"""sip='column_wise' with 2 SIPs: each SIP gets K//2 columns."""
dp = DPPolicy(sip="column_wise", cube="replicate", pe="column_wise")
shards = resolve_dp_policy( shards = resolve_dp_policy(
dp, shape=(128, 256), itemsize=2, dp, shape=(128, 256), itemsize=2,
num_pe=8, num_cubes=1, num_sips=2, num_pe=8, num_cubes=1, target_sip=3,
) )
# 2 SIPs × 1 cube × 8 PEs = 16 shards assert all(s.sip == 3 for s in shards)
assert len(shards) == 16 assert all(0 <= s.pe < 8 for s in shards)
assert all(s.cube == 0 for s in shards)
# SIP0 shards: first half of K (0 to K//2)
# SIP1 shards: second half of K (K//2 to K)
total_bytes = 128 * 256 * 2 # 64KB
sip0_shards = [s for s in shards if s.pe_index < 8]
sip1_shards = [s for s in shards if s.pe_index >= 8]
# SIP0 offsets start at 0
assert sip0_shards[0].offset_bytes == 0
# SIP1 offsets start at half
assert sip1_shards[0].offset_bytes == total_bytes // 2
# Total coverage
assert sum(s.nbytes for s in sip0_shards) == total_bytes // 2
assert sum(s.nbytes for s in sip1_shards) == total_bytes // 2
# ── SP3. sip="row_wise" ────────────────────────────────────────────── # ── SP2. column-wise placement composed across two SIPs ─────────────
def test_sip_row_wise_splits_across_sips(): def test_compose_two_sips_column_wise_covers_tensor():
"""sip='row_wise' with 2 SIPs: each SIP gets M//2 rows.""" """Bench splits K-axis across 2 SIPs by calling resolve twice and
dp = DPPolicy(sip="row_wise", cube="replicate", pe="column_wise") giving each SIP half of the tensor (half-shape + offset). Shards
shards = resolve_dp_policy( from both SIPs together cover the whole K axis."""
full_shape = (128, 256)
itemsize = 2
# Per-SIP half-shape (K split across SIPs).
half_shape = (128, 128)
dp = DPPolicy(cube="replicate", pe="column_wise")
shards_sip0 = resolve_dp_policy(
dp, shape=half_shape, itemsize=itemsize,
num_pe=8, num_cubes=1, target_sip=0,
)
shards_sip1 = resolve_dp_policy(
dp, shape=half_shape, itemsize=itemsize,
num_pe=8, num_cubes=1, target_sip=1,
)
total_bytes = full_shape[0] * full_shape[1] * itemsize
sip0_bytes = sum(s.nbytes for s in shards_sip0)
sip1_bytes = sum(s.nbytes for s in shards_sip1)
assert sip0_bytes + sip1_bytes == total_bytes
assert all(s.sip == 0 for s in shards_sip0)
assert all(s.sip == 1 for s in shards_sip1)
# ── SP3. row-wise placement composed across two SIPs ────────────────
def test_compose_two_sips_row_wise_covers_tensor():
full_shape = (128, 256)
itemsize = 2
half_shape = (64, 256) # per-SIP half of M
dp = DPPolicy(cube="replicate", pe="column_wise")
shards_sip0 = resolve_dp_policy(
dp, shape=half_shape, itemsize=itemsize,
num_pe=8, num_cubes=1, target_sip=0,
)
shards_sip1 = resolve_dp_policy(
dp, shape=half_shape, itemsize=itemsize,
num_pe=8, num_cubes=1, target_sip=1,
)
total_bytes = full_shape[0] * full_shape[1] * itemsize
assert sum(s.nbytes for s in shards_sip0) + sum(s.nbytes for s in shards_sip1) == total_bytes
# ── SP4. cube × PE sharding is independent per SIP ──────────────────
def test_cube_pe_sharding_independent_per_sip():
"""Intra-SIP cube + PE layout matches across SIPs; only sip field differs."""
dp = DPPolicy(cube="column_wise", pe="column_wise")
s0 = resolve_dp_policy(
dp, shape=(128, 256), itemsize=2, dp, shape=(128, 256), itemsize=2,
num_pe=8, num_cubes=1, num_sips=2, num_pe=4, num_cubes=2, target_sip=0,
) )
assert len(shards) == 16 s1 = resolve_dp_policy(
sip0_shards = [s for s in shards if s.pe_index < 8]
sip1_shards = [s for s in shards if s.pe_index >= 8]
# SIP0: rows 0..63, SIP1: rows 64..127
total_bytes = 128 * 256 * 2
assert sip0_shards[0].offset_bytes == 0
assert sip1_shards[0].offset_bytes == total_bytes // 2
# ── SP4. 3-level resolve ─────────────────────────────────────────────
def test_3level_resolve_flat_index():
"""3-level: sip × cube × pe produces correct flat indices."""
dp = DPPolicy(sip="column_wise", cube="replicate", pe="column_wise")
shards = resolve_dp_policy(
dp, shape=(128, 256), itemsize=2, dp, shape=(128, 256), itemsize=2,
num_pe=8, num_cubes=2, num_sips=2, num_pe=4, num_cubes=2, target_sip=1,
) )
# 2 SIPs × 2 cubes × 8 PEs = 32 shards assert len(s0) == len(s1) == 2 * 4
assert len(shards) == 32 for a, b in zip(s0, s1):
assert a.sip == 0 and b.sip == 1
# Flat index: sip_id * cubes_per_sip * num_pe + cube_id * num_pe + pe_id assert (a.cube, a.pe, a.offset_bytes, a.nbytes) == (
indices = [s.pe_index for s in shards] b.cube, b.pe, b.offset_bytes, b.nbytes
# SIP0: 0..15, SIP1: 16..31
assert min(indices) == 0
assert max(indices) == 31
assert len(set(indices)) == 32 # all unique
def test_3level_offsets_cover_full_tensor():
"""3-level sharding covers the entire tensor with no gaps."""
dp = DPPolicy(sip="column_wise", cube="replicate", pe="column_wise")
shards = resolve_dp_policy(
dp, shape=(128, 256), itemsize=2,
num_pe=4, num_cubes=1, num_sips=2,
) )
# 2 SIPs × 1 cube × 4 PEs = 8 shards
# sip="column_wise": K=128 per SIP, pe="column_wise": 32 cols per PE
total = 128 * 256 * 2
# For non-replicate, total shard bytes == tensor bytes
# (replicate within cube means cube shards overlap, but sip shards don't)
sip0_bytes = sum(s.nbytes for s in shards if s.pe_index < 4)
sip1_bytes = sum(s.nbytes for s in shards if s.pe_index >= 4)
assert sip0_bytes + sip1_bytes == total
# ── SP5. sip="replicate" backward compat ───────────────────────────── # ── SP5. PE_CPU num_programs (contract unchanged) ───────────────────
def test_sip_replicate_backward_compat():
"""sip='replicate' produces same result as before (2-level)."""
dp_old = DPPolicy(cube="replicate", pe="column_wise")
dp_new = DPPolicy(sip="replicate", cube="replicate", pe="column_wise")
shards_old = resolve_dp_policy(
dp_old, shape=(128, 256), itemsize=2,
num_pe=8, num_cubes=2, num_sips=2,
)
shards_new = resolve_dp_policy(
dp_new, shape=(128, 256), itemsize=2,
num_pe=8, num_cubes=2, num_sips=2,
)
assert len(shards_old) == len(shards_new)
for a, b in zip(shards_old, shards_new):
assert a.pe_index == b.pe_index
assert a.offset_bytes == b.offset_bytes
assert a.nbytes == b.nbytes
# ── SP6. PE_CPU num_programs ──────────────────────────────────────────
def test_pe_cpu_sets_num_programs(): def test_pe_cpu_sets_num_programs():
"""PE_CPU should create TLContext with num_programs = PEs per cube.""" """TLContext reports num_programs from its initializer — used by PE_CPU
# This test validates the interface contract. when it launches a kernel on behalf of its shards."""
# After implementation, PE_CPU should derive num_programs from the
# number of PE shards in the kernel launch's target cube.
from kernbench.triton_emu.tl_context import TLContext from kernbench.triton_emu.tl_context import TLContext
# With 8 PEs per cube, num_programs should be 8
tl = TLContext(pe_id=3, num_programs=8) tl = TLContext(pe_id=3, num_programs=8)
assert tl.program_id(0) == 3 assert tl.program_id(0) == 3
assert tl.num_programs(0) == 8 assert tl.num_programs(0) == 8
+23 -17
View File
@@ -2,11 +2,13 @@ import pytest
from kernbench.policy.address.allocator import AddressConfig, AllocationError, PEMemAllocator from kernbench.policy.address.allocator import AddressConfig, AllocationError, PEMemAllocator
from kernbench.policy.placement.dp import ( from kernbench.policy.placement.dp import (
DPPolicy,
ShardSpec, ShardSpec,
column_wise, column_wise,
tiled_column_major,
replicate, replicate,
resolve_dp_policy,
row_wise, row_wise,
tiled_column_major,
tiled_row_major, tiled_row_major,
) )
from kernbench.runtime_api.kernel import ( from kernbench.runtime_api.kernel import (
@@ -40,9 +42,9 @@ _CFG = AddressConfig(
) )
def _make_allocators(num_pe: int = 8) -> dict[int, PEMemAllocator]: def _make_allocators(num_pe: int = 8) -> dict[tuple[int, int, int], PEMemAllocator]:
return { return {
i: PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=i, cfg=_CFG) (0, 0, i): PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=i, cfg=_CFG)
for i in range(num_pe) for i in range(num_pe)
} }
@@ -133,7 +135,7 @@ def test_column_wise_placement():
assert len(shards) == 8 assert len(shards) == 8
expected_nbytes = 1024 * 64 * 2 # 128 KB expected_nbytes = 1024 * 64 * 2 # 128 KB
for i, s in enumerate(shards): for i, s in enumerate(shards):
assert s.pe_index == i assert s.local_pe == i
assert s.nbytes == expected_nbytes assert s.nbytes == expected_nbytes
# offsets are contiguous # offsets are contiguous
assert shards[0].offset_bytes == 0 assert shards[0].offset_bytes == 0
@@ -151,7 +153,7 @@ def test_row_wise_placement():
assert len(shards) == 8 assert len(shards) == 8
expected_nbytes = 128 * 512 * 2 # 128 KB expected_nbytes = 128 * 512 * 2 # 128 KB
for i, s in enumerate(shards): for i, s in enumerate(shards):
assert s.pe_index == i assert s.local_pe == i
assert s.nbytes == expected_nbytes assert s.nbytes == expected_nbytes
assert shards[0].offset_bytes == 0 assert shards[0].offset_bytes == 0
assert sum(s.nbytes for s in shards) == 1024 * 512 * 2 assert sum(s.nbytes for s in shards) == 1024 * 512 * 2
@@ -166,7 +168,7 @@ def test_replicate_placement():
assert len(shards) == 8 assert len(shards) == 8
full_nbytes = 1024 * 512 * 2 # 1 MB full_nbytes = 1024 * 512 * 2 # 1 MB
for i, s in enumerate(shards): for i, s in enumerate(shards):
assert s.pe_index == i assert s.local_pe == i
assert s.nbytes == full_nbytes assert s.nbytes == full_nbytes
assert s.offset_bytes == 0 # each is a full copy assert s.offset_bytes == 0 # each is a full copy
@@ -188,10 +190,10 @@ def test_tiled_column_major():
# tile (m=0,k=0) → PE0, tile (m=0,k=1) → PE1, ..., (m=0,k=3) → PE3 # tile (m=0,k=0) → PE0, tile (m=0,k=1) → PE1, ..., (m=0,k=3) → PE3
# tile (m=1,k=0) → PE4, tile (m=1,k=1) → PE5, ..., (m=1,k=3) → PE7 # tile (m=1,k=0) → PE4, tile (m=1,k=1) → PE5, ..., (m=1,k=3) → PE7
# tile (m=2,k=0) → PE0, ... # tile (m=2,k=0) → PE0, ...
assert shards[0].pe_index == 0 assert shards[0].local_pe == 0
assert shards[1].pe_index == 1 assert shards[1].local_pe == 1
assert shards[7].pe_index == 7 assert shards[7].local_pe == 7
assert shards[8].pe_index == 0 # wraps around assert shards[8].local_pe == 0 # wraps around
# total coverage # total coverage
assert sum(s.nbytes for s in shards) == 1024 * 512 * 2 assert sum(s.nbytes for s in shards) == 1024 * 512 * 2
@@ -212,10 +214,10 @@ def test_tiled_row_major():
# tile (m=0,k=0) → PE0, tile (m=1,k=0) → PE1, ..., (m=3,k=0) → PE3 # tile (m=0,k=0) → PE0, tile (m=1,k=0) → PE1, ..., (m=3,k=0) → PE3
# tile (m=0,k=1) → PE4, tile (m=1,k=1) → PE5, ..., (m=3,k=1) → PE7 # tile (m=0,k=1) → PE4, tile (m=1,k=1) → PE5, ..., (m=3,k=1) → PE7
# tile (m=0,k=2) → PE0, ... # tile (m=0,k=2) → PE0, ...
assert shards[0].pe_index == 0 assert shards[0].local_pe == 0
assert shards[1].pe_index == 1 assert shards[1].local_pe == 1
assert shards[7].pe_index == 7 assert shards[7].local_pe == 7
assert shards[8].pe_index == 0 # wraps around assert shards[8].local_pe == 0 # wraps around
# total coverage # total coverage
assert sum(s.nbytes for s in shards) == 1024 * 512 * 2 assert sum(s.nbytes for s in shards) == 1024 * 512 * 2
@@ -226,7 +228,11 @@ def test_tiled_row_major():
def test_deploy_tensor_hbm(): def test_deploy_tensor_hbm():
"""Deploy with column_wise placement → TensorHandle with valid PA shards.""" """Deploy with column_wise placement → TensorHandle with valid PA shards."""
allocs = _make_allocators() allocs = _make_allocators()
placement = column_wise(shape=(1024, 512), itemsize=2, num_pe=8) placement = resolve_dp_policy(
DPPolicy(cube="replicate", pe="column_wise"),
shape=(1024, 512), itemsize=2,
num_pe=8, num_cubes=1, target_sip=0,
)
th = deploy_tensor( th = deploy_tensor(
name="W", name="W",
shape=(1024, 512), shape=(1024, 512),
@@ -253,7 +259,7 @@ def test_deploy_tensor_hbm():
def test_deploy_tensor_tcm(): def test_deploy_tensor_tcm():
"""Deploy with TCM → uses pe_tcm_addr allocation.""" """Deploy with TCM → uses pe_tcm_addr allocation."""
allocs = _make_allocators() allocs = _make_allocators()
placement = [ShardSpec(pe_index=0, offset_bytes=0, nbytes=256)] placement = [ShardSpec(sip=0, cube=0, pe=0, offset_bytes=0, nbytes=256)]
th = deploy_tensor( th = deploy_tensor(
name="small", name="small",
shape=(128,), shape=(128,),
@@ -271,7 +277,7 @@ def test_deploy_tensor_overflow():
"""Allocation exceeding PE HBM capacity raises AllocationError.""" """Allocation exceeding PE HBM capacity raises AllocationError."""
allocs = _make_allocators() allocs = _make_allocators()
# 6 GB per PE slice, try to allocate 7 GB # 6 GB per PE slice, try to allocate 7 GB
big_shard = ShardSpec(pe_index=0, offset_bytes=0, nbytes=7 * _GB) big_shard = ShardSpec(sip=0, cube=0, pe=0, offset_bytes=0, nbytes=7 * _GB)
with pytest.raises(AllocationError): with pytest.raises(AllocationError):
deploy_tensor( deploy_tensor(
name="toobig", name="toobig",
+2 -2
View File
@@ -75,8 +75,8 @@ def test_recv_async_simpy_runner():
(1, 8 * 8), (1, 8 * 8),
dtype="f16", dtype="f16",
dp=DPPolicy( dp=DPPolicy(
sip="replicate", cube="replicate", pe="column_wise", cube="replicate", pe="column_wise",
num_sips=1, num_cubes=1, num_cubes=1,
), ),
name="async_in", name="async_in",
) )
+19 -7
View File
@@ -12,7 +12,7 @@ import pytest
from kernbench.policy.address.allocator import AddressConfig, PEMemAllocator from kernbench.policy.address.allocator import AddressConfig, PEMemAllocator
from kernbench.policy.address.pe_mmu import PeMMU from kernbench.policy.address.pe_mmu import PeMMU
from kernbench.policy.address.va_allocator import VirtualAllocator from kernbench.policy.address.va_allocator import VirtualAllocator
from kernbench.policy.placement.dp import column_wise, ShardSpec from kernbench.policy.placement.dp import DPPolicy, ShardSpec, resolve_dp_policy
from kernbench.runtime_api.tensor import ( from kernbench.runtime_api.tensor import (
TensorHandle, TensorHandle,
TensorShard, TensorShard,
@@ -37,9 +37,9 @@ _CFG = AddressConfig(
) )
def _make_allocators(num_pe: int = 8) -> dict[int, PEMemAllocator]: def _make_allocators(num_pe: int = 8) -> dict[tuple[int, int, int], PEMemAllocator]:
return { return {
i: PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=i, cfg=_CFG) (0, 0, i): PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=i, cfg=_CFG)
for i in range(num_pe) for i in range(num_pe)
} }
@@ -88,7 +88,11 @@ def test_deploy_tensor_assigns_va_base():
"""deploy_tensor with VA allocator assigns va_base to TensorHandle.""" """deploy_tensor with VA allocator assigns va_base to TensorHandle."""
allocs = _make_allocators() allocs = _make_allocators()
va_alloc = _make_va_allocator() va_alloc = _make_va_allocator()
placement = column_wise(shape=(1024, 512), itemsize=2, num_pe=8) placement = resolve_dp_policy(
DPPolicy(cube="replicate", pe="column_wise"),
shape=(1024, 512), itemsize=2,
num_pe=8, num_cubes=1, target_sip=0,
)
th = deploy_tensor( th = deploy_tensor(
name="W", name="W",
@@ -107,7 +111,11 @@ def test_deploy_tensor_va_covers_all_shards():
"""VA allocation covers the entire tensor; each shard is at va_base + offset.""" """VA allocation covers the entire tensor; each shard is at va_base + offset."""
allocs = _make_allocators() allocs = _make_allocators()
va_alloc = _make_va_allocator() va_alloc = _make_va_allocator()
placement = column_wise(shape=(1024, 512), itemsize=2, num_pe=8) placement = resolve_dp_policy(
DPPolicy(cube="replicate", pe="column_wise"),
shape=(1024, 512), itemsize=2,
num_pe=8, num_cubes=1, target_sip=0,
)
th = deploy_tensor( th = deploy_tensor(
name="W", name="W",
@@ -128,7 +136,11 @@ def test_deploy_tensor_does_not_install_mmu_mappings():
allocs = _make_allocators() allocs = _make_allocators()
va_alloc = _make_va_allocator() va_alloc = _make_va_allocator()
mmus = _make_mmus() mmus = _make_mmus()
placement = column_wise(shape=(1024, 512), itemsize=2, num_pe=8) placement = resolve_dp_policy(
DPPolicy(cube="replicate", pe="column_wise"),
shape=(1024, 512), itemsize=2,
num_pe=8, num_cubes=1, target_sip=0,
)
deploy_tensor( deploy_tensor(
name="W", name="W",
@@ -153,7 +165,7 @@ def test_tensor_va_property():
allocs = _make_allocators(1) allocs = _make_allocators(1)
va_alloc = _make_va_allocator() va_alloc = _make_va_allocator()
placement = [ShardSpec(pe_index=0, offset_bytes=0, nbytes=4096)] placement = [ShardSpec(sip=0, cube=0, pe=0, offset_bytes=0, nbytes=4096)]
t = Tensor(shape=(2048,), dtype="f16", name="test") t = Tensor(shape=(2048,), dtype="f16", name="test")
t._handle = deploy_tensor( t._handle = deploy_tensor(
+15 -5
View File
@@ -20,7 +20,7 @@ from kernbench.policy.address.allocator import AddressConfig, PEMemAllocator
from kernbench.policy.address.pe_mmu import PeMMU from kernbench.policy.address.pe_mmu import PeMMU
from kernbench.policy.address.phyaddr import PhysAddr from kernbench.policy.address.phyaddr import PhysAddr
from kernbench.policy.address.va_allocator import VirtualAllocator from kernbench.policy.address.va_allocator import VirtualAllocator
from kernbench.policy.placement.dp import DPPolicy, column_wise from kernbench.policy.placement.dp import DPPolicy, resolve_dp_policy
from kernbench.runtime_api.tensor import deploy_tensor from kernbench.runtime_api.tensor import deploy_tensor
from kernbench.sim_engine.engine import GraphEngine from kernbench.sim_engine.engine import GraphEngine
from kernbench.runtime_api.context import RuntimeContext from kernbench.runtime_api.context import RuntimeContext
@@ -70,7 +70,7 @@ def _make_standalone(shape, num_pe=NUM_PE):
sram_bytes_per_cube=32 * _MB, sram_bytes_per_cube=32 * _MB,
) )
allocators = { allocators = {
i: PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=i, cfg=cfg) (0, 0, i): PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=i, cfg=cfg)
for i in range(num_pe) for i in range(num_pe)
} }
va_alloc = VirtualAllocator(va_base=0x1_0000_0000, va_size=64 * _GB, page_size=4096) va_alloc = VirtualAllocator(va_base=0x1_0000_0000, va_size=64 * _GB, page_size=4096)
@@ -110,7 +110,11 @@ def test_2d_va_translates_to_local_hbm():
cols_per_pe = K // NUM_PE cols_per_pe = K // NUM_PE
block_bytes = M * cols_per_pe * ELEM_BYTES block_bytes = M * cols_per_pe * ELEM_BYTES
placement = column_wise(shape=(M, K), itemsize=ELEM_BYTES, num_pe=NUM_PE) placement = resolve_dp_policy(
DPPolicy(cube="replicate", pe="column_wise"),
shape=(M, K), itemsize=ELEM_BYTES,
num_pe=NUM_PE, num_cubes=1, target_sip=0,
)
handle = deploy_tensor( handle = deploy_tensor(
name="src", shape=(M, K), dtype="fp16", name="src", shape=(M, K), dtype="fp16",
placement=placement, allocators=allocators, va_allocator=va_alloc, placement=placement, allocators=allocators, va_allocator=va_alloc,
@@ -178,7 +182,11 @@ def test_1d_va_translates_to_local_hbm():
elems_per_pe = N_1D // NUM_PE elems_per_pe = N_1D // NUM_PE
block_bytes = elems_per_pe * ELEM_BYTES block_bytes = elems_per_pe * ELEM_BYTES
placement = column_wise(shape=(1, N_1D), itemsize=ELEM_BYTES, num_pe=NUM_PE) placement = resolve_dp_policy(
DPPolicy(cube="replicate", pe="column_wise"),
shape=(1, N_1D), itemsize=ELEM_BYTES,
num_pe=NUM_PE, num_cubes=1, target_sip=0,
)
handle = deploy_tensor( handle = deploy_tensor(
name="src_1d", shape=(N_1D,), dtype="fp16", name="src_1d", shape=(N_1D,), dtype="fp16",
placement=placement, allocators=allocators, va_allocator=va_alloc, placement=placement, allocators=allocators, va_allocator=va_alloc,
@@ -207,7 +215,9 @@ def test_1d_e2e_completes():
correlation_id="vo6", spec=graph.spec, correlation_id="vo6", spec=graph.spec,
) )
dp = DPPolicy(sip="column_wise", cube="column_wise", pe="column_wise") # ADR-0026: DPPolicy is intra-device only; SIP scoping comes from the
# RuntimeContext's target_device. This 1D e2e runs on a single SIP.
dp = DPPolicy(cube="column_wise", pe="column_wise")
src = ctx.zeros((N_1D,), dtype=DTYPE, dp=dp, name="src_1d") src = ctx.zeros((N_1D,), dtype=DTYPE, dp=dp, name="src_1d")
dst = ctx.empty((N_1D,), dtype=DTYPE, dp=dp, name="dst_1d") dst = ctx.empty((N_1D,), dtype=DTYPE, dp=dp, name="dst_1d")