mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 12:56:07 +00:00
remove some View add logic [pr] (#12584)
no longer simplify the case of v0+v1 where v0 has a mask
This commit is contained in:
chenyu
@@ -154,21 +154,6 @@ class TestRealStrides(unittest.TestCase):
|
||||
))
|
||||
self.assertEqual(st.is_expanded(), (False, False, False, True, False))
|
||||
|
||||
class TestRealSimplifies(unittest.TestCase):
|
||||
def tearDown(self):
|
||||
self.st = self.st.simplify()
|
||||
assert len(self.st.views) == 1
|
||||
|
||||
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)))
|
||||
|
||||
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)))
|
||||
|
||||
class TestIndexExpressions2d(unittest.TestCase):
|
||||
def setUp(self):
|
||||
shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5
|
||||
|
||||
@@ -62,6 +62,7 @@ class TestShapeTrackerAdd(unittest.TestCase):
|
||||
b = ShapeTracker.from_shape((100,))
|
||||
assert a+b == b
|
||||
|
||||
@unittest.skip("no longer simplifies")
|
||||
def test_simple_add_permute(self):
|
||||
a = ShapeTracker.from_shape((10, 10))
|
||||
a = a.permute((1,0))
|
||||
|
||||
@@ -28,7 +28,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_merge_view_recursion_err(self):
|
||||
vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
|
||||
vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
|
||||
self.assertEqual(vm2+vm1, vm1)
|
||||
self.assertEqual(vm2+vm1, None)
|
||||
|
||||
def test_merge_view_recursion_err2(self):
|
||||
vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)
|
||||
|
||||
@@ -69,161 +69,5 @@ class TestMergeDims(unittest.TestCase):
|
||||
# print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
|
||||
self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
|
||||
|
||||
class TestMergeViews(unittest.TestCase):
|
||||
def test_with_mask_0(self):
|
||||
# from test/test_ops.py::TestOps::test_pad_reflect_mode
|
||||
v0 = View(shape=(1, 1, 5, 8), strides=(0, 0, 5, 1), offset=-3, mask=((0, 1), (0, 1), (0, 5), (3, 8)), contiguous=False)
|
||||
v1 = View(shape=(1, 1, 2, 2), strides=(0, 0, 8, 1), offset=3, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(1, 1, 2, 2), strides=(0, 0, 5, 1), offset=0, mask=None, contiguous=False))
|
||||
|
||||
def test_with_mask_1(self):
|
||||
# from test/test_ops.py::TestOps::test_pad_reflect_mode
|
||||
v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, 3, 1), offset=-6, mask=((0, 3), (0, 3), (2, 4), (1, 3)), contiguous=False)
|
||||
v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, 3, 1), offset=7, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=1, mask=None, contiguous=False))
|
||||
|
||||
def test_with_mask_2(self):
|
||||
# from test/test_ops.py::TestOps::test_pad_reflect_mode
|
||||
v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (0, 2)), contiguous=False)
|
||||
v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, -3, 1), offset=3, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=3, mask=None, contiguous=False))
|
||||
|
||||
def test_with_mask_3(self):
|
||||
# from test/test_ops.py::TestOps::test_pad_reflect_mode
|
||||
# has a mask in the final view
|
||||
v0 = View(shape=(3, 3, 4, 4), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False)
|
||||
v1 = View(shape=(3, 3, 4, 2), strides=(48, 16, 4, 1), offset=0, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(3, 3, 4, 2), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False))
|
||||
|
||||
def test_with_mask_4(self):
|
||||
# from test/test_ops.py::TestOps::test_pad_reflect_mode
|
||||
# has a mask in the final view
|
||||
v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (1, 3)), contiguous=False)
|
||||
v1 = View(shape=(3, 3, 3, 3), strides=(45, 15, 3, 1), offset=6, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(3, 3, 3, 3), strides=(0, 0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0), (0, 0)), contiguous=False))
|
||||
|
||||
def test_with_mask_5(self):
|
||||
# from test/test_ops.py::TestOps::test_pad_reflect_mode
|
||||
# has a mask in the final view
|
||||
v0 = View(shape=(1, 1, 6, 5), strides=(0, 0, 5, 1), offset=-5, mask=((0, 1), (0, 1), (1, 6), (0, 5)), contiguous=False)
|
||||
v1 = View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=3, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=-2, mask=((0, 1), (0, 1), (1, 6), (0, 3)), contiguous=False))
|
||||
|
||||
@unittest.expectedFailure # TODO: fix these
|
||||
def test_merges_from_fuzzer1(self):
|
||||
v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
|
||||
v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
|
||||
target = View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, target)
|
||||
|
||||
@unittest.expectedFailure # TODO: fix these
|
||||
def test_merges_from_fuzzer2(self):
|
||||
v0 = View(shape=(5, 10, 12), strides=(100, 1, 10), offset=-20, mask=((0, 5), (0, 10), (2, 12)), contiguous=False)
|
||||
v1 = View(shape=(10, 6, 5, 2, 2), strides=(12, 2, 120, 1, 0), offset=0, mask=None, contiguous=False)
|
||||
target = View(shape=(10, 6, 5, 2, 2), strides=(1, 20, 100, 10, 0), offset=-20, mask=((0, 10), (1, 6), (0, 5), (0, 2), (0, 2)), contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, target)
|
||||
|
||||
@unittest.expectedFailure # TODO: fix these
|
||||
def test_merges_from_fuzzer3(self):
|
||||
v0 = View(shape=(8, 7, 3), strides=(1, 12, -4), offset=6, mask=((2, 6), (0, 7), (0, 3)), contiguous=False)
|
||||
v1 = View(shape=(4, 2, 6, 2, 1), strides=(42, 21, 3, 1, 0), offset=4, mask=None, contiguous=False)
|
||||
target = View(shape=(4, 2, 6, 2, 1), strides=(2, 1, 12, -4, 0), offset=14, mask=((1, 3), (0, 2), (0, 6), (0, 2), (0, 1)), contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, target)
|
||||
|
||||
@unittest.expectedFailure # TODO: fix these
|
||||
def test_merges_from_fuzzer4(self):
|
||||
v0 = View(shape=(7, 21, 3), strides=(54, 3, 1), offset=-9, mask=((0, 6), (3, 21), (0, 3)), contiguous=False)
|
||||
v1 = View(shape=(5, 3, 3, 7), strides=(63, 1, 3, 9), offset=63, mask=None, contiguous=False)
|
||||
target = View(shape=(5, 3, 3, 7), strides=(54, 1, 3, 9), offset=45, mask=((0, 5), (0, 3), (0, 3), (1, 7)), contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, target)
|
||||
|
||||
@unittest.expectedFailure # TODO: fix these
|
||||
def test_merges_from_fuzzer5(self):
|
||||
v0 = View(shape=(5, 1, 24), strides=(20, 0, 1), offset=-2, mask=((0, 5), (0, 1), (2, 22)), contiguous=False)
|
||||
v1 = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 24, 0, 0), offset=0, mask=None, contiguous=False)
|
||||
target = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 20, 0, 0), offset=-2, mask=((1, 11), (0, 2), (0, 5), (0, 2), (0, 1)), contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, target)
|
||||
|
||||
def test_merge_views_variable(self):
|
||||
from tinygrad import Variable
|
||||
N = 100
|
||||
start_pos = Variable("start_pos", 1, N-1)
|
||||
v0 = View(shape=(N, 32, 2), strides=(32, 1, 0), offset=0, mask=((0, N), (0, 32), (0, 1)), contiguous=False)
|
||||
v1 = View(shape=(1, 8, 1, 32), strides=(0, 0, 0, 2), offset=start_pos*64, mask=None, contiguous=False)
|
||||
target = View(shape=(1, 8, 1, 32), strides=(0,0,0,1), offset=start_pos*32, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, target)
|
||||
|
||||
def test_view_padded_area1(self):
|
||||
# test_multinomial
|
||||
v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
|
||||
v1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(1,), strides=(0,), offset=0, mask=((0, 0),), contiguous=False))
|
||||
|
||||
def test_view_padded_area2(self):
|
||||
# test_pad_reflect_mode
|
||||
v0 = View(shape=(1, 1, 10, 7), strides=(0, 0, 5, 1), offset=-15, mask=((0, 1), (0, 1), (3, 8), (0, 5)), contiguous=False)
|
||||
v1 = View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True))
|
||||
|
||||
def test_view_padded_area3(self):
|
||||
# test_roll
|
||||
v0 = View(shape=(2, 4), strides=(0, 1), offset=4, mask=((0, 1), (0, 4)), contiguous=False)
|
||||
v1 = View(shape=(1, 4), strides=(0, 1), offset=4, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(1, 4), strides=(0, 0), offset=0, mask=((0, 0), (0, 0)), contiguous=False))
|
||||
|
||||
def test_view_padded_area4(self):
|
||||
# test_std_mean
|
||||
v0 = View(shape=(2,), strides=(0,), offset=0, mask=((0, 1),), contiguous=False)
|
||||
v1 = View(shape=(1, 1, 1), strides=(0, 0, 0), offset=1, mask=None, contiguous=False)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0)), contiguous=False))
|
||||
|
||||
def test_empty_shape_view1(self):
|
||||
# test_stack_slice
|
||||
v0 = View(shape=(3, 5), strides=(0, 1), offset=0, mask=((0, 1), (0, 5)), contiguous=False)
|
||||
v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
|
||||
v = v0 + v1
|
||||
self.assertIsNotNone(v)
|
||||
self.assertEqual(v, View(shape=(), strides=(), offset=0, mask=None, contiguous=True))
|
||||
|
||||
def test_empty_shape_view2(self):
|
||||
# test_std_mean
|
||||
v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
|
||||
v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
|
||||
v = v0 + v1
|
||||
# TODO: why is this different?
|
||||
self.assertIsNone(v)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
+2
-43
@@ -4,7 +4,7 @@ from dataclasses import dataclass
|
||||
from typing import cast, Sequence
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import resolve, UOp, Variable, sint, smax, smin, sint_to_uop, Ops, ssimplify
|
||||
from tinygrad.helpers import prod, all_int, flatten, ceildiv
|
||||
from tinygrad.helpers import prod, all_int, flatten
|
||||
|
||||
# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
|
||||
def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
|
||||
@@ -171,7 +171,6 @@ class View:
|
||||
if not all_int(vm1.shape):
|
||||
# if all strides are 0 and vm2 is unmasked, return vm1
|
||||
if all(x == 0 for x in vm2.strides+vm1.strides) and vm2.mask is None: return vm1
|
||||
# TODO: handle more cases
|
||||
return None
|
||||
|
||||
# Project vm1's offset and strides on to vm2.
|
||||
@@ -184,47 +183,7 @@ class View:
|
||||
if not resolve((s1 := s1 - o)!=0): continue # if s1 can possibly be 0
|
||||
terms[d2].append((d1, s1))
|
||||
strides[d1] += ssimplify(s1 * vm2.strides[d2])
|
||||
|
||||
# Merge dimensions in vm2 if required.
|
||||
# NB: Merging too many dimensions can make it difficult to project vm2's mask, hence only combining when required.
|
||||
idxs: list[UOp] = [UOp.variable(f"idx{i}", 0, s-1, dtypes.index) for i,s in enumerate(vm1.shape)]
|
||||
merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
|
||||
extents: list[tuple[sint, UOp]] = []
|
||||
for term, s, o in zip(reversed(terms), reversed(vm2.shape), reversed(origin)):
|
||||
merged_term += (sum([idxs[d1] * s1 for d1, s1 in term]) + o) * merged_size
|
||||
merged_size *= s
|
||||
if resolve(merged_term < merged_size, False) and resolve(0 <= merged_term, False):
|
||||
extents.append((merged_size, merged_term))
|
||||
merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
|
||||
if resolve(merged_term != 0): return None
|
||||
if (vm2_shape := tuple(s for s,_ in reversed(extents))) != vm2.shape:
|
||||
if (reshaped_vm2 := vm2.reshape(vm2_shape)) is None: return None
|
||||
# NOTE: this != to prevent infinite loop
|
||||
if reshaped_vm2.shape != vm2.shape: return reshaped_vm2 + vm1
|
||||
|
||||
if vm2.mask:
|
||||
# Try to project vm2's mask on to vm1.
|
||||
newb, newe, bad = [0] * len(vm1.shape), list(vm1.shape), False
|
||||
for (b, e), o, term, (_, t) in zip(vm2.mask, origin, terms, reversed(extents)):
|
||||
if resolve(b <= (t := t.simplify()).vmin and t.vmax < e, False): continue
|
||||
if len(term) != 1:
|
||||
if not term and newe:
|
||||
# t should be a constant if no terms contribute to this dimension, but it might not be simplified
|
||||
if t.vmin != t.vmax: return None
|
||||
newe[0] = 0
|
||||
else: bad = True
|
||||
continue
|
||||
d1, s1 = term[0]
|
||||
newb[d1] = smax(newb[d1], ceildiv(b - o if s1 > 0 else e - o - 1, s1))
|
||||
newe[d1] = smin(newe[d1], (b - o if s1 < 0 else e - o - 1) // s1 + 1)
|
||||
|
||||
# If any of vm1 was masked off, try again with that mask in place.
|
||||
if any((b, e) != (0, s) for b, e, s in zip(newb, newe, vm1.shape)):
|
||||
return vm2 + View.create(vm1.shape, vm1.strides, vm1.offset, tuple(zip(newb, newe)))
|
||||
# Otherwise if vm2's mask was violated, then cannot merge.
|
||||
if bad: return None
|
||||
|
||||
return View.create(vm1.shape, tuple(strides), ssimplify(sum(o * s for o, s in zip(origin, vm2.strides)) + vm2.offset))
|
||||
return None
|
||||
|
||||
def __unsafe_resize(self, arg: tuple[tuple[sint, sint], ...], mask=None) -> View:
|
||||
offset = sum([s * x[0] for s, x in zip(self.strides,arg)])
|
||||
|
||||
Reference in New Issue
Block a user