16 seconds

examples/hlb_cifar10.py

@@ -3,7 +3,7 @@
 # tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 # https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
 # https://siboehm.com/articles/22/CUDA-MMM
-import random, time
+import random, time, math
 import numpy as np
 from typing import Optional
 from extra.lr_scheduler import OneCycleLR
@@ -220,13 +220,10 @@ def train_cifar():
     return X_cutmix, Y_cutmix
 
   def random_permutation(rows:int, cols:int) -> Tensor:
-    # Alternating modular additions are bijections on the rows x cols domain.
-    idx = Tensor.arange(rows * cols)
-    row, col = idx // cols, idx % cols
-    for _ in range(4):
-      row = (row + col * col + random.randrange(rows) * col + random.randrange(rows)) % rows
-      col = (col + row * row + random.randrange(cols) * row + random.randrange(cols)) % cols
-    return row * cols + col
+    size = rows * cols
+    # An affine map is a permutation when its stride is coprime to the domain size.
+    while math.gcd(stride:=random.randrange(1, size), size) != 1: pass
+    return (Tensor.arange(size) * stride + random.randrange(size)) % size
 
   def shuffled_augmentations(X:Tensor, Y:Tensor):
     perms = random_permutation(X.shape[0] // BS, BS)
@@ -398,10 +395,9 @@ def train_cifar():
     while i <= STEPS:
       if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
         # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
-        corrects = []
-        corrects_ema = []
-        losses = []
-        losses_ema = []
+        correct_sum = loss_sum = None
+        correct_sum_ema = loss_sum_ema = None
+        correct_len = correct_len_ema = eval_batches = 0
         for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False):
           if len(GPUS) > 1:
             Xt.shard_(GPUS, axis=0)
@@ -411,23 +407,31 @@ def train_cifar():
           out = eval_forward_jitted(model, Xt_contiguous).clone().realize()
           out_flipped = eval_forward_jitted(model, Xt_contiguous[..., ::-1].contiguous().realize())
           correct, loss = eval_step_jitted(out, out_flipped, Yt)
-          losses.append(loss.numpy().tolist())
-          corrects.extend(correct.numpy().tolist())
+          batch_correct, batch_loss = correct.sum().realize(), loss.clone().realize()
+          correct_sum = batch_correct if correct_sum is None else correct_sum + batch_correct
+          loss_sum = batch_loss if loss_sum is None else loss_sum + batch_loss
+          correct_len += correct.numel()
+          eval_batches += 1
           if model_ema:
             out_ema = eval_forward_ema_jitted(model_ema.net_ema, Xt_contiguous).clone().realize()
             out_flipped_ema = eval_forward_ema_jitted(model_ema.net_ema, Xt_contiguous[..., ::-1].contiguous().realize())
             correct_ema, loss_ema = eval_step_ema_jitted(out_ema, out_flipped_ema, Yt)
-            losses_ema.append(loss_ema.numpy().tolist())
-            corrects_ema.extend(correct_ema.numpy().tolist())
+            batch_correct_ema, batch_loss_ema = correct_ema.sum().realize(), loss_ema.clone().realize()
+            correct_sum_ema = batch_correct_ema if correct_sum_ema is None else correct_sum_ema + batch_correct_ema
+            loss_sum_ema = batch_loss_ema if loss_sum_ema is None else loss_sum_ema + batch_loss_ema
+            correct_len_ema += correct_ema.numel()
 
         # collect accuracy across ranks
-        correct_sum, correct_len = sum(corrects), len(corrects)
-        if model_ema: correct_sum_ema, correct_len_ema = sum(corrects_ema), len(corrects_ema)
+        assert correct_sum is not None and loss_sum is not None
+        correct_count, eval_loss = correct_sum.item(), (loss_sum / eval_batches).item()
+        if model_ema:
+          assert correct_sum_ema is not None and loss_sum_ema is not None
+          correct_count_ema, eval_loss_ema = correct_sum_ema.item(), (loss_sum_ema / eval_batches).item()
 
