diff --git a/test/unit/test_allreduce.py b/test/unit/test_allreduce.py index 2404408fe8..c45b28ab7f 100644 --- a/test/unit/test_allreduce.py +++ b/test/unit/test_allreduce.py @@ -6,7 +6,7 @@ from tinygrad.ops import Ops class TestRingAllReduce(unittest.TestCase): def test_schedule_ring(self): with Context(RING=2): - N = 6 + N = 4 ds = tuple(f"CPU:{i}" for i in range(N)) t = Tensor.empty(N, N*100).shard(ds, axis=0).realize() schedules = t.sum(0).schedule_with_vars()[0] @@ -17,5 +17,13 @@ class TestRingAllReduce(unittest.TestCase): # copy topology forms a ring self.assertEqual(len(set(pairs)), N) + def test_correct_ring(self): + with Context(RING=2): + N = 4 + ds = tuple(f"CPU:{i}" for i in range(N)) + t = Tensor.ones(N, N*100).contiguous().shard(ds, axis=0).realize() + out = t.sum(0) + self.assertListEqual(out.tolist(), [4]*N*100) + if __name__ == '__main__': unittest.main() diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 273545bb08..cb6dfd95de 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -114,7 +114,7 @@ class Ops(FastEnum): VALID = auto(); SPECIAL = auto(); NOOP = auto() # noqa: E702 # reduce - REDUCE_AXIS = auto(); REDUCE = auto() # noqa: E702 + REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto() # noqa: E702 # helper ops GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702 diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index fbdea504c0..c93977410b 100644 --- a/tinygrad/shape/view.py +++ b/tinygrad/shape/view.py @@ -3,7 +3,7 @@ import functools, operator, itertools from dataclasses import dataclass from typing import Optional, cast, Sequence from tinygrad.dtype import dtypes -from tinygrad.ops import resolve, UOp, Variable, sint, sym_infer, smax, smin, sint_to_uop +from tinygrad.ops import resolve, UOp, Variable, sint, sym_infer, smax, smin, sint_to_uop, Ops from tinygrad.helpers import prod, all_int, argsort, flatten, ceildiv @functools.cache @@ -107,7 +107,8 @@ class View: @staticmethod @functools.cache def create(shape:tuple[sint, ...], strides:Optional[tuple[sint, ...]]=None, offset:sint=0, mask:Optional[tuple[tuple[sint, sint], ...]]=None): - if not all(s >= 0 for s in shape): raise ValueError(f"Trying to create View with negative dimension: {shape=}") + # TODO: resolve shouldn't be needed here + if not all(resolve(s >= 0) for s in shape): raise ValueError(f"Trying to create View with negative dimension: {shape=}") strides = canonicalize_strides(shape, strides) if strides else strides_for_shape(shape) # canonicalize 0 in shape if 0 in shape: return View(shape, (0,) * len(shape), offset=0, mask=None, contiguous=True) @@ -138,7 +139,7 @@ class View: @functools.cache # pylint: disable=method-cache-max-size-none def unbind(self) -> tuple[View, dict[Variable, int]]: - var_unboundvar_val = [(v, v.unbind()) for v in self.vars()] + var_unboundvar_val = [(v, v.unbind()) for v in self.vars() if v.op is Ops.BIND] unbound_vars = {v:uv for v,(uv,_) in var_unboundvar_val} def substitute(x:sint): return x if isinstance(x, int) else x.substitute(unbound_vars) new_shape = tuple(map(substitute, self.shape)) @@ -196,8 +197,8 @@ class View: else: bad = True continue d1, s1 = term[0] - newb[d1] = max(newb[d1], ceildiv(b - o if s1 > 0 else e - o - 1, s1)) - newe[d1] = min(newe[d1], (b - o if s1 < 0 else e - o - 1) // s1 + 1) + newb[d1] = smax(newb[d1], ceildiv(b - o if s1 > 0 else e - o - 1, s1)) + newe[d1] = smin(newe[d1], (b - o if s1 < 0 else e - o - 1) // s1 + 1) # If any of vm1 was masked off, try again with that mask in place. if any((b, e) != (0, s) for b, e, s in zip(newb, newe, vm1.shape)): diff --git a/tinygrad/spec.py b/tinygrad/spec.py index 56a57f309d..0369651690 100644 --- a/tinygrad/spec.py +++ b/tinygrad/spec.py @@ -32,7 +32,8 @@ except (ImportError, AttributeError): z3_imported = False buffer_spec = PatternMatcher([ (UPat(Ops.UNIQUE, dtypes.void, ()), lambda: True), - (UPat(Ops.DEVICE, dtypes.void, (), name="device"), lambda device: isinstance(device.arg, str)), + (UPat(Ops.DEVICE, dtypes.void, (), name="d"), lambda d: + isinstance(d.arg, str) or (isinstance(d.arg, tuple) and all(isinstance(s, str) for s in d.arg))), (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), allow_any_len=True, name="buf"), lambda buf: isinstance(buf.arg, int) and isinstance(buf.dtype, (DType, ImageDType))), (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),), name="buf_view"), @@ -75,8 +76,9 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([ (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="root", src=(UPat.var("x"),), arg=None), lambda root,x: root.dtype == x.dtype), - # COPY + # COPY/ALLREDUCE (UPat(Ops.COPY, name="copy", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda copy,x: copy.dtype == x.dtype), + (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)), ]) # ***** uop type spec ***** diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 58521f22bf..9305fca8bd 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -14,7 +14,8 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.IGNORE: "#00C000", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF", - Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500"} + Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.FUSE: "#FFa500", + Ops.ALLREDUCE: "#ff40a0"} # VIZ API