gate MULTIOUTPUT

test/test_schedule.py

@@ -2414,5 +2414,16 @@ class TestUOpBecome(unittest.TestCase):
     b.shrink(((0,4),)).assign(a_view).realize()
     self.assertListEqual(b.tolist(), [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])
 
+class TestScheduleMultioutput(unittest.TestCase):
+  def test_simplest_multioutput(self):
+    with Context(MULTIOUTPUT=1):
+      a = Tensor.ones(256, 256).contiguous().realize()
+      r = a.sum(axis=1)
+      b = r+1
+      c = r+2
+      run_schedule(check_schedule([b, c], 1))
+      np.testing.assert_allclose(b.numpy(), 257)
+      np.testing.assert_allclose(c.numpy(), 258)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

tinygrad/codegen/late/control_flow.py

@@ -86,7 +86,7 @@ def do_merge_ends(s:UOp):
   replaces = {}
   for k,v in stacked.items():
     if len(v) == 1: continue
-    rep = UOp(v[0].op, src=tuple([k] + [y for x in v for y in x.src[1:]]), arg=x[0].arg)
+    rep = UOp(v[0].op, src=tuple([k] + [y for x in v for y in x.src[1:]]), arg=v[0].arg)
     for x in v: replaces[x] = rep
   if not len(replaces) and not len(dangling_ifs): return None
   ret = s.substitute(replaces)

tinygrad/helpers.py

@@ -170,6 +170,8 @@ SPEC = ContextVar("SPEC", 0)
 # TODO: disable by default due to speed
 IGNORE_OOB = ContextVar("IGNORE_OOB", 1)
 PCONTIG = ContextVar("PCONTIG", 0)  # partial contiguous in rangeify
+MULTIOUTPUT = ContextVar("MULTIOUTPUT", 0)
+DEBUG_RANGEIFY = ContextVar("DEBUG_RANGEIFY", 0)
 
 @dataclass(frozen=True)
 class Metadata:

tinygrad/schedule/indexing.py

@@ -1,10 +1,10 @@
-from typing import Iterator
+from typing import Iterator, cast
 import functools, operator, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored
+from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, MULTIOUTPUT
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.BUFFER_VIEW,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
@@ -141,6 +141,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
 
 @profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
+  if debug: print("**************************")
   rctx = IndexingContext()
 
   # get ops to realize
@@ -252,5 +253,28 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op.
     rctx.range_map[x] = (rngs, out_rngs)
 
+  if MULTIOUTPUT:
+    # second forward pass to fuse children
+    replaced_ranges = {}
+    for x in tsink.toposort():
+      if x not in rctx.realize_map: continue
+      out_rngs = rctx.range_map[x][1]
+      _realize_axis = cast(list[int], rctx.realize_map[x])
+      consumers = [rctx.range_map[u][0] for u in consumer_map[x] if u in rctx.range_map]
+      if len(consumers) < 2: continue
+      assert all(len(out_rngs) == len(rr) for rr in consumers)
+      for i,c in enumerate(zip(*consumers)):
+        out_rng = out_rngs[i]
+        # check if they are all simple ranges
+        if not all(y.op is Ops.RANGE and y.vmax == out_rng.vmax for y in c): continue
+        for r in c: replaced_ranges[r] = out_rngs[i]
+        _realize_axis.remove(i)
+      if len(_realize_axis) == 0: del rctx.realize_map[x]
+      else: rctx.realize_map[x] = _realize_axis
+
+    # do all the replaces
+    for k,(v0,v1) in rctx.range_map.items():
+      rctx.range_map[k] = (tuple(x.substitute(replaced_ranges) for x in v0), tuple(x.substitute(replaced_ranges) for x in v1))
+
   tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   return tsink, rctx

tinygrad/schedule/rangeify.py

@@ -4,9 +4,11 @@ from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
 from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate
 from tinygrad.uop.symbolic import symbolic_flat
-from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, Metadata
+from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, \
+  Metadata, DEBUG_RANGEIFY, MULTIOUTPUT
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented
 from tinygrad.codegen.opt import Opt
+from tinygrad.codegen.late.control_flow import pm_merge_ends
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTIGUOUS, IndexingContext, apply_movement_op
 
 # creation can recurse a lot
@@ -477,7 +479,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(tsink, earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites")
 
   # convert movement ops to ranges
-  tsink, rctx = run_rangeify(tsink, getenv("DEBUG_RANGEIFY", 0))
+  tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
 
   # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
   tsink = graph_rewrite(tsink, symbolic_flat+pm_reduce_unparented, name="symbolic")  # this supports const folding
@@ -494,6 +496,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
 
   # bufferize -> store
   tsink = graph_rewrite(tsink, pm_add_buffers, bottom_up=True, name="bufferize to store")
+  if MULTIOUTPUT: tsink = graph_rewrite(tsink, pm_merge_ends, name="merge end ranges")
   tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels")
 
   # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign

tinygrad/uop/ops.py

@@ -362,7 +362,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
   def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), (self,)+src, **kwargs)
   def end(self, *src:UOp, ends:Sequence[UOp]):
-    if len(ends) == 0: return self
+    if len(ends) == 0:
+      if len(src): return UOp(Ops.NOOP, src=(self,*src))
+      return self
     return UOp(Ops.END, src=(*ends, self, *src), arg=len(ends))
   def after(self, *src:UOp): return UOp(Ops.AFTER, self.dtype, (self,)+src)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))