remove MSTACK from allreduce

test/unit/test_allreduce.py

@@ -26,7 +26,8 @@ class TestRingAllReduce(unittest.TestCase):
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
       self.assertEqual(len(copies), 24)
-      self.assertEqual(len(sinks), 26)
+      # 4 fewer than MSTACK: the reduce-scatter kernel stores straight into the scratch MSELECT, no redundant diagonal copies
+      self.assertEqual(len(sinks), 22)
 
   @Context(RING=0, ALL2ALL=0)
   def test_schedule_naive(self):

tinygrad/schedule/allreduce.py

@@ -21,19 +21,30 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
 
   # naive: copy to all devices. if you shrink later, that'll be handled
   if not use_ring and not use_all2all:
-    return functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)])
+    if isinstance(device, str):
+      return functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)])
+    # copies are gathered into per-source scratch buffers with MSELECT stores, a device's own shard is read directly
+    dnum = UOp.variable("_device_num", 0, ndev-1)
+    terms:list[UOp] = []
+    for i in range(ndev):
+      scratch = UOp.new_buffer(device, numel, buf.dtype)
+      state = scratch.after(*[scratch.mselect(j).store(buf.mselect(i).copy_to_device(device[j])) for j in range(ndev) if j != i])
+      terms.append(dnum.eq(i).where(buf, state.reshape(shape)))
+    return functools.reduce(lambda x,y: x.alu(op, y), terms)
 
   # chunk data into ndev pieces
   factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
   base, left = divmod(numel // factor,  ndev)
   chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0)))
 
-  # reduce-scatter
+  # reduce-scatter: with all2all chunk i is reduced on buf.device[i], with ring it ends up on buf.device[(i-1)%ndev]
   reduced_chunks:list[UOp] = []
+  starts:list[int] = []
   for i,(s,e) in enumerate(chunks):
     if use_all2all:
       chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(ndev)]
       reduced_chunks.append(functools.reduce(lambda x,y: x.alu(op, y), chunks_on_i))
+      starts.append(i)
     else:
       chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),))
       for step in range(ndev-1):
@@ -41,20 +52,28 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
         cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None)
         reduced = cp.alu(op, chunk.copy_to_device(buf.device[dest], dest))
       reduced_chunks.append(reduced)
+      starts.append((i+ndev-1)%ndev)
 
-  # allgather
-  copied_chunks:list[UOp] = []
-  for i,rc in enumerate(reduced_chunks):
-    if isinstance(device, str): copied_chunks.append(rc.copy_to_device(device))
-    elif use_all2all: copied_chunks.append(UOp.mstack(*(rc.copy_to_device(buf.device[j]) for j in range(ndev))))
+  # single device output: copy the reduced chunks straight to the output device
+  if isinstance(device, str):
+    return UOp.usum(*[rc.copy_to_device(device).pad(((s,numel-e),)) for (s,e),rc in zip(chunks, reduced_chunks)]).reshape(shape)
+
+  # allgather: store each chunk into an MSELECT of a scratch buffer, then copy it to the other devices
+  gathered:list[UOp] = []
+  for (s,e),rc,start in zip(chunks, reduced_chunks, starts):
+    scratch = UOp.new_buffer(device, e-s, buf.dtype)
+    state = scratch.after(scratch.mselect(start).store(rc))
+    if use_all2all:
+      state = state.after(*[scratch.mselect(j).store(state.mselect(start).copy_to_device(device[j])) for j in range(ndev) if j != start])
     else:
-      chain:list[UOp] = [rc]
-      for step in range(ndev-1):
-        chain.append(rc := rc.copy_to_device(buf.device[(i+step)%ndev]))
-      copied_chunks.append(UOp.mstack(*(chain[(j-i+1)%ndev] for j in range(ndev))))
+      # forward the chunk around the ring
+      for step in range(1, ndev):
+        prev, dest = (start+step-1)%ndev, (start+step)%ndev
+        state = state.after(scratch.mselect(dest).store(state.mselect(prev).copy_to_device(device[dest])))
+    gathered.append(state)
 
   # reassemble
-  return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
+  return UOp.usum(*[g.pad(((s,numel-e),)) for (s,e),g in zip(chunks, gathered)]).reshape(shape)
 
 def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
   if output is None: output = UOp.const(red.dtype, Invalid, shape=red.shape).clone(device=red.device)