sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

tighter idiv bound for symbolic denominator (#10226)

Browse Source
This commit is contained in:
chenyu
2025-05-08 22:38:56 -04:00
committed by GitHub
parent 82a6160ff7
commit 99f6d89dfb
2 changed files with 25 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_uop_vmin_vmax.py
+21
View File
@@ -147,6 +147,27 @@ class TestVminVmaxDivMod(unittest.TestCase):
self.assertEqual(uop.vmin, -3)
self.assertEqual(uop.vmax, 3)
def test_vmin_vmax_div_symbolic(self):
x = UOp.variable('x', 1, 10)
y = UOp.variable('y', 3, 5)
self.assertEqual((x//y).vmin, 0)
self.assertEqual((x//y).vmax, 3)
self.assertEqual(((-x)//y).vmin, -3)
self.assertEqual(((-x)//y).vmax, 0)
self.assertEqual((x//(-y)).vmin, -3)
self.assertEqual((x//(-y)).vmax, 0)
self.assertEqual(((-x)//(-y)).vmin, 0)
self.assertEqual(((-x)//(-y)).vmax, 3)
self.assertEqual((100//y).vmin, 20)
self.assertEqual((100//y).vmax, 33)
self.assertEqual(((-100)//y).vmin, -33)
self.assertEqual(((-100)//y).vmax, -20)
self.assertEqual((100//(-y)).vmin, -33)
self.assertEqual((100//(-y)).vmax, -20)
self.assertEqual(((-100)//(-y)).vmin, 20)
self.assertEqual(((-100)//(-y)).vmax, 33)
def test_vmin_vmax_mod_positive(self):
# vmin and vmax for modulo of a variable by a positive constant
x = UOp.variable('x', 10, 20)
tinygrad/ops.py
+4 -3
View File
@@ -634,9 +634,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
if (c:=s1_vmin) == s1_vmax: # s1 is a const
if c > 0: return cdiv(s0_vmin, c), cdiv(s0_vmax, c)
if c < 0: return cdiv(s0_vmax, c), cdiv(s0_vmin, c)
# don't know exact bounds, but know the sign
if (s0_vmax <= 0 and s1_vmax < 0) or (s0_vmin >= 0 and s1_vmin > 0): return 0, dtypes.max(self.dtype)
if (s0_vmax <= 0 and s1_vmin > 0) or (s0_vmin >= 0 and s1_vmax < 0): return dtypes.min(self.dtype), 0
if (s0_vmax <= 0 and s1_vmax < 0): return cdiv(s0_vmax, s1_vmin), cdiv(s0_vmin, s1_vmax)
if (s0_vmin >= 0 and s1_vmin > 0): return cdiv(s0_vmin, s1_vmax), cdiv(s0_vmax, s1_vmin)
if (s0_vmax <= 0 and s1_vmin > 0): return cdiv(s0_vmin, s1_vmin), cdiv(s0_vmax, s1_vmax)
if (s0_vmin >= 0 and s1_vmax < 0): return cdiv(s0_vmax, s1_vmax), cdiv(s0_vmin, s1_vmin)
if self.op is Ops.MAX: return max(s0_vmin, s1_vmin), max(s0_vmax, s1_vmax)
if self.op is Ops.CMPLT: return (s0_vmax<s1_vmin, s0_vmin<s1_vmax)
if self.op is Ops.CMPNE: return ((s0_vmax < s1_vmin) or (s1_vmax < s0_vmin), not (s0_vmin == s0_vmax == s1_vmin == s1_vmax))