from __future__ import annotations
import itertools
from dataclasses import dataclass
from typing import Any
from modelexpress.refit.reshard.types import (
OpChain,
RecordedCopy,
UnsupportedReshard,
)
@dataclass
class Shard:
shard_offset: tuple shape: tuple session: str addr: int elsize: int digest: str | None = None
@dataclass
class PullSegment:
session: str
src_addr: int
param_name: str
dst_byte: int
nbytes: int
def _row_major_strides(shape) -> list:
strides = [1] * len(shape)
for d in range(len(shape) - 2, -1, -1):
strides[d] = strides[d + 1] * shape[d + 1]
return strides
def op_chain_to_box(op_chain: OpChain, global_shape) -> list:
box = [[0, int(g)] for g in global_shape]
for op_name, args, _kw in op_chain:
if op_name == "narrow":
dim, start, size = args
box[dim][0] += int(start)
box[dim][1] = box[dim][0] + int(size)
elif op_name == "contiguous":
pass elif op_name == "__getitem__":
key = args[0] if args else slice(None)
keys = key if isinstance(key, tuple) else (key,)
dim = 0
for k in keys:
if k is Ellipsis:
raise UnsupportedReshard("ellipsis indexing is not box-derivable")
if not isinstance(k, slice):
raise UnsupportedReshard(
f"integer index {k!r} collapses a dim; not box-derivable"
)
if k.step not in (None, 1):
raise UnsupportedReshard(
f"strided slice step={k.step} is not box-derivable"
)
lo, hi = box[dim]
extent = hi - lo
start = 0 if k.start is None else int(k.start)
stop = extent if k.stop is None else int(k.stop)
box[dim][0] = lo + start
box[dim][1] = lo + stop
dim += 1
else:
raise UnsupportedReshard(
f"op {op_name!r} is not box-derivable (transpose/reshape/permute/chunk)"
)
return [tuple(b) for b in box]
def _replay_view(op_chain: OpChain, global_shape) -> tuple:
import torch
src = torch.empty(tuple(int(g) for g in global_shape), device="meta")
t: Any = src
for op_name, args, kw in op_chain:
t = getattr(t, op_name)(*args, **dict(kw))
if not isinstance(t, torch.Tensor):
raise UnsupportedReshard(f"op-chain resolved to non-tensor {type(t).__name__}")
if t is not src and t._base is not src:
raise UnsupportedReshard("op-chain includes a copy (not a pure view)")
return (
int(t.storage_offset()),
tuple(int(s) for s in t.shape),
tuple(int(s) for s in t.stride()),
)
def _view_to_box_perm(
v_offset: int, v_shape: tuple, v_stride: tuple, global_shape
) -> tuple:
ndim = len(global_shape)
vdim = len(v_shape)
if vdim > ndim:
raise UnsupportedReshard(
f"rank increase {vdim}>{ndim} (unsqueeze) not box-derivable"
)
gstrides = _row_major_strides(global_shape)
stride_to_dim: dict = {}
for s in range(ndim):
stride_to_dim.setdefault(gstrides[s], s)
perm: list = [None] * vdim
for d in range(vdim):
if v_shape[d] == 1:
continue sdim = stride_to_dim.get(v_stride[d])
if sdim is None:
raise UnsupportedReshard(
f"non-permutation stride {v_stride[d]} (dim-merging reshape)"
)
perm[d] = sdim
used = {s for s in perm if s is not None}
remaining = [s for s in range(ndim) if s not in used]
ri = 0
for d in range(vdim):
if perm[d] is None:
perm[d] = remaining[ri]
ri += 1
if len(set(perm)) != vdim:
raise UnsupportedReshard("view strides are not a permutation of source dims")
box_lo = [(v_offset // gstrides[s]) % global_shape[s] for s in range(ndim)]
box = [[box_lo[s], box_lo[s] + 1] for s in range(ndim)]
for d in range(vdim):
box[perm[d]][1] = box[perm[d]][0] + v_shape[d]
if any(box[s][0] < 0 or box[s][1] > global_shape[s] for s in range(ndim)):
raise UnsupportedReshard("resolved box exceeds source bounds")
if sum(box[s][0] * gstrides[s] for s in range(ndim)) != v_offset:
raise UnsupportedReshard("view offset does not decompose to a box start")
return [tuple(b) for b in box], perm
def resolve_slice(copy: RecordedCopy, global_shape) -> tuple:
v_off, v_shape, v_stride = _replay_view(copy.op_chain, global_shape)
box, perm = _view_to_box_perm(v_off, v_shape, v_stride, global_shape)
if tuple(copy.dest_shape) != tuple(v_shape):
raise UnsupportedReshard(
f"{copy.param_name}: dest {tuple(copy.dest_shape)} "
f"!= resolved view {tuple(v_shape)}"
)
dest_stride_src = [0] * len(global_shape)
for d, s in enumerate(perm):
dest_stride_src[s] = copy.dest_stride[d]
return box, dest_stride_src
def intersect(box_a: list, box_b: list):
out = []
for (a0, a1), (b0, b1) in zip(box_a, box_b, strict=True):
lo, hi = max(a0, b0), min(a1, b1)
if lo >= hi:
return None
out.append((lo, hi))
return out
def paired_runs(overlap: list, src_origin, src_strides, dst_origin, dst_strides):
ndim = len(overlap)
sizes = [hi - lo for lo, hi in overlap]
src_start = [lo - o for (lo, _), o in zip(overlap, src_origin, strict=True)]
dst_start = [lo - o for (lo, _), o in zip(overlap, dst_origin, strict=True)]
run_len = 1
p = ndim
while (
p - 1 >= 0 and src_strides[p - 1] == run_len and dst_strides[p - 1] == run_len
):
run_len *= sizes[p - 1]
p -= 1
src_base = sum(src_start[d] * src_strides[d] for d in range(ndim))
dst_base = sum(dst_start[d] * dst_strides[d] for d in range(ndim))
if p == 0:
yield src_base, dst_base, run_len
return
for idx in itertools.product(*[range(sizes[d]) for d in range(p)]):
s_off = src_base + sum(idx[d] * src_strides[d] for d in range(p))
d_off = dst_base + sum(idx[d] * dst_strides[d] for d in range(p))
yield s_off, d_off, run_len
def plan_pull(
copy: RecordedCopy,
global_shape,
src_dtype: Any,
elsize: int,
shards: list,
) -> list:
if src_dtype != copy.dest_dtype:
raise UnsupportedReshard(
f"{copy.param_name}: source dtype {src_dtype} != dest dtype {copy.dest_dtype} "
"-> conversion required, not a raw byte copy"
)
needed, dest_strides = resolve_slice(copy, global_shape)
need_origin = [lo for lo, _ in needed]
segments: list = []
for sh in shards:
shard_box = [(o, o + s) for o, s in zip(sh.shard_offset, sh.shape, strict=True)]
overlap = intersect(needed, shard_box)
if overlap is None:
continue
src_strides = _row_major_strides(sh.shape)
for s_off, d_off, n in paired_runs(
overlap, sh.shard_offset, src_strides, need_origin, dest_strides
):
segments.append(
PullSegment(
session=sh.session,
src_addr=sh.addr + s_off * sh.elsize,
param_name=copy.param_name,
dst_byte=(copy.dest_offset + d_off) * elsize,
nbytes=n * elsize,
)
)
return segments