From 1757067fa65627ce8317f7cce507b259ba409b5b Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:39:34 -0700 Subject: [PATCH] add device range as src[1] to multi (kimi) (#17264) * add device range as src[1] to multi (kimi) * cleanup * fix tests --- test/backend/test_multitensor.py | 8 +++++--- test/unit/test_realize_is_realize.py | 2 +- tinygrad/callify.py | 2 +- tinygrad/engine/jit.py | 2 +- tinygrad/schedule/multi.py | 29 ++++++++++++++-------------- tinygrad/uop/ops.py | 14 +++++++++----- tinygrad/uop/spec.py | 4 +++- 7 files changed, 35 insertions(+), 26 deletions(-) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index ae63a1ce97..8b58f4c554 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -1,6 +1,6 @@ import unittest, random from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable -from tinygrad.uop.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp, AxisType from tinygrad.helpers import getenv, prod, Context from tinygrad.nn.state import get_parameters from tinygrad.engine.realize import run_linear, compile_linear @@ -52,8 +52,10 @@ class TestMultiTensor(unittest.TestCase): def test_shard(self): X = Tensor.ones(256).contiguous().realize() X.shard_(devices_2, 0) - for lb in X.uop.src: - assert lb.shape == (128,) + assert X.uop.src[0].shape == (128,) + # the MULTI carries and ends the DEVICE range as its second src + assert X.uop.src[1].op is Ops.RANGE and X.uop.src[1].arg[-1] is AxisType.DEVICE + assert X.uop.ended_ranges == X.uop.src[1:] (X + X).realize() @unittest.expectedFailure # TODO: fix diff --git a/test/unit/test_realize_is_realize.py b/test/unit/test_realize_is_realize.py index b2e141a6cd..63fb266545 100644 --- a/test/unit/test_realize_is_realize.py +++ b/test/unit/test_realize_is_realize.py @@ -30,7 +30,7 @@ class TestRealizeIsRealized(unittest.TestCase): def test_multi(self): d = Device.DEFAULT t = Tensor.ones(8).contiguous().shard((d, d), axis=0).realize() - assert all(u.is_realized for u in t.uop.src) + assert t.uop.src[0].is_realized def test_empty(self): t = Tensor.empty(4, 4).realize() diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 4a5211d4f2..bcf26ec15a 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -83,7 +83,7 @@ def contiguous_mops_to_view(c:UOp, src:UOp): resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view") if resolved.op is not Ops.MULTI: return None if (view := _make_buffer_view(resolved.src[0])) is None: return None - return view.reshape(resolved.src[0].shape).multi(resolved.arg).contiguous(tag=c.tag) + return view.reshape(resolved.src[0].shape).multi(resolved.arg, resolved.src[1]).contiguous(tag=c.tag) return None diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index d0ab840bd0..e65c742552 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -206,7 +206,7 @@ def _prepare_jit_inputs(args, kwargs): for x in args + tuple(kwargs.values()): it = x if isinstance(x, (tuple,list)) else x.values() if isinstance(x, dict) else [] tensors += [t for t in it if t.__class__ is Tensor and not any(t is y for y in tensors)] - def get_input_uops() -> list[UOp]: return flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors]) + def get_input_uops() -> list[UOp]: return flatten([[t.uop.src[0]] if t.uop.op is Ops.MULTI else [t.uop] for t in tensors]) if any(u.is_virtual for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()") if len(unrealized_tensors := [x for x in tensors if not x.uop.is_realized]): Tensor.realize(*unrealized_tensors) input_uops = get_input_uops() diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index dba7c9c571..b35c254ff0 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -68,7 +68,7 @@ def alu_multi(root:UOp): axis = root.axis assert axis is not None srcs = shard_srcs(root.src, axis) - return srcs[0].alu(root.op, *srcs[1:]).multi(axis) + return srcs[0].alu(root.op, *srcs[1:]).multi(axis, next(m.src[1] for m in root.src if m.op is Ops.MULTI)) def reduce_multi(root:UOp, multi:UOp): op, num_axes = root.arg @@ -81,25 +81,25 @@ def reduce_multi(root:UOp, multi:UOp): return local.allreduce(op, multi.device) # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct new_axis = multi.axis - num_axes if multi.axis is not None else None - return multi.src[0]._rop(op, tuple(range(num_axes))).multi(axis=new_axis) + return multi.src[0]._rop(op, tuple(range(num_axes))).multi(new_axis, multi.src[1]) 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)) - return multi.src[0].reshape(new_shape).multi(new_axis) + return multi.src[0].reshape(new_shape).multi(new_axis, multi.src[1]) def expand_multi(root:UOp, multi:UOp): new_axis = None if multi.axis is None else multi.axis + len(root.marg) - return multi.src[0]._mop(Ops.EXPAND, arg=root.marg).multi(new_axis) + return multi.src[0]._mop(Ops.EXPAND, arg=root.marg).multi(new_axis, multi.src[1]) def pad_multi(root:UOp, multi:UOp): assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]), f"padding not supported for {root.marg=}" local_pad = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg)) - return multi.src[0]._mop(Ops.PAD, local_pad).multi(multi.axis) + return multi.src[0]._mop(Ops.PAD, local_pad).multi(multi.axis, multi.src[1]) def permute_multi(root:UOp, multi:UOp): # all permutes supported! - return multi.src[0].permute(root.marg).multi(root.axis) + return multi.src[0].permute(root.marg).multi(root.axis, multi.src[1]) def shrink_multi(root:UOp, multi:UOp): shard_bounds = tuple((s,e-s) for s,e in multi.bounds) if multi.axis is not None else () @@ -111,17 +111,17 @@ def shrink_multi(root:UOp, multi:UOp): non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else s for i, s in enumerate(root.marg)) return multi.src[0].copy_to_device(multi.device, arg=shard_bounds.index(root.marg[multi.axis]))._mop(Ops.SHRINK, non_shard_shrink) local_shrink = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg)) - return multi.src[0]._mop(Ops.SHRINK, local_shrink).multi(multi.axis) + return multi.src[0]._mop(Ops.SHRINK, local_shrink).multi(multi.axis, multi.src[1]) def flip_multi(root:UOp, multi:UOp): assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis" - return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).multi(multi.axis) + return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).multi(multi.axis, multi.src[1]) def stack_multi(root:UOp): # STACK adds a leading axis: srcs are sharded one axis below the output axis = root.axis assert axis is not None - return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).multi(axis) + return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).multi(axis, next(m.src[1] for m in root.src if m.op is Ops.MULTI)) def copy_multi(multi:UOp, device:str | tuple[str, ...]): assert multi.axis is not None, "all multi ops have axis" @@ -130,11 +130,11 @@ def copy_multi(multi:UOp, device:str | tuple[str, ...]): return pieces[0].cat(*pieces[1:], dim=multi.axis) return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device) -def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).multi(src.axis) +def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).multi(src.axis, src.src[1]) def passthrough_multi(root:UOp, multi:UOp): new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src[1:]) - return UOp(root.op, root.dtype, src=new_src, arg=root.arg).multi(multi.axis) + return UOp(root.op, root.dtype, src=new_src, arg=root.arg).multi(multi.axis, multi.src[1]) def rewrite_into_function(call:UOp): if call.arg.precompile: return None @@ -144,7 +144,8 @@ def rewrite_into_function(call:UOp): assert new_body.op is Ops.TUPLE if any(s.op is Ops.MULTI for s in new_body.src): shard_call = call.replace(src=(UOp.maketuple(*[s.src[0] if s.op is Ops.MULTI else s for s in new_body.src]),)+new_args) - return UOp.maketuple(*[shard_call.gettuple(i).multi(s.axis) if s.op is Ops.MULTI else shard_call.gettuple(i) for i, s in enumerate(new_body.src)]) + return UOp.maketuple(*[shard_call.gettuple(i).multi(s.axis, s.src[1]) if s.op is Ops.MULTI else shard_call.gettuple(i) + for i, s in enumerate(new_body.src)]) return call.replace(src=(new_body,)+new_args) def param_to_multi(p:UOp): @@ -166,13 +167,13 @@ multi_pm = PatternMatcher([ (UPat(Ops.AFTER, src=(UPat(Ops.MULTI), UPat(Ops.STORE, src=(UPat(Ops.MULTI, name="dest"), UPat(Ops.MULTI, name="src"))))), store_after_multi), (UPat(Ops.COPY, src=(UPat(Ops.MULTI, name="multi"),), name="copy"), lambda multi,copy: copy_multi(multi, copy.arg)), (UPat(Ops.ALLREDUCE, src=(UPat(Ops.MULTI, name="multi"),), name="red"), - lambda multi,red: multi.src[0].allreduce(*red.arg).multi(axis=multi.axis)), + lambda multi,red: multi.src[0].allreduce(*red.arg).multi(multi.axis, multi.src[1])), # resolve TUPLE+GETTUPLE (needed in multi) (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]), # GETTUPLE on MULTI: passthrough MULTI (e.g. when FUNCTION was replaced by MULTI(GETTUPLE(...))) (UPat(Ops.GETTUPLE, src=(UPat(Ops.MULTI, name="multi"),), name="g"), - lambda g, multi: multi.src[0].gettuple(g.arg).multi(multi.axis) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} + lambda g, multi: multi.src[0].gettuple(g.arg).multi(multi.axis, multi.src[1]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} else multi), # rewrite into FUNCTION calls explicitly for MULTI (value-producing) (UPat(Ops.FUNCTION, name="call"), rewrite_into_function), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 9137c89d4a..d19c7b818c 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -438,7 +438,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.FLIP: if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}") return ps - case Ops.MULTI: return tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(ps)) + case Ops.MULTI: return tuple(s*(int(self.src[1].vmax)+1) if a == self.axis else s for a,s in enumerate(ps)) case Ops.REDUCE: num_axes = self.arg[1] if not isinstance(num_axes, int) or num_axes < 0 or num_axes > len(ps): @@ -473,7 +473,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @property def shard_shape(self) -> tuple[sint, ...]: if not isinstance(self.device, tuple) or self.axis is None: return self.shape - return tuple(x//len(self.device) if i == self.axis else x for i,x in enumerate(self.shape)) + dcount = int(self.src[1].vmax)+1 if self.op is Ops.MULTI else len(self.device) + return tuple(x//dcount if i == self.axis else x for i,x in enumerate(self.shape)) @property def max_shard_shape(self) -> tuple[int, ...]: return to_max_shape(self.shard_shape) @@ -483,7 +484,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op in range_start: return self.src[range_start[self.op]:] if self.op is Ops.AFTER: return tuple(flatten([x.ended_ranges for x in self.src[1:]])) # MULTI ends the DEVICE range: its src is per-device index math, the device axis is carried by the axis metadata - if self.op is Ops.MULTI: return tuple(r for r in self.src[0].ranges if r.arg[-1] is AxisType.DEVICE) + if self.op is Ops.MULTI: return self.src[1:] return () # determine what ranges this is in @@ -663,10 +664,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # *** multi-device helpers *** - def multi(self, axis:int|None): + def multi(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" - return UOp(Ops.MULTI, src=(self,), arg=axis) + # a MULTI 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 + return UOp(Ops.MULTI, src=(self, device_range), arg=axis) @property def bounds(self): diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 49baafcac2..99616a6b3f 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -175,7 +175,9 @@ spec_tensor = PatternMatcher([ len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])), # MULTI/MSELECT/MSTACK - (UPat(Ops.MULTI, name="multi"), lambda multi: all(matches_dtype(x, multi.dtype) for x in multi.src) and isinstance(multi.arg, int)), + # a MULTI always has two srcs: the value and the DEVICE range it ends + (UPat(Ops.MULTI, 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), (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)),