From d1cce7a4766fcab73eb2e7ea17a94d5d9d084929 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 16 Apr 2026 19:06:40 +0800 Subject: [PATCH] put the ranges on store instead of after (#15759) * put the ranges on store instead of after * better assert * fix stuff * comment out slow rules i don't understand * simpler rule * closer * return false for store * fix loop * only a few schedule failures remain * remove stores to self * all tests pass locally * remove junk * regression test and fix * better test, bump broken torch count * bugfix with regression test * new fusion is better --- extra/torch_backend/example.py | 2 +- extra/torch_backend/test_kernel_fusion.py | 26 ++++++------ test/backend/test_randomness.py | 1 + test/backend/test_schedule.py | 4 +- test/mockgpu/amd/amdgpu.py | 2 +- test/unit/test_assign.py | 35 ++++++++++----- tinygrad/schedule/indexing.py | 29 +++++++------ tinygrad/schedule/rangeify.py | 52 ++++++++++------------- tinygrad/uop/symbolic.py | 4 +- 9 files changed, 84 insertions(+), 71 deletions(-) diff --git a/extra/torch_backend/example.py b/extra/torch_backend/example.py index 7a609b7a17..506d3bde56 100644 --- a/extra/torch_backend/example.py +++ b/extra/torch_backend/example.py @@ -24,7 +24,7 @@ if __name__ == "__main__": kernel_count = GlobalCounters.kernel_count assert kernel_count > 0, "No kernels, test failed" # NOTE: this is 124 on torch 2.10.0 - expected_kernels = 334 + expected_kernels = 355 expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected." if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning) assert kernel_count <= expected_kernels, f"{expectation}" diff --git a/extra/torch_backend/test_kernel_fusion.py b/extra/torch_backend/test_kernel_fusion.py index f79733fd37..0a748fe4ad 100644 --- a/extra/torch_backend/test_kernel_fusion.py +++ b/extra/torch_backend/test_kernel_fusion.py @@ -23,7 +23,7 @@ class TestKernelFusionRegression(unittest.TestCase): def fn(): x = torch.randn(128, 128, device=device) return (x + 1.0) * 2.0 - 0.5 - self._check_kernel_count(fn, 7) + self._check_kernel_count(fn, 6) def test_relu_fusion(self): def fn(): @@ -31,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase): conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device) with torch.no_grad(): return torch.nn.functional.relu(conv(x)) - self._check_kernel_count(fn, 8) + self._check_kernel_count(fn, 7) def test_batchnorm_fusion(self): def fn(): @@ -41,26 +41,26 @@ class TestKernelFusionRegression(unittest.TestCase): bn.eval() with torch.no_grad(): return torch.nn.functional.relu(bn(conv(x))) - self._check_kernel_count(fn, 12) + self._check_kernel_count(fn, 11) def test_reduce_fusion(self): def fn(): x = torch.randn(64, 64, device=device) return (x * 2.0).sum() - self._check_kernel_count(fn, 7) + self._check_kernel_count(fn, 6) def test_matmul_elementwise_fusion(self): def fn(): x = torch.randn(32, 32, device=device) w = torch.randn(32, 32, device=device) return torch.nn.functional.relu(x @ w + 1.0) - self._check_kernel_count(fn, 9) + self._check_kernel_count(fn, 8) def test_pooling_fusion(self): def fn(): x = torch.randn(1, 8, 16, 16, device=device) return torch.nn.functional.max_pool2d(x * 2.0, 2) - self._check_kernel_count(fn, 7) + self._check_kernel_count(fn, 6) def test_residual_add_relu_fusion(self): def fn(): @@ -68,7 +68,7 @@ class TestKernelFusionRegression(unittest.TestCase): identity = torch.randn(1, 8, 16, 16, device=device) out = x + identity return torch.nn.functional.relu(out) - self._check_kernel_count(fn, 9) + self._check_kernel_count(fn, 8) def test_inplace_add_relu_fusion(self): def fn(): @@ -76,7 +76,7 @@ class TestKernelFusionRegression(unittest.TestCase): y = torch.randn(1, 16, 32, 32, device=device) x += y return torch.nn.functional.relu(x) - self._check_kernel_count(fn, 9) + self._check_kernel_count(fn, 8) def test_conv_bn_add_relu_fusion(self): def fn(): @@ -89,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase): out = bn(conv(x)) out += identity return torch.nn.functional.relu(out) - self._check_kernel_count(fn, 14) + self._check_kernel_count(fn, 13) def test_multiple_inplace_ops_fusion(self): def fn(): @@ -97,7 +97,7 @@ class TestKernelFusionRegression(unittest.TestCase): x += 1.0 x *= 2.0 return torch.nn.functional.relu(x) - self._check_kernel_count(fn, 6) + self._check_kernel_count(fn, 5) def test_view_inplace_no_fusion_break(self): def fn(): @@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase): view = x[1:3] view += 1.0 return x.sum() - self._check_kernel_count(fn, 10) + self._check_kernel_count(fn, 8) def test_batchnorm_running_stats_update(self): def fn(): @@ -114,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase): bn.train() with torch.no_grad(): return bn(x) - self._check_kernel_count(fn, 10) + self._check_kernel_count(fn, 9) # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer def test_mnist_training_fusion(self): @@ -135,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase): loss.backward() optimizer.step() return loss - self._check_kernel_count(fn, 26) + self._check_kernel_count(fn, 25) if __name__ == "__main__": unittest.main() diff --git a/test/backend/test_randomness.py b/test/backend/test_randomness.py index bcc6bf72e1..fc874b5aaf 100644 --- a/test/backend/test_randomness.py +++ b/test/backend/test_randomness.py @@ -360,6 +360,7 @@ class TestRandomness(unittest.TestCase): torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)] self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples))) + @unittest.skip("this test is flaky") def test_multinomial_counterexample(self): tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True) torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True) diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index e39ae28c9a..2eede24a50 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -754,7 +754,7 @@ class TestSchedule(unittest.TestCase): p = P[0] p = p.pad(((1, 0), )) p = p.repeat([2]) - run_schedule(check_schedule(p, 4)) # TODO: this is high + run_schedule(check_schedule(p, 3)) tiny_ret = p.numpy() P = np.ones((3, 3), dtype=np.float32) @@ -1075,7 +1075,7 @@ class TestSchedule(unittest.TestCase): idx = Tensor([1,2,5,6], dtype=dtypes.int32) flat_base[idx] = Tensor([99,99,99,99]) base.assign(flat_base.reshape(4, 4)) - sched = check_schedule(base, 6) # TODO: this is high + sched = check_schedule(base, 4) run_schedule(sched) expected = list(range(16)) for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v diff --git a/test/mockgpu/amd/amdgpu.py b/test/mockgpu/amd/amdgpu.py index 92031cfd87..7035dfd1f8 100644 --- a/test/mockgpu/amd/amdgpu.py +++ b/test/mockgpu/amd/amdgpu.py @@ -7,7 +7,7 @@ from tinygrad.runtime.support import c MOCKGPU_ARCH = "cdna4" if DEV.arch == "gfx950" else "rdna4" if DEV.arch.startswith("gfx12") else "rdna3" assert (ma:=getenv("MOCKGPU_ARCH", "")) == "", "MOCKGPU_ARCH is deprecated, use DEV=" + \ - str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100"))) + str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100"))) # type: ignore GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000, "cdna4": 90500}[MOCKGPU_ARCH] import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4 diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index c83158cf61..3671f583d4 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -621,6 +621,23 @@ class TestAssign(unittest.TestCase): # N matmuls + N assigns + 1 final read = 2*N+1 (AFTER embedding allows full graph scheduling with shared contiguous reuse) self.assertEqual(GlobalCounters.kernel_count, 2*N+1) + def test_double_assign_from_const(self): + a = Tensor.empty(2) + a.assign(Tensor.ones(2)) + a.assign(Tensor.ones(2)) + GlobalCounters.reset() + a.realize() + self.assertEqual(GlobalCounters.kernel_count, 1) + self.assertEqual(a.tolist(), [1.,1.]) + + def test_nested_after_contiguous_store(self): + # Mirrors the nested contiguous-write-then-assign-back shape from torch backend view updates. + base = Tensor.empty(3, dtype=dtypes.int64) + base.assign(Tensor([1, 2, 3], dtype=dtypes.int64)) + contig = base.contiguous() + contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64)) + base.assign(contig).realize() + self.assertEqual(base.tolist(), [1,4,3]) class TestAssignOrdering(unittest.TestCase): """Tests for complex assign orderings that could differ between lazy and eager execution. @@ -958,11 +975,10 @@ class TestAfterCachePatterns(unittest.TestCase): a_store = a.uop.store(c.uop) b_store = b.uop.store(c.uop) - with self.assertRaises(AssertionError): - a = Tensor(a.uop.after(a_store, b_store)) - a.realize() - np.testing.assert_array_equal(a.numpy(), 1) - np.testing.assert_array_equal(b.numpy(), 1) + a = Tensor(a.uop.after(a_store, b_store)) + a.realize() + np.testing.assert_array_equal(a.numpy(), 1) + np.testing.assert_array_equal(b.numpy(), 1) def test_double_store_after_different_sizes(self): full = Tensor.zeros(2).contiguous() @@ -974,11 +990,10 @@ class TestAfterCachePatterns(unittest.TestCase): full_store = full.uop.store(full_src.uop) head_store = head.uop.store(head_src.uop) - with self.assertRaises(AssertionError): - head = Tensor(head.uop.after(head_store, full_store)) - head.realize() - np.testing.assert_array_equal(head.numpy(), [3]) - np.testing.assert_array_equal(full.numpy(), [1, 2]) + head = Tensor(head.uop.after(head_store, full_store)) + head.realize() + np.testing.assert_array_equal(head.numpy(), [3]) + np.testing.assert_array_equal(full.numpy(), [1, 2]) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 25aa9551d2..f98dbbc676 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -27,13 +27,11 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp): pm_generate_realize_map = PatternMatcher([ # always realize - (UPat({Ops.COPY, Ops.CONTIGUOUS}, name="tr"), realize), - # realize AFTER of STORE+AFTER - (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE)), allow_any_len=True, name="tr"), realize), + (UPat({Ops.COPY, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize), # realize srcs of these (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs), - # sometimes realize/unrealize src of store+after - (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))))), realize_store_after_src), + # sometimes we need to realize the src of STORE if there's a self-access + (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src), ]) @dataclass(frozen=True) @@ -60,8 +58,7 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp): new_srcs = [] for s in x.src: new_src = s - if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT} or \ - (s.op is Ops.AFTER and not any(c.op in {Ops.STORE, Ops.END} for c in s.src[1:])): + if s.op in {Ops.PARAM, Ops.BUFFER_VIEW, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}: if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0]) elif s in ctx.realize_map: realized_ranges = ctx.realize_map[s] @@ -166,8 +163,8 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # no ranges on kernels, they are internal if x.op in {Ops.CALL, Ops.FUNCTION, Ops.LINEAR}: continue - # only STORE+AFTER has range - if x.op is Ops.AFTER and all(s.op is not Ops.STORE for s in x.src[1:]): continue + # AFTER doesn't have range + if x.op is Ops.AFTER: continue # treat MSTACK/MSELECT like SINK if x.op in {Ops.MSTACK, Ops.MSELECT}: continue @@ -180,15 +177,21 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # 2. from the single consumer if this op only has one consumer # 3. potentially new if this op has 2+ consumers + shape = x._shape + if x.op is Ops.STORE: + # TODO: TestTensorVariable.test_symbolic_var_sum_alt_name fails with this, fix canonicalize on variables. + #assert x.src[0].shape == x.src[1].shape, f"STORE must have matching input shapes, {x.src[0].shape} != {x.src[1].shape}" + shape = x.src[0].shape + consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map] if x in rctx.realize_map: # if this is in the realize_map, we create new ranges (at the output) - out_rngs = tuple(rctx.new_range(s) for s in x.shape) + out_rngs = tuple(rctx.new_range(s) for s in shape) # all ranges are ended now ending_ranges[x] = [] # mark all ranges as ended assert rctx.realize_map[x] is None - rctx.realize_map[x] = list(range(len(x.shape))) + rctx.realize_map[x] = list(range(len(shape))) elif len(consumer_rngs) == 0: # if no consumers have ranges and this isn't realized, this doesn't have ranges either. continue @@ -214,7 +217,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: minimum_valid = UOp.const(dtypes.bool, False).usum(valids) _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid")) else: - _out_rngs.append(rctx.new_range(x.shape[i])) + _out_rngs.append(rctx.new_range(shape[i])) _realize_axis.append(i) out_rngs = tuple(_out_rngs) @@ -231,7 +234,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: ending_ranges[x] = [] if len(_realize_axis): rctx.realize_map[x] = _realize_axis - out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)]) + out_rngs = tuple([(rctx.new_range(shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)]) # TODO: some ops don't have shape, enable this after the `.st` property is removed #assert len(out_rngs) == len(x.shape), \ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 84af8c1d0d..e8acc31e7b 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -71,8 +71,7 @@ pm_mops = PatternMatcher([ if r.src[0]._shape is not None and len(idx.src[1:]) == len(r.shape) else None), # move movement ops and INDEX after AFTER (but not when AFTER has a raw STORE with shaped children — from replace_contig_with_store_after) (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True), - lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg) - if a.src[0]._shape is not None and not any(s.op is Ops.STORE and s.src[0]._shape is not None for s in a.src[1:]) else None), + lambda r,a: UOp(r.op, r.dtype, (a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], r.arg)), (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])), # lower SHAPED_WMMA to WMMA with CONTRACT/UNROLL (UPat(Ops.SHAPED_WMMA, name="x"), lower_shaped_wmma), @@ -81,21 +80,14 @@ pm_mops = PatternMatcher([ # ***************** # 0. do some cleanup rewrites, mostly copied from the old stuff -def fix_store_after_hazard(after:UOp, target:UOp, src:UOp): +def fix_store_hazard(target:UOp, src:UOp): # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set()) base = target.base reaches_base: dict[UOp, bool] = {} for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS): reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src) - if reaches_base[s] and s.op in unsafe: return after.replace(src=(after.src[0], target.store(src.contiguous()))) - -def normalize_store_after_target_chain(after:UOp, target:UOp, src:UOp): - root_target = target - while root_target.op is Ops.AFTER: root_target = root_target.src[0] - # when RHS depends on the previous assign result, break with contiguous - if target in src.toposort(): src = src.contiguous() - return after.replace(src=(root_target, root_target.store(src))) + if reaches_base[s] and s.op in unsafe: return target.store(src.contiguous()) def split_reduceop(reduce:UOp, x:UOp): if prod(reduce.shape) == 0: return None @@ -173,6 +165,9 @@ earliest_rewrites = mop_cleanup+PatternMatcher([ lambda x,copy: x.replace(src=(copy.replace(src=(x.src[0],)+copy.src[1:]),)+x.src[1:]) \ if isinstance(x.device, str) and x.device.startswith("DISK") else None), + # SINK only ever references the base + (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.base for y in x.src))), + # ** copy rules ** # COPY and source size need to match @@ -182,22 +177,17 @@ earliest_rewrites = mop_cleanup+PatternMatcher([ # copy only to different device (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP) if x.device == copy.device else None), - # ** assign rules (STORE+AFTER) ** + # ** store rules ** - # move bitcast from store+after target to source - (UPat(Ops.AFTER, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(Ops.STORE, src=(UPat(Ops.BITCAST), UPat(name="src"))))), - lambda target, src: target.after(target.store(src.bitcast(target.dtype)))), + # fix store hazard (dest is in used in src) by adding contiguous: TestAssign.test_post_flipped_assignment + (UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src"))), fix_store_hazard), - # wrap STORE in inner AFTER when target is a view — gives the STORE its own ranges from the view shape - (UPat(Ops.AFTER, src=(UPat(name="buf"), UPat(Ops.STORE, src=(UPat(name="target"), UPat()))), name="after"), - lambda after, buf, target: after.replace(src=(buf, target.after(after.src[1]))) if target.shape != buf.shape else None), + # remove two STOREs that store the same thing to the same place: TestSchedule.test_dedup_assign + (UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1").after(UPat.var("a1").store(UPat.var("src"))), lambda buf,src,a1:a1), - # make source contiguous if it has hazardous movement ops on the dest buffer - (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src")))), name="after"), fix_store_after_hazard), - - # normalize target chain: walk through AFTERs to root, insert contiguous if needed - (UPat(Ops.AFTER, src=(UPat(Ops.AFTER, name="target"), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="after"), - lambda after, target, src: normalize_store_after_target_chain(after, target, src)), + # move bitcast from store dest to source: TestAssign.test_assign_bitcast + (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))), + lambda target, src: target.store(src.bitcast(target.dtype))), # ** size 0 ** @@ -257,6 +247,9 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp): if (x.op is Ops.BUFFERIZE and x.arg.addrspace == AddrSpace.GLOBAL) or x.op is Ops.MSTACK: accessed_buffers.append(x) return False + if x.op is Ops.STORE: + # don't look inside stores, this doesn't count toward buffer accesses + return False if x.op is Ops.PARAM: accessed_buffers.append(x) if x.op is Ops.INDEX: @@ -326,6 +319,10 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([ pm_remove_bufferize = PatternMatcher([ # remove reindexing with cost function (UPat.var("src").f(Ops.BUFFERIZE, allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize), + # STORE to self is NOOP + (UPat.var("x").store(UPat.var("x")), lambda x: UOp(Ops.NOOP)), + # END on NOOP is NOOP + (UPat(Ops.END, src=(UPat(Ops.NOOP, name="x"),), allow_any_len=True), lambda x: x), ]) def late_buffer_view(t:UOp, b:UOp): @@ -391,9 +388,6 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True): if (after:=x.src[0]).op is Ops.AFTER: buf = after.src[0].buf_uop.base if not (stores := [s for s in after.src[1:] if s.op is Ops.STORE and s.src[0].op is Ops.INDEX]): return buf - # the ranges are created on the AFTER, and the stores might be different sizes - # so we block all multi store AFTERs - assert len(stores) <= 1, "rangeify doesn't support multiple stores on one after" # BUFFERIZE(INDEX(...)); store through the underlying global index instead. ended_stores = [] for store in stores: @@ -482,8 +476,8 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp): buf = after.buf_uop # HACK to put the buffer in the MAP instead of MSTACK/MSELECT if buf.op in {Ops.MSTACK, Ops.MSELECT}: buf = buf.src[0] - assert buf not in ctx.map - ctx.map[buf] = after + # NOTE: this is bottom up, so we only add it once + if buf not in ctx.map: ctx.map[buf] = after return buf def renumber_range(ctx:LocalAddBufferContext, r:UOp): diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 9f391ae9b5..f1b7d0e66e 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -285,8 +285,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)), # only RANGE/IF/STORE/KERNEL have side effects (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+ - tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.BUFFERIZE} - else y.src for y in x.src[1:]])))), + tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.BUFFERIZE} + else y.src for y in x.src[1:]]))))), # after with 1 src is just src[0] (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s), # VECTORIZE/CONST