improve vmin/vmax for IDIV (tinygrad#9678)

Author: chenyuxyz

test/external/fuzz_symbolic.py

@@ -70,7 +70,7 @@ def gt(expr, rng=None):
       v = [v1,v2,v3]
       rn = 0
       for t,r in zip(tape, rngs): rn, _ = t(rn, r)
-      num = eval(expr.render())
+      num = eval(expr.render(simplify=False))
       if num != rn:
         unsimplified_num = eval(expr.render(simplify=False))
         assert unsimplified_num == rn, "UNSIMPLIFIED MISMATCH!"

test/unit/test_uop_symbolic.py

@@ -323,8 +323,8 @@ def test_mul_div(self):
 
   def test_add_div(self):
     # careful about the lower bounds and upper bounds
-    self.helper_test_variable((Variable("a", 0, 5)-2)//4, -1, 0, "((a+-2)//4)")
-    self.helper_test_variable((Variable("a", 0, 5)-1)//4, -1, 1, "((a+-1)//4)")
+    self.helper_test_variable((Variable("a", 0, 5)-2)//4, 0, 0, "0")
+    self.helper_test_variable((Variable("a", 0, 5)-1)//4, 0, 1, "((a+-1)//4)")
     self.helper_test_variable((Variable("a", 0, 5))//4, 0, 1, "(a//4)")
     self.helper_test_variable((Variable("a", 0, 5)+1)//4, 0, 1, "((a+1)//4)")
     self.helper_test_variable((Variable("a", 0, 5)+2)//4, 0, 1, "((a+2)//4)")

test/unit/test_uop_vmin_vmax.py

@@ -120,10 +120,32 @@ def test_vmin_vmax_division_positive(self):
 
   def test_vmin_vmax_division_negative(self):
     # vmin and vmax for division of a variable by a negative constant
+    # always positive
     x = UOp.variable('x', 10, 20)
     uop = x // -2
     self.assertEqual(uop.vmin, -10)
     self.assertEqual(uop.vmax, -5)
+    uop = x // -3
+    self.assertEqual(uop.vmin, -6)
+    self.assertEqual(uop.vmax, -3)
+
+    # always negative
+    x = UOp.variable('x', -20, -10)
+    uop = x // -2
+    self.assertEqual(uop.vmin, 5)
+    self.assertEqual(uop.vmax, 10)
+    uop = x // -3
+    self.assertEqual(uop.vmin, 3)
+    self.assertEqual(uop.vmax, 6)
+
+    # cross 0
+    x = UOp.variable('x', -10, 10)
+    uop = x // -2
+    self.assertEqual(uop.vmin, -5)
+    self.assertEqual(uop.vmax, 5)
+    uop = x // -3
+    self.assertEqual(uop.vmin, -3)
+    self.assertEqual(uop.vmax, 3)
 
   def test_vmin_vmax_mod_positive(self):
     # vmin and vmax for modulo of a variable by a positive constant
@@ -144,7 +166,7 @@ def test_vmin_vmax_division_with_mixed_range(self):
     # vmin and vmax for division of a variable with a range crossing zero
     x = UOp.variable('x', -10, 10)
     uop = x // 3
-    self.assertEqual(uop.vmin, -4)  # -10//3 = -4
+    self.assertEqual(uop.vmin, -3)  # -10//3 = -3 (in C)
     self.assertEqual(uop.vmax, 3)   # 10//3 = 3
 
   def test_vmin_vmax_mod_with_mixed_range(self):

tinygrad/ops.py

@@ -616,12 +616,13 @@ def _min_max(self) -> tuple[ConstType, ConstType]:
       if self.op is Ops.MOD and s1_vmin > 0:
         return (0, s1_vmax-1) if s0_vmin >= 0 else (-(s1_vmax-1), s1_vmax-1)
       if self.op is Ops.IDIV:
-        if s1_vmin == s1_vmax:  # min/max are equal in a CONST
-          if s1_vmin > 0: return s0_vmin//s1_vmin, s0_vmax//s1_vmin
-          if s1_vmin < 0 and s0_vmin >= 0: return -(s0_vmax//-s1_vmin), -(s0_vmin//-s1_vmin)
+        assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
+        if s1_vmin == s1_vmax:  # s1 is a const
+          if s1_vmin > 0: return cdiv(s0_vmin, s1_vmin), cdiv(s0_vmax, s1_vmin)
+          if s1_vmin < 0: return cdiv(s0_vmax, s1_vmin), cdiv(s0_vmin, s1_vmin)
         # don't know exact bounds, but know the sign
-        if (s0_vmax <= 0 and s1_vmin < 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_vmin < 0): return dtypes.min(self.dtype), 0
+        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 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))