track size in shapetracker (#3026)

* track size in shapetracker

* shapetracker adapter

* size is an int

* create Buffer with st.size

* only compare the views for the jit

* fix webgpu

extra/optimization/helpers.py

@@ -1,7 +1,14 @@
 # stuff needed to unpack a kernel
 from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
 from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
+#from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.helpers import prod
+from tinygrad.shape.shapetracker import ShapeTracker as RealShapeTracker
+class ShapeTracker:
+  def __new__(cls, views, size=None):
+    return RealShapeTracker(views, prod([s if st != 0 else 1 for s,st in zip(views[0].shape, views[0].strides)]) if size is None else size)
+  @staticmethod
+  def from_shape(shape): return RealShapeTracker.from_shape(shape)
 from tinygrad.shape.view import View
 from tinygrad.shape.symbolic import Variable, NumNode
 inf, nan = float('inf'), float('nan')

extra/to_movement_ops.py

@@ -2,10 +2,9 @@ from tqdm import tqdm
 import itertools
 from collections import defaultdict
 from typing import List, Tuple, DefaultDict
-from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from extra.optimization.helpers import load_worlds, ast_str_to_ast, ShapeTracker
 from tinygrad.ops import MovementOps, BufferOps, LazyOp
 from tinygrad.helpers import prod
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.symbolic import sym_infer, Node
 
 def apply_mop(st: ShapeTracker, mop_arg: Tuple[MovementOps, Tuple]) -> ShapeTracker:

test/imported/test_indexing.py

@@ -9,7 +9,7 @@ from tinygrad import Tensor, dtypes, Device
 from tinygrad.lazy import LazyBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.helpers import CI
+from tinygrad.helpers import CI, prod
 
 random.seed(42)
 
@@ -27,7 +27,7 @@ def set_(reference: Tensor, shape, strides, offset):
   assert reference.lazydata.base.realized, "base has to be realized before setting it to strided's base"
   # TODO: this shouldn't directly create a LazyBuffer
   strided = Tensor(LazyBuffer(device=reference.device,
-                              st=ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),)),
+                              st=ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),), prod(shape)),
                               op=None, dtype=reference.dtype, srcs=(), base=reference.lazydata.base))
   assert strided.lazydata.st.real_strides() == strides, "real_strides should equal strides for strided"
   return strided

test/test_linearizer.py

@@ -117,9 +117,9 @@ class TestLinearizer(unittest.TestCase):
   def test_simplify_uop(self):
     def helper_test_simplify(uop, dtype, vin, arg=None):
       ast = LazyOp(BufferOps.CONST, (),
-                   ConstBuffer(42, dtypes.float, ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),))))
+                   ConstBuffer(42, dtypes.float, ShapeTracker((View(shape=(), strides=(), offset=0, mask=None, contiguous=True),), 1)))
       ast = LazyOp(BufferOps.STORE, (ast,),
-                   MemBuffer(0, dtypes.float, ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),))))
+                   MemBuffer(0, dtypes.float, ShapeTracker((View(shape=(), strides=(), offset=0, mask=None, contiguous=True),), 1)))
       lin = Linearizer(ast=ast) # this is a dummy ast
 
       lin.uops = []

test/test_linearizer_failures.py

@@ -7,8 +7,8 @@ from tinygrad.helpers import OSX, CI
 from test.external.fuzz_linearizer import run_linearizer
 
 # stuff needed to unpack a kernel
+from extra.optimization.helpers import ShapeTracker
 from tinygrad.ops import LazyOp, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer, get_lazyop_info
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 inf, nan = float('inf'), float('nan')
 

test/test_search.py

@@ -12,7 +12,7 @@ class TestTimeLinearizer(unittest.TestCase):
 
   def test_reasonable_time(self):
     si = [si for si in Tensor([1,2,3,4]).add(1).lazydata.schedule() if si.ast.op not in LoadOps][0]
