forked from tinygrad/tinygrad
enable alloc_fragment support with UNSHARD (kimi) (#17272)
* enable alloc_fragment support with UNSHARD (kimi) * cleaner with implicit barrier * cleanups * cleaner * strongloop * dcount cleanups
This commit is contained in:
George Hotz
@@ -0,0 +1,176 @@
|
||||
"""
|
||||
tilelang-style matmul_relu written with tinygrad UOp APIs.
|
||||
|
||||
Demonstrates that tilelang's T.alloc_fragment is expressible with existing
|
||||
tinygrad primitives: a per-thread REG buffer, wrapped in Ops.UNSHARD over a
|
||||
LOCAL thread range to form the full logical tile. The kernel is written
|
||||
against the full-tile UNSHARD view, and multi_pm (the same pass that lowers
|
||||
multi-device UNSHARDs) resolves it into per-thread shard code.
|
||||
|
||||
Reference tilelang kernel:
|
||||
|
||||
@tilelang.jit
|
||||
def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
|
||||
dtype=T.float16, accum_dtype=T.float32):
|
||||
M, N, K = T.const('M, N, K')
|
||||
C = T.empty([M, N], dtype)
|
||||
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
||||
A_shared = T.alloc_shared((block_M, block_K), dtype)
|
||||
B_shared = T.alloc_shared((block_K, block_N), dtype)
|
||||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||||
T.clear(C_local)
|
||||
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
|
||||
T.copy(A[by * block_M, ko * block_K], A_shared)
|
||||
T.copy(B[ko * block_K, bx * block_N], B_shared)
|
||||
T.gemm(A_shared, B_shared, C_local)
|
||||
for i, j in T.Parallel(block_M, block_N):
|
||||
C_local[i, j] = T.max(C_local[i, j], 0)
|
||||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||||
return C
|
||||
|
||||
API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py):
|
||||
|
||||
T.Kernel(gx, gy, threads=T) -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL range (threads)
|
||||
T.alloc_shared(shape, dtype) -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL)
|
||||
T.alloc_fragment(shape, dt) -> per-thread REG placeholder, wrapped in Ops.UNSHARD over the AxisType.LOCAL
|
||||
range: fragment.unshard(axis, tid). The full logical tile is shard x threads,
|
||||
exactly like device sharding, but the sharding axis is the thread axis
|
||||
carried by the RANGE metadata instead of a device tuple. C_local[i, j] with
|
||||
i in this thread's shard is INDEX on the UNSHARD, which multi_pm resolves
|
||||
into INDEX on the per-thread REG shard.
|
||||
T.copy(gmem_slice, smem) -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the
|
||||
smem tile AFTER the copy; the implicit-barrier pass turns the store->load
|
||||
dependency of the loop that consumes it into a workgroup barrier
|
||||
T.gemm (no WMMA) -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
|
||||
with k a loop-carried STRONGLOOP range (codegen builds the register accumulator
|
||||
from this self-referential store automatically)
|
||||
T.copy(fragment, gmem) -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
|
||||
UNSHARD lowering -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
|
||||
full-tile view become per-thread shard ops, no UNSHARD survives into the program.
|
||||
"""
|
||||
|
||||
from tinygrad.dtype import dtypes, AddrSpace, DType
|
||||
from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
|
||||
from tinygrad.helpers import cdiv
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# tilelang builtins, expressed with tinygrad UOp APIs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
|
||||
"""T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
|
||||
return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)
|
||||
|
||||
def alloc_fragment(shape:tuple[int, ...], dtype:DType, axis:int, tnum:UOp) -> UOp:
|
||||
"""T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL axis.
|
||||
|
||||
Each thread privately owns shape[axis]//threads elements along `axis` in a
|
||||
REG buffer. The UNSHARD over the LOCAL thread range presents the full
|
||||
logical tile: full_shape = shard_shape with `axis` multiplied by the number
|
||||
of threads. This is exactly how UNSHARD carries the DEVICE axis today,
|
||||
except the sharding axis is the thread axis carried by the RANGE metadata.
|
||||
"""
|
||||
nthreads = int(tnum.vmax) + 1
|
||||
assert tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL, "fragment must shard over a LOCAL range"
|
||||
assert shape[axis] % nthreads == 0
|
||||
shard_shape = tuple(s // nthreads if i == axis else s for i, s in enumerate(shape))
|
||||
fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG)
|
||||
return fragment.unshard(axis, tnum)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# 64x64 output tile per block, 64 threads (tilelang uses 128 with a 2D per-thread fragment layout;
|
||||
# an UNSHARD splits one axis, so here each thread owns BLOCK_M//THREADS full tile rows of BLOCK_N fp32)
|
||||
BLOCK_M = BLOCK_N = BLOCK_K = 64
|
||||
THREADS = 64
|
||||
ROWS = BLOCK_M // THREADS # fragment rows per thread
|
||||
ROWS_B = BLOCK_N // THREADS # B tile columns copied per thread
|
||||
|
||||
def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
|
||||
"""C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + fragments."""
|
||||
M, K = a.shape
|
||||
K2, N = b.shape
|
||||
assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype)
|
||||
assert not (K % BLOCK_K or M % BLOCK_M or N % BLOCK_N), "test sizes must be multiples of the block sizes"
|
||||
|
||||
# with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=64) as (bx, by):
|
||||
bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
|
||||
by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
|
||||
tid = UOp.range(THREADS, 2, AxisType.LOCAL)
|
||||
|
||||
# A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
|
||||
# B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
|
||||
A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype)
|
||||
B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype)
|
||||
|
||||
# C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype)
|
||||
C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, axis=0, tnum=tid)
|
||||
|
||||
# T.clear(C_local) -- each thread zeroes its own fragment rows
|
||||
ir0, j0 = UOp.range(ROWS, 4, AxisType.STRONGLOOP), UOp.range(BLOCK_N, 5, AxisType.STRONGLOOP)
|
||||
C_loc = C_local[tid*ROWS + ir0, j0].set(0.0, end=(ir0, j0))
|
||||
|
||||
# for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
|
||||
# (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
|
||||
ko = UOp.range(cdiv(K, BLOCK_K), 3, AxisType.STRONGLOOP)
|
||||
|
||||
# T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies ROWS row(s)
|
||||
# set returns the tile AFTER the copy; codegen turns that store->load dependency into a workgroup barrier
|
||||
iar, ka = UOp.range(ROWS, 6, AxisType.STRONGLOOP), UOp.range(BLOCK_K, 7, AxisType.STRONGLOOP)
|
||||
A_shared = A_shared[tid*ROWS + iar, ka].set(a[by*BLOCK_M + tid*ROWS + iar, ko*BLOCK_K + ka], end=(iar, ka))
|
||||
|
||||
# T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared) -- each thread copies ROWS_B column(s)
|
||||
kb, ibr = UOp.range(BLOCK_K, 8, AxisType.STRONGLOOP), UOp.range(ROWS_B, 9, AxisType.STRONGLOOP)
|
||||
B_shared = B_shared[kb, tid*ROWS_B + ibr].set(b[ko*BLOCK_K + kb, bx*BLOCK_N + tid*ROWS_B + ibr], end=(kb, ibr))
|
||||
|
||||
# T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment rows.
|
||||
# identical to custom_gemm: a self-referential store over the loop-carried kk range,
|
||||
# which codegen turns into a register accumulator
|
||||
# kk nests outside the fragment row/col loops (lower ids nest outer), so B loads are contiguous and
|
||||
# the A row value is loaded once per kk
|
||||
ir, kk = UOp.range(ROWS, 10, AxisType.STRONGLOOP), UOp.range(BLOCK_K, 11, AxisType.STRONGLOOP)
|
||||
jj = UOp.range(BLOCK_N, 12, AxisType.STRONGLOOP)
|
||||
acc = C_loc.after(kk)[tid*ROWS + ir, jj] + A_shared[tid*ROWS + ir, kk].cast(dtypes.float32) * B_shared[kk, jj].cast(dtypes.float32)
|
||||
# closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
|
||||
C_loc = C_loc[tid*ROWS + ir, jj].set(acc, end=(kk, ir, jj, ko))
|
||||
|
||||
# for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
|
||||
# T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
|
||||
# STRONGLOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
|
||||
ie, je = UOp.range(ROWS, 13, AxisType.STRONGLOOP), UOp.range(BLOCK_N, 14, AxisType.STRONGLOOP)
|
||||
c_st = c[by*BLOCK_M + tid*ROWS + ie, bx*BLOCK_N + je].store(C_loc[tid*ROWS + ie, je].relu().cast(c.dtype))
|
||||
|
||||
# all open ranges are closed at the final store (ko was closed above).
|
||||
# the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code
|
||||
return c_st.end(je, ie, tid, by, bx).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# python wrapper: same signature as the tilelang function
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def matmul_relu(a:Tensor, b:Tensor) -> Tensor:
|
||||
"""C = relu(A @ B), fp16 in/out with an fp32 fragment accumulator."""
|
||||
c = Tensor.empty(a.shape[0], b.shape[1], dtype=a.dtype, device=a.device)
|
||||
return c.custom_kernel(a, b, fxn=matmul_relu_kernel)[0]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
from tinygrad import Device
|
||||
assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)"
|
||||
M = K = N = 256 # 4x4 grid of 64x64 tiles, 4 K chunks
|
||||
|
||||
a = Tensor.randn(M, K, dtype=dtypes.float16).contiguous()
|
||||
b = Tensor.randn(K, N, dtype=dtypes.float16).contiguous()
|
||||
ref = (a.float() @ b.float()).relu().realize()
|
||||
|
||||
out = matmul_relu(a, b).realize()
|
||||
|
||||
import numpy as np
|
||||
np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
|
||||
print("matmul_relu passed!")
|
||||
@@ -21,6 +21,7 @@ from tinygrad.codegen.late.coalesce import indexing_simplify
|
||||
from tinygrad.codegen.opt.postrange import apply_opts
|
||||
from tinygrad.codegen.late.gater import pm_move_gates_from_index
|
||||
from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
|
||||
from tinygrad.schedule.multi import multi_pm
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
|
||||
from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
|
||||
@@ -267,7 +268,7 @@ def add_raw_barrier(after:UOp):
|
||||
|
||||
def add_war_barrier(end:UOp):
|
||||
# a LOCAL buffer stored and loaded in the same loop needs a barrier at the end of the loop body
|
||||
rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.LOOP) and r.vmax > 0]
|
||||
rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.LOOP, AxisType.STRONGLOOP) and r.vmax > 0]
|
||||
if not rngs or end.src[0].op is Ops.BARRIER: return None
|
||||
sl = end.src[0].backward_slice_with_self
|
||||
# only stores that are inside this loop body (not in the backward slice through AFTER chains from other loops)
|
||||
@@ -286,8 +287,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
|
||||
if DEBUG >= 5: print(pyrender(ast))
|
||||
if SPEC: type_verify(ast, spec_tensor)
|
||||
|
||||
# resolve UNSHARDs (multi-device UNSHARDs are already resolved by the scheduler; this handles in-kernel shards, e.g. fragments)
|
||||
sink = graph_rewrite(ast, multi_pm, name="multi_pm")
|
||||
|
||||
# preprocess
|
||||
sink = graph_rewrite(ast, pm_mops, name="early movement ops", bottom_up=True)
|
||||
sink = graph_rewrite(sink, pm_mops, name="early movement ops", bottom_up=True)
|
||||
|
||||
# first we optimize
|
||||
if optimize:
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
|
||||
from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape
|
||||
from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape, sint_to_uop
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.schedule.allreduce import handle_allreduce
|
||||
|
||||
def shard_count(multi:UOp) -> int:
|
||||
# the shard count: the device count for multi-device UNSHARDs, the range size otherwise (e.g. threads for LOCAL fragments)
|
||||
return len(multi.device) if isinstance(multi.device, tuple) else int(multi.src[1].vmax)+1
|
||||
|
||||
# ***** multi rewrite MSELECT/MSTACK *****
|
||||
|
||||
def _apply_shrink(marg, s:UOp, i:int) -> UOp:
|
||||
@@ -50,7 +54,12 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
|
||||
# normalize srcs to local shards on axis
|
||||
devices = [x.device for x in msrcs if x.device is not None]
|
||||
assert all_same(devices), f"all buffers must have the same device {devices}"
|
||||
dcount = len(devices[0])
|
||||
# without devices the sharding range comes from the UNSHARD itself (e.g. a LOCAL thread range);
|
||||
# device shards range over the devices instead
|
||||
if len(devices): sharding_rng = UOp.range(len(devices[0]), -1, AxisType.DEVICE)
|
||||
else:
|
||||
sharding_rng = next((m.src[1] for m in msrcs if m.op is Ops.UNSHARD), None)
|
||||
assert sharding_rng is not None, "shard_srcs requires a device or a sharding range"
|
||||
|
||||
out_shape = _broadcast_shape(*[x.shape for x in msrcs])
|
||||
srcs:list[UOp] = []
|
||||
@@ -60,9 +69,9 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
|
||||
# same axis, just copy through
|
||||
srcs.append(mlb.src[0])
|
||||
else:
|
||||
# otherwise every device gets the full copy, sharded iff this src has the axis (broadcast srcs stay whole)
|
||||
# otherwise every shard gets the full copy, sharded iff this src has the axis (broadcast srcs stay whole)
|
||||
full = mlb if mlb.axis is None else copy_multi(mlb, mlb.device)
|
||||
srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, dcount))
|
||||
srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, sharding_rng))
|
||||
return srcs
|
||||
|
||||
def alu_multi(root:UOp):
|
||||
@@ -86,7 +95,7 @@ def reduce_multi(root:UOp, multi:UOp):
|
||||
|
||||
def reshape_multi(root:UOp, multi:UOp):
|
||||
if prod(multi.shape) != prod(new_shape:=root.marg): raise RuntimeError("reshape must maintain prod(shape)")
|
||||
if (new_axis:=root.axis) is not None: new_shape = tuple(s//len(multi.device) if a==new_axis else s for a,s in enumerate(new_shape))
|
||||
if (new_axis:=root.axis) is not None: new_shape = tuple(s//shard_count(multi) if a==new_axis else s for a,s in enumerate(new_shape))
|
||||
return multi.src[0].reshape(new_shape).unshard(new_axis, multi.src[1])
|
||||
|
||||
def expand_multi(root:UOp, multi:UOp):
|
||||
@@ -103,6 +112,14 @@ def permute_multi(root:UOp, multi:UOp):
|
||||
return multi.src[0].permute(root.marg).unshard(root.axis, multi.src[1])
|
||||
|
||||
def shrink_multi(root:UOp, multi:UOp):
|
||||
if multi.axis is not None:
|
||||
shard_sz = multi.src[0].shape[multi.axis]
|
||||
s, l = root.marg[multi.axis] # SHRINK marg is (start, length)
|
||||
# shrink to exactly this range's own shard: this resolves the UNSHARD into its per-shard view
|
||||
# (e.g. a fragment indexed by its LOCAL thread range becomes that thread's REG shard, no copy needed)
|
||||
if sint_to_uop(l).ssimplify() == shard_sz and (sint_to_uop(s)-multi.src[1]*shard_sz).ssimplify() == 0:
|
||||
non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else t for i, t in enumerate(root.marg))
|
||||
return multi.src[0]._mop(Ops.SHRINK, non_shard_shrink)
|
||||
shard_bounds = tuple((s,e-s) for s,e in multi.bounds) if multi.axis is not None else ()
|
||||
assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in shard_bounds, \
|
||||
f"shrinking not supported for {root.marg=}"
|
||||
@@ -124,6 +141,15 @@ def stack_multi(root:UOp):
|
||||
assert axis is not None
|
||||
return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).unshard(axis, next(m.src[1] for m in root.src if m.op is Ops.UNSHARD))
|
||||
|
||||
def index_multi(root:UOp, multi:UOp):
|
||||
# INDEX on UNSHARD: resolve the sharded axis into this range's own shard (idx - rng*shard_sz)
|
||||
if multi.axis is None: return None
|
||||
shard_sz = multi.src[0].shape[multi.axis]
|
||||
local = (root.src[1+multi.axis] - multi.src[1]*shard_sz).simplify()
|
||||
# the index along the sharded axis must be provably inside this shard
|
||||
if local.vmin < 0 or local.vmax >= shard_sz: return None
|
||||
return multi.src[0].index(*root.src[1:1+multi.axis], local, *root.src[2+multi.axis:])
|
||||
|
||||
def copy_multi(multi:UOp, device:str | tuple[str, ...]):
|
||||
assert multi.axis is not None, "all multi ops have axis"
|
||||
if isinstance(device, str):
|
||||
@@ -165,6 +191,7 @@ multi_pm = PatternMatcher([
|
||||
(UPat(Ops.PERMUTE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), permute_multi),
|
||||
(UPat(Ops.FLIP, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), flip_multi),
|
||||
(UPat(Ops.STACK, name="root", custom_early_reject=set([Ops.UNSHARD])), stack_multi),
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.UNSHARD, name="multi"),), name="root", allow_any_len=True), index_multi),
|
||||
(UPat(Ops.AFTER, src=(UPat(Ops.UNSHARD), UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="dest"), UPat(Ops.UNSHARD, name="src"))))), store_after_multi),
|
||||
(UPat(Ops.COPY, src=(UPat(Ops.UNSHARD, name="multi"),), name="copy"), lambda multi,copy: copy_multi(multi, copy.arg)),
|
||||
(UPat(Ops.ALLREDUCE, src=(UPat(Ops.UNSHARD, name="multi"),), name="red"),
|
||||
|
||||
+23
-16
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
|
||||
class AxisType(Enum):
|
||||
def __repr__(self): return str(self)
|
||||
DEVICE = auto(); GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto() # noqa: E702
|
||||
UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto() # noqa: E702
|
||||
UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto(); STRONGLOOP = auto() # noqa: E702
|
||||
|
||||
@dataclass(frozen=True, order=True)
|
||||
class ParamArg:
|
||||
@@ -36,13 +36,14 @@ class ParamArg:
|
||||
args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
|
||||
return f"ParamArg({', '.join(args)})"
|
||||
axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L",
|
||||
AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
|
||||
AxisType.STRONGLOOP: "S", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
|
||||
axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN",
|
||||
AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
|
||||
AxisType.LOOP: "WHITE", AxisType.STRONGLOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red",
|
||||
AxisType.UNROLL: "magenta"}
|
||||
|
||||
# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
|
||||
axis_to_pos = {AxisType.DEVICE: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
|
||||
AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
|
||||
axis_to_pos = {AxisType.DEVICE: -2, AxisType.LOOP: -1, AxisType.STRONGLOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
|
||||
AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
|
||||
|
||||
range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
|
||||
Ops.SLICE: 2, Ops.LINEAR: 0}
|
||||
@@ -665,17 +666,20 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
# *** multi-device helpers ***
|
||||
|
||||
def unshard(self, axis:int|None, device_range:UOp|None=None):
|
||||
assert isinstance(self.device, tuple), f"multi device must be tuple, {self.device} isn't"
|
||||
assert axis is not None, "multi None is no longer supported"
|
||||
# an UNSHARD always has two srcs: the value and the DEVICE range it ends (defaults to a DEVICE range over the devices)
|
||||
if device_range is None: device_range = UOp.range(len(self.device), -1, AxisType.DEVICE)
|
||||
assert device_range.op is Ops.RANGE and device_range.arg[-1] is AxisType.DEVICE
|
||||
# an UNSHARD always has two srcs: the value and the range it ends (defaults to a DEVICE range over the devices)
|
||||
# the range need not be DEVICE, e.g. a LOCAL range shards a kernel tile into per-thread fragments
|
||||
if device_range is None:
|
||||
assert isinstance(self.device, tuple), f"multi device must be tuple, {self.device} isn't"
|
||||
device_range = UOp.range(len(self.device), -1, AxisType.DEVICE)
|
||||
assert device_range.op is Ops.RANGE
|
||||
return UOp(Ops.UNSHARD, src=(self, device_range), arg=axis)
|
||||
|
||||
@property
|
||||
def bounds(self):
|
||||
if self.axis is None: raise RuntimeError("bounds is not defined when axis is None")
|
||||
return tuple(itertools.pairwise(itertools.accumulate([self.src[0].shape[self.axis] for _ in self.device], initial=0)))
|
||||
dcount = int(self.src[1].vmax)+1 if self.op is Ops.UNSHARD else len(self.device)
|
||||
return tuple(itertools.pairwise(itertools.accumulate([self.src[0].shape[self.axis] for _ in range(dcount)], initial=0)))
|
||||
|
||||
@functools.cached_property
|
||||
def axis(self) -> int|None:
|
||||
@@ -705,7 +709,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
target = ssimplify(prod(self.src[0].shape[:src_axis]))
|
||||
if target not in arg_acc: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards")
|
||||
new_axis = len(arg_acc) - arg_acc[::-1].index(target) - 1
|
||||
if self.shape[new_axis] % len(self.device) != 0: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards")
|
||||
dcount = len(self.device) if isinstance(self.device, tuple) else \
|
||||
int(next(u.src[1] for u in self.src[0].toposort() if u.op is Ops.UNSHARD).vmax)+1
|
||||
if self.shape[new_axis] % dcount != 0: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards")
|
||||
return new_axis
|
||||
if self.op is Ops.PERMUTE: return self.marg.index(src_axis) if src_axis is not None else None
|
||||
if self.op is Ops.EXPAND: return src_axis + len(self.marg) if src_axis is not None else None
|
||||
@@ -716,15 +722,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
dnum = UOp.range(dcount, -1, AxisType.DEVICE)
|
||||
return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))
|
||||
|
||||
def _shard(self, axis:int, dcount:int) -> UOp:
|
||||
def _shard(self, axis:int, rng:UOp) -> UOp:
|
||||
assert rng.op is Ops.RANGE, f"_shard requires a RANGE, got {rng.op}"
|
||||
if len(self.shape) == 0: return self # scalars broadcast, no sharding needed
|
||||
dnum = UOp.range(dcount, -1, AxisType.DEVICE)
|
||||
dcount = int(rng.vmax)+1
|
||||
if self.shape[axis] % dcount != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {dcount=}")
|
||||
sz = self.shape[axis] // dcount
|
||||
return self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
|
||||
return self.shrink(tuple((0,s) if i != axis else (rng*sz,rng*sz+sz) for i,s in enumerate(self.shape)))
|
||||
def shard(self, devices:tuple[str, ...], axis:int|None=None) -> UOp:
|
||||
copied = self.copy_to_device(devices)
|
||||
return copied if axis is None else copied._shard(axis, len(devices)).unshard(axis)
|
||||
return copied if axis is None else copied._shard(axis, UOp.range(len(devices), -1, AxisType.DEVICE)).unshard(axis)
|
||||
|
||||
def copy_to_device(self, device:str|tuple[str, ...], arg=None):
|
||||
assert arg is None or isinstance(self.device, tuple)
|
||||
@@ -844,7 +851,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
|
||||
if self.op is Ops.BUFFER: return self.arg.addrspace
|
||||
if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU
|
||||
if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
|
||||
if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END}:
|
||||
if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END, Ops.UNSHARD}:
|
||||
return self.src[0].addrspace
|
||||
if self.op in GroupOp.Movement: return self.src[0].addrspace
|
||||
if self.op in {Ops.STACK, Ops.WMMA, Ops.GROUP} or self.op in GroupOp.Elementwise:
|
||||
|
||||
@@ -175,9 +175,9 @@ spec_tensor = PatternMatcher([
|
||||
len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])),
|
||||
|
||||
# UNSHARD/MSELECT/MSTACK
|
||||
# an UNSHARD always has two srcs: the value and the DEVICE range it ends
|
||||
# an UNSHARD always has two srcs: the value and the range it ends (usually DEVICE, but any typed range can be sharded over)
|
||||
(UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 2 and matches_dtype(multi.src[0], multi.dtype)
|
||||
and isinstance(multi.arg, int) and multi.src[1].op is Ops.RANGE and multi.src[1].arg[-1] is AxisType.DEVICE),
|
||||
and isinstance(multi.arg, int) and multi.src[1].op is Ops.RANGE and isinstance(multi.src[1].arg[-1], AxisType)),
|
||||
(UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
|
||||
(UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
|
||||
|
||||
|
||||
Reference in New Issue
Block a user