Merge branch 'master' into invalid_try_3

multi_migration.md

@@ -0,0 +1,86 @@
+# Multi-device op migration: MULTI/MSELECT/MSTACK → PAD / WHERE / STACK+INDEX
+
+## Status (updated)
+
+- **Stage 0 — DONE.** Internal `Ops.PAD` fills **Invalid** (`schedule/indexing.py:104`, bool keeps 0-fill); external `Tensor.pad`/`pad_to` always emit an explicit fill mask (`mixin/op.py:289`, `mixin/movement.py:267`) — required because a bare Invalid-pad leaks through elementwise ALU (`pad(x)+1` would read 0 instead of 1 in pad regions). REDUCE inputs with Invalid contribute the reduce identity (`pm_invalid_reduce_identity` in `uop/symbolic.py`, run in `get_kernel_graph` after gate lifting in `schedule/rangeify.py`) — only WHERE-alt gates whose condition involves a reduce range are rewritten, so gather-with-Invalid-index still poisons whole lanes. Same-condition nested where collapse rule added (`c?(c?t:f):f2 -> c?t:f2`) so the mask form folds to a single gate. All suites green (`test/unit`, `test/null`, `test/backend`, `test/external/external_test_schedule_scaling.py`, mypy, ruff).
+- **Stage 1 — representation in place behind `SYMBOLIC_MULTI`.** `symbolic_multi_pm` (`schedule/multi.py`) converts `MULTI→_unshard` (raw Invalid pad), `MSELECT→dnum.eq(i).where(x, Invalid)`, `MSTACK→STACK.index(dnum)` + INDEX(STACK,var)→nested-where lowering. `_unshard` uses the raw Invalid pad; `_unshard_fill` (0-fill) is used for the ALU allreduce in `copy_multi` because gated stores leave stale pad regions (the ALU-sum path can't use the identity rule). Basic shard ops work; full parity is Stage 2.
+- **Stage 2 — remaining.** Buffer level, reduce/allreduce split for shard-axis reduces, API surface.
+- **Stage 3 — remaining.**
+
+Notes: `test_schedule.py:test_pad_reduce_unsafe_multiview_st` went 4→5 kernels (pad now materializes an explicit mask; the mask form is also what makes the previously-wrong masked-pad+hazard case correct). `test_jit_footguns.py:test_symbolic_pad_view_frozen` went 2→4: the explicit mask recomputes from the symbolic shape, fixing the frozen-pad footgun. Also fixed a latent infinite loop: `(x+y) !=/< c → x !=/< c-y` collapse rules in `codegen/simplify.py` now only fire when the remaining side still contains the range (they previously shuffled constants forever when both sides were range-free).
+
+## Goal
+
+Replace the three multi-device UOps with a symbolic `_device_num` representation:
+
+| Old op | New form |
+|---|---|
+| `MULTI(x, axis)` | `x._unshard(axis)` — PAD with `_device_num`-dependent bounds back to full shape (helper already exists at `tinygrad/uop/ops.py:704-707`) |
+| `MSELECT(x, i)` | `dnum.eq(i).where(x, x.const_like(Invalid))` |
+| `MSTACK(s0..sn)` | `UOp(Ops.STACK, src=srcs).index(dnum)` — leading device axis, indexed per-device |
+
+where `dnum = UOp.variable("_device_num", 0, ndev-1)`. The per-device specialization mechanism already exists: `unwrap_multi` (`tinygrad/engine/realize.py:148-153`) binds `_device_num` per device at exec time.
+
+**Key semantic decision (approved):** internal `Ops.PAD` produces **Invalid** in padded regions; external `Tensor.pad` API still pads with 0. Staged migration: introduce the new representation first, keep old ops working, migrate call sites incrementally, delete old ops last.
+
+## Background: current design
+
+- `Ops.MULTI(src, axis)` (`tinygrad/uop/__init__.py:100`) — per-shard graph marker. Eliminated by `multi_pm` (`tinygrad/schedule/multi.py:162-195`) as the first step of `get_kernel_graph` (`tinygrad/schedule/rangeify.py:548`). Shape/axis tracking: `UOp.axis`/`UOp.bounds` (`tinygrad/uop/ops.py:667-702`).
+- `Ops.MSELECT(x, i)` / `Ops.MSTACK(srcs)` (`__init__.py:96`) — buffer-level ops. Spec at `tinygrad/uop/spec.py:181-184`; device prop `ops.py:816-819`; per-kernel PARAMs via debuf (`rangeify.py:474`); per-device dependency states (`tinygrad/schedule/__init__.py:11-17`); `MultiBuffer` (`ops.py:904-930`, `tinygrad/device.py:88-99`); only MSTACK can be `realized` (`ops.py:920-930`).
+- `_shard`/`_unshard` (`ops.py:704-714`) already emit symbolic SHRINK/PAD bounds with `_device_num`.
+- Naive allreduce already uses the target pattern: `dnum.eq(i).where(buf, state)` (`tinygrad/schedule/allreduce.py:27-33`).
+
+## Existing Invalid machinery (rely on this)
+
+- `pm_data_invalid` (`tinygrad/uop/symbolic.py:71-92`): Invalid poisons ALU (ops move inside the gate); gated LOAD folds to alt/0, gated STORE folds to NOOP.
+- `pm_remove_invalid` (`symbolic.py:94-96`): leftover Invalid → 0 in final codegen (`codegen/__init__.py:345`). Spec forbids Invalid in final programs (`spec.py:217`), so materialized Invalid regions read as 0.
+- STORE of CONST(Invalid) → NOOP (`rangeify.py:423-424`).
+- `identity_element(op, dtype)` exists (`ops.py:51`): ADD→0, MUL→1, MAX→dtype.min.
+- `found_after` (`rangeify.py:26`) already matches `WHERE(cond, PAD(x), Invalid)`.
+
+## Stage 0 — internal PAD = Invalid; external pad = explicit 0
+
+1. `tinygrad/schedule/indexing.py:100-104` (`convert_pad_to_where_to_keep_behavior_local`): fill value `0` → `UOp.const(x.dtype, Invalid)`, **except `dtypes.bool` keeps 0-fill** (False is the bool-reduce identity, and the external-pad mask below needs it).
+2. `tinygrad/mixin/op.py:282-290` (`_pad_constant`): **remove the `if value == 0: return base` shortcut** — always emit `pad(bool_ones).where(base, value)`. Required because bare Invalid-pad leaks through elementwise ALU: `pad(x)+1` gate-lifts to `where(valid, x+1, Invalid)` and reads 0 instead of 1 in pad regions. The mask lowers to a pure index expression (`valid.where(1,0)`), no extra kernel. External behavior unchanged for all `value`.
+3. **New rule**: `REDUCE(where(c, x, Invalid), op)` → `REDUCE(where(c, x, identity_element(op, dtype)), op)`. Must fire in rangeify/symbolic *before* codegen builds the accumulator loop — otherwise `pm_data_invalid` gate-lifts `acc + where(c,x,Invalid)` into `where(c, acc+x, Invalid)` and one invalid lane poisons the whole reduction. Placement (symbolic.py vs the reduce path in indexing.py) TBD at implementation; verify with `Tensor.pad(...).sum()/max()` tests.
+4. Audit: schedule tests with kernel counts involving pads; circular/reflect/replicate pads don't use PAD fill (verified, `op.py:292-312`) — unaffected; `allreduce.py:59,76` usum-of-padded-chunks gets *more* correct (disjoint regions).
+
+## Stage 1 — new representation behind env flag
+
+New `symbolic_multi_pm` PatternMatcher (in `schedule/multi.py` or new file), gated by env (e.g. `SYMBOLIC_MULTI`), run in `get_kernel_graph` right after `multi_pm`:
+
+- `MULTI(x, axis)` → `x._unshard(axis)`
+- `MSELECT(x, i)` → `dnum.eq(i).where(x, x.const_like(Invalid))` (Invalid from `tinygrad.dtype`)
+- `MSTACK(srcs)` → `STACK(*srcs).index(dnum)`, plus new lowering `INDEX(STACK(vals), var)` → nested `var.eq(k).where(src_k, Invalid)` (analogous to `convert_stack_to_where`, `indexing.py:113-121`; must fire before `validate_index` spec, `spec.py:118-122`)
+
+Flag off = zero behavior change; flag on = new forms flow through rangeify and specialize per device at exec.
+
+## Stage 2 — migrate producers/consumers (one commit each, independently testable)
+
+1. `UOp.shard` (`ops.py:715-717`): emit symbolic `_shard`+`_unshard` full-shape form directly instead of `.multi(axis)`; delete movement-op `multi_pm` rules that PAD subsumes (`pad_multi`, `permute_multi`, `expand_multi`, `reshape_multi`, `flip_multi`, `shrink_multi` — `multi.py:93-125`).
+2. ALU/STACK: `alu_multi`/`shard_srcs`/`stack_multi` (`multi.py:55-78,127-131`) become plain elementwise on full-shape padded tensors. `reduce_multi` (`multi.py:80-91`) keeps the shard-axis → local-reduce + ALLREDUCE split; Invalid-pad + identity rule replaces neutral-pad-value reasoning.
+3. allreduce (`schedule/allreduce.py`): naive path already matches; migrate ring/all2all MSELECT/MSTACK scratch-buffer assembly (lines 35-76) to WHERE/STACK+INDEX forms.
+4. Buffer level: debuf (`rangeify.py:474`), `_states`/`_unwrap_src` (`schedule/__init__.py:11-17`), `_collect_bufs` (`schedule/memory.py:9`), `unwrap_multi` (`realize.py:148-153`), JIT (`jit.py:127-130, 237`), callify (`callify.py:52-95`), `buffer`/`realized`/`buf_uop`/`has_buffer_identity` (`ops.py:841-930`).
+5. API surface: `UOp.multi/mselect/mstack` (`ops.py:662-725`), `Tensor.shard` (`tensor.py:333-347`), gradient (`mixin/gradient.py:72`), `_multi_like` (`mixin/creation.py:16-20`), embedding backward (`nn/__init__.py:309-354`), `copy_to_device(arg=)` MSELECT path (`ops.py:719-723`).
+
+## Stage 3 — removal
+
+Delete `Ops.MULTI/MSELECT/MSTACK` from the enum (`uop/__init__.py:96,100`), spec rules, viz colors (`viz/serve.py:51,56`), `UOp.axis`/`bounds` machinery (`ops.py:667-702`), remaining `multi_pm` rules, and `MultiBuffer` if fully subsumed. Flip flag default-on, then delete the flag.
+
+## Open implementation details
+
+- REDUCE-identity rule placement (must precede codegen accumulator construction).
+- INDEX(STACK, var) spec timing — the value-STACK INDEX violates the pointer-INDEX spec until lowered.
+- Whether `MultiBuffer`/tuple-`device` survives as the runtime container, or buffers become single-device with the device axis explicit in shape — decides how much of Stage 2.4 is rewrite vs delete.
+- Bool carve-out in Stage 0.1: verify no internal consumer needs Invalid-filled bool pads.
+
+## Verification (run at each stage)
+
+```bash
+python -m pytest test/unit/test_multitensor.py test/unit/test_allreduce.py test/null/test_multitensor.py test/unit/test_call.py -x -q -n12
+python -m pytest test/external/external_test_schedule_scaling.py -x -q   # test_concat_scaling
+python -m mypy tinygrad/
+python -m ruff check .
+```
+
+Also pad/reduce numeric tests after Stage 0 (`test_ops` pad tests, `Tensor.pad(...).sum()/max()`).

test/null/test_schedule.py

@@ -600,7 +600,8 @@ class TestSchedule(unittest.TestCase):
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
     # TODO: this should be 3 if fix store hazard worked correctly
-    check_schedule(p, 4)
+    # NOTE: pad now always has an explicit fill mask (internal PAD is Invalid-filled), which materializes here
+    check_schedule(p, 5)
 
   def test_conv2d(self, allowed=4, dtype=dtypes.float):
     old_default_float, dtypes.default_float = dtypes.default_float, dtype

test/null/test_uops.py

@@ -8,6 +8,7 @@ from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
+from tinygrad.mixin.movement import MovementMixin
 from test.helpers import eval_uop, to_uops_list
 
 class TestDTypeFromUOp(unittest.TestCase):
@@ -457,7 +458,7 @@ class TestContiguousViewOffset(unittest.TestCase):
   def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7)
   def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1)
   def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None)
