tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 10:56:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

add device range as src[1] to multi (kimi) (#17264)

Browse Source 
* add device range as src[1] to multi (kimi)

* cleanup

* fix tests
This commit is contained in:
George Hotz
2026-07-28 16:39:34 -07:00
committed by GitHub
parent f11f884f49
commit 1757067fa6
7 changed files with 35 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_multitensor.py
+5 -3
View File
@@ -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
test/unit/test_realize_is_realize.py
+1 -1
View File
@@ -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()
tinygrad/callify.py
+1 -1
View File
@@ -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
tinygrad/engine/jit.py
+1 -1
View File
@@ -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()
tinygrad/schedule/multi.py
+15 -14
View File
@@ -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),
tinygrad/uop/ops.py
+9 -5
View File
@@ -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):
tinygrad/uop/spec.py
+3 -1
View File
@@ -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)),