-    rawbufs = [Buffer(Device.DEFAULT, si.out.st.size(), si.out.dtype)] + [Buffer(Device.DEFAULT, x.st.size(), x.dtype) for x in si.inputs]
+    rawbufs = [Buffer(Device.DEFAULT, si.out.st.size, si.out.dtype)] + [Buffer(Device.DEFAULT, x.st.size, x.dtype) for x in si.inputs]
     tm = time_linearizer(Linearizer(si.ast), rawbufs, allow_test_size=False, cnt=10)
     assert tm > 0 and tm != float('inf')
 

test/test_symbolic_shapetracker.py

@@ -1,4 +1,5 @@
 import unittest
+from tinygrad.helpers import prod
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.symbolic import Variable, NumNode
 from tinygrad.tensor import Tensor
@@ -61,7 +62,7 @@ class TestSymbolicVarVals(unittest.TestCase):
   def test_var_vals_mask(self):
     x = Variable("x", 1, 100).bind(3)
     view = View.create(shape=(3,4), strides=(4,1), offset=0, mask=((0, x), (0, 4)))
-    st = ShapeTracker(views=(view,))
+    st = ShapeTracker(views=(view,), size=3*4)
     assert st.var_vals == {Variable("x", 1, 100): 3}
 
   def test_var_vals_complex(self):
@@ -88,13 +89,13 @@ class TestShapeTrackerUnbind(unittest.TestCase):
     v = Variable("v", 1, 100)
     bv = Variable("v", 1, 100).bind(3)
     t = Tensor.rand(3, 4).reshape(bv, 4)
-    assert t.lazydata.st.unbind() == ShapeTracker((View.create(shape=(v, 4)),))
+    assert t.lazydata.st.unbind() == ShapeTracker((View.create(shape=(v, 4)),), 3*4)
 
   def test_shrink_unbind(self):
     v = Variable("v", 1, 100)
     bv = Variable("v", 1, 100).bind(2)
     t = Tensor.rand(3, 4).shrink(((bv, bv+1), (0, 4)))
-    assert t.lazydata.st.unbind() == ShapeTracker((View.create(shape=(1, 4), offset=4*v),))
+    assert t.lazydata.st.unbind() == ShapeTracker((View.create(shape=(1, 4), offset=4*v),), 3*4)
 
 class TestSymbolicReshape(unittest.TestCase):
   def test_reshape_into_symbols_simple(self):
@@ -176,7 +177,7 @@ class TestSymbolicShapeExpr(unittest.TestCase):
     idx = (gidx0, lidx1, NumNode(1))
     shape = (i+1, 8, 4)
     strides = (1, (i*4)+4, i+1)
-    st = ShapeTracker((View.create(shape, strides), ))
+    st = ShapeTracker((View.create(shape, strides), ), prod(shape))
     idx, _valid = st.expr_idxs(idx)
     assert idx.render() == "((lidx1*((i*4)+4))+1+gidx0+i)"
 

test/unit/test_shapetracker.py

@@ -108,20 +108,20 @@ class TestRealDoesntSimplify(unittest.TestCase):
   def test_1(self):
     self.st = ShapeTracker((
       View.create((8, 3, 1, 2, 11, 1), (33, 11, 0, 0, 1, 0), 0, None),
-      View.create((8, 6, 11), (66, 11, 1), 0, None)))
+      View.create((8, 6, 11), (66, 11, 1), 0, None)), prod((8, 3, 1, 2, 11, 1)))
     assert self.st.real_strides() == (33, None, 1)
 
   def test_2(self):
     self.st = ShapeTracker((
       View.create((2, 2, 4, 3, 3), (72, 9, 18, -3, -1), 8, None),
-      View.create((4, 4, 3, 3), (36, 9, 3, 1), 0, None)))
+      View.create((4, 4, 3, 3), (36, 9, 3, 1), 0, None)), prod((2, 2, 4, 3, 3)))
     assert self.st.real_strides() == (None, 18, -3, -1)
 
 class TestRealStrides(unittest.TestCase):
   def test_1(self):
     self.st = ShapeTracker((
       View.create((2048,), (1,), 0, ((0, 512),)),
-      View.create((16, 32, 4), (128, 4, 1), 0, None)))
+      View.create((16, 32, 4), (128, 4, 1), 0, None)), prod((2048,)))
     st = self.st.real_strides()
     print(self.st, st)
     assert st == (None, 4, 1)
@@ -137,12 +137,12 @@ class TestRealSimplifies(unittest.TestCase):
   def test_1(self):
     self.st = ShapeTracker((
       View.create((1, 3, 2, 11, 4, 28), (0, 308, 0, 28, 0, 1), 0, None),
-      View.create((1, 3, 2, 11, 26, 1, 1, 3), (0, 2464, 0, 112, 1, 0, 0, 29), 0, None)))
+      View.create((1, 3, 2, 11, 26, 1, 1, 3), (0, 2464, 0, 112, 1, 0, 0, 29), 0, None)), prod((1, 3, 2, 11, 4, 28)))
 
   def test_2(self):
     self.st = ShapeTracker((
       View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
-      View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
+      View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)), prod((8, 3, 3, 11, 2, 28)))
 
 class TestIndexExpressions2d(unittest.TestCase):
 
@@ -797,35 +797,38 @@ class TestGetContraction(unittest.TestCase):
 class TestShapeTrackerSize(unittest.TestCase):
   def test_simple_size(self):
     st = ShapeTracker.from_shape((100, 100))
-    self.assertEqual(st.size(), 100*100)
+    self.assertEqual(st.size, 100*100)
 
   def test_expand_size(self):
     st = ShapeTracker.from_shape((100, 100))
     st = st.reshape((100, 100, 1))
     st = st.expand((100, 100, 100))
-    self.assertEqual(st.size(), 100*100)
+    self.assertEqual(st.size, 100*100)
 
   def test_expand_size_flatten(self):
     st = ShapeTracker.from_shape((100, 100))
     st = st.reshape((100, 100, 1))
     st = st.expand((100, 100, 100))
     st = st.reshape((100*100*100,))
-    self.assertEqual(st.size(), 100*100)
+    self.assertEqual(st.size, 100*100)
 
+  @unittest.skip("size is now the buffer size")
   def test_shrink_size_axis_0(self):
     st = ShapeTracker.from_shape((100, 100))
     st = st.shrink(((0, 50), (0, 100)))
-    self.assertEqual(st.size(), 50*100)
+    self.assertEqual(st.size, 50*100)
 
+  @unittest.skip("size is now the buffer size")
   def test_shrink_size_axis_0_variable(self):
     st = ShapeTracker.from_shape((100, 100))
     st = st.shrink(((0, Variable("a", 0, 50)), (0, 100)))
-    self.assertEqual(st.size(), 50*100)
+    self.assertEqual(st.size, 50*100)
 
+  @unittest.skip("size is now the buffer size")
   def test_shrink_size_axis_1(self):
     st = ShapeTracker.from_shape((100, 100))
     st = st.shrink(((0, 100), (0, 50)))
-    self.assertEqual(st.size(), 9950)    # careful here
+    self.assertEqual(st.size, 9950)    # careful here
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_shapetracker_math.py

@@ -57,7 +57,7 @@ class TestShapeTrackerBasics(unittest.TestCase):
 
   def test_simplify_is_correct(self):
     multiv = ShapeTracker(views=(View(shape=(15, 3), strides=(9, 1), offset=6, mask=None, contiguous=False),
-                                 View(shape=(4, 3), strides=(12, 4), offset=0, mask=None, contiguous=False)))
+                                 View(shape=(4, 3), strides=(12, 4), offset=0, mask=None, contiguous=False)), size=15*3)
     assert st_equal(multiv, multiv.simplify())
 
 class TestShapeTrackerAdd(unittest.TestCase):
