From 2adedf5ccb95d3d05aaf2fad24df4e3b468e3c6f Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 14 Jun 2026 17:15:52 -0400 Subject: [PATCH] clean up fold_divmod_general [pr] (#16622) genralized fold_binary_numerator in fold_divmod_congruence --- tinygrad/uop/divandmod.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index e4669c36f2..d64b3f2dec 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -1,7 +1,7 @@ import functools, itertools, math from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp from tinygrad.dtype import dtypes -from tinygrad.helpers import floordiv, floormod, unwrap +from tinygrad.helpers import unwrap # NOTE: this cache is only on index UOps @functools.cache @@ -19,10 +19,8 @@ def fold_divmod_general(d: UOp) -> UOp|None: # ** Constant Denominator Rules ** # these rules strictly require y to be a scalar constant > 0 if y.op is Ops.CONST and (c := y.arg) > 0: - # nested_div_mod: (x%(k*c))//c -> (x//c)%k (requires k>0), and (x%(k*c))%c -> x%c - if x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None: - if d.op is Ops.FLOORMOD: return x.src[0] % y - if k > 0: return x.src[0] // y % k + # nested_div: (x%(k*c))//c -> (x//c)%k (requires k>0); the mod case is handled by remove_nested_mod below + if d.op is Ops.FLOORDIV and x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None and k > 0: return x.src[0] // y % k # remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2 if d.op is Ops.FLOORMOD: @@ -38,15 +36,10 @@ def fold_divmod_general(d: UOp) -> UOp|None: decomp = [(u.divides(f:=u.const_factor()),f) for u in uops_no_const] terms, factors = zip(*decomp) - # fold_binary_numerator: fold if expression has one non-constant term that takes on two values - if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1: - y1 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmin+const, c) - y2 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmax+const, c) - return (y2-y1)*(v-v.vmin) + y1 - # fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c - # when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period - rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors] + # try both signs of the remainder for a lone term (covers a binary numerator that crosses one period) + # or on an exact f%c == c//2 tie; otherwise pick the smaller to keep the product over terms small + rem_choices = [(r, r-c) if (r:=f%c)*2 == c or len(terms)==1 else (min(r, r-c, key=abs),) for f in factors] for rems in itertools.product(*rem_choices): if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c: if d.op is Ops.FLOORMOD: return rem - rem.vmin//c*c