diff --git a/tinygrad/ops.py b/tinygrad/ops.py index c0e0571c15..0d49080f22 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -951,7 +951,7 @@ def div_and_mod_folding(x: UOp, c: int, which: Literal[Ops.MOD, Ops.IDIV], split def lt_folding(x:UOp, c:int) -> Optional[UOp]: p, np = partition(split_uop(x, Ops.ADD), lambda u: u.const_factor() == 1) if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d: - return cast(UOp, functools.reduce(operator.add, np).divides(d)).lt(c//d) + return cast(UOp, functools.reduce(operator.add, np).divides(d))<(c//d) return None def fold_unrolled_divs(divs:UOp): @@ -1132,14 +1132,14 @@ symbolic = symbolic_simple+PatternMatcher([ ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2)), # (x//c1)//c2 -> x//(c1*c2) # ** lt ** # c0*x 0 and c1.arg > 0 else None), + ((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.ints)) 0 and c1.arg > 0 else None), # c0*x 0 else None), + ((UPat.var("x", dtype=dtypes.ints)//UPat.cvar("c0", vec=False)) 0 else None), # ** move add/mul consts to end (NOTE: this is still happening before constant folding) ** (UPat(Ops.ADD, src=(UPat.var("x"), UPat.cvar("c1"))) + UPat.var("y"), lambda x,c1,y: (x+y)+c1), (UPat(Ops.MUL, src=(UPat.var("x"), UPat.cvar("c1"))) * UPat.var("y"), lambda x,c1,y: (x*y)*c1), @@ -1147,10 +1147,10 @@ symbolic = symbolic_simple+PatternMatcher([ # unrolled arange div folding (UPat(Ops.ADD, name="divs", src=[UPat(), UPat(Ops.IDIV)]), fold_unrolled_divs), # generic lt folding - (UPat.var("x", dtypes.sints).lt(UPat.cvar("c", vec=False)), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None), + (UPat.var("x", dtypes.sints) 0 # not x < 1 -> X > 0 - (UPat.var("x", dtypes.ints).lt(1).ne(True), lambda x: newx.lt(1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None), + ((UPat.var("x", dtypes.ints)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None), # ** div ** # # div folding (UPat.var("x", dtypes.sints) // UPat.cvar("c", vec=False), lambda x,c: div_and_mod_folding(x,c.arg,Ops.IDIV) if 0 < c.arg else None),