@@ -104,7 +104,7 @@ class TestShapeTrackerAddVariable(unittest.TestCase):
     var_j = Variable('i', 1, 10)
     vm1 = View(shape=(var_i, var_j, 3), strides=(3, 0, 1), offset=0, mask=None, contiguous=False)
     vm2 = View(shape=(var_i, var_j, 3), strides=(var_j*3, 3, 1), offset=0, mask=None, contiguous=True)
-    ShapeTracker((vm1,)) + ShapeTracker((vm2,))
+    ShapeTracker((vm1,), prod((var_i, var_j, 3))) + ShapeTracker((vm2,), prod((var_i, var_j, 3)))
 
   @unittest.skip("two vars not supported")
   def test_merge_symbolic_views_2(self):

tinygrad/codegen/kernel.py

@@ -325,8 +325,8 @@ class Kernel:
           stride[j] = bst
           bst *= shp[j]
 
-    self.sts.append(ShapeTracker((View.create(tuple(shp), tuple(stride)),)))
-    self.bufs.append(LocalBuffer(name=f"ldata{i}", size=self.sts[-1].size()))
+    self.sts.append(ShapeTracker((View.create(tuple(shp), tuple(stride)),), prod(shp)))
+    self.bufs.append(LocalBuffer(name=f"ldata{i}", size=self.sts[-1].size))
     if DEBUG >= 4: print("aliasing buffer", self.sts[i])
     self.local_alias[i] = cast(LocalBuffer, self.bufs[-1])
 

tinygrad/codegen/linearizer.py

@@ -206,14 +206,14 @@ class Linearizer(Kernel):
       self.loop_uops[var.expr] = self.uop(UOps.DEFINE_GLOBAL, dtypes.int32, (), var.expr)
     # define local buffers
     for lb in self.local_alias.values():
-      self.buf_uops[self.bufs.index(lb)] = self.uop(UOps.DEFINE_LOCAL, PtrDType(dtypes.float32), (), (lb.name, self.sts[self.bufs.index(lb)].size()))
+      self.buf_uops[self.bufs.index(lb)] = self.uop(UOps.DEFINE_LOCAL, PtrDType(dtypes.float32), (), (lb.name, self.sts[self.bufs.index(lb)].size))
     # add a local buffer for multistage reduce. # TODO: use local alias
     if self.group_for_reduce:
       # TODO: the strides of this can be controlled
       self.sts.append(ShapeTracker.from_shape(tuple([1] * self.global_dims + list(self.full_shape[self.global_dims:self.global_dims+self.local_dims+len(self.group_for_reduce)]) + [1] * (self.shape_len - self.upcasted - len(self.group_for_reduce) - self.first_reduce) + [x[0] for x in self.upcasted_axis(0)])))  # noqa: E501
       temp_dtype = self.get_base_dtype(get_lazyop_info(self.reduceop).dtype)
-      self.bufs.append(LocalBuffer("temp", self.sts[-1].size(), temp_dtype))
-      self.buf_uops.append(self.uop(UOps.DEFINE_LOCAL, PtrDType(temp_dtype), (), ("temp", self.sts[-1].size())))
+      self.bufs.append(LocalBuffer("temp", self.sts[-1].size, temp_dtype))
+      self.buf_uops.append(self.uop(UOps.DEFINE_LOCAL, PtrDType(temp_dtype), (), ("temp", self.sts[-1].size)))
 
     # kernel name (before late upcast)
     self.name = ("r_" if self.reduceop else "E_") + colored('_', 'BLACK').join([colored(str(x), c) for x,c in zip(self.full_shape, self.colors())])

tinygrad/features/search.py

@@ -64,7 +64,8 @@ def bufs_from_lin(lin:Linearizer) -> List[Buffer]:
   for x in lin.membufs: bufsts[x.idx].append(x)
   rawbufs:List[Optional[Buffer]] = [None]*len(bufsts)
   for k,lx in bufsts.items():
