no range tags in kernels (#12962)

This commit is contained in:
George Hotz
2025-10-28 12:33:48 +08:00
committed by GitHub
parent 39c2117dea
commit bbe0bebbf3
+7 -3
View File
@@ -384,8 +384,8 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
return buf
def renumber_range(ctx:LocalAddBufferContext, r:UOp):
if r.tag is not None: return None
ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
if r.tag != (): return None
ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
ctx.range += 1
return ret
@@ -443,6 +443,10 @@ pm_remove_tags = PatternMatcher([
(UPat(GroupOp.All, name="x"), remove_metadata_tags),
])
pm_add_range_tags = PatternMatcher([
(UPat(Ops.RANGE, name="x"), lambda x: x.rtag(()))
])
@dataclass(frozen=True)
class Kernel:
ast: UOp
@@ -532,7 +536,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Tagged Rangeify")
# bufferize -> store
tsink = graph_rewrite(tsink, pm_add_buffers, bottom_up=True, name="bufferize to store")
tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_range_tags, bottom_up=True, name="bufferize to store")
tsink = graph_rewrite(tsink, split_kernels, ctx=uop_list, name="split kernels")
# if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign