tsink_base wasn't needed

2025-10-08 22:29:37 +08:00
1 changed files with 3 additions and 5 deletions
@@ -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])