From bd296a73593498e451b0e4f3a3bf9f2f0ca3ab14 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:46:45 -0700 Subject: [PATCH] enable alloc_fragment support with UNSHARD (kimi) (#17272) * enable alloc_fragment support with UNSHARD (kimi) * cleaner with implicit barrier * cleanups * cleaner * strongloop * dcount cleanups --- extra/gemm_fragment.py | 176 +++++++++++++++++++++++++++++++++++ tinygrad/codegen/__init__.py | 8 +- tinygrad/schedule/multi.py | 37 +++++++- tinygrad/uop/ops.py | 39 ++++---- tinygrad/uop/spec.py | 4 +- 5 files changed, 239 insertions(+), 25 deletions(-) create mode 100644 extra/gemm_fragment.py diff --git a/extra/gemm_fragment.py b/extra/gemm_fragment.py new file mode 100644 index 0000000000..6f1e2f5c8f --- /dev/null +++ b/extra/gemm_fragment.py @@ -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!") diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index ddc3884932..520b93dfb2 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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: diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index ab01eb3778..b85ebe4639 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -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"), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 12cd821f81..a19a9eee7e 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index c6f89c98de..952e55e666 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -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)),