-        eval_acc_pct = correct_sum/correct_len*100.0
-        if model_ema: acc_ema = correct_sum_ema/correct_len_ema*100.0
-        print(f"eval     {correct_sum}/{correct_len} {eval_acc_pct:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
-        if model_ema: print(f"eval ema {correct_sum_ema}/{correct_len_ema} {acc_ema:.2f}%, {(sum(losses_ema)/len(losses_ema)):7.2f} val_loss STEP={i}")
+        eval_acc_pct = correct_count/correct_len*100.0
+        if model_ema: acc_ema = correct_count_ema/correct_len_ema*100.0
+        print(f"eval     {correct_count}/{correct_len} {eval_acc_pct:.2f}%, {eval_loss:7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
+        if model_ema: print(f"eval ema {correct_count_ema}/{correct_len_ema} {acc_ema:.2f}%, {eval_loss_ema:7.2f} val_loss STEP={i}")
 
       if STEPS == 0 or i == STEPS: break
 

tinygrad/codegen/__init__.py

@@ -8,11 +8,11 @@ from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_simplify_valid, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
@@ -22,6 +22,7 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
 from tinygrad.schedule.rangeify import pm_mops
+from tinygrad.schedule.indexing import apply_movement_op
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalese import memory_coalesing, pm_simplify_add_image
@@ -29,6 +30,8 @@ from tinygrad.helpers import all_same, flatten, argsort, partition
 from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts
 
+empty_matcher = PatternMatcher([])
+
 def do_number_param(ctx:list[int], x:UOp):
   if x.arg.slot != -1: return None
   ctx[0] += 1
@@ -133,16 +136,40 @@ unbroadcast = pm_wmma_add+PatternMatcher([
 ])
 
 def do_devectorize(ctx, b:UOp):
-  ren = ctx[-1] if isinstance(ctx, tuple) else ctx
+  ren = ctx.ren if isinstance(ctx, DevectorizeContext) else (ctx[-1] if isinstance(ctx, tuple) else ctx)
   if b.op in GroupOp.Elementwise and b.dtype in dtypes.floats and ren.supports_float4: return None
-  if (shape:=b.shape) == (): return None
+  if (shape:=b._shape) is None or shape == (): return None
   # broadcasting needs to be already unpacked
-  if not all_same([x.shape for x in b.src]): return None
+  if any(x._shape != shape for x in b.src): return None
   src = []
-  for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(dtypes.index, i) for i in idx]
-    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
-  return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
+  for idx_c in shape_indexes(shape):
+    new_src = tuple(index_lane(ctx, x, idx_c) if isinstance(ctx, DevectorizeContext) else
+                    (x.src[idx_c[0].arg] if len(idx_c) == 1 and x.op is Ops.STACK else UOp(Ops.INDEX, x.dtype, (x,)+idx_c)) for x in b.src)
+    src.append(UOp(b.op, b.dtype, new_src, b.arg, b.tag))
+  return UOp(Ops.STACK, b.dtype, tuple(src)).reshape(shape) if b.op is not Ops.STORE else UOp.group(*src)
+
+@dataclass
+class DevectorizeContext:
+  ren: Renderer
+  lanes: dict[tuple[UOp, tuple[UOp, ...]], UOp]
+  @property
+  def rewrite_cache_key(self): return (type(self.ren), self.ren.target)
+
+devectorize_caches: list[dict[tuple[UOp, tuple[UOp, ...]], UOp]] = []
+class scoped_devectorize_cache:
+  def __enter__(self): devectorize_caches.append({})
+  def __exit__(self, *args): devectorize_caches.pop()
+
+def index_lane(ctx:DevectorizeContext, x:UOp, idxs:tuple[UOp, ...]) -> UOp:
+  key = (x, idxs)
+  if (ret:=ctx.lanes.get(key)) is not None: return ret
+  if x.op is Ops.STACK and idxs and idxs[0].op is Ops.CONST:
+    ret = index_lane(ctx, x.src[idxs[0].arg], idxs[1:]) if len(idxs) > 1 else x.src[idxs[0].arg]
+  elif x.op in GroupOp.Movement and len(idxs) == len(x.shape):
+    ret = index_lane(ctx, x.src[0], apply_movement_op(x.op, x.src[0].shape, x.marg, idxs))
+  else: ret = UOp(Ops.INDEX, x.dtype, (x,)+idxs)
+  ctx.lanes[key] = ret
+  return ret
 
 def index_elementwise(x:UOp, idx:UOp):
   indexes = idx.src[1:]
@@ -317,14 +344,18 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+unbroadcast+pm_add_loads, name="*** unbroadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2,
+                       ctx=DevectorizeContext(ren, devectorize_caches[-1] if devectorize_caches else {}), name="devectorize2")
 
   # some coalesing misses without this
   sink = graph_rewrite(sink, sym, name="early symbolic")
 
   # do memory coalesing (late)
   sink = memory_coalesing(sink, ren)
-  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  if IMAGE: sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+
+  has_invalid = (sink.op is Ops.CONST and sink.arg is Invalid) or any(u.op is Ops.CONST and u.arg is Invalid for u in sink.backward_slice)
+  if has_invalid: sink = graph_rewrite(sink, pm_simplify_valid, name="simplify valid after coalescing")
 
   # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
@@ -343,24 +374,24 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   # late decomps + move gates from unrenderable INVALID where
   candidate_dtypes = {*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong}
   emulated_dtypes = set(EMULATED_DTYPES.tolist(dtypes)) | (candidate_dtypes - ren.supported_dtypes())
