From 7fa94ffa4a90be0b8b485316207c284323de4aed Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 8 Oct 2025 22:29:37 +0800 Subject: [PATCH] tsink_base wasn't needed --- tinygrad/schedule/indexing.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 54b0a0bb9a..91664d253e 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -25,11 +25,11 @@ def realize_assign(ctx:dict[UOp, None], a:UOp) -> None: pm_generate_realize_map = PatternMatcher([ # always realize SINK src (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)), - # always realize ASSIGN/COPY/BUFFER_VIEW/CONTIGUOUS + # always realize COPY/BUFFER_VIEW/CONTIGUOUS (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize), # realize srcs of COPY, MSELECT, MSTACK (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs), - # realize input to assign (might be optimized out) + # realize ASSIGN and input to assign (might be optimized out) (UPat(Ops.ASSIGN, name="a"), realize_assign), ]) @@ -104,8 +104,6 @@ pm_apply_rangeify = PatternMatcher([ ]) def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: - tsink_base = UOp.sink(*[x.base for x in tsink.src]) - rctx = IndexingContext() # get ops to realize @@ -113,7 +111,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # explicit rangeify ending_ranges: dict[UOp, bool] = {} - for x in tsink_base.reverse_toposort(consumer_map:=tsink_base.get_consumer_map()): + for x in tsink.reverse_toposort(consumer_map:=tsink.get_consumer_map()): if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue ending_ranges[x] = any(ending_ranges[u] for u in consumer_map[x])