mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

simpler cancel_divmod [PR] (#16616)

Browse source
This commit is contained in:
chenyu 2026-06-14 11:41:31 -04:00 • committed by GitHub
commit 2e181f4259
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

9
tinygrad/uop/divandmod.py
View file

@ -8,12 +8,9 @@ from tinygrad.helpers import floordiv, floormod, unwrap
def fold_divmod_general(d: UOp) -> UOp|None:
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.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
if y_min*y_max > 0 and (qv:=floordiv(x_min,y_min)) == floordiv(x_min,y_max) == floordiv(x_max,y_min) == floordiv(x_max,y_max):
return x - qv*y if d.op is Ops.FLOORMOD else d.const_like(qv)
if y.vmin==y.vmax==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
# x//y is constant
if (xdiv:=x//y).vmin == xdiv.vmax: return x - xdiv.vmin*y if d.op is Ops.FLOORMOD else xdiv.const_like(xdiv.vmin)
# split uops for the rest of the processing
x_peeled, const = x.pop_const()