From b4c799b81032b02b66fa9763d0e9101b5f704a2d Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 6 Feb 2026 17:44:09 +0800 Subject: [PATCH] better --- tinygrad/codegen/__init__.py | 8 +------- tinygrad/schedule/rangeify.py | 8 +++++++- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 306128bb68..0b9e864e8e 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -18,15 +18,9 @@ from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_in ReduceContext, correct_load_store, pm_render, pm_add_loads from tinygrad.codegen.opt.postrange import apply_opts, make_images from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse -from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops +from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize -pm_syntactic_sugar = PatternMatcher([ - # INDEX on ptr INDEX concats them - (UPat(Ops.INDEX, name="i1").f(Ops.INDEX, name="i2", allow_any_len=True), - lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None), -]) - def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -> UOp: if ren is None: ren = Renderer() diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index a8e5c978bc..ff5ef67324 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -14,6 +14,12 @@ from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, ALWAYS_CONTI import sys sys.setrecursionlimit(10000) +pm_syntactic_sugar = PatternMatcher([ + # INDEX on ptr INDEX concats them + (UPat(Ops.INDEX, name="i1").f(Ops.INDEX, name="i2", allow_any_len=True), + lambda i1,i2: i2.replace(src=i1.src+i2.src[1:]) if isinstance(i1.dtype, PtrDType) and not isinstance(i2.dtype, PtrDType) else None), +]) + # movement op on INDEX as a PatternMatcher pm_mops = PatternMatcher([ (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), @@ -563,7 +569,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: uop_list: list[UOp] = [] tsink = graph_rewrite(sink, add_tags, ctx=(uop_list, set()), bottom_up=True, name="number the uops") - tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites+replace_contiguous, ctx={}, name="earliest rewrites") + tsink = graph_rewrite(tsink, pm_syntactic_sugar+pm_mops+earliest_rewrites+replace_contiguous, ctx={}, bottom_up=True, name="earliest rewrites") # convert movement ops to ranges tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))