factorizing shapetracker from chonker
MODIFY
test/test_shapetracker.py
test/test_shapetracker.py
@@ -66,14 +66,40 @@ class TestComplexShapeTracker(unittest.TestCase):
self.st.permute(1,0)
assert self.st.contiguous
- @unittest.skip("reshape is even more complex")
- def test_super_complex(self):
+ def test_permute_reshape(self):
+ self.st = ShapeTracker((4, 4))
+ self.st.permute(1,0)
+ self.st.reshape(2, 2, 2, 2)
+ # TODO: should also be tested by test_super_complex
+ assert len(self.st.views) == 1
+
+ def test_factorize_split(self):
self.st = ShapeTracker((4, 4))
self.st.permute(1,0)
self.st.reshape(2, 2, 2, 2)
self.st.permute(2,3,0,1)
assert self.st.contiguous
+ def test_factorize_combine(self):
+ self.st = ShapeTracker((4, 4, 4))
+ self.st.permute(2, 0, 1)
+ self.st.reshape(4, 16)
+ self.st.permute(1, 0)
+ assert self.st.contiguous
+
+ def test_factorize_combine_add_ones(self):
+ self.st = ShapeTracker((4, 4, 4))
+ self.st.permute(2, 0, 1)
+ self.st.reshape(4, 16, 1, 1)
+ self.st.permute(1, 0, 2, 3)
+ assert self.st.contiguous
+
+ def test_super_complex_2_fail(self):
+ self.st = ShapeTracker((4, 4, 4))
+ self.st.permute(2, 0, 1)
+ self.st.reshape(16, 4)
+ assert len(self.st.views) != 1
+
def test_work(self):
self.st = ShapeTracker((64, 1024, 4))
self.st.reshape(1, 64, 128, 32)

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
@@ -53,6 +53,8 @@ class ZeroView:
acc *= self.shape[0]
self.expr = 'valid=' + ' && '.join(expr)
+ def __repr__(self): return f"ZeroView<{self.old_shape}, {self.arg}>"
+
ViewTypes = Union[View, ZeroView]
@functools.lru_cache(maxsize=None)
@@ -85,9 +87,7 @@ class ShapeTracker:
def offset(self): return self.views[-1].offset
def expr(self): return ';'.join([v.expr for v in self.views[::-1] if v.expr != 'idx=idx' and v.expr != 'valid=valid'])
- def movement_op(self, op, arg):
- getattr(self, str(op).split(".")[1].lower())(*arg)
- return self
+ def movement_op(self, op, arg): return getattr(self, str(op).split(".")[1].lower())(*arg)
def needs_valid(self): return any(isinstance(v, ZeroView) for v in self.views)
# TODO: do we really need this for conv?
@@ -99,6 +99,7 @@ class ShapeTracker:
self.views[-1] = view
else:
self.views.append(view)
+ return self
def reshape(self, *new_shape):
assert all(isinstance(x, int) for x in new_shape)
@@ -109,22 +110,45 @@ class ShapeTracker:
old_strides = [y for x,y in zip(self.shape, self.strides) if x != 1]
new_strides = [0 if x == 1 else old_strides.pop(0) for x in new_shape]
self.views[-1] = View(new_shape, new_strides, self.offset)
- return
+ return self
+
+ # check if the new dimensions factorize from the old ones
+ # NOTE: if you don't make a copy here, the list is popped in the lrucache
+ min_shape_strides = to_shape_strides(self.shape, self.strides)[:]
+ curr_dim, curr_stride = min_shape_strides.pop(0)
+ new_strides = []
+ for s in new_shape:
+ if curr_dim%s == 0:
+ curr_dim //= s
+ new_strides.append(curr_stride * curr_dim)
+ if curr_dim == 1:
+ if len(min_shape_strides) == 0:
+ # there might still be 1s in the shape
+ while len(new_strides) != len(new_shape):
+ assert new_shape[len(new_strides)] == 1
+ new_strides.append(1)
+ self.views[-1] = View(new_shape, new_strides, self.offset)
+ return self # early return, it factorized!
+ curr_dim, curr_stride = min_shape_strides.pop(0)
+ else:
+ break # didn't factorize
view = View(new_shape, strides_for_shape(new_shape))
if self.contiguous:
self.views[-1] = view # NOTE: if it's contiguous it can't have an offset
else:
self.views.append(view)
+ return self
def permute(self, *axis):
assert all(isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis)
assert len(set(axis)) == len(axis) and len(axis) == len(self.shape), f"can't permute {self.shape} with {axis}"
self.views[-1] = View([self.shape[a] for a in axis], [self.strides[a] for a in axis], self.offset)
+ return self
# TODO: this is a special case of slice with strides, remove it
# though it's nice that it can't change size
- def flip(self, *axis): self.stride(*[-1 if i in axis else 1 for i in range(len((self.shape)))])
+ def flip(self, *axis): return self.stride(*[-1 if i in axis else 1 for i in range(len((self.shape)))])
# *** under this line are not invertible ***
@@ -142,12 +166,14 @@ class ShapeTracker:
if zeroview.expr != "valid=valid":
# if we add a ZeroView, we add another (stock) view also for modding
self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))]
+ return self
def expand(self, *new_shape):
assert all(isinstance(x, int) for x in new_shape)
assert all(x == y or x == 1 for x,y in zip(self.shape, new_shape)), f"can't expand {self.shape} into {new_shape}"
strides = [s if x == y else 0 for s,(x,y) in zip(self.strides, zip(self.shape, new_shape))]
self.views[-1] = View(new_shape, strides, self.offset)
+ return self
# TODO: combine with slice? this doesn't require a ZeroView, though slice shouldn't always either
def stride(self, *mul):
@@ -156,3 +182,4 @@ class ShapeTracker:
new_shape = [(s+(abs(m)-1))//abs(m) for s,m in zip(self.shape, mul)]
offset = sum([(s-1)*z for s,z,m in zip(self.shape, self.strides, mul) if m < 0])
self.views[-1] = View(new_shape, strides, self.offset + offset)
+ return self

factorizing shapetracker from chonker