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

fix divmod recombine for floordiv (#16062)

Browse source
This commit is contained in:
chenyu 2026-05-06 14:22:42 -04:00 • committed by GitHub
commit af4140f3be
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
test/null/test_uop_symbolic.py
View file

@ -745,6 +745,11 @@ class TestSymbolic(unittest.TestCase):
a = Variable("a", 0, 10)
self.helper_test_variable((25*a+3)%10 + ((25*a+3)//10)*10, 3, 253, "((a*25)+3)")
def test_div_mod_recombine_negative_div_unsound(self):
# ((b//d)%div)*mul + (b//(d*div))*(div*mul) only equals (b//d)*mul when div>0
b = Variable("b", -100, 100)
self.helper_test_variable(((b//(-3))%(-2)) + (b//6)*(-2), -33, 34, "(b//6*-2+b//-3%-2)")
def test_mod_nest_by_factor(self):
# (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for FLOORMOD
gidx0 = Variable("gidx0", 0, 15)

4
tinygrad/uop/symbolic.py
View file

@ -38,8 +38,8 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None:
q, exact = v.src[0], False
# (base%div)*mul + (base//div)*(div*mul) -> base*mul
if q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
# ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul
if not exact and base.op is Ops.FLOORDIV and base.src[1].op is Ops.CONST:
# ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul if div>0
if not exact and div > 0 and base.op is Ops.FLOORDIV and base.src[1].op is Ops.CONST:
exact = q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
if exact: return (base*mul).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
# ((base//div)%d)*div + base%div -> base%(div*d)