-  def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
+  def test_shrink_invalid(self): self._check(MovementMixin.pad(UOp.empty(4), ((2,2),))[0], None)
   def test_strided(self): self._check(UOp.empty(4)[::2], None)
 
 if __name__ == '__main__':

test/unit/test_jit_footguns.py

@@ -137,11 +137,11 @@ class TestJitFootguns(unittest.TestCase):
     from tinygrad import Variable
     a = Tensor.rand(3, 10).realize()
 
-    # broken: pad is a view, BIND values frozen at capture (i=2)
+    # fixed: pad now has an explicit fill mask (internal PAD is Invalid-filled), which recomputes from the symbolic shape
     @TinyJit
     def f_broken(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
     for i in range(1, 5): f_broken(a[:, :Variable("i", 1, 10).bind(i)])
-    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2)  # should be 4!
+    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 4)
 
     # workaround: contiguous fuses pad into kernel
     @TinyJit

tinygrad/codegen/simplify.py

@@ -95,11 +95,12 @@ pm_reduce_unparented = PatternMatcher([
 ])
 
 pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
-  # lift x+y out of reduce on lt
-  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  # lift x+y out of reduce on lt. only fire if x still has the range: with both sides range-free it just shuffles constants
+  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"),
+   lambda x,y,c: (x < (c.cast(y.dtype)-y)) if not no_range(x) and no_range(y) and no_range(c) else None),
   # lift x*y out of reduce
   ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
+    lambda x,y,c: (x < ((c+y-1) // y)) if not no_range(x) and no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
   # sum over r in [0,N) of [lower<=r<upper]*val -> clamp(min(upper,N) - max(lower,0), 0, N) * val
   (UPat.any(
     (UPat(Ops.RANGE, name="r") < UPat.var("upper")).where(UPat.var("val"), 0),
@@ -119,8 +120,9 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
 ])+symbolic
 
 pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
-  # lift x+y out of reduce on ne
-  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
+  # lift x+y out of reduce on ne (same range guard as the lt version)
+  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"),
+   lambda x,y,c: (x != (c.cast(y.dtype)-y)) if not no_range(x) and no_range(y) and no_range(c) else None),
   # reduce on gated load becomes can substitute the range and remove the reduce
   ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
    lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),

tinygrad/mixin/movement.py

@@ -265,7 +265,8 @@ class MovementMixin:
     return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))
 
   def pad_to(self, shape, *args) -> Self:
-    return self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)))
+    # NOTE: this calls the overridden pad (OpMixin.pad when available) so the fill is an explicit 0, not Invalid
+    return self.pad(tuple((0, 0) if ns is None else (0, ns-s) for s, ns in zip(self.shape, argfix(shape, *args), strict=True)))
 
   def view(self, shape, *args) -> Self:
     """`.view` is an alias for `.reshape`."""

