diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 9f2e985efaa9c..cdb4219dd6e70 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -1355,11 +1355,12 @@ def test_invalid_gate_simplifies_vectorize(self): idx0 = (r0 + uconst(-1)) // uconst(3) idx1 = r0 % uconst(3) - idx:UOp = (r0 < 3).where(UOp(Ops.STACK, dtypes.weakint, (idx0, idx1)), UOp.invalid()) + # raw WHERE: the carrier keeps a bare Invalid; mixin where would broadcast the ()-shaped cond/Invalid to the STACK's shape + idx:UOp = (r0 < 3)._select(UOp(Ops.STACK, dtypes.weakint, (idx0, idx1)), UOp.invalid()) idx = graph_rewrite(idx, pm_simplify_valid) # independent simplification: (r0-1)//3 -> (r0+2)//3 - 1, and r0%3 -> r0 when r0 in [0,2] expected_vec = UOp(Ops.STACK, dtypes.weakint, ((r0 + uconst(2)) // uconst(3) + uconst(-1), r0)) - self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid())) + self.assertEqual(idx, (r0 < 3)._select(expected_vec, UOp.invalid())) class TestRangeSplitting(unittest.TestCase): def test_range_split_on_mod(self): diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 40e18af9a7b19..317c9ff4a9a81 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -197,7 +197,7 @@ def apply_opt(self, opt:Opt, append_opt:bool=True): for b in self.bufs: if rng in (i:=b.src[1].get_idx()).backward_slice_with_self: nb = b.replace(src=(b.src[0], i.valid(valid&b.src[1].get_valid()))) - replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(b.dtype, Invalid)) + replaces[b] = nb if b in store_targets else valid._select(nb, UOp.const(b.dtype, Invalid)) self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}") elif opt.op is OptOps.SWAP: try: diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 8f134735c8d85..3186db9d2d94b 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -111,7 +111,7 @@ def reduce_unparented(red:UOp) -> UOp|None: lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)), # AND on WHERE ((UPat(Ops.PARAM, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"), - lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)), + lambda x,y,c,r: y._select(c, c.const_like(0)).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)), # MUL casted bool ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)), ])+symbolic @@ -121,7 +121,8 @@ def reduce_unparented(red:UOp) -> UOp|None: ((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), # 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)), + lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0]))._select( + expr.substitute({r:idx.cast(r.dtype).valid(v)}), expr.const_like(0))), ]) def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp|None: diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index db865d964ceba..abe116c25fe1c 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -6,7 +6,7 @@ from tinygrad.mixin.creation import CreationMixin if TYPE_CHECKING: - from tinygrad.uop.ops import UOp + from tinygrad.uop.ops import UOp, sint class ElementwiseMixin(CreationMixin): @@ -19,7 +19,7 @@ def ufix(self, x: 'Self|ConstType|UOp') -> Self: return x if isinstance(x, type(self)) else self._wrap_uop(self._uop.ufix(x)) # implemented in OpMixin, broadcasting needs the movement ops - def _broadcasted(self, y: 'Self|ConstType|UOp', reverse: bool = False) -> tuple[Self, Self]: + def _broadcasted(self, y: 'Self|ConstType|UOp', reverse: bool = False, match_dtype: bool = True) -> tuple[Self, Self]: raise NotImplementedError def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self: @@ -411,10 +411,30 @@ def logaddexp(self, other: Self | ConstType) -> Self: m = a.maximum(b) return ((a-m).exp() + (b-m).exp()).log() + m - def where(self, x: Self | ConstType, y: Self | ConstType) -> Self: - ref: Self = x if isinstance(x, type(self)) else y if isinstance(y, type(self)) else \ - self.cast(least_upper_dtype(dtypes.from_py(x), dtypes.from_py(y))) - return self.alu(Ops.WHERE, ref.ufix(x), ref.ufix(y)) + def where(self, x: 'Self|ConstType|sint', y: 'Self|ConstType|sint') -> Self: + """ + Returns a tensor of elements selected from either `x` or `y`, depending on `self`. + `output_i = x_i if self_i else y_i`. + + ```python exec="true" source="above" session="tensor" result="python" + cond = Tensor([[True, True, False], [True, False, False]]) + print(cond.where(1, 3).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + Tensor.manual_seed(42) + cond = Tensor.randn(2, 3) + print(cond.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print((cond > 0).where(cond, -float("inf")).numpy()) + ``` + """ + if isinstance(x, type(self)): x, y = x._broadcasted(y) + elif isinstance(y, type(self)): y, x = y._broadcasted(x) + else: x, y = self.ufix(x)._broadcasted(y) + cond, x = self.cast(dtypes.bool)._broadcasted(x, match_dtype=False) + cond, y = cond._broadcasted(y, match_dtype=False) + return cond.alu(Ops.WHERE, x, y) def masked_fill(self, mask:Self, value:Self|PyConst) -> Self: """ diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index a6ed3ff06bb10..24aefae6bf870 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -357,7 +357,7 @@ def pad(self, padding:Sequence[sint]|Sequence[tuple[sint, sint]|None], mode:str= if mode in {"reflect", "replicate"}: return self._pad_reflect_replicate(pX, mode) raise NotImplementedError(f"{mode=} is not supported") - def _broadcasted(self, y:Self|ConstType|UOp, reverse:bool=False) -> tuple[Self, Self]: + def _broadcasted(self, y:Self|ConstType|UOp, reverse:bool=False, match_dtype:bool=True) -> tuple[Self, Self]: if not isinstance(y, type(self)): y = self.ufix(y) x, y = (self, y) if not reverse else (y, self) # ValueError: unsized ptr has shape (-1,) which can't broadcast; RuntimeError: shape mismatch @@ -365,7 +365,7 @@ def _broadcasted(self, y:Self|ConstType|UOp, reverse:bool=False) -> tuple[Self, out_shape = _broadcast_shape(x.shape, y.shape) x, y = x._broadcast_to(out_shape), y._broadcast_to(out_shape) except (RuntimeError, ValueError): pass - if x.dtype == y.dtype: return x, y + if x.dtype == y.dtype or not match_dtype: return x, y return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype) def dot(self, w:Self, dtype:DTypeLike|None=None) -> Self: diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 0d6a82630bb71..887d8775e5d69 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -84,7 +84,7 @@ 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)) + return valid._select(bx.src[0], UOp.const(x.dtype, 0)) def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp): if x.arg[1] == 0: return None diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 5ada5b299f112..21066d41948d5 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -6,7 +6,7 @@ from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, _from_np_dtype, _to_np_dtype, PyConst from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc -from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape +from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable from tinygrad.mixin.rand import RandMixin from tinygrad.schedule import create_linear_with_vars from tinygrad.device import Buffer, canonicalize_device @@ -473,32 +473,6 @@ def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None: def __delitem__(self, indices) -> None: raise TypeError("Tensor does not support deleting items") - # ***** broadcasted elementwise ops ***** - - def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor: - """ - Returns a tensor of elements selected from either `x` or `y`, depending on `self`. - `output_i = x_i if self_i else y_i`. - - ```python exec="true" source="above" session="tensor" result="python" - cond = Tensor([[True, True, False], [True, False, False]]) - print(cond.where(1, 3).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - Tensor.manual_seed(42) - cond = Tensor.randn(2, 3) - print(cond.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print((cond > 0).where(cond, -float("inf")).numpy()) - ``` - """ - if isinstance(x, Tensor): x, y = x._broadcasted(y) - elif isinstance(y, Tensor): y, x = y._broadcasted(x) - else: x, y = self.ufix(x)._broadcasted(y) - out_shape = _broadcast_shape(self.shape, x.shape) - return self.cast(dtypes.bool)._broadcast_to(out_shape)._apply_uop(UOp.where, x._broadcast_to(out_shape), y._broadcast_to(out_shape)) - # ***** op wrappers ***** # unlike Tensors, UOps are immutable, so these don't go in mixin diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index aae7622ca39bf..9c86ed878b511 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -616,8 +616,9 @@ def _rop(self, op:Ops, axis:tuple[int, ...]): return ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis)) if axis != reduce_axis else ret @staticmethod def invalid(): return UOp.const(dtypes.weakint, Invalid) - def valid(self, cond): - return cond.where(self, self.const_like(Invalid)) + # TODO: we currently use WHERE without broadcasting and violates spec + def _select(self, x, y): return self.alu(Ops.WHERE, x, y) + def valid(self, cond): return cond._select(self, self.const_like(Invalid)) def get_idx(self) -> UOp: assert dtypes.is_int(self.dtype), "Can only call get_idx on index dtype" if self.op is Ops.STACK: return UOp.vectorize(*(x.get_idx() for x in self.src)) @@ -1308,7 +1309,7 @@ def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.match_dtype, ( def end(self, *src:UPat, **kwargs): return UPat(Ops.END, src=(self,)+src, **kwargs) def const_like(self, b:ConstLike): return UPat.const(self.match_dtype, cast(ConstType, b)) - def _broadcasted(self, y, reverse=False) -> tuple[UPat, UPat]: + def _broadcasted(self, y, reverse=False, match_dtype=True) -> tuple[UPat, UPat]: y = self.ufix(y) return (y, self) if reverse else (self, y) def ufix(self, x): return self.const_like(x) if not isinstance(x, UPat) else x @@ -1318,6 +1319,10 @@ def mod(self, x, reverse=False): return self._binop(Ops.FLOORMOD, x, reverse) def alu(self, op:Ops, *src:UPat): asrc = (self,)+src return UPat(op, dtypes.bool if op in {Ops.CMPLT, Ops.CMPNE} else asrc[-1].match_dtype, list(asrc) if op in GroupOp.Commutative else asrc) + # a pattern describes the WHERE node structurally, tensor-level broadcast/cast semantics don't apply + def where(self, x, y): + ref = x if isinstance(x, UPat) else y if isinstance(y, UPat) else self.cast(least_upper_dtype(dtypes.from_py(x), dtypes.from_py(y))) + return self.alu(Ops.WHERE, ref.ufix(x), ref.ufix(y)) def match(self:UPat, uop:UOp, store:dict[str, UOp]) -> list[dict[str, UOp]]: if self.is_any: return flatten([x.match(uop, store.copy()) for x in self.src[0]]) @@ -1672,7 +1677,7 @@ def select_dtype(u): return dtypes.long if u.overflows(dtypes.int32) else dtypes x.cast(dt:=least_upper_dtype(select_dtype(u), x.dtype, y.dtype)).alu(u.op, y.cast(dt)).cast(u.dtype)), (UPat(Ops.CONST, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None), (UPat(Ops.WHERE, dtypes.weakint, src=(UPat.var("cond"), UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint))), lambda cond,x,y: - cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.weakint)), + cond._select(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt)).cast(dtypes.weakint)), (UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.weakint)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.weakint)), (UPat(Ops.STACK, src=UPat().cast(dtypes.weakint), name="v"), lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt) for s in v.src)).cast(dtypes.weakint)), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 4ffbeb6e036f6..6124e0a67d8ef 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -79,18 +79,18 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat) pm_data_invalid = PatternMatcher([ (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_pat,), name="op"), lambda i,op: i.cast(op.dtype)), - (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_gate,), name="op"), lambda cond,x,op,i: cond.where(op.replace(src=(x,)), i.cast(op.dtype))), + (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_gate,), name="op"), lambda cond,x,op,i: cond._select(op.replace(src=(x,)), i.cast(op.dtype))), # binary ops move inside the gate, with Invalid cast to the result dtype (bool for comparisons) - (UPat(GroupOp.Binary, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i.cast(alu.dtype))), - (UPat(GroupOp.Binary, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i.cast(alu.dtype))), + (UPat(GroupOp.Binary, src=(invalid_gate, UPat.var("y")), name="alu"), lambda cond,x,y,alu,i: cond._select(x.alu(alu.op,y), i.cast(alu.dtype))), + (UPat(GroupOp.Binary, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond._select(y.alu(alu.op,x), i.cast(alu.dtype))), (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i), # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid) - (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i), + (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not()._select(val, i) if val.arg != Invalid else i), # lift Invalid out: a.where(cond.where(x, Invalid), c) -> (~a|cond).where(a.where(x, c), Invalid) # when a is cond, ~a|cond is True and would drop the Invalid gate (losing the valid), so keep cond as the gate (UPat.var("a").where(invalid_gate, UPat.var("c")), lambda cond,i,x,a,c: - (cond if a is cond else (a.logical_not()|cond)).where(a.where(x,c), i) if c.arg != Invalid else None), - (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if b.arg != Invalid else None), + (cond if a is cond else (a.logical_not()|cond))._select(a._select(x,c), i) if c.arg != Invalid else None), + (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond)._select(a._select(b, x), i) if b.arg != Invalid else None), # fold gated LOAD/STORE (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(), UPat())), lambda i: UOp(Ops.NOOP)), (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(),), allow_any_len=True, name="x"), @@ -230,13 +230,13 @@ def canonicalize_simplex(X:UOp) -> UOp|None: (UPat.cvar("y") * (UPat.var("x", dtype=dtypes.weakint) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)), # y*(x+c) -> y*x + y*c # ** where folding ** (UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), - lambda cond, t, f: cond.where(f,t) if f.arg is not Invalid else None), + lambda cond, t, f: cond._select(f,t) if f.arg is not Invalid else None), # alu of two where with same conds can combine, only do if true branch or false branch is const (UPat(GroupOp.Binary, name="alu", src=(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("c").where(UPat.var("tt"), UPat.var("ff")))), \ - lambda alu,c,t,tt,f,ff: c.where(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), + lambda alu,c,t,tt,f,ff: c._select(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), # if its a plus we add the associative variation too ((UPat.var("y")+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \ - lambda y,c,t,tt,f,ff: y+c.where(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), + lambda y,c,t,tt,f,ff: y+c._select(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), # ALU/variable min==max -> CONST (UPat({Ops.CMPLT, Ops.CMPNE, Ops.FLOORDIV, Ops.FLOORMOD, Ops.PARAM, Ops.BIND, Ops.SPECIAL}, name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None), @@ -372,7 +372,7 @@ def reduce_mul_chain(r:UOp) -> UOp|None: def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None: keep, drop = partition(cond.split_uop(Ops.AND), lambda c: any(r in x.ranges for r in c.ranges)) - return UOp.const(dtypes.bool, True).uprod(*keep).where(x, i) if drop else None + return UOp.const(dtypes.bool, True).uprod(*keep)._select(x, i) if drop else None pm_drop_and_clauses = PatternMatcher([(invalid_gate, drop_and_clauses)]) # move conditions from where to load's valid, drop clauses already in load @@ -387,7 +387,7 @@ def can_move(c:UOp) -> bool: if len(keep) == len(where_clauses): return None idx = buf.index(idx.get_idx().valid(load_valid.uprod(*moved))) ret_idx = idx.cast(or_cast.dtype) if or_cast.op is Ops.CAST else idx - return UOp.const(dtypes.bool, True).uprod(*keep).where(ret_idx, ret_idx.const_like(0)) + return UOp.const(dtypes.bool, True).uprod(*keep)._select(ret_idx, ret_idx.const_like(0)) # where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer pm_move_where_on_load = PatternMatcher([ @@ -400,7 +400,7 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None: if x.dtype is not dtypes.weakint: return None # Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD): return None - return cond.where(uop_given_valid(cond, x, try_simplex=False), i) + return cond._select(uop_given_valid(cond, x, try_simplex=False), i) # TODO: this is O(number of WHERE * number of node) # def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None: