From f64908b6c505c69b152ea439fbecae936bf07ec3 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 25 Jun 2026 14:47:27 -0700 Subject: [PATCH] late lowering for weakint --- tinygrad/codegen/__init__.py | 12 ++++++------ tinygrad/uop/symbolic.py | 3 --- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 567c2c4067..15d1681394 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -109,12 +109,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch) # devectorize - sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store+load_store_indexing, - ctx=ren, name="devectorize") - - # lower the index dtype to a concrete int - sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower all index dtypes") - sink = graph_rewrite(sink, symbolic, name="post index symbolic") + sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store, ctx=ren, name="devectorize") + sink = graph_rewrite(sink, gep_pushing, name="gep pushing") # optional pre matcher if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher") @@ -127,6 +123,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # do memory coalesing (late) sink = memory_coalesing(sink, ren) + # lower the index dtype to a concrete int + sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing, name="lower all index dtypes") + sink = graph_rewrite(sink, symbolic, name="post index symbolic") + # floordiv+mod / dtype decomp (early) supported_ops = tuple(ren.code_for_op.keys()) pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 62f98d788f..3d533514a1 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -320,9 +320,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ 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 - (UPat(Ops.STACK, src=UPat(Ops.CONST), name="vec"), - lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src)) if len(vec.src) > 0 else None), ])+div_and_mod_symbolic+gep_pushing # ******** we take a small aside to "simplify_valid" to rewrite valids ********