From aa79f6781cbe05e6eef10d35d6729576809e61fb Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 19 Nov 2025 12:01:36 -0800 Subject: [PATCH] one more merge --- tinygrad/uop/divandmod.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index e25d7d0232..4933c40e79 100644 --- a/tinygrad/uop/divandmod.py +++ b/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