-    rawbufs[k] = Buffer(Device.DEFAULT, prod(lx[0].dtype.shape) if isinstance(lx[0].dtype, ImageDType) else max(y.st.size() for y in lx), lx[0].dtype)
+    buffer_size = prod(lx[0].dtype.shape) if isinstance(lx[0].dtype, ImageDType) else max(y.st.size for y in lx)
+    rawbufs[k] = Buffer(Device.DEFAULT, buffer_size, lx[0].dtype)
   assert all(r is not None for r in rawbufs)
   return cast(List[Buffer], rawbufs)
 

tinygrad/jit.py

@@ -66,8 +66,9 @@ class TinyJit(Generic[ReturnType]):
 
     if self.cnt >= 2:
       # jit exec
-      assert self.expected_vals == expected_vals, "mismatch of var_vals"
-      assert self.expected_name_sts_dtype == expected_name_sts_dtype, f"mismatch of sts, expected {self.expected_name_sts_dtype} got {expected_name_sts_dtype}"  # noqa: E501
+      assert self.expected_vals == expected_vals and self.expected_name_sts_dtype is not None, "missing/mismatch of var_vals"
+      assert all(x[0] == y[0] and x[1].views == y[1].views and x[2] == x[2] for x,y in zip(self.expected_name_sts_dtype, expected_name_sts_dtype)), \
+        f"mismatch of sts, expected {self.expected_name_sts_dtype} got {expected_name_sts_dtype}"
       for (j,i),input_idx in self.input_replace.items(): self.jit_cache[j].rawbufs[i] = input_rawbuffers[input_idx]
       for ji in self.jit_cache: ji.prg(cast(List[Buffer], ji.rawbufs), var_vals, wait=DEBUG>=2, jit=True)
     elif self.cnt == 1:

tinygrad/lazy.py

@@ -63,7 +63,7 @@ class LazyBuffer:
     return LazyBuffer.loadop(LoadOps.CONST, tuple(), self.dtype, self.device, arg=val).reshape((1,)*len(self.shape)).expand(self.shape)
 
   def contiguous(self):
-    if not self.st.contiguous or self.st.size() != self.base.st.size() or self.is_unrealized_const():
+    if not self.st.contiguous or self.is_unrealized_const():
       ret = self.e(LoadOps.CONTIGUOUS)
       sti = self.st.invert(self.base.shape)
       if sti is not None: self.base.contiguous_child = ref(ret), sti
@@ -91,7 +91,7 @@ class LazyBuffer:
     if self.device == device: return self
 
     # double COPY = one COPY
-    if self.st.contiguous and self.st.size() == self.base.st.size() and not self.base.realized and self.base.op == LoadOps.COPY:
+    if self.st.contiguous and not self.base.realized and self.base.op == LoadOps.COPY:
       return self.base.srcs[0].copy_to_device(device).reshape(self.st.shape)
 
     # const doesn't have to be copied (issues with disk tensor)