tinygrad/mixin/op.py

@@ -286,8 +286,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     X = self.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, self.shape))) if has_neg else self
     pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
     base = MovementMixin.pad(X, pads)
-    if value == 0: return base
-    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, value)
+    if base is X: return X  # no padding, nothing to fill
+    # the fill is always explicit: internal PAD fills with Invalid, so the mask is required for every value (incl. 0)
+    # for 0 use a literal that loses every promotion: a Python 0.0 would promote int/bool tensors to float, int 0 bool to int
+    fill = (False if X.dtype == dtypes.bool else 0) if value == 0 else value
+    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, fill)
 
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder

tinygrad/schedule/indexing.py

@@ -1,7 +1,7 @@
 from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes
 from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
@@ -101,7 +101,8 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
   if x not in ctx.range_map: return None
   bx = create_bufferize_and_index_based_on_ranges(ctx, x)
   valid: UOp = UOp.const(dtypes.bool, True).uprod([r.get_valid() for r in ctx.range_map[x][0]])
-  return valid.where(bx.src[0], UOp.const(x.dtype, 0))
+  # internal PAD fills with Invalid. bool keeps 0-fill: False is the bool reduce identity and external pad masks need it
+  return valid.where(bx.src[0], UOp.const(x.dtype, 0 if x.dtype == dtypes.bool else Invalid))
 
 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   if x.arg[1] == 0: return None

