From 33d80db64cc93629cbf355e16b9e005b39971e18 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 22 Jul 2025 19:12:37 -0700
Subject: [PATCH] this can happen later

---
 tinygrad/codegen/devectorizer.py |  4 ++--
 tinygrad/codegen/linearize.py    | 11 ++++++++++-
 tinygrad/runtime/ops_python.py   |  2 +-
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py
index a94381a74a..84eff10b2e 100644
--- a/tinygrad/codegen/devectorizer.py
+++ b/tinygrad/codegen/devectorizer.py
@@ -310,8 +310,8 @@ pm_render = PatternMatcher([
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op is Ops.CUSTOM else None),
   # gate any stores that aren't gated with ifs
-  #(UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True),
-  #  lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:])),
+  (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat()), name="store"),
+    lambda gate,store: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(gate,)),)+store.src[2:])),
 ])
 
 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***
diff --git a/tinygrad/codegen/linearize.py b/tinygrad/codegen/linearize.py
index 0581e7b386..5d19c4ae55 100644
--- a/tinygrad/codegen/linearize.py
+++ b/tinygrad/codegen/linearize.py
@@ -92,7 +92,16 @@ class BlockContext:
       # RANGE/IF add to the next ctx
       # STORE/ASSIGN subtract from the next ctx
       if u.op in {Ops.RANGE, Ops.IF}: ctx.child_ctxs[u] = _sort_ctx(ctx.block_ctxs[u] + (u,))
-      elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([x for x in ctx.block_ctxs[u] if x not in u.src])
+      elif u.op is Ops.STORE:
+        if len(definereg:=[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_REG]):
+          # old assign logic
+          ctx.child_ctxs[u] = tuple([y for y in ctx.last_ctx(u.src[1]) if y not in definereg[0].src[1:]])
+        elif any(x.op is Ops.DEFINE_LOCAL for x in u.src[0].toposort()):
+          # deal with non-reduce locals. probably wrong
+          idx_context, store_context = ctx.last_ctx(u.src[0]), ctx.last_ctx(u.src[1])
+          ctx.child_ctxs[u] = tuple([y for y in store_context if y not in idx_context and y.op is Ops.RANGE])
+        else: ctx.child_ctxs[u] = ()
+      #elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([x for x in ctx.block_ctxs[u] if x not in u.src])
     return ctx
 
 # ***** make blocks *****
diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py
index 487fb5b92e..43a263910e 100644
--- a/tinygrad/runtime/ops_python.py
+++ b/tinygrad/runtime/ops_python.py
@@ -56,7 +56,7 @@ class PythonProgram:
         assert dtype is not None, f"{uop} is missing a dtype"
         dl[i] = dtype
         if uop is Ops.STORE:
-          assert len(inp) == 2, "expected store is ([(memory, offset, gate)], [value])"
+          #assert len(inp) == 2, "expected store is ([(memory, offset, gate)], [value])"
           for j,val in enumerate(inp[1] if dtp[1].count > 1 else [inp[1]]):
             for (m,o,g),v in zip(inp[0], val):
               if g: _store(m, o+j, v)