From e0956c518c895cf0e9afc387e992f671e880c496 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 13 Dec 2024 19:56:15 -0500 Subject: [PATCH] move some ifs from merge_dims to reshape [pr] (#8229) the third return value is only used in reshape --- test/unit/test_view.py | 20 ++++++++++++++++---- tinygrad/shape/view.py | 6 +++--- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/test/unit/test_view.py b/test/unit/test_view.py index d745db6bf832..16be4b3b94db 100644 --- a/test/unit/test_view.py +++ b/test/unit/test_view.py @@ -50,7 +50,7 @@ def test_0_in_strides(self): shape = (2, 3, 4) self.assertEqual(merge_dims(shape, (0, 4, 1)), ((2, 0, 0), (12, 1, 12))) self.assertEqual(merge_dims(shape, (0, 0, 1)), ((6, 0, 0), (4, 1, 4))) - self.assertEqual(merge_dims(shape, (3, 1, 0)), ((6, 1, 6), (4, 0, 0))) + self.assertEqual(merge_dims(shape, (3, 1, 0)), ((6, 1, 6), (4, 0, 4))) self.assertEqual(merge_dims(shape, (0, 0, 0)), ((24, 0, 0),)) def test_pad_reshape(self): @@ -59,16 +59,28 @@ def test_pad_reshape(self): # shift mask on stride 0 self.assertEqual(merge_dims((2, 3), (0, 1), ((0, 1), (0, 2))), ((6, 1, 3),)) # permute 0 / 1 - self.assertEqual(merge_dims((3, 2), (1, 0), ((0, 2), (1, 2))), ((3, 1, 3), (2, 0, 0))) + self.assertEqual(merge_dims((3, 2), (1, 0), ((0, 2), (1, 2))), ((3, 1, 3), (2, 0, 2))) + + # st = ShapeTracker.from_shape((1, 1, 2)).pad(((1, 0), (1, 0), (0, 1))) + # print(f"{st.views[-1]}") + self.assertEqual(merge_dims((2, 2, 3), (0, 0, 1), ((1, 2), (1, 2), (0, 2))), ((12, 1, 3),)) + + # st = ShapeTracker.from_shape((1, 1, 2, 2)).pad(((1, 0), (1, 0), (0, 1), (0, 1))) + # print(f"{st.views[-1]}") + self.assertEqual(merge_dims((2, 2, 3, 3), (0, 0, 2, 1), ((1, 2), (1, 2), (0, 2), (0, 2))), ((12, 2, 3), (3, 1, 3))) + + # st = ShapeTracker.from_shape((1, 1, 1, 1)).pad(((0, 2), (0, 0), (0, 1), (0, 3))) + # print(f"{st.views[-1]}") + # self.assertEqual(merge_dims((3, 1, 2, 4), (0, 0, 0, 0), ((0, 1), (0, 1), (0, 1), (0, 1))), ((24, 0, 0),)) def test_different_1_pad(self): # st = ShapeTracker.from_shape((2, 2, 1)).pad(((0, 0), (0, 0), (0, 1))) # print(f"{st.views[-1]}") - self.assertEqual(merge_dims((2, 2, 2), (2, 1, 0), ((0, 2), (0, 2), (0, 1))), ((4, 1, 4), (2, 0, 0))) + self.assertEqual(merge_dims((2, 2, 2), (2, 1, 0), ((0, 2), (0, 2), (0, 1))), ((4, 1, 4), (2, 0, 2))) # st = ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))) # print(f"{st.views[-1]}") - self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 0))) + self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4))) if __name__ == '__main__': unittest.main() diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index f5250a519d5e..d0a476ea4838 100644 --- a/tinygrad/shape/view.py +++ b/tinygrad/shape/view.py @@ -29,8 +29,8 @@ def merge_dims(shape:Tuple[int, ...], strides:Tuple[int, ...], mask:Optional[Tup if s == 1: continue last_s, last_st, last_pre_expand_s = ret[-1] # merge last dim with this dim if merging or strides matched - if merging or last_st == s * st: ret[-1] = (last_s * s, st, (s if merging else last_pre_expand_s * s) if st != 0 else 0) - else: ret.append((s, st, s if st != 0 else 0)) + if merging or last_st == s * st: ret[-1] = (last_s * s, st, (s if merging else last_pre_expand_s * s)) + else: ret.append((s, st, s)) # merge this dim to next dim if size is 1 merging = (mask[i][1] - mask[i][0] == 1) if mask is not None else s == 1 return tuple(ret) @@ -321,13 +321,13 @@ def reshape(self, new_shape: Tuple[sint, ...]) -> Optional[View]: r_strides, r_new_shape = [], reversed(new_shape) for merged_size, new_stride, real_size in reversed(merge_dims(self.shape, self.strides, self.mask)): + if new_stride == 0: real_size = 0 # TODO: write with get_contraction acc = 1 # TODO: third resolve shouldn't be needed while resolve(acc <= merged_size) and resolve(acc != merged_size) and resolve((new_dim := next(r_new_shape, 0)) > 0): r_strides.append(new_stride * acc) acc = acc * new_dim - # merge dim merges if (1) previous_stride = stride * dim, (2) dim = 1, stride = 0, either padded or not if not resolve(acc < real_size): new_stride = 0 if resolve(acc != merged_size): return None