one more merge

tinygrad/uop/divandmod.py

@@ -2,16 +2,18 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
 
-def cancel_divmod(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # simple cancel div/mod case when the range of the numerator lies within a single denominator interval
+def fold_divmod_general(d: UOp) -> UOp|None:
+  # get srcs
+  x,y = d.src
+
+  # cancel_divmod: simple cancel div/mod case when the range of the numerator lies within a single denominator interval
   x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
   assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
   if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
   if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
     return x - q*y if d.op is Ops.MOD else d.const_like(q)
-  return None
 
-def fold_divmod_general(d: UOp, x: UOp, y: UOp) -> UOp|None:
+  # split uops for the rest of the rules
   x_peeled, const = x.pop_const()
   uops_no_const = list(x_peeled.split_uop(Ops.ADD))
 
@@ -86,7 +88,7 @@ def nest_div_by_smallest_factor(d: UOp, x: UOp, y: UOp) -> UOp|None:
   if ((c := y.arg) < 0): return None
   factors = [u.const_factor() for u in x.split_uop(Ops.ADD) if u.op not in (Ops.CONST, Ops.VCONST)]
   div = min([y.arg]+[abs(f) for f in factors if abs(f) > 1 and (c%f)==0])
-  newxs = fold_divmod_general(newx:=(x//div), x, y.const_like(div))
+  newxs = fold_divmod_general(newx:=(x//y.const_like(div)))
   if div==y.arg or newxs is None or x.vmin<0 or newx.vmin<0: return None
   return newxs//(c//div)
 
@@ -102,9 +104,7 @@ div_and_mod_symbolic = PatternMatcher([
     lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
 
   # ** 2. Slow Rules **
-  # NOTE: if you move this one below `fold_divmod_general` you get more uops in test/external/external_benchmark_schedule.py
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.var("y"))), cancel_divmod),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.var("y"))), fold_divmod_general),
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d"), fold_divmod_general),
   (UPat(Ops.IDIV, dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), nest_div_by_smallest_factor),
 
   # NOTE: these have to go at the bottom or TestSymbolicOps.test_var loops