tinygrad/schedule/multi.py

@@ -1,6 +1,6 @@
 from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, graph_rewrite, broadcast_axes, _broadcast_shape
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, Invalid
 from tinygrad.schedule.allreduce import handle_allreduce
 
 # ***** multi rewrite MSELECT/MSTACK *****
@@ -43,6 +43,34 @@ _early_allreduce = PatternMatcher([
 ])
 if not getenv("LATE_ALLREDUCE", 1): replace_allreduce = _early_allreduce + replace_allreduce
 
+# ***** symbolic multi rewrite (SYMBOLIC_MULTI) *****
+# replaces MULTI/MSELECT/MSTACK with the symbolic _device_num representation:
+#   MULTI(x, axis) -> x._unshard(axis): PAD with _device_num-dependent bounds back to full shape, other shards Invalid
+#   MSELECT(x, i) -> dnum==i ? x : Invalid
+#   MSTACK(srcs)  -> STACK(srcs).index(dnum), lowered to nested dnum==k ? src_k : Invalid
+# the per-device specialization binds _device_num at exec time (unwrap_multi in engine/realize.py)
+
+def _dnum(ndev:int) -> UOp: return UOp.variable("_device_num", 0, ndev-1)
+
+def mselect_to_where(ms:UOp) -> UOp:
+  return _dnum(len(ms.src[0].device)).eq(ms.arg).where(ms.src[0], ms.src[0].const_like(Invalid))
+
+def mstack_to_stack_index(ms:UOp) -> UOp:
+  return UOp(Ops.STACK, src=ms.src).index(_dnum(len(ms.src)))
+
+def index_stack_to_where(stack:UOp, var:UOp) -> UOp:
+  ret = stack.src[0].const_like(Invalid)
+  for k in range(len(stack.src)-1, -1, -1): ret = var.eq(k).where(stack.src[k], ret)
+  return ret
+
+symbolic_multi_pm = PatternMatcher([
+  (UPat(Ops.MULTI, src=(UPat(),), name="multi"), lambda multi: multi.src[0]._unshard(multi.arg)),
+  (UPat(Ops.MSELECT, src=(UPat(),), name="ms"), mselect_to_where),
+  (UPat(Ops.MSTACK, name="ms"), mstack_to_stack_index),
+  # lower INDEX into a value-STACK to nested selects (must fire before the pointer-INDEX spec in codegen)
+  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="stack"), UPat.var("var"))), index_stack_to_where),
+])
+
 # ***** multi functions *****
 
 def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
@@ -128,7 +156,7 @@ def copy_multi(multi:UOp, device:str | tuple[str, ...]):
   if isinstance(device, str):
     pieces = [multi.src[0].mselect(i).copy_to_device(device) for i in range(len(multi.device))]
     return pieces[0].cat(*pieces[1:], dim=multi.axis)
-  return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device)
+  return multi.src[0]._unshard_fill(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)
 

tinygrad/schedule/rangeify.py

@@ -4,14 +4,14 @@ import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, pm_invalid_reduce_identity
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
-from tinygrad.schedule.multi import multi_pm
+from tinygrad.schedule.multi import multi_pm, symbolic_multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function
 
 # creation can recurse a lot
@@ -545,7 +545,8 @@ pm_copy_to_store = PatternMatcher([
 
 @profile_matches
 def get_kernel_graph(sink:UOp) -> UOp:
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
+  # SYMBOLIC_MULTI replaces the per-shard multi_pm rules with the symbolic _device_num representation (Stage 1)
+  tsink = graph_rewrite(sink, symbolic_multi_pm if getenv("SYMBOLIC_MULTI") else multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
 
@@ -555,6 +556,8 @@ def get_kernel_graph(sink:UOp) -> UOp:
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
   tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, name="symbolic+reduce_collapse+debuf")
+  # Invalid (from internal PAD) in reduce inputs contributes the reduce identity, must run after gate lifting
+  tsink = graph_rewrite(tsink, pm_invalid_reduce_identity, name="reduce invalid to identity")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")

tinygrad/uop/ops.py

@@ -228,6 +228,7 @@ class recursive_property(property):
 
 # we import this late so we can use resolve/smax in mixins
 from tinygrad.mixin.op import OpMixin
+from tinygrad.mixin.movement import MovementMixin
 from tinygrad.mixin.rand import RandMixin
 
 # NOTE: this should be frozen, but frozen is slower
@@ -704,6 +705,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return src_axis
 
   def _unshard(self, axis:int) -> UOp:
+    bsz, dcount = self.shape[axis], len(self.device)
+    dnum = UOp.variable("_device_num", 0, dcount-1)
+    # raw PAD with _device_num-dependent bounds: the other shards' regions are Invalid ("no data"), never a fill value
+    return MovementMixin.pad(self, tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))
+
+  def _unshard_fill(self, axis:int) -> UOp:
+    # 0-filled variant of _unshard for ALU allreduce: materialized pad regions must be explicit 0, not gated-stale Invalid
     bsz, dcount = self.shape[axis], len(self.device)
     dnum = UOp.variable("_device_num", 0, dcount-1)
     return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))