@@ -234,7 +234,7 @@ def _recurse_lb(buf:LazyBuffer, realizes:Set[LazyBuffer], allbufs:Dict[LazyBuffe
   allbufs[buf] = None
   if buf.op in LoadOps: realizes.add(buf.base)
   if buf.op == LoadOps.COPY:
-    assert buf.srcs[0].st.contiguous and buf.srcs[0].st.size() == buf.srcs[0].base.st.size(), "can only copy contig"
+    assert buf.srcs[0].st.contiguous, "can only copy contig"
     realizes.add(buf.srcs[0].base)
   for x in buf.srcs: _recurse_lb(x, realizes, allbufs, simple_pads)
 
@@ -278,7 +278,7 @@ def create_schedule(outs:List[LazyBuffer], seen:Optional[Set[LazyBuffer]]=None)
           # can only have one output buffer
           # can only reduce contiguous
           # max one reduceop per kernel
-          if len(realized_children) > 1 or not st.contiguous or st.size() != r.st.size() or (tr in reduce_for_op and reduce_for_op[tr] != r):
+          if len(realized_children) > 1 or not st.contiguous or (tr in reduce_for_op and reduce_for_op[tr] != r):
             can_chase = tr not in reduce_for_op or reduce_for_op[tr] == r
             forced_realize = True
             break
@@ -304,7 +304,7 @@ def create_schedule(outs:List[LazyBuffer], seen:Optional[Set[LazyBuffer]]=None)
           tr_next = next(iter(tr.children))
           st_childs = dedup([s for s in tr_next.srcs if s.base == tr])
           if len(st_childs) > 1: break
-          if st.size() != st_childs[0].st.size(): break
+          if st.size != st_childs[0].st.size: break
           st = st + st_childs[0].st
           if not st.contiguous or tr_next.op in ReduceOps: break
           tr = tr_next

tinygrad/ops.py

@@ -83,9 +83,9 @@ class FlopCounter:
     return ret
 
 InterpretedFlopCounter: Dict[Op, Callable] = {
-  BufferOps.LOAD: lambda arg: FlopCounter(arg.st.shape, arg.dtype, 0, {arg.idx: arg.dtype.itemsize*arg.st.size()}),
+  BufferOps.LOAD: lambda arg: FlopCounter(arg.st.shape, arg.dtype, 0, {arg.idx: arg.dtype.itemsize*arg.st.size}),
   BufferOps.CONST: lambda arg: FlopCounter(arg.st.shape, arg.dtype, 0, {}),
-  BufferOps.STORE: lambda self,arg: FlopCounter(arg.st.shape, arg.dtype, self.consume_flops(), {**self.mem, arg.idx: arg.dtype.itemsize*arg.st.size()}),  # noqa: E501
+  BufferOps.STORE: lambda self,arg: FlopCounter(arg.st.shape, arg.dtype, self.consume_flops(), {**self.mem, arg.idx: arg.dtype.itemsize*arg.st.size}),  # noqa: E501
   UnaryOps.CAST: lambda self,arg: FlopCounter(self.shape, arg[0], self.consume_flops(), self.mem),   # cast uses no flops
   **{op:lambda self: FlopCounter(self.shape, self.dtype, self.consume_flops() + prod(self.shape), self.mem) for op in UnaryOps if op != UnaryOps.CAST},  # noqa: E501
   **{op:lambda self,y,op=op: FlopCounter(self.shape,  dtypes.bool if op in (BinaryOps.CMPLT, BinaryOps.CMPEQ) else self.dtype, self.consume_flops() + y.consume_flops() + prod(self.shape), {**self.mem, **y.mem}) for op in BinaryOps},  # noqa: E501

tinygrad/realize.py

@@ -2,7 +2,7 @@ from typing import List, Dict, Optional, cast
 from tinygrad.ops import LoadOps, ScheduleItem, BufferOps, GlobalCounters
 from tinygrad.device import Device, Buffer, BufferCopy, JITRunner, update_stats, InterpretedASTRunner
 from tinygrad.graph import print_tree, realized_lazybuffer
-from tinygrad.helpers import prod, colored, getenv
+from tinygrad.helpers import colored, getenv
 from tinygrad.shape.symbolic import Variable
 
 # *** schedule running ***
@@ -39,12 +39,12 @@ def run_schedule(schedule:List[ScheduleItem]):
 
     # we don't have an output buffer, we have to create it, and create to max size if it has symbolic shape
     si.out.realized = si.out.output_buffer if si.out.output_buffer is not None else \
-      Buffer(si.out.device, prod((s if isinstance(s, int) else s.max for s in si.out.shape)), si.out.dtype,
+      Buffer(si.out.device, si.out.st.size, si.out.dtype,
              "PLACEHOLDER" if isinstance(prg, InterpretedASTRunner) else None)
     del si.out.srcs
 
     # run the function (put it in JIT)
     assert all(x.realized is not None for x in si.inputs), f"can't run, some inputs aren't realized {[x for x in si.inputs if x.realized is None]}"
     if prg: prg.exec([si.out.realized] + [cast(Buffer, x.realized) for x in si.inputs], si.var_vals)
-    else: update_stats(colored(f"empty {si.out.st.size():10d} {si.out.dtype}", "yellow"), 0, 0, {}, None, 1, device=si.out.device)
+    else: update_stats(colored(f"empty {si.out.st.size:10d} {si.out.dtype}", "yellow"), 0, 0, {}, None, 1, device=si.out.device)
     realized_lazybuffer(si.out, GlobalCounters.kernel_count)

tinygrad/shape/shapetracker.py

@@ -8,6 +8,8 @@ from tinygrad.helpers import prod, DEBUG, merge_dicts, getenv
 from tinygrad.shape.symbolic import Variable, MulNode, Node, SumNode, NumNode, sint
 from tinygrad.shape.view import View, _merge_dims
 
+def shape_to_size(shape:Tuple[sint, ...]) -> int: return prod([x.max if isinstance(x, Node) else x for x in shape])
+
 def expr_node_mask(view:View, idx:Node, valid:Optional[Node]=None) -> Node:
   expr = [valid] if valid is not None else []
   if view.mask is not None:
@@ -39,7 +41,7 @@ def merge_views(vm2:View, vm1:View) -> Optional[View]:
   if vm1.contiguous and vm1.shape == vm2.shape: return vm2
   if vm2.contiguous: return vm1
   if vm2.mask or vm1.offset != 0: return None  # this isn't supported yet
-  if None in (strides := ShapeTracker((vm2, vm1)).real_strides()): return None
+  if None in (strides := ShapeTracker((vm2, vm1), shape_to_size(vm2.shape)).real_strides()): return None
   return View.create(vm1.shape, cast(Tuple[sint, ...], strides), vm2.offset, vm1.mask)
 
 @functools.lru_cache(maxsize=None)
@@ -54,41 +56,33 @@ def idxs_to_idx(shape:Tuple[int, ...], idxs:Tuple[Node, ...]) -> Node:
 @dataclass(frozen=True)
 class ShapeTracker:
   views: Tuple[View, ...]
-  def __post_init__(self):
-    assert isinstance(self.views, tuple) and all(isinstance(v, View) for v in self.views), "ShapeTracker must be created with a tuple of Views"
+  size: int
 
   def __add__(self, st:ShapeTracker) -> ShapeTracker:
-    base = ShapeTracker(self.views)
-    for v in st.views: base = ShapeTracker(base.views + (v,)).simplify() # one view at a time = better simplification
+    base = ShapeTracker(self.views, self.size)
+    for v in st.views: base = ShapeTracker(base.views + (v,), self.size).simplify() # one view at a time = better simplification
     return base
 
   def invert(self, out_shape:Tuple[sint, ...]) -> Optional[ShapeTracker]:
     ret = tuple(v.invert(s) for v,s in zip(self.views[::-1], [x.shape for x in self.views[::-1][1:]]+[out_shape]))
-    return ShapeTracker(cast(Tuple[View, ...], ret)).reshape(out_shape) if all(x is not None for x in ret) else None
+    return ShapeTracker(cast(Tuple[View, ...], ret), shape_to_size(out_shape)).reshape(out_shape) if all(x is not None for x in ret) else None
 
   @staticmethod
-  def from_shape(shape:Tuple[sint, ...]): return ShapeTracker((View.create(shape),))
+  def from_shape(shape:Tuple[sint, ...]): return ShapeTracker((View.create(shape),), shape_to_size(shape))
 
   @property
-  def contiguous(self) -> bool: return len(self.views) == 1 and self.views[0].contiguous
+  def contiguous(self) -> bool:
+    return len(self.views) == 1 and self.views[0].contiguous and shape_to_size(self.views[0].shape) == self.size
 
   @property
   def shape(self) -> Tuple[sint, ...]: return self.views[-1].shape
 
-  @functools.lru_cache(maxsize=None)   # NOTE: this keeps all ShapeTrackers alive
-  def size(self) -> int:
-    if 0 in self.shape: return 0
-    ret = self.expr_idxs()[0].max
-    while not isinstance(ret, int): ret = ret.max    # TODO: this is a while loop?!? it should be more clear what max does
-    assert isinstance(ret, int), f"ret must be integer, {ret=} isn't"
-    return ret+1
-
   def vars(self) -> Set[Variable]: return set.union(*[v.vars() for v in self.views], set())
 
   @property
   def var_vals(self) -> Dict[Variable, int]: return merge_dicts([dict([v.unbind()]) for v in self.vars()])
 
-  def unbind(self) -> ShapeTracker: return ShapeTracker(tuple(v.unbind() for v in self.views))
+  def unbind(self) -> ShapeTracker: return ShapeTracker(tuple(v.unbind() for v in self.views), self.size)
 
   def to_movement_ops(self) -> List[Tuple[MovementOps, Tuple]]:
     to_apply:List[Tuple[MovementOps, Tuple]] = []
@@ -156,20 +150,26 @@ class ShapeTracker:
   def simplify(self) -> ShapeTracker:
     if len(self.views) >= 2 and (new_view := merge_views(self.views[-2], self.views[-1])) is not None:
       if DEBUG >= 5: print(f"st simplify : {self.views[-2]} + {self.views[-1]} = {new_view}")
-      return ShapeTracker(self.views[:-2] + (new_view,)).simplify()
+      return ShapeTracker(self.views[:-2] + (new_view,), self.size).simplify()
     return self
 
   # *** under this line are the movement ops ***
 
-  def pad(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].pad(arg), ))
-  def shrink(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].shrink(arg), ))
-  def expand(self, new_shape: Tuple[sint, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].expand(new_shape), ))
-  def permute(self, axis: Tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].permute(axis), ))
-  def stride(self, mul: Tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].stride(mul), ))
+  def pad(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker:
+    return ShapeTracker(self.views[0:-1] + (self.views[-1].pad(arg), ), self.size)
+  def shrink(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker:
+    return ShapeTracker(self.views[0:-1] + (self.views[-1].shrink(arg), ), self.size)
+  def expand(self, new_shape: Tuple[sint, ...]) -> ShapeTracker:
+    return ShapeTracker(self.views[0:-1] + (self.views[-1].expand(new_shape), ), self.size)
+  def permute(self, axis: Tuple[int, ...]) -> ShapeTracker:
+    return ShapeTracker(self.views[0:-1] + (self.views[-1].permute(axis), ), self.size)
+  def stride(self, mul: Tuple[int, ...]) -> ShapeTracker:
+    return ShapeTracker(self.views[0:-1] + (self.views[-1].stride(mul), ), self.size)
 
   def reshape(self, new_shape: Tuple[sint, ...]) -> ShapeTracker:
-    if getenv("MERGE_VIEW", 1) and (new_view := self.views[-1].reshape(new_shape)) is not None: return ShapeTracker(self.views[0:-1] + (new_view,))
-    return ShapeTracker(self.views + (View.create(new_shape), ))
+    if getenv("MERGE_VIEW", 1) and (new_view := self.views[-1].reshape(new_shape)) is not None:
+      return ShapeTracker(self.views[0:-1] + (new_view,), self.size)
+    return ShapeTracker(self.views + (View.create(new_shape), ), self.size)
 
 # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
 # TODO: if we remove movementops from lazy.py we can delete this

tinygrad/shape/symbolic.py

@@ -308,7 +308,7 @@ def sym_render(a: Union[Node, int], ops=None, ctx=None) -> str: return str(a) if
 def sym_infer(a: Union[Node, int], var_vals: Dict[Variable, int]) -> int:
   if isinstance(a, (int, float)): return a
   ret = a.substitute({k:NumNode(v) for k, v in var_vals.items()})
-  assert isinstance(ret, NumNode), f"sym_infer didn't produce NumNode from {a} with {var_vals}"
+  assert isinstance(ret, NumNode), f"sym_infer didn't produce NumNode from {a} with {var_vals}, it produced {ret}"
   return ret.b
 
 # symbolic int