From ca36ef0186178accbe442a082e96c5bb3d4036e9 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 12 Mar 2026 15:49:43 +0800 Subject: [PATCH] found_assign --- tinygrad/schedule/rangeify.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 7191177c34..a3e49d0f7c 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -33,7 +33,7 @@ def found_assign(ctx:dict[UOp, UOp], assign:UOp, src:UOp): # *** fold moved ASSIGNs (hack for openpilot) *** pm_fold_moved_assign = PatternMatcher([ - (UPat(Ops.ASSIGN, src=(UPat(), UPat((*GroupOp.Movement, Ops.CAST), name="src")), name="assign"), found_assign), + (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement, Ops.CAST), name="src")))), name="assign"), found_assign), # replace ALU sources with assign versions found above (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None), ])