tinygrad/uop/symbolic.py

@@ -1,7 +1,7 @@
 # all of symbolic lives here now
 import math, struct
 from collections import defaultdict
-from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu, identity_element
 from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid
 from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
 from tinygrad.uop.divandmod import div_and_mod_symbolic
@@ -98,6 +98,32 @@ pm_remove_invalid = PatternMatcher([
    if any(x.arg is Invalid for x in s.src) else None),
 ])
 
+# Invalid in a reduce input means "no data": those lanes must contribute the reduce identity, not poison the accumulator.
+# this must run after pm_data_invalid gate lifting (gates float to the top of the reduce input) and before codegen builds
+# the accumulator loop, otherwise acc+where(c,x,Invalid) gate-lifts and one invalid lane poisons the whole reduction.
+# only WHERE alt gates whose condition involves a reduce range are rewritten: those are "this element has no data".
+# an Invalid behind a reduce-range-independent gate (e.g. gather with an Invalid index) poisons the whole lane: keep it.
+def _invalid_to_identity(u:UOp, ident:UOp, red_ranges:frozenset[UOp]) -> UOp|None:
+  if u.op is Ops.WHERE:
+    if u.src[2].base.op is Ops.CONST and u.src[2].base.arg is Invalid and u.src[2].dtype == ident.dtype:
+      alt = ident if not red_ranges.isdisjoint(u.src[0].ranges) else None
+    else: alt = _invalid_to_identity(u.src[2], ident, red_ranges)
+    then = _invalid_to_identity(u.src[1], ident, red_ranges)
+    if alt is None and then is None: return None
+    return u.replace(src=(u.src[0], u.src[1] if then is None else then, u.src[2] if alt is None else alt))
+  if u.op in GroupOp.Elementwise-{Ops.WHERE}:
+    new_srcs = tuple(_invalid_to_identity(s, ident, red_ranges) for s in u.src)
+    if all(n is None for n in new_srcs): return None
+    return u.replace(src=tuple(s if n is None else n for s,n in zip(u.src, new_srcs)))
+  return None
+
+def reduce_invalid_identity(r:UOp) -> UOp|None:
+  red_ranges = frozenset(x for x in r.src[1:] if x.op is Ops.RANGE)
+  new_src = _invalid_to_identity(r.src[0], r.const_like(identity_element(r.arg[0], r.dtype)), red_ranges)
+  return r.replace(src=(new_src,)+r.src[1:]) if new_src is not None else None
+
+pm_invalid_reduce_identity = PatternMatcher([(UPat(Ops.REDUCE, name="r"), reduce_invalid_identity)])
+
 symbolic_simple = pm_data_invalid + PatternMatcher([
   # ** self folding **
   (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
@@ -180,6 +206,8 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
   # a.where(b.where(c, d), d) -> (a & b).where(c, d)
   (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
+  # nested where with the same condition in the then position: c ? (c ? t : f) : f2 -> c ? t : f2
+  (UPat.var("c").where(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("f2")), lambda c,t,f,f2: c.where(t, f2)),
   # a.where(c, b.where(c, d)) -> (a | b).where(c, d)
   (UPat.var("a").where(UPat.var("c"), UPat.var("b").where(UPat.var("c"), UPat.var("d"))), lambda a,b,c,d: (a|b).where(c,d)),
 ])+mop_cleanup