diff --git a/test/test_const_folding.py b/test/test_const_folding.py new file mode 100644 index 0000000000..333a686f4c --- /dev/null +++ b/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),))) diff --git a/test/test_jit.py b/test/test_jit.py index e9ca4b68b2..7f584a5533 100644 --- a/test/test_jit.py +++ b/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): diff --git a/tinygrad/features/multi.py b/tinygrad/features/multi.py index 5276a6f4af..4b446d670a 100644 --- a/tinygrad/features/multi.py +++ b/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]) diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 726eac7718..d2b3dd99e4 100644 --- a/tinygrad/lazy.py +++ b/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): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index f9ce5a697b..2de6835df6 100644 --- a/tinygrad/tensor.py +++ b/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: