fix _to_const_val and const folding around it (#4017)

* fix _to_const_val and const folding around it

is_unrealized_contiguous_const is too strict and almost never hit if const is expanded.
suffice to check if there's no pad

* that test is folded

* test_const_folding

test/test_const_folding.py

@@ -0,0 +1,64 @@
+import unittest
+from tinygrad import Tensor
+from tinygrad.ops import BufferOps
+from tinygrad.engine.schedule import create_schedule
+
+def _check_ast_count(desired_count:int, t:Tensor):
+  # NOTE: this has side effect because everything can be scheduled only once
+  asts = [s for s in create_schedule([t.lazydata]) if s.ast[0].op is BufferOps.STORE]
+  assert len(asts) == desired_count
+
+class TestSimpleConstFolding(unittest.TestCase):
+  def test_add_literal_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) + 0)
+  def test_add_tensor_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) + Tensor.zeros(4))
+
+  def test_sub_literal_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) - 0)
+  def test_sub_tensor_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) - Tensor.zeros(4))
+
+  def test_mul_literal_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) * 0)
+  def test_mul_tensor_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) * Tensor.zeros(4))
+
+  def test_mul_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) * 1)
+  def test_mul_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) * Tensor.ones(4))
+
+  def test_div_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) / 1)
+  def test_div_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) / Tensor.ones(4))
+
+  # TODO: fix pow const folding
+  @unittest.expectedFailure
+  def test_pow_literal_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) ** 0)
+  @unittest.expectedFailure
+  def test_pow_tensor_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) ** Tensor.zeros(4))
+
+  @unittest.expectedFailure
+  def test_pow_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) ** 1)
+  @unittest.expectedFailure
+  def test_pow_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) ** Tensor.ones(4))
+
+class TestMovedConstFolding(unittest.TestCase):
+  def test_add_shrunk_zero(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))
+
+  def test_add_padded_zero(self):
+    # TODO: it's 1 now, this might be possible to fold
+    _check_ast_count(1, Tensor([1, 2, 3, 4]) + Tensor.zeros(2).pad(((1, 1),)))
+
+  def test_mul_shrunk_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) * Tensor.ones(6).shrink(((1, 5),)))
+
+  def test_add_padded_one(self):
+    _check_ast_count(1, Tensor([1, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))

test/test_jit.py

@@ -338,11 +338,6 @@ class TestJit(unittest.TestCase):
     assert isinstance(jf.jit_cache[1].prg, graph_t)
 
   def test_jit_const_inputs(self):
-    @TinyJit
-    def f(x,y): return (x+y).realize()
-    for _ in range(5):
-      np.testing.assert_equal(f(Tensor.ones(3), Tensor.zeros(3)).numpy(), np.ones(3))
-
     @TinyJit
     def g(x,y,z): return (x+y+z).realize()
     for i in range(5):

tinygrad/features/multi.py

@@ -83,7 +83,7 @@ class MultiLazyBuffer:
     return functools.reduce(lambda x,y: x.e(BinaryOps.ADD, y), llbs)
 
   # TODO: fix this
-  def is_unrealized_contiguous_const(self): return False
+  def is_unrealized_unpadded_const(self): return False
 
   # passthroughs
   def is_realized(self) -> bool: return all([lb.base.realized is not None for lb, r in zip(self.lbs, self.real) if r is True])

tinygrad/lazy.py

@@ -91,7 +91,7 @@ class LazyBuffer:
     return create_lazybuffer(self.device, ShapeTracker.from_shape(new_shape), dtype, UnaryOps.CAST, (dtype, bitcast), (self,))
 
   def is_unrealized_const(self): return self.base.realized is None and self.base.op is LoadOps.CONST
-  def is_unrealized_contiguous_const(self): return self.base == self and self.base.realized is None and self.op is LoadOps.CONST
+  def is_unrealized_unpadded_const(self): return self.is_unrealized_const() and all(v.mask is None for v in self.st.views)
 
   def _copy(self, device:str) -> LazyBuffer:
     if (dstart:=self.device.split(":")[0]) in {"EXT", "DISK"} or (dstart in {"HSA", "CUDA"} and device.split(":")[0] == dstart):

tinygrad/tensor.py

@@ -885,7 +885,7 @@ class Tensor:
 
   def _to_const_val(self, x:Union[Tensor, ConstType]) -> Union[Tensor, ConstType]:
     # TODO: update with multi
-    return x.lazydata.base.arg if isinstance(x, Tensor) and isinstance(x.lazydata, LazyBuffer) and x.lazydata.is_unrealized_contiguous_const() \
+    return x.lazydata.base.arg if isinstance(x, Tensor) and isinstance(x.lazydata, LazyBuffer) and x.lazydata.is_unrealized_unpadded_const() \
       and not x.requires_grad and self._broadcasted(x)[0].shape == self.shape else x
 
   def add(self, x:Union[Tensor, ConstType], reverse=False) -> Tensor: