commit - release 1

src/kernbench/policy/placement/dp.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from math import ceil
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class DPPolicy:
+    """Two-level data-parallel policy: cube-level + pe-level."""
+
+    cube: Literal["replicate", "shard_m", "shard_k"] = "replicate"
+    pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
+
+
+def resolve_dp_policy(
+    policy: DPPolicy,
+    *,
+    shape: tuple[int, int],
+    itemsize: int,
+    num_pe: int,
+    num_cubes: int = 1,
+) -> list[ShardSpec]:
+    """Resolve a DPPolicy into a list[ShardSpec] with two-level resolution.
+
+    Cube-level policy distributes across cubes, pe-level distributes within
+    each cube. ShardSpec.pe_index uses flat indexing: cube_id * num_pe + pe_id.
+    """
+    _PE_RESOLVERS = {
+        "replicate": replicate,
+        "column_wise": column_wise,
+        "row_wise": row_wise,
+    }
+    resolver = _PE_RESOLVERS.get(policy.pe)
+    if resolver is None:
+        raise ValueError(f"Unknown pe-level policy: {policy.pe}")
+
+    if num_cubes <= 1:
+        return resolver(shape=shape, itemsize=itemsize, num_pe=num_pe)
+
+    # Two-level resolution: cube-level → pe-level
+    M, K = shape
+    all_shards: list[ShardSpec] = []
+
+    for cube_id in range(num_cubes):
+        # Determine per-cube shape based on cube-level policy
+        if policy.cube == "replicate":
+            cube_shape = (M, K)
+            cube_offset = 0
+        elif policy.cube == "shard_m":
+            chunk_m = M // num_cubes
+            cube_shape = (chunk_m, K)
+            cube_offset = cube_id * chunk_m * K * itemsize
+        elif policy.cube == "shard_k":
+            chunk_k = K // num_cubes
+            cube_shape = (M, chunk_k)
+            cube_offset = cube_id * M * chunk_k * itemsize
+        else:
+            raise ValueError(f"Unknown cube-level policy: {policy.cube}")
+
+        # Resolve pe-level within this cube's shape
+        pe_shards = resolver(shape=cube_shape, itemsize=itemsize, num_pe=num_pe)
+
+        # Remap pe_index to flat index and adjust offset
+        for ps in pe_shards:
+            flat_idx = cube_id * num_pe + ps.pe_index
+            all_shards.append(ShardSpec(
+                pe_index=flat_idx,
+                offset_bytes=cube_offset + ps.offset_bytes,
+                nbytes=ps.nbytes,
+            ))
+
+    return all_shards
+
+
+@dataclass(frozen=True)
+class ShardSpec:
+    pe_index: int
+    offset_bytes: int
+    nbytes: int
+
+
+def column_wise(
+    *, shape: tuple[int, int], itemsize: int, num_pe: int,
+) -> list[ShardSpec]:
+    """Split K axis into num_pe equal parts. Each PE gets (M, K/P)."""
+    M, K = shape
+    chunk_k = K // num_pe
+    chunk_bytes = M * chunk_k * itemsize
+    shards = []
+    for i in range(num_pe):
+        shards.append(ShardSpec(
+            pe_index=i,
+            offset_bytes=i * chunk_bytes,
+            nbytes=chunk_bytes,
+        ))
+    return shards
+
+
+def row_wise(
+    *, shape: tuple[int, int], itemsize: int, num_pe: int,
+) -> list[ShardSpec]:
+    """Split M axis into num_pe equal parts. Each PE gets (M/P, K)."""
+    M, K = shape
+    chunk_m = M // num_pe
+    chunk_bytes = chunk_m * K * itemsize
+    shards = []
+    for i in range(num_pe):
+        shards.append(ShardSpec(
+            pe_index=i,
+            offset_bytes=i * chunk_bytes,
+            nbytes=chunk_bytes,
+        ))
+    return shards
+
+
+def replicate(
+    *, shape: tuple[int, int], itemsize: int, num_pe: int,
+) -> list[ShardSpec]:
+    """Full copy per PE. Each PE gets (M, K)."""
+    M, K = shape
+    full_bytes = M * K * itemsize
+    return [
+        ShardSpec(pe_index=i, offset_bytes=0, nbytes=full_bytes)
+        for i in range(num_pe)
+    ]
+
+
+def tiled_column_major(
+    *, shape: tuple[int, int], itemsize: int, num_pe: int,
+    tile_m: int, tile_k: int,
+) -> list[ShardSpec]:
+    """2D tiling, column-major order (K axis first), round-robin across PEs."""
+    M, K = shape
+    tiles_m = ceil(M / tile_m)
+    tiles_k = ceil(K / tile_k)
+    tile_bytes = tile_m * tile_k * itemsize
+    row_bytes = K * itemsize
+    shards = []
+    idx = 0
+    for mi in range(tiles_m):
+        for ki in range(tiles_k):
+            offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
+            shards.append(ShardSpec(
+                pe_index=idx % num_pe,
+                offset_bytes=offset,
+                nbytes=tile_bytes,
+            ))
+            idx += 1
+    return shards
+
+
+def tiled_row_major(
+    *, shape: tuple[int, int], itemsize: int, num_pe: int,
+    tile_m: int, tile_k: int,
+) -> list[ShardSpec]:
+    """2D tiling, row-major order (M axis first), round-robin across PEs."""
+    M, K = shape
+    tiles_m = ceil(M / tile_m)
+    tiles_k = ceil(K / tile_k)
+    tile_bytes = tile_m * tile_k * itemsize
+    row_bytes = K * itemsize
+    shards = []
+    idx = 0
+    for ki in range(tiles_k):
+        for mi in range(tiles_m):
+            offset = (mi * tile_m * row_bytes) + (ki * tile_k * itemsize)
+            shards.append(ShardSpec(
+                pe_index=idx % num_pe,
+                offset_bytes=offset,
+                nbytes=tile_bytes,
+            ))
+            idx += 1
+    return shards