-  if any(u.dtype in emulated_dtypes for u in sink.toposort()):
+  if sink.dtype in emulated_dtypes or any(u.dtype in emulated_dtypes for u in sink.backward_slice):
     sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
   pm_decomp = pm_decomp+\
     get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\
     get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
   sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="late decompositions")
-  sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")
+  if has_invalid: sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")
 
   # final rules for the renderer (without sym)
-  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends+pm_no_index
+  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else empty_matcher
+  pm_final_rewrite = symbolic_simple+extra_matcher+pm_split_ends+pm_no_index
   sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
   # this was the linearizer
   sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
 
   # put unnumbered variable PARAMs in slots
-  num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
+  num_params = sum(x.op is Ops.PARAM and x.arg.slot != -1 for x in sink.backward_slice_with_self)
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
@@ -434,7 +465,7 @@ pm_to_program = PatternMatcher([
 
 @track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
 @Context(ALLOW_DEVICE_USAGE=0)
-def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
+def do_to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp:
   """
   Transform an AST into a compiled PROGRAM. May trigger BEAM search.
 
@@ -467,14 +498,14 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   if prg.src[0].arg.estimates is None and (estimated:=do_estimates(prg, prg.src[0], prg.src[1])) is not None: prg = estimated
   if len(prg.src) == 2:
     prg = do_assemble(renderer, prg, prg.src[1]) if isinstance(renderer, ISARenderer) else do_render(renderer, prg, prg.src[1])
-  if len(prg.src) == 3 and (compiled:=do_compile(renderer, prg, prg.src[2])) is not None: prg = compiled
+  if compile_binary and len(prg.src) == 3 and (compiled:=do_compile(renderer, prg, prg.src[2])) is not None: prg = compiled
   return prg
 
 to_program_cache: dict[tuple, UOp] = {}
-def to_program(ast:UOp, renderer:Renderer) -> UOp:
+def to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp:
   config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
   # UOps are structurally interned, so identity is already a collision-free structural
   # cache key within this process and avoids recursively hashing every kernel graph.
-  key = (ast, type(renderer), renderer.target, *[x.value for x in config])
-  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  key = (ast, type(renderer), renderer.target, compile_binary, *[x.value for x in config])
+  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer, compile_binary=compile_binary)
   return prg

tinygrad/device.py

@@ -284,6 +284,8 @@ class Compiler:
       lib = self.compile(src)
       if self.cachekey is not None: diskcache_put(self.cachekey, src, lib)
     return lib
+  def compile_cached_batch(self, srcs:list[tuple[str, str]]) -> list[tuple[str, str, bytes]]:
+    return [(name, src, self.compile_cached(src)) for name,src in srcs]
   def disassemble(self, lib:bytes): pass
 
 class Compiled:

tinygrad/engine/realize.py

@@ -4,10 +4,11 @@ import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo
-from tinygrad.device import Device, Buffer, MultiBuffer
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite, ProgramInfo, scoped_rewrite_cache
+from tinygrad.device import Device, Buffer, MultiBuffer, Compiler
 from tinygrad.renderer import Estimates
-from tinygrad.codegen import to_program
+from tinygrad.codegen import to_program, scoped_devectorize_cache
+from tinygrad.uop.spec import scoped_type_verify_cache
 from tinygrad.codegen.opt.postrange import bufs_from_ast
 
 # **************** Helpers ****************
@@ -240,11 +241,33 @@ pm_beam = PatternMatcher([
    lambda ctx,call,sink: call.replace(src=(sink.replace(arg=replace(sink.arg, beam=ctx)), *call.src[1:])) if sink.arg.beam == 0 else None),
 ])
 
+def compile_call(ctx, call:UOp, ast:UOp):
+  dev = call.device if isinstance(call.device, str) else call.device[0]
+  return call.replace(src=(to_program(ast, Device[dev].renderer, compile_binary=not ctx), *call.src[1:]))
+
 pm_compile = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), lambda call,ast:
-    call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))),
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), compile_call),
 ])
 
+def batch_compile(linear:UOp) -> UOp:
+  groups:dict[Compiler, list[UOp]] = {}
+  replacements = {}
+  for call in linear.toposort():
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or len(call.src[0].src) != 3: continue
+    dev = call.device if isinstance(call.device, str) else call.device[0]
+    compiler = Device[dev].compiler
+    cache = call.src[0].__dict__.setdefault('_compiled_programs', {})
+    if (cached:=cache.get(compiler)) is not None: replacements[call.src[0]] = cached
+    else: groups.setdefault(compiler, []).append(call.src[0])
+  for compiler,prgs in groups.items():
+    prgs = list(dict.fromkeys(prgs))
+    compiled = compiler.compile_cached_batch([(p.arg.function_name, p.src[2].arg) for p in prgs])
+    for prg,(name,src,lib) in zip(prgs, compiled):
+      replacements[prg] = compiled_prg = prg.replace(src=prg.src[:2]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)),
+                                                      arg=replace(prg.arg, name=name))
+      prg.__dict__['_compiled_programs'][compiler] = compiled_prg
+  return linear.substitute(replacements, walk=True, enter_calls=True) if replacements else linear
+
 pm_optimize_local_size = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size),
 ])
@@ -262,7 +285,9 @@ pm_exec = PatternMatcher([
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
-  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
+  with scoped_rewrite_cache(), scoped_devectorize_cache(), scoped_type_verify_cache():
+    linear = graph_rewrite(linear, pm_compile, ctx=True, name="render kernels", walk=True)
+  linear = batch_compile(linear)
   if getenv("HCQ2"):
     from extra.hcq2.hcq2 import hcq_compile
     linear = hcq_compile(linear, input_uops, jit=jit)

tinygrad/helpers.py

@@ -413,6 +413,13 @@ def diskcache_get(table:str, key:dict|str|int) -> Any:
   if (val:=res.fetchone()) is not None: return pickle.loads(val[0])
   return None
 
+def diskcache_get_batch(table:str, keys:list[str]) -> dict[str, Any]:
+  if CACHELEVEL < 1 or not keys: return {}
+  cur = db_connection().cursor()
+  try: rows = cur.execute(f"SELECT key, val FROM '{table}_{VERSION}' WHERE key IN ({','.join('?' for _ in keys)})", keys).fetchall()
+  except sqlite3.OperationalError: return {}
+  return {key:pickle.loads(val) for key,val in rows}
+
 _db_tables = set()
 def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False):
   if CACHELEVEL < 1: return val
@@ -430,6 +437,16 @@ def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False):
   cur.close()
   return val
 
+def diskcache_put_batch(table:str, items:list[tuple[str, Any]]) -> None:
+  if CACHELEVEL < 1 or not items: return
+  conn, cur = db_connection(), db_connection().cursor()
+  if table not in _db_tables:
+    cur.execute(f"CREATE TABLE IF NOT EXISTS '{table}_{VERSION}' (key text, val blob, PRIMARY KEY (key))")
+    _db_tables.add(table)
+  cur.executemany(f"REPLACE INTO '{table}_{VERSION}' (key, val) VALUES (?, ?)", [(key, pickle.dumps(val)) for key,val in items])
+  conn.commit()
+  cur.close()
+
 def diskcache(func:Callable[..., T]):
   def wrapper(*args, **kwargs) -> T:
     table, key = f"cache_{func.__name__}", hashlib.sha256(pickle.dumps((args, kwargs))).hexdigest()

tinygrad/renderer/__init__.py

@@ -75,6 +75,8 @@ class Renderer:
   compiler: Compiler = Compiler()
 
   def __init__(self, target:Target): self.target = target
+  @property
+  def rewrite_cache_key(self): return (type(self), self.target)
   def __reduce__(self): return self.__class__, (self.target,)
   def render(self, uops:list[UOp]) -> str: raise NotImplementedError("needs a renderer")
   def asm(self, prg:UOp, lin:UOp) -> bytes: raise NotImplementedError("needs an assembler")

tinygrad/runtime/ops_amd.py

@@ -13,7 +13,7 @@ from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.elf import elf_loader
+from tinygrad.runtime.support.elf import elf_loader, elf_symbols
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
@@ -559,24 +559,37 @@ class AMDCopyQueue(HWQueue):
 
     sdma_queue.signal_doorbell(dev)
 
+class AMDProgramModule:
+  __slots__ = ("lib_gpu", "image", "symbols", "rodata_entry", "__weakref__")
+  def __init__(self, dev:AMDDevice, lib_gpu:HCQBuffer, image:bytes, spec:BufferSpec, symbols:dict[str, int], rodata_entry:int):
+    self.lib_gpu, self.image, self.symbols, self.rodata_entry = lib_gpu, image, symbols, rodata_entry
+    weakref.finalize(self, HCQProgram._fini, dev, lib_gpu, spec)
+
 class AMDProgram(HCQProgram):
   def __init__(self, dev:AMDDevice, name:str, lib:bytes, **kwargs):
     # TODO; this API needs the type signature of the function and global_size/local_size
     self.dev, self.name, self.lib = dev, name, lib
 
-    image, sections, relocs = elf_loader(self.lib)
+    module_key = hashlib.sha256(self.lib).digest()
+    if (module:=self.dev.program_modules.get(module_key)) is None:
+      image, sections, relocs = elf_loader(self.lib)
+      symbols = elf_symbols(self.lib)
 
-    rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
+      for apply_image_offset, rel_sym_offset, typ, addent in relocs:
+        if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
+        else: raise RuntimeError(f"unknown AMD reloc {typ}")
+
+      lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True))
+      self.dev.allocator._copyin(lib_gpu, image)
+      self.dev.synchronize()
+      rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
+      self.dev.program_modules[module_key] = module = AMDProgramModule(self.dev, lib_gpu, bytes(image), buf_spec, symbols, rodata_entry)
+    self.module = module
+    self.lib_gpu, image, symbols, rodata_entry = module.lib_gpu, module.image, module.symbols, module.rodata_entry
+
+    rodata_entry = symbols.get(f"{name}.kd", rodata_entry)
     assert rodata_entry >= 0, ".rodata section not found"
 
-    for apply_image_offset, rel_sym_offset, typ, addent in relocs:
-      if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
-      else: raise RuntimeError(f"unknown AMD reloc {typ}")
-
-    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True))
-    self.dev.allocator._copyin(self.lib_gpu, image)
-    self.dev.synchronize()
-
     desc_sz = ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)
     desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+desc_sz]))
     self.group_segment_size = desc.group_segment_fixed_size
@@ -606,7 +619,6 @@ class AMDProgram(HCQProgram):
 
     super().__init__(CLikeArgsState, self.dev, self.name, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz, lib=self.lib,
                      base=self.lib_gpu.va_addr)
-    weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
                wait=False, timeout:int|None=None):
@@ -1004,6 +1016,7 @@ class AMDDevice(HCQCompiled):
                      functools.partial(AMDCopyQueue, self, max_copy_size=self.max_copy_size) if self.has_sdma_queue else None,
                      kernargs_size=(8 << 10) if self.is_usb() else (16 << 20), sigalloc_size=0x100 if self.is_usb() else 0x1000,
                      can_recover=self.is_am(), arch=self.arch)
+    self.program_modules:weakref.WeakValueDictionary[bytes, AMDProgramModule] = weakref.WeakValueDictionary()
 
     # Scratch setup
     self.max_private_segment_size = 0

tinygrad/runtime/support/compiler_amd.py

@@ -1,5 +1,5 @@
 import ctypes, hashlib, tempfile, subprocess, pathlib, shutil
-from tinygrad.helpers import system, getenv
+from tinygrad.helpers import system, getenv, diskcache_get, diskcache_get_batch, diskcache_put_batch
 from tinygrad.runtime.autogen import comgr
 try:
   comgr.amd_comgr_get_version(ctypes.byref(major:=ctypes.c_uint64()), ctypes.byref(minor:=ctypes.c_uint64()))
@@ -45,7 +45,7 @@ def set_options(action_info, options:bytes):
   return amd_comgr_action_info_set_option_list(action_info, to_char_p_p(options_list:=options.split(b' ')), len(options_list))
 
 # AMD_COMGR_SAVE_TEMPS=1 AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1
-def compile_hip(prg:str, arch="gfx1100", asm=False, use_device_libs=True) -> bytes:
+def compile_hip(prg:str, arch="gfx1100", asm=False, use_device_libs=True, backend_opt=3) -> bytes:
   check(comgr.amd_comgr_create_action_info(ctypes.byref(action_info := comgr.amd_comgr_action_info_t())))
   check(comgr.amd_comgr_action_info_set_language(action_info, comgr.AMD_COMGR_LANGUAGE_HIP))
   check(comgr.amd_comgr_action_info_set_isa_name(action_info, b"amdgcn-amd-amdhsa--" + arch.encode()))
@@ -80,7 +80,8 @@ def compile_hip(prg:str, arch="gfx1100", asm=False, use_device_libs=True) -> byt
     if status != 0:
       print(_get_comgr_data(data_set_bc, comgr.AMD_COMGR_DATA_KIND_LOG).decode())
       raise RuntimeError("compile failed")
-    check(set_options(action_info, b"-O3 -mllvm -amdgpu-internalize-symbols"))
+    check(set_options(action_info, f"-O{backend_opt} -mllvm -vectorize-loops=false "
+                                     "-mllvm -vectorize-slp=false -mllvm -unroll-threshold=0".encode()))
     check(comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_CODEGEN_BC_TO_RELOCATABLE, action_info, data_set_bc, data_set_reloc))
 
   check(set_options(action_info, b""))
@@ -99,6 +100,58 @@ class HIPCompiler(Compiler):
   def compile(self, src:str) -> bytes:
     try: return compile_hip(src, self.arch, src.split('\n', 1)[0].strip() == '.text', use_device_libs="__ocml_" in src)
     except RuntimeError as e: raise CompileError(e) from e
+  def compile_cached_batch(self, srcs:list[tuple[str, str]]) -> list[tuple[str, str, bytes]]:
+    batch_size = getenv("AMD_COMPILE_BATCH_SIZE", 256)
+    if batch_size <= 1 or any(src.split('\n', 1)[0].strip() == '.text' for _,src in srcs): return super().compile_cached_batch(srcs)
+
+    renamed = []
+    for name,src in srcs:
+      new_name = f"{name}_{hashlib.sha256(src.encode()).hexdigest()[:8]}"
+      renamed.append((new_name, src.replace(f" {name}(", f" {new_name}(", 1)))
+
+    ret:list[tuple[str, str, bytes]|None] = [None] * len(srcs)
+    missing:dict[tuple[int, bool], list[int]] = {}
+    cache_writes:list[tuple[str, tuple[str, str]]] = []
+    module_writes:list[tuple[str, bytes]] = []
+    cache_srcs = [f"batchO1\n{src}" for _,src in renamed]
+    cached_sources = diskcache_get_batch(self.cachekey, cache_srcs) if self.cachekey is not None else {}
+    refs = {v[1] for v in cached_sources.values() if isinstance(v, tuple) and v[0] == "batch"}
+    module_keys = {key:f"__batch__{key}" for key in refs}
+    same_table_modules = diskcache_get_batch(self.cachekey, list(module_keys.values())) if self.cachekey is not None else {}
+    cached_modules = {key:same_table_modules[module_keys[key]] for key in refs if module_keys[key] in same_table_modules}
+    if self.cachekey is not None and (old_refs:=refs-cached_modules.keys()):
+      cached_modules.update(diskcache_get_batch(f"{self.cachekey}_batch", list(old_refs)))
+    for i,(name,src) in enumerate(renamed):
+      opt = 1
+      cached = cached_sources.get(cache_srcs[i])
+      if isinstance(cached, tuple) and cached[0] == "batch": cached = cached_modules.get(cached[1])
+      if cached is not None: ret[i] = (name, src, cached)
+      else: missing.setdefault((opt, "__ocml_" in src), []).append(i)
+    for (opt,use_device_libs),indices in missing.items():
+      for start in range(0, len(indices), batch_size):
+        batch = indices[start:start+batch_size]
+        preamble = dict.fromkeys(line for i in batch for line in
+                                 renamed[i][1].split('extern "C" __attribute__((global))', 1)[0].splitlines())
+        kernels = []
+        for i in batch:
+          _, kernel = renamed[i][1].split('extern "C" __attribute__((global))', 1)
+          kernels.append('extern "C" __attribute__((global))'+kernel)
+        combined = '\n'.join((*preamble, *kernels))
+        module_key = hashlib.sha256(f"O{opt}\n{combined}".encode()).hexdigest()
+        lib = diskcache_get(self.cachekey, f"__batch__{module_key}") if self.cachekey is not None else None
+        if lib is None and self.cachekey is not None: lib = diskcache_get(f"{self.cachekey}_batch", module_key)
+        if lib is None:
+          try: lib = compile_hip(combined, self.arch, use_device_libs=use_device_libs, backend_opt=opt)
+          except RuntimeError as e: raise CompileError(e) from e
+          if self.cachekey is not None: module_writes.append((f"__batch__{module_key}", lib))
+        for i in batch:
+          name,src = renamed[i]
+          if self.cachekey is not None: cache_writes.append((f"batchO{opt}\n{src}", ("batch", module_key)))
+          ret[i] = (name, src, lib)
+    if self.cachekey is not None:
+      diskcache_put_batch(self.cachekey, module_writes+cache_writes)
+    assert all(x is not None for x in ret)
+    return [x for x in ret if x is not None]
   def disassemble(self, lib:bytes): amdgpu_disassemble(lib)
 
 class HIPCCCompiler(Compiler):

tinygrad/runtime/support/elf.py

@@ -6,6 +6,22 @@ from tinygrad.runtime.autogen import libc
 @dataclass(frozen=True)
 class ElfSection: name:str; header:libc.Elf64_Shdr|libc.Elf32_Shdr; content:bytes # noqa: E702
 
+def elf_symbols(blob:bytes) -> dict[str, int]:
+  ecls = {libc.ELFCLASS32: "Elf32", libc.ELFCLASS64: "Elf64"}[blob[libc.EI_CLASS]]
+  header = getattr(libc, f"{ecls}_Ehdr").from_buffer_copy(blob)
+  section_headers = (getattr(libc, f"{ecls}_Shdr") * header.e_shnum).from_buffer_copy(blob[header.e_shoff:])
+  sym_t = getattr(libc, f"{ecls}_Sym")
+  symbols = {}
+  for sh in section_headers:
+    if sh.sh_type not in (libc.SHT_SYMTAB, libc.SHT_DYNSYM): continue
+    str_sh = section_headers[sh.sh_link]
+    strtab = blob[str_sh.sh_offset:str_sh.sh_offset+str_sh.sh_size]
+    for sym in (sym_t * (sh.sh_size // sh.sh_entsize)).from_buffer_copy(blob[sh.sh_offset:]):
+      symbols[strtab[sym.st_name:strtab.find(b'\0', sym.st_name)].decode()] = sym.st_value
+  return symbols
+
+def elf_symbol_address(blob:bytes, name:str) -> int|None: return elf_symbols(blob).get(name)
+
 def link_sym(sym:str, libs:list[str]) -> int:
   for lib in libs:
     try: return unwrap(ctypes.cast(getattr(ctypes.CDLL(ctypes.util.find_library(lib)), sym), ctypes.c_void_p).value)

tinygrad/schedule/indexing.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
 from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
-from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC
+from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC, prod
 
 ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.SLICE,
                       Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
@@ -167,6 +167,18 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
       rngs = tuple(r if (sz == sh and off == 0) else (r-off).valid(graph_rewrite((r >= off) & (r < (sh+off)),
         symbolic+pm_simplify_valid, name="pad")) for r,sh,(off,sz) in zip(rngs, in_shape, arg))
     case Ops.RESHAPE:
+      if all(isinstance(s, int) for s in in_shape+arg) and tuple(s for s in in_shape if s != 1) == tuple(s for s in arg if s != 1):
+        out_rngs = iter(r for s,r in zip(arg, rngs) if s != 1)
+        return tuple(UOp.const(dtypes.index, 0) if s == 1 else next(out_rngs) for s in in_shape)
+      if all(isinstance(s, int) and s > 0 for s in in_shape+arg) and all(r.op is Ops.CONST and isinstance(r.arg, int) for r in rngs):
+        flat_idx = sum(int(r.arg)*prod(arg[i+1:]) for i,r in enumerate(rngs))
+        ret = []
+        for s in cast(tuple[int, ...], in_shape)[::-1]:
+          ret.append(UOp.const(dtypes.index, flat_idx % s))
+          flat_idx //= s
+        return tuple(ret[::-1])
+      if len(in_shape) == 1:
+        return (UOp.const(dtypes.index, 0).usum([r*prod(arg[i+1:]) for i,r in enumerate(rngs)]),)
       sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early
       sub_array = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER, dtype=r.dtype) for i,r in enumerate(sink.ranges)}
       rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src

tinygrad/uop/ops.py

@@ -1589,15 +1589,21 @@ if TRACK_MATCH_STATS or PROFILE:
 # A pure Python sentinel, but *typed* as UOp so it fits all the dict annotations
 SENTINEL: Final[UOp] = cast(UOp, object())
 class BottomUpGate(Exception): pass
+rewrite_caches: list[dict[Any, dict[UOp, UOp]]] = []
+class scoped_rewrite_cache:
+  def __enter__(self): rewrite_caches.append({})
+  def __exit__(self, *args): rewrite_caches.pop()
+
 class RewriteContext:
-  __slots__ = ("pm", "bpm", "bpm_cache", "ctx", "replace", "enter_calls")
-  def __init__(self, pm, bpm, ctx=None, enter_calls=False):
+  __slots__ = ("pm", "bpm", "bpm_cache", "ctx", "replace", "enter_calls", "shared_cache")
+  def __init__(self, pm, bpm, ctx=None, enter_calls=False, shared_cache=None):
     self.pm: PatternMatcher|None = pm
     self.bpm: PatternMatcher|None = bpm
     self.bpm_cache: dict[UOp, UOp|None] = {}
     self.ctx = ctx
     self.replace: dict[UOp, UOp] = {}
     self.enter_calls = enter_calls
+    self.shared_cache: dict[UOp, UOp]|None = shared_cache
 
   # no cache needed: pm_rewrite is called at most once per UOp due to the replace dict check in unified_rewrite
   def pm_rewrite(self, x:UOp) -> UOp|None: return unwrap(self.pm).rewrite(x, self.ctx)
@@ -1636,7 +1642,7 @@ class RewriteContext:
     stack: list[tuple[UOp, int, UOp]] = [(root, 0, root)]
     on_stack = {root}  # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
     waitlist: dict[UOp, list[tuple[UOp, int, UOp]]] = {}  # UOps waiting on a dependency to be in self.replace
-    replace, ctx, stack_limit, enter_calls = self.replace, self.ctx, REWRITE_STACK_LIMIT.value, self.enter_calls
+    replace, ctx, stack_limit, enter_calls, shared_cache = self.replace, self.ctx, REWRITE_STACK_LIMIT.value, self.enter_calls, self.shared_cache
     bpm, cached_bpm_rewrite = self.bpm, self.cached_bpm_rewrite
     pm_rewrite = self.pm.rewrite if self.pm is not None else None
     while stack:
@@ -1644,6 +1650,10 @@ class RewriteContext:
       n, stage, new_n = stack.pop()
       if n in replace: continue  # skip any nodes we have seen
       if stage == 0:
+        if shared_cache is not None and (cached:=shared_cache.get(n, SENTINEL)) is not SENTINEL:
+          replace[n] = cached
+          if n in waitlist: stack.extend(waitlist.pop(n))
+          continue
         # if bottom up, we rewrite this node early. in both cases, we add its srcs to the stack
         if bpm is not None:
           # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match
@@ -1682,6 +1692,7 @@ class RewriteContext:
             # if top down, do the rewrite. if no rewrite or bottom up, we are done rewriting this node so we add it to the dict
             if pm_rewrite is None or (new_src_n:=pm_rewrite(new_n, ctx)) is None:
               replace[n] = new_n
+              if shared_cache is not None: shared_cache[n] = new_n
               if n in waitlist: stack.extend(waitlist.pop(n))
               continue
           else:
@@ -1698,12 +1709,17 @@ class RewriteContext:
         else:
           # otherwise we are done
           replace[n] = replaced_new_n
+          if shared_cache is not None: shared_cache[n] = replaced_new_n
           if n in waitlist: stack.extend(waitlist.pop(n))
     return replace[root]
 
 @profile_matches
 def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp:
-  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
+  cache_ctx = None if ctx is None else getattr(ctx, 'rewrite_cache_key', SENTINEL)
+  cache_key = pm if ctx is None else (pm, cache_ctx)
+  shared_cache = rewrite_caches[-1].setdefault(cache_key, {}) if rewrite_caches and cache_ctx is not SENTINEL and \
+    not bottom_up and bpm is None and not walk else None
+  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls, shared_cache)
   return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
 
 def sint_to_uop(x:sint, dtype=dtypes.index) -> UOp: return UOp.const(dtype, x) if isinstance(x, int) else x.cast(dtype)

tinygrad/uop/spec.py

@@ -33,15 +33,23 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
   return validate_index_with_z3(sz, idx, gate)
 
 def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
-  lst = list(ast.toposort()) if isinstance(ast, UOp) else ast
+  lst = [*ast.backward_slice, ast] if isinstance(ast, UOp) else ast
   if SPEC > 1: test_pyrender(lst[-1])  # assume this is the sink
 
+  verified = type_verify_caches[-1].setdefault(check_spec, set()) if type_verify_caches else None
   with Context(TRACK_MATCH_STATS=0):
     for i,u in enumerate(lst):
+      if verified is not None and u in verified: continue
       ret: bool|None = check_spec.rewrite(u)
       if ret is not True:
         if DEBUG >= 3: print_uops(lst)
         raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}")
+      if verified is not None: verified.add(u)
+
+type_verify_caches: list[dict[PatternMatcher, set[UOp]]] = []
+class scoped_type_verify_cache:
+  def __enter__(self): type_verify_caches.append({})
+  def __exit__(self, *args): type_verify_caches.pop()
 
 # ***** new specs *****
 

tinygrad/uop/symbolic.py

@@ -51,6 +51,7 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None:
   #   q == b//div     -> b*mul              (full recombine)
   #   q == (b//div)%d -> (b%(div*d))*mul    (partial recombine into a wider mod, needs d>0)
   terms = list(x.split_uop(Ops.ADD))
+  if not any(u.op is Ops.FLOORMOD or (u.op is Ops.MUL and any(s.op is Ops.FLOORMOD for s in u.src)) for u in terms): return None
   for i,u in enumerate(terms):
     mod, mul = u.pop_const(Ops.MUL)
     if mod.op is not Ops.FLOORMOD or mod.src[1].op is not Ops.CONST: continue
@@ -315,6 +316,10 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=False) -> UOp:
     expr, is_upper, c = res
     bounds[expr][int(is_upper)] = c
   if not bounds: return uop
+  if not try_simplex:
+    uop_nodes = uop.backward_slice_with_self
+    bounds = defaultdict(lambda: [None, None], ((expr, v) for expr,v in bounds.items() if expr in uop_nodes))
+    if not bounds: return uop
 
   # simplify uop given that valid is True
   all_candidates = []
@@ -347,7 +352,7 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=False) -> UOp:
 
 def _valid_priority(v: UOp, valids:list[UOp]) -> int:
   # we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified
-  return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
+  return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.backward_slice_with_self else 0 for other in valids)
 
 def simplify_valid(valid:UOp) -> UOp|None:
   if valid.op_in_backward_slice_with_self(Ops.INDEX): return None  # this should only be for indexing, skip if there's a INDEX
@@ -402,7 +407,7 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None:
   if x.dtype is not dtypes.index: return None
   # Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot
   if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD): return None
-  return cond.where(uop_given_valid(cond, x, try_simplex=False), i)
+  return cond.where(new_x, i) if (new_x:=uop_given_valid(cond, x, try_simplex=False)) is not x else None
 
 # TODO: this is O(number of WHERE * number of node)
 # def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None: