remove unmeasured CIFAR micro-optimizations

# Conflicts:
#	tinygrad/uop/ops.py

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
@@ -49,14 +49,13 @@ class UnsyncedBatchNorm:
       # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
       # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
       batch_mean = x.mean(axis=(1,3,4))
-      y = (x - batch_mean.detach().reshape(shape=[batch_mean.shape[0], 1, -1, 1, 1]))  # d(var)/d(mean) = 0
-      batch_var = (y*y).mean(axis=(1,3,4))
+      batch_var = (x*x).mean(axis=(1,3,4)) - batch_mean*batch_mean
       batch_invstd = batch_var.add(self.eps).pow(-0.5)
 
       # NOTE: wow, this is done all throughout training in most PyTorch models
       if self.track_running_stats:
         self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
-        batch_var_adjust = prod(y.shape[1:])/(prod(y.shape[1:])-y.shape[2])
+        batch_var_adjust = prod(x.shape[1:])/(prod(x.shape[1:])-x.shape[2])
         self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * batch_var_adjust * batch_var.detach().cast(self.running_var.dtype))
         self.num_batches_tracked += 1
     else:
@@ -70,25 +69,37 @@ class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
     super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
     self.weight.is_param_(False)
 
+class MatmulConv2d(nn.Conv2d):
+  def __call__(self, x:Tensor) -> Tensor:
+    if not getenv("MATMUL_CONV", 1): return super().__call__(x)
+    assert self.groups == 1 and self.stride == self.dilation == self.padding == 1 and self.bias is None
+    bs, cin, _, _ = x.shape
+    cout, _, ky, kx = self.weight.shape
+    patches = x.pad((1, 1, 1, 1))._pool((ky, kx), 1, 1)
+    oy, ox = patches.shape[2:4]
+    patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(bs*oy*ox, cin*ky*kx).contiguous().contiguous_backward()
+    out = (patches @ self.weight.reshape(cout, cin*ky*kx).T).contiguous().contiguous_backward()
+    return out.reshape(bs, oy, ox, cout).permute(0, 3, 1, 2)
+
 class ConvGroup:
   def __init__(self, channels_in, channels_out):
-    self.conv1 = nn.Conv2d(channels_in,  channels_out, kernel_size=3, padding=1, bias=False)
-    self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
+    self.conv1 = MatmulConv2d(channels_in,  channels_out, kernel_size=3, padding=1, bias=False)
+    self.conv2 = MatmulConv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
 
     self.norm1 = BatchNorm(channels_out)
     self.norm2 = BatchNorm(channels_out)
 
   def __call__(self, x):
-    x = self.conv1(x)
+    x = self.conv1(x).contiguous()
     x = x.max_pool2d(2)
     x = x.float()
-    x = self.norm1(x)
+    x = self.norm1(x).contiguous()
     x = x.cast(dtypes.default_float)
-    x = x.quick_gelu()
+    x = x.quick_gelu().contiguous()
     residual = x
-    x = self.conv2(x)
+    x = self.conv2(x).contiguous()
     x = x.float()
-    x = self.norm2(x)
+    x = self.norm2(x).contiguous()
     x = x.cast(dtypes.default_float)
     x = x.quick_gelu()
 
@@ -111,7 +122,10 @@ class SpeedyResNet:
   def __call__(self, x, training=True):
     # pad to 32x32 because whitening conv creates 31x31 images that are awfully slow to compute with
     # TODO: remove the pad but instead let the kernel optimize itself
-    forward = lambda x: x.conv2d(self.whitening).pad((1,0,0,1)).sequential(self.net)
+    def forward(x):
+      x = x.conv2d(self.whitening).pad((1,0,0,1)).contiguous()
+      for layer in self.net: x = layer(x).contiguous()
+      return x
     return forward(x) if training else (forward(x) + forward(x[..., ::-1])) / 2.
 
 # hyper-parameters were exactly the same as the original repo
@@ -216,15 +230,31 @@ def train_cifar():
     Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
     return X_cutmix, Y_cutmix
 
-  @TinyJit
-  def augmentations(X:Tensor, Y:Tensor):
-    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensive to generate
+  def random_permutation(rows:int, cols:int) -> Tensor:
+    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)
+    X, Y = X[:perms.shape[0]], Y[:perms.shape[0]]
     if getenv("RANDOM_CROP", 1):
       X = random_crop(X, crop_size=32)
     if getenv("RANDOM_FLIP", 1):
       # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
       X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
     X, Y = X[perms], Y[perms]
+    return X, Y, perms
+
+  @TinyJit
+  def augmentations(X:Tensor, Y:Tensor):
+    X, Y, _ = shuffled_augmentations(X, Y)
+    return X, Y
+
+  @TinyJit
+  def augmentations_cutmix(X:Tensor, Y:Tensor):
+    X, Y, perms = shuffled_augmentations(X, Y)
     return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
 
   # the operations that remain inside batch fetcher is the ones that involves random operations
@@ -234,8 +264,10 @@ def train_cifar():
       st = time.monotonic()
       X, Y = X_in, Y_in
       if is_train:
-        X, Y, X_cm, Y_cm = augmentations(X, Y)
-        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
+        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']:
+          _, _, X, Y = augmentations_cutmix(X, Y)
+        else:
+          X, Y = augmentations(X, Y)
       et = time.monotonic()
       print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
 
@@ -255,7 +287,6 @@ def train_cifar():
 
   class modelEMA():
     def __init__(self, w, net):
-      # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
       self.net_ema = SpeedyResNet(w)
       for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
         net_ema_param.assign(net_param.numpy())
@@ -280,6 +311,15 @@ def train_cifar():
 
   # initialize model weights
   model = SpeedyResNet(W)
+  model_state = get_state_dict(model)
+  random_params = [x for name, x in model_state.items() if x.is_param and "bias" not in name]
+  Tensor.manual_seed(getenv('SEED', hyp['seed']))
+  random_values = Tensor.rand(sum(x.numel() for x in random_params))
+  offset = 0
+  for param in random_params:
+    bound = prod(param.shape[1:]) ** -0.5
+    param.replace(((random_values[offset:offset+param.numel()] * (2 * bound)) - bound).reshape(param.shape).cast(param.dtype))
+    offset += param.numel()
 
   # padding is not timed in the original repo since it can be done all at once
   X_train = pad_reflect(X_train, size=hyp['net']['pad_amount'])
@@ -296,7 +336,7 @@ def train_cifar():
         x.to_(GPUS)
 
   # parse the training params into bias and non-bias
-  params_dict = get_state_dict(model)
+  params_dict = model_state
   params_bias = []
   params_non_bias = []
   for params in params_dict:
@@ -320,7 +360,7 @@ def train_cifar():
   lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
 
   def train_step(model, optimizer, lr_scheduler, X, Y):
-    out = model(X)
+    out = model(X).contiguous()
     loss_batchsize_scaler = 512/BS
     loss = cross_entropy(out, Y, reduction='none', label_smoothing=hyp['opt']['label_smoothing']).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
 
@@ -331,15 +371,19 @@ def train_cifar():
       return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
     return loss.realize()
 
-  train_step_jitted = TinyJit(train_step)
+  train_step_jitted = TinyJit(train_step, warmup=False)
 
-  def eval_step(model, X, Y):
-    out = model(X, training=False)
-    loss = cross_entropy(out, Y, reduction='mean')
+  def eval_forward(model, X):
+    return model(X).realize()
+
+  def eval_step(out, out_flipped, Y):
+    out = (out + out_flipped) / 2.
     correct = out.argmax(axis=1) == Y.argmax(axis=1)
-    return correct.realize(), loss.realize()
-  eval_step_jitted     = TinyJit(eval_step)
-  eval_step_ema_jitted = TinyJit(eval_step)
+    return correct.sum().realize()
+  eval_forward_jitted = TinyJit(eval_forward, warmup=False)
+  eval_forward_ema_jitted = TinyJit(eval_forward, warmup=False)
+  eval_step_jitted = TinyJit(eval_step, warmup=False)
+  eval_step_ema_jitted = TinyJit(eval_step, warmup=False)
 
   # 97 steps in 2 seconds = 20ms / step
   # step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
@@ -360,31 +404,37 @@ 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 = correct_sum_ema = None
+        correct_len = correct_len_ema = 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)
             Yt.shard_(GPUS, axis=0)
 
-          correct, loss = eval_step_jitted(model, Xt, Yt)
-          losses.append(loss.numpy().tolist())
-          corrects.extend(correct.numpy().tolist())
+          Xt_contiguous = Xt.contiguous().realize()
+          out = eval_forward_jitted(model, Xt_contiguous).clone().realize()
+          out_flipped = eval_forward_jitted(model, Xt_contiguous[..., ::-1].contiguous().realize())
+          batch_correct = eval_step_jitted(out, out_flipped, Yt)
+          correct_sum = batch_correct if correct_sum is None else correct_sum + batch_correct
+          correct_len += Yt.shape[0]
           if model_ema:
-            correct_ema, loss_ema = eval_step_ema_jitted(model_ema.net_ema, Xt, Yt)
-            losses_ema.append(loss_ema.numpy().tolist())
-            corrects_ema.extend(correct_ema.numpy().tolist())
+            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())
+            batch_correct_ema = eval_step_ema_jitted(out_ema, out_flipped_ema, Yt)
+            correct_sum_ema = batch_correct_ema if correct_sum_ema is None else correct_sum_ema + batch_correct_ema
+            correct_len_ema += Yt.shape[0]
 
         # 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
+        correct_count = correct_sum.item()
+        if model_ema:
+          assert correct_sum_ema is not None
+          correct_count_ema = correct_sum_ema.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}% 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}% STEP={i}")
 
       if STEPS == 0 or i == STEPS: break
 

test/null/test_graph_rewrite.py

@@ -348,6 +348,19 @@ class TestStopEarly(unittest.TestCase):
     ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
     assert ret == cn+d
 
+class TestFastSubstitute(unittest.TestCase):
+  def test_replacement_tree_is_substituted(self):
+    a, b, c, d = [UOp.variable(x, 0, 10) for x in "abcd"]
+    self.assertIs((a+4).substitute({a:b+c, b:d}), (d+c)+4)
+
+  def test_rebuilt_node_is_substituted(self):
+    a, b, c, d = [UOp.variable(x, 0, 10) for x in "abcd"]
+    self.assertIs(((a+b)*2).substitute({a:c, c+b:d}), d*2)
+
+  def test_mapping_cycle(self):
+    a, b = [UOp.variable(x, 0, 10) for x in "ab"]
+    with self.assertRaises(RuntimeError): (a+1).substitute({a:b, b:a})
+
 class TestWalkRewrite(unittest.TestCase):
   """Tests for graph_rewrite with walk=True (MLIR Walk Pattern Rewrite Driver semantics).
   walk=True gives a single-pass traversal that does NOT revisit or re-traverse into rewritten subtrees.

tinygrad/codegen/__init__.py

@@ -1,18 +1,19 @@
 from dataclasses import replace, dataclass
 import itertools, functools
-from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
+from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, PROFILE, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
+from tinygrad.uop.ops import TRACK_MATCH_STATS
 from tinygrad.uop.ops import AxisType
 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 +23,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 +31,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
@@ -113,6 +117,10 @@ def broadcast_and_devec_wmma(b:UOp):
     src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
   return UOp.stack(*src).reshape(b.shape)
 
+@functools.cache
+def shape_indexes(shape:tuple[int, ...]) -> tuple[tuple[UOp, ...], ...]:
+  return tuple(tuple(UOp.const(dtypes.index, i) for i in idx) for idx in itertools.product(*map(range, shape)))
+
 pm_wmma_add = PatternMatcher([
   (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
    lambda add, wmma: UOp(wmma.op, src=(wmma.src[0], wmma.src[1], wmma.src[2]+add), arg=wmma.arg)),
@@ -128,15 +136,48 @@ unbroadcast = pm_wmma_add+PatternMatcher([
   (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
 ])
 
-def do_devectorize(b:UOp):
-  if b.shape == (): return None
+def do_devectorize(ctx, b:UOp):
+  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) 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)
+
+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))
+  elif x.op is Ops.INDEX:
+    ret = index_lane(ctx, x.src[0], x.src[1:]+idxs)
+  elif x.op in GroupOp.Elementwise:
+    ret = UOp(x.op, x.dtype, tuple(index_lane(ctx, s, idxs) if s._shape else s for s in x.src), x.arg, x.tag)
+  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:]
+  return UOp(x.op, x.dtype, tuple(UOp(Ops.INDEX, s.dtype, (s,)+indexes) if s._shape else s for s in x.src), x.arg, x.tag)
+
+def index_elementwise_lane(ctx:DevectorizeContext, x:UOp, idx:UOp):
+  indexes = idx.src[1:]
+  return UOp(x.op, x.dtype, tuple(index_lane(ctx, s, indexes) if s._shape else s for s in x.src), x.arg, x.tag)
 
 def do_stack_wmma(u:UOp):
   if all(x.op in (Ops.STACK, Ops.WMMA) for x in u.src): return None
@@ -152,11 +193,13 @@ def do_stack_wmma(u:UOp):
 ew_devectorizer = PatternMatcher([
   # unpack broadcasting
   (UPat(GroupOp.Elementwise, name="b"), do_devectorize),
+  (UPat(GroupOp.Elementwise, name="x").f(Ops.INDEX, allow_any_len=True, name="idx"), index_elementwise),
 ])
 
 devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # unpack broadcasting
   (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize),
+  (UPat(GroupOp.Elementwise, name="x").f(Ops.INDEX, allow_any_len=True, name="idx"), index_elementwise_lane),
   # INDEX without src is nothing (TODO: this should be in mop_cleanup)
   (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
   # unpack WMMA
@@ -298,53 +341,54 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
 
   # add gpu dims (late). this works after devectorize, but it's faster here
-  sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
+  if VIZ or PROFILE or TRACK_MATCH_STATS: sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
+  elif (gpu_sink:=pm_add_gpudims.rewrite(sink, ren)) is not None: sink = gpu_sink
 
   # **** optimizations are done, now we lower to actual code ****
 
   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")
-
-  # simplify indexing
-  sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=DevectorizeContext(ren, {}), 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)
 
-  # extra symbolic before decomp. crashes without this?
-  sink = graph_rewrite(sink, sym, name="extra symbolic")
+  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")
+    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)
 
   # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
   sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, name="lower all index dtypes")
 
-  # final symbolic before decomp
-  sink = graph_rewrite(sink, symbolic, name="final symbolic")
-
   # **** decomps ****
 
-  # floordiv+mod / dtype decomp (early)
+  # final symbolic + floordiv/mod + dtype decomp
   supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
-  sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
+  pm_decomp = symbolic+get_simplifying_rewrite_patterns(supported_ops)
 
   # late decomps + move gates from unrenderable INVALID where
-  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
+  candidate_dtypes = {*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong}
+  emulated_dtypes = set(EMULATED_DTYPES.tolist(dtypes)) | (candidate_dtypes - ren.supported_dtypes())
+  needs_dtype_decomp = sink.dtype in emulated_dtypes or any(u.dtype in emulated_dtypes for u in sink.backward_slice)
+  if needs_dtype_decomp:
+    sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
+    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
@@ -425,7 +469,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.
 
@@ -436,25 +480,50 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   Returns:
     The Ops.PROGRAM with SINK/LINEAR/SOURCE/BINARY.
   """
+  from tinygrad.codegen.opt.gemm import cooperative_gemm_program, direct_conv_bwd_activation_program
+  from tinygrad.codegen.opt.reduce import activation_var_grad_program, bn_grad_512_program, channel_reduce_program, col2im_program
+  from tinygrad.codegen.opt.reduce import im2col_program, moments_512_program
+  from tinygrad.codegen.opt.reduce import maxpool_backward_program, maxpool_program
+  if ast.op is Ops.SINK and (prg:=direct_conv_bwd_activation_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=cooperative_gemm_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=activation_var_grad_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=moments_512_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=bn_grad_512_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=channel_reduce_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=col2im_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=im2col_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=maxpool_backward_program(ast, renderer, compile_binary)) is not None: return prg
+  if ast.op is Ops.SINK and (prg:=maxpool_program(ast, renderer, compile_binary)) is not None: return prg
   if ast.op is Ops.PROGRAM: prg = ast
   elif ast.op is Ops.SINK:
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
-    prog_info = ProgramInfo.from_sink(full_sink)
     # instruction selection
     if isinstance(renderer, ISARenderer):
       full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True)
       full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True)
-    prg = UOp(Ops.PROGRAM, src=(full_sink,), arg=prog_info)
+    prg = UOp(Ops.PROGRAM, src=(full_sink,))
   else: raise RuntimeError(f"can't call to_program on {ast.op}")
+  if VIZ:
+    if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
+    prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
+    graph_rewrite(prg, PatternMatcher([]), name="View Program")
+    return prg
+  # PROGRAM lowering is a linear root-only pipeline. Driving it through graph_rewrite
+  # needlessly walks the full SINK and LINEAR graphs between each stage.
+  if len(prg.src) == 1: prg = do_linearize(renderer, prg, prg.src[0])
   if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
-  prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
-  if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
+  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 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)
-  key = (ast.key, 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)
+  # 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, 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/codegen/decomp/op.py

@@ -43,19 +43,22 @@ def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
 
 # ***** threefry *****
 
-def threefry2x32(x: UOp, key: UOp):
+def threefry2x32(x: UOp, key: UOp, ctx:Renderer|None=None):
+  native_shifts = ctx is not None and Ops.SHL in ctx.code_for_op and Ops.SHR in ctx.code_for_op
+  shl = (lambda v, s: v << s) if native_shifts else (lambda v, s: v * 2**s)
+  shr = (lambda v, s: v >> s) if native_shifts else (lambda v, s: v // 2**s)
   # split x and key from uint64 to two uint32
-  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
-  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
+  x0, x1 = x.cast(dtypes.uint32), shr(x, 32).cast(dtypes.uint32)
+  key0, key1 = key.cast(dtypes.uint32), shr(key, 32).cast(dtypes.uint32)
 
   rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
   ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
   xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
   for i in range(5):
-    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
+    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ (shl(xr[1], r) + shr(xr[1], 32-r))
     xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
 
-  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
+  return shl(xr[1].cast(dtypes.uint64), 32) | xr[0].cast(dtypes.uint64)
 
 # ***** decomposition patterns *****
 

tinygrad/codegen/late/coalese.py

@@ -42,7 +42,7 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N
   if not is_image_shape(buf._shape): return None
   if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int)
   start_idx = idx_x.stack(idx_y)
-  idx = uop_given_valid(valid, start_idx)
+  idx = uop_given_valid(valid, start_idx, try_simplex=True)
   drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1])
 
   if not drop_stmt and idx is start_idx: return None
@@ -74,7 +74,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
   # search for dims that drop the most valid statements
   best_drop, cands = -1, []
   for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
-    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)))
+    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)), try_simplex=True)
     dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
     if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
     elif dropped == best_drop: cands.append((ch, cw, cidx))

tinygrad/codegen/opt/gemm.py

@@ -0,0 +1,461 @@
+from dataclasses import dataclass, replace
+from math import prod
+from tinygrad.dtype import dtypes
+from tinygrad.renderer import Estimates, Renderer
+from tinygrad.uop.ops import AxisType, Ops, ProgramInfo, UOp, ssimplify
+
+WMMA_M = WMMA_N = WMMA_K = 16
+BLOCK_M = BLOCK_N = 128
+BLOCK_K = 32
+THREADS = 128
+
+@dataclass(frozen=True)
+class GemmMatch:
+  c:UOp
+  a:UOp
+  b:UOp
+  m:int
+  n:int
+  k:int
+  old:UOp|None = None
+  scale:float = 0.0
+  a_kxm:bool = False
+  b_kxn:bool = False
+
+@dataclass(frozen=True)
+class BatchedGemmMatch:
+  c:UOp
+  a:UOp
+  b:UOp
+  m:int
+  n:int
+  k:int
+  batch:int
+
+@dataclass(frozen=True)
+class DirectConvBwdActivationMatch:
+  m:int
+  cin:int
+  cout:int
+  spatial:int
+  residual:bool = False
+
+@dataclass(frozen=True)
+class GemmOutputNCHW:
+  spatial:int
+
+def _match_gemm(ast:UOp, device:str, arch:str) -> GemmMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) < 3 or end.src[0].op is not Ops.STORE: return None
+  store, ranges = end.src[0], end.src[1:]
+  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.LOOP for x in ranges) or store.src[0].op is not Ops.INDEX: return None
+  value = store.src[1]
+  old, old_index, scale = None, None, 0.0
+  if value.op is Ops.CAST and value.dtype == dtypes.half and value.src[0].op is Ops.REDUCE: reduce = value.src[0]
+  elif value.op is Ops.CAST and value.dtype == dtypes.float and value.src[0].op is Ops.ADD:
+    add_lhs, add_rhs = value.src[0].src
+    if add_lhs.op is not Ops.CAST or add_lhs.dtype != dtypes.half or add_lhs.src[0].op is not Ops.REDUCE:
+      add_lhs, add_rhs = add_rhs, add_lhs
+    if add_lhs.op is not Ops.CAST or add_lhs.dtype != dtypes.half or add_lhs.src[0].op is not Ops.REDUCE: return None
+    if add_rhs.op is not Ops.MUL: return None
+    old_idx = next((x for x in add_rhs.src if x.op is Ops.INDEX and x.dtype == dtypes.half), None)
+    scale_uop = next((x for x in add_rhs.src if x.op is Ops.CONST and x.dtype == dtypes.half), None)
+    if old_idx is None or scale_uop is None: return None
+    reduce, old, old_index, scale = add_lhs.src[0], old_idx.src[0], old_idx.src[1], float(scale_uop.arg)
+  else: return None
+  if reduce.arg != (Ops.ADD, 0) or len(reduce.src) != 2 or reduce.src[1].op is not Ops.RANGE: return None
+  k, product = reduce.src[1], reduce.src[0]
+  if product.op is not Ops.CAST or product.dtype != dtypes.float or product.src[0].op is not Ops.MUL: return None
+  lhs, rhs = product.src[0].src
+  if lhs.op is not Ops.INDEX or rhs.op is not Ops.INDEX: return None
+  if any(x.dtype != dtypes.half for x in (lhs.src[0], rhs.src[0])): return None
+  if store.src[0].src[0].dtype not in (dtypes.half, dtypes.float): return None
+
+  out_idx = ssimplify(store.src[0].src[1].get_idx())
+  if old_index is not None and ssimplify(old_index.get_idx()) is not out_idx: return None
+  k_size, total_size = int(k.vmax)+1, prod(int(x.vmax)+1 for x in ranges)
+  lhs_idx, rhs_idx = ssimplify(lhs.src[1].get_idx()), ssimplify(rhs.src[1].get_idx())
+  for m in ranges:
+    m_size, n_size = int(m.vmax)+1, total_size//(int(m.vmax)+1)
+    n = ssimplify(out_idx%n_size)
+    if ssimplify(out_idx//n_size) is not m or int(n.vmin) != 0 or int(n.vmax) != n_size-1: continue
+    n64 = old is None and n_size == 64
+    a_mxk_idx, a_kxm_idx = ssimplify(m*k_size+k), ssimplify(k*m_size+m)
+    b_nxk_idx, b_kxn_idx = ssimplify(n*k_size+k), ssimplify(k*n_size+n)
+    for a_idx, a_buf, b_idx, b_buf in ((lhs_idx, lhs.src[0], rhs_idx, rhs.src[0]), (rhs_idx, rhs.src[0], lhs_idx, lhs.src[0])):
+      if not (a_idx is a_mxk_idx or a_idx is a_kxm_idx): continue
+      if not (b_idx is b_nxk_idx or b_idx is b_kxn_idx): continue
+      if n64 and (a_idx is not a_mxk_idx or b_idx is not b_nxk_idx): continue
+      g = GemmMatch(store.src[0].src[0], a_buf, b_buf, m_size, n_size, k_size, old, scale,
+                    a_idx is a_kxm_idx, b_idx is b_kxn_idx)
+      bm, bn = _gemm_block_m(g), _gemm_block_n(g)
+      if m_size % bm or n_size % bn or k_size % BLOCK_K: continue
+      if (m_size//bm)*(n_size//bn) < (32 if old is not None else 512): continue
+      return g
+  return None
+
+def _match_batched_gemm(ast:UOp, device:str, arch:str) -> BatchedGemmMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 4 or end.src[0].op is not Ops.STORE: return None
+  store, ranges = end.src[0], end.src[1:]
+  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.LOOP for x in ranges) or store.src[0].op is not Ops.INDEX: return None
+  reduce = store.src[1]
+  if reduce.op is not Ops.REDUCE or reduce.dtype != dtypes.float or reduce.arg != (Ops.ADD, 0) or len(reduce.src) != 2: return None
+  k, product = reduce.src[1], reduce.src[0]
+  if k.op is not Ops.RANGE or product.op is not Ops.CAST or product.dtype != dtypes.float or product.src[0].op is not Ops.MUL: return None
+  lhs, rhs = product.src[0].src
+  if lhs.op is not Ops.INDEX or rhs.op is not Ops.INDEX or lhs.dtype != dtypes.half or rhs.dtype != dtypes.half: return None
+  out_idx, lhs_idx, rhs_idx = (ssimplify(x.src[1].get_idx()) for x in (store.src[0], lhs, rhs))
+  k_size = int(k.vmax)+1
+  for m in ranges:
+    for n in ranges:
+      if n is m: continue
+      batch = next(x for x in ranges if x is not m and x is not n)
+      m_size, n_size, batch_size = int(m.vmax)+1, int(n.vmax)+1, int(batch.vmax)+1
+      if ssimplify(out_idx//(n_size*batch_size)) is not m or ssimplify((out_idx//batch_size)%n_size) is not n or \
+         ssimplify(out_idx%batch_size) is not batch: continue
+      if m_size % 64 or n_size % 32 or k_size % BLOCK_K: continue
+      a_idx, b_idx = ssimplify((batch*k_size+k)*m_size+m), ssimplify((batch*k_size+k)*n_size+n)
+      if lhs_idx is a_idx and rhs_idx is b_idx: a, b = lhs.src[0], rhs.src[0]
+      elif rhs_idx is a_idx and lhs_idx is b_idx: a, b = rhs.src[0], lhs.src[0]
+      else: continue
+      return BatchedGemmMatch(store.src[0].src[0], a, b, m_size, n_size, k_size, batch_size)
+  return None
+
+def _gemm_block_m(g:GemmMatch) -> int:
+  if g.old is None and (g.m,g.n,g.k) in ((393216,576,256), (24576,512,4608)): return 64
+  return 32 if g.old is not None and (g.m < 512 or (g.a_kxm and g.b_kxn and g.k >= 65536)) else \
+    64 if g.old is not None else BLOCK_M
+def _gemm_block_n(g:GemmMatch) -> int:
+  if g.old is not None and g.a_kxm and g.b_kxn and (g.m,g.n,g.k) == (256,2304,65536): return 128
+  if g.old is not None and g.a_kxm and g.b_kxn and (g.m,g.n,g.k) == (256,2304,98304): return 64
+  if g.n == 288 and g.old is None and not g.a_kxm and g.b_kxn: return 96
+  if g.n == 576 and g.old is None and not g.a_kxm and g.b_kxn: return 192
+  if g.old is not None and g.a_kxm and g.b_kxn and g.m == 256 and g.k >= 65536: return BLOCK_K
+  return 64 if g.n == 64 and g.old is None and not g.a_kxm and not g.b_kxn else BLOCK_N
+def _gemm_block_k(g:GemmMatch) -> int:
+  if g.old is not None and g.a_kxm and g.b_kxn and g.m == 256 and g.k >= 65536: return 64
+  if g.old is not None and (g.m < 512 or (g.a_kxm and g.b_kxn and g.k >= 65536)) and g.k % 64 == 0: return 64
+  return 144 if _gemm_block_n(g) == 64 and g.k in (288, 576) else BLOCK_K
+def _batched_block_k(g:BatchedGemmMatch) -> int:
+  if (g.batch, g.m, g.n, g.k) == (192, 64, 288, 8192): return 16
+  if (g.batch, g.m, g.n, g.k) == (96, 64, 576, 4096): return BLOCK_K
+  return 128 if g.m == 64 and g.batch >= 96 and g.k % 128 == 0 else BLOCK_K
+def _batched_block_n(g:BatchedGemmMatch) -> int:
+  if (g.batch, g.m, g.n, g.k) == (192, 64, 288, 8192): return 288
+  return 192 if g.n == 576 else BLOCK_N
+def _batched_threads(g:BatchedGemmMatch) -> int:
+  if _batched_block_n(g) == 288: return 192
+  return _batched_block_n(g)
+def _batched_family_name(g:BatchedGemmMatch) -> str:
+  return f"coop_bgemm_bm64_bn{_batched_block_n(g)}_bk{_batched_block_k(g)}_t{_batched_threads(g)}" \
+    f"{'_partial_n' if g.n % _batched_block_n(g) else ''}"
+def _gemm_threads(g:GemmMatch) -> int:
+  if _gemm_block_n(g) == 192: return 192
+  return THREADS
+
+def _gemm_family_name(g:GemmMatch, output_nchw:GemmOutputNCHW|None) -> str:
+  return f"coop_gemm_bm{_gemm_block_m(g)}_bn{_gemm_block_n(g)}_bk{_gemm_block_k(g)}_t{_gemm_threads(g)}" \
+    f"{'_kxm' if g.a_kxm else ''}{'_kxn' if g.b_kxn else ''}{'_acc' if g.old is not None else ''}" \
+    f"{'_nchw' if output_nchw is not None else ''}"
+
+def _render_gemm(g:GemmMatch, name:str, output_nchw:GemmOutputNCHW|None=None) -> str:
+  bm, bn_size, bk, threads = _gemm_block_m(g), _gemm_block_n(g), _gemm_block_k(g), _gemm_threads(g)
+  pack_b = g.b_kxn and not g.a_kxm
+  as_stride, bs_stride = (bm if g.a_kxm else bk+8), (bn_size if g.b_kxn else bk+8)
+  waves_m, waves_n = (2, 3) if threads == 192 else (1, 2) if threads == 64 else \
+                     (((1, 8) if bm == 32 else (2, 4) if bm == 64 else (4, 2)) if threads == 256 else
+                      ((1, 4) if bm <= 64 else (2, 2)))
+  tiles_m, tiles_n = bm//(waves_m*WMMA_M), bn_size//(waves_n*WMMA_N)
+  cslot, aslot, bslot = g.c.arg.slot, g.a.arg.slot, g.b.arg.slot
+  buffers = (g.c, g.a, g.b) + ((g.old,) if g.old is not None else ())
+  params = ', '.join(f'{"float" if x.dtype == dtypes.float else "half"}* p{x.arg.slot}' for x in sorted(buffers, key=lambda x:x.arg.slot))
+  params += ', int M, int N, int K' + (', int S' if output_nchw is not None else '')
+  lines = [
+    '#define half _Float16',
+    'typedef half half16 __attribute__((ext_vector_type(16)));',
+    'typedef float float8 __attribute__((ext_vector_type(8)));',
+    'typedef unsigned uint8 __attribute__((ext_vector_type(8)));',
+    'typedef unsigned uint16 __attribute__((ext_vector_type(16)));',
+    'typedef unsigned short ushort16 __attribute__((ext_vector_type(16)));',
+    '#define HALF_BITS(x) (unsigned short)(x), (unsigned short)((x)>>16)',
+    '#define WMMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32',
+    f'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({threads}, {threads}))) {name}({params}) {{',
+    f'  __attribute__((shared, aligned(32))) half As[{bk*bm if g.a_kxm else bm*as_stride}], '
+    f'Bs[{bk*bn_size if g.b_kxn else bn_size*bs_stride}];',
+    '  int bn=__builtin_amdgcn_workgroup_id_x(), bm=__builtin_amdgcn_workgroup_id_y(), tid=__builtin_amdgcn_workitem_id_x();',
+    f'  int wave=tid>>5, lane=tid&31, wm=wave/{waves_n}, wn=wave%{waves_n}, row=lane&15, halfrow=lane>>4;',
+  ]
+  lines.append(f'  float8 c[{tiles_m}][{tiles_n}]={{}};')
+  lines += [f'  for (int kt=0; kt<K/{bk}; kt++) {{']
+  if g.a_kxm:
+    aseg_count = bm*bk//16
+    lines += ['    #pragma unroll', f'    for (int q=0; q<{(aseg_count+threads-1)//threads}; q++) {{',
+              f'      int aseg=tid+q*{threads};']
+    if aseg_count % threads: lines.append(f'      if (aseg<{aseg_count}) {{')
+    lines += [f'      int ak=aseg/{bm//16}, am=(aseg%{bm//16})*16;',
+              f'      *((half16*)(As+ak*{bm}+am))=*((half16*)(p{aslot}+((long)(kt*{bk}+ak)*M)+bm*{bm}+am));']
+    if aseg_count % threads: lines.append('      }')
+    lines.append('    }')
+  else:
+    prefix, suffix = (f'    if (tid<{bm}) {{ ', ' }') if threads != bm else ('    ', '')
+    lines += [f'{prefix}long ao=((long)(bm*{bm}+tid)*K)+kt*{bk};']
+    lines += ['      #pragma unroll', f'      for (int q=0; q<{bk//16}; q++) '
+              f'*((half16*)(As+tid*{as_stride}+q*16))=*((half16*)(p{aslot}+ao+q*16));']
+    lines[-1] += suffix
+  if pack_b:
+    bsegs = bn_size//16
+    if threads != bsegs*16: lines.append(f'    if (tid<{bsegs*16}) {{')
+    lines += ['    #pragma unroll', f'    for (int q=0; q<{bk//32}; q++) {{',
+              f'      int bkp=tid/{bsegs}+q*16, bn0=(tid%{bsegs})*16;',
+              f'      half16 bv0=*((half16*)(p{bslot}+((long)(kt*{bk}+bkp*2)*N)+bn*{bn_size}+bn0));',
+              f'      half16 bv1=*((half16*)(p{bslot}+((long)(kt*{bk}+bkp*2+1)*N)+bn*{bn_size}+bn0));',
+              '      ushort16 pb0=__builtin_bit_cast(ushort16,bv0), pb1=__builtin_bit_cast(ushort16,bv1);',
+              f'      *((uint16*)(((unsigned*)Bs)+bkp*{bn_size}+bn0))=__builtin_convertvector(pb0,uint16)|'
+              '(__builtin_convertvector(pb1,uint16)<<16);', '    }']
+    if threads != bsegs*16: lines.append('    }')
+  elif g.b_kxn:
+    bsegs = bn_size//16
+    lines += ['    #pragma unroll', f'    for (int q=0; q<{bn_size*bk//(threads*16)}; q++) {{',
+              f'      int bseg=tid+q*{threads}, bki=bseg/{bsegs}, bni=(bseg%{bsegs})*16;',
+              f'      *((half16*)(Bs+bki*{bn_size}+bni))=*((half16*)(p{bslot}+'
+              f'((long)(kt*{bk}+bki)*N)+bn*{bn_size}+bni));', '    }']
+  else:
+    prefix, suffix = (f'    if (tid<{bn_size}) {{ ', ' }') if threads != bn_size else ('    ', '')
+    lines += [f'{prefix}long bo=((long)(bn*{bn_size}+tid)*K)+kt*{bk};']
+    lines += ['      #pragma unroll', f'      for (int q=0; q<{bk//16}; q++) '
+              f'*((half16*)(Bs+tid*{bs_stride}+q*16))=*((half16*)(p{bslot}+bo+q*16));']
+    lines[-1] += suffix
+  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
+  lines += ['    #pragma unroll', f'    for (int ki=0; ki<{bk//WMMA_K}; ki++) {{']
+  lines += [f'      half16 av[{tiles_m}], bv[{tiles_n}];', '      #pragma unroll', f'      for (int im=0; im<{tiles_m}; im++) {{']
+  if g.a_kxm:
+    lines += ['        half16 v;', '        #pragma unroll',
+              f'        for (int e=0; e<16; e++) v[e]=As[(ki*16+e)*{bm}+(wm*{tiles_m}+im)*16+row];', '        av[im]=v;']
+  else:
+    lines += [f'        av[im]=*((half16*)(As+(((wm*{tiles_m}+im)*16+row)*{as_stride}+ki*16)));']
+  lines += ['      }', '      #pragma unroll', f'      for (int jn=0; jn<{tiles_n}; jn++) {{']
+  if pack_b:
+    half_bits = ','.join(f'HALF_BITS(bp[{e}])' for e in range(8))
+    lines += ['        uint8 bp;', '        #pragma unroll',
+              f'        for (int e=0; e<8; e++) bp[e]=((unsigned*)Bs)[(ki*8+e)*{bn_size}+(wn*{tiles_n}+jn)*16+row];',
+              f'        ushort16 bb=(ushort16){{{half_bits}}};', '        bv[jn]=__builtin_bit_cast(half16,bb);']
+  elif g.b_kxn:
+    lines += ['        half16 v;', '        #pragma unroll',
+              f'        for (int e=0; e<16; e++) v[e]=Bs[(ki*16+e)*{bn_size}+(wn*{tiles_n}+jn)*16+row];', '        bv[jn]=v;']
+  else:
+    lines += [f'        bv[jn]=*((half16*)(Bs+(((wn*{tiles_n}+jn)*16+row)*{bs_stride}+ki*16)));']
+  lines += ['      }', '      #pragma unroll', f'      for (int im=0; im<{tiles_m}; im++) {{', '        #pragma unroll',
+            f'        for (int jn=0; jn<{tiles_n}; jn++) c[im][jn]=WMMA(av[im],bv[jn],c[im][jn]);', '      }', '    }']
+  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");', '  }']
+  lines += ['  #pragma unroll', f'  for (int im=0; im<{tiles_m}; im++) {{', '    #pragma unroll',
+            f'    for (int jn=0; jn<{tiles_n}; jn++) {{', '      #pragma unroll', '      for (int e=0; e<8; e++) {',
+            f'        int om=bm*{bm}+(wm*{tiles_m}+im)*16+e*2+halfrow, on=bn*{bn_size}+(wn*{tiles_n}+jn)*16+row;']
+  if output_nchw is not None:
+    lines.append('        int spatial_size=S*S;')
+    lines.append('        long oi=((long)(om/spatial_size)*N+on)*spatial_size+om%spatial_size;')
+  else: lines.append('        long oi=(long)om*N+on;')
+  if g.old is None: lines.append(f'        p{cslot}[oi]=(half)c[im][jn][e];')
+  else: lines.append(f'        p{cslot}[oi]=(float)((half)c[im][jn][e]+(half){g.scale}*p{g.old.arg.slot}[oi]);')
+  lines += ['      }', '    }', '  }']
+  lines += ['}']
+  return '\n'.join(lines)
+
+def _render_batched_gemm(g:BatchedGemmMatch, name:str) -> tuple[str, int]:
+  bm, bn_size, bk, threads, waves_m = 64, _batched_block_n(g), _batched_block_k(g), _batched_threads(g), 1
+  waves_n = threads//32
+  exact_n = g.n % bn_size == 0
+  tiles_m, tiles_n = bm//(waves_m*WMMA_M), bn_size//(waves_n*WMMA_N)
+  cslot, aslot, bslot = g.c.arg.slot, g.a.arg.slot, g.b.arg.slot
+  params = ', '.join(f'{"float" if x.dtype == dtypes.float else "half"}* p{x.arg.slot}'
+                     for x in sorted((g.c, g.a, g.b), key=lambda x:x.arg.slot))
+  params += ', int B, int M, int N, int K'
+  lines = [
+    '#define half _Float16',
+    'typedef half half16 __attribute__((ext_vector_type(16)));',
+    'typedef float float8 __attribute__((ext_vector_type(8)));',
+    '#define WMMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32',
+    f'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({threads}, {threads}))) {name}({params}) {{',
+    f'  __attribute__((shared, aligned(32))) half As[{bm*bk}], Bs[{bn_size*bk}];',
+    '  int bn=__builtin_amdgcn_workgroup_id_x(), bm=__builtin_amdgcn_workgroup_id_y(), '
+    'batch=__builtin_amdgcn_workgroup_id_z(), tid=__builtin_amdgcn_workitem_id_x();',
+    f'  int wave=tid>>5, lane=tid&31, wm=wave/{waves_n}, wn=wave%{waves_n}, row=lane&15, halfrow=lane>>4;',
+  ]
+  lines.append(f'  float8 c[{tiles_m}][{tiles_n}]={{}};')
+  lines.append(f'  for (int kt=0; kt<K/{bk}; kt++) {{')
+  aseg_count = bm*bk//16
+  lines += ['    #pragma unroll', f'    for (int q=0; q<{(aseg_count+threads-1)//threads}; q++) {{',
+            f'      int aseg=tid+q*{threads};']
+  if aseg_count % threads: lines.append(f'      if (aseg<{aseg_count}) {{')
+  lines += [f'      int ak=aseg/{bm//16}, am=(aseg%{bm//16})*16;',
+            f'      *((half16*)(As+ak*{bm}+am))=*((half16*)(p{aslot}+'
+            f'((long)(batch*K+kt*{bk}+ak)*M)+bm*{bm}+am));']
+  if aseg_count % threads: lines.append('      }')
+  lines.append('    }')
+  bsegs = bn_size//16
+  lines += ['    #pragma unroll', f'    for (int q=0; q<{bn_size*bk//(threads*16)}; q++) {{',
+            f'      int bseg=tid+q*{threads}, bki=bseg/{bsegs}, bni=(bseg%{bsegs})*16;',
+            f'      *((half16*)(Bs+bki*{bn_size}+bni))=' + ('' if exact_n else f'bn*{bn_size}+bni<N ? ') +
+            f'*((half16*)(p{bslot}+((long)(batch*K+kt*{bk}+bki)*N)+bn*{bn_size}+bni))' +
+            (';' if exact_n else ' : (half16){0};'), '    }']
+  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
+  lines += ['    #pragma unroll', f'    for (int ki=0; ki<{bk//WMMA_K}; ki++) {{',
+            f'      half16 av[{tiles_m}], bv[{tiles_n}];', '      #pragma unroll',
+            f'      for (int im=0; im<{tiles_m}; im++) {{', '        half16 v;', '        #pragma unroll',
+            f'        for (int e=0; e<16; e++) v[e]=As[(ki*16+e)*{bm}+(wm*{tiles_m}+im)*16+row];',
+            '        av[im]=v;', '      }', '      #pragma unroll', f'      for (int jn=0; jn<{tiles_n}; jn++) {{',
+            '        half16 v;', '        #pragma unroll',
+            f'        for (int e=0; e<16; e++) v[e]=Bs[(ki*16+e)*{bn_size}+(wn*{tiles_n}+jn)*16+row];',
+            '        bv[jn]=v;', '      }', '      #pragma unroll', f'      for (int im=0; im<{tiles_m}; im++) {{',
+            '        #pragma unroll', f'        for (int jn=0; jn<{tiles_n}; jn++) c[im][jn]=WMMA(av[im],bv[jn],c[im][jn]);',
+            '      }', '    }']
+  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");', '  }']
+  lines += ['  #pragma unroll', f'  for (int im=0; im<{tiles_m}; im++) {{', '    #pragma unroll',
+            f'    for (int jn=0; jn<{tiles_n}; jn++) {{',
+            f'      int om=bm*{bm}+(wm*{tiles_m}+im)*16+halfrow, on=bn*{bn_size}+(wn*{tiles_n}+jn)*16+row;',
+            ('      {' if exact_n else '      if (on<N) {'), '        #pragma unroll',
+            f'        for (int e=0; e<8; e++) p{cslot}[((long)(om+e*2)*N+on)*B+batch]=c[im][jn][e];',
+            '      }', '    }', '  }']
+  lines += ['}']
+  return '\n'.join(lines), bm
+
+def _render_direct_conv_bwd_activation(g:DirectConvBwdActivationMatch, name:str) -> str:
+  bm, bn, bk, threads = 128, min(g.cin, 64), 32, 128
+  persist_weights = g.cin <= 64
+  tiles_m = 2
+  tiles_n, aslot, bslot = bn//16, (3 if g.residual else 2), (4 if g.residual else 3)
+  params = 'float* p0, float* p1, half* p2, half* p3, half* p4' if g.residual else 'float* p0, half* p1, half* p2, half* p3'
+  lines = [
+    '#define half _Float16',
+    'typedef half half16 __attribute__((ext_vector_type(16)));',
+    'typedef half half8 __attribute__((ext_vector_type(8)));',
+    'typedef float float8 __attribute__((ext_vector_type(8)));',
+    'typedef unsigned uint8 __attribute__((ext_vector_type(8)));',
+    'typedef unsigned short ushort16 __attribute__((ext_vector_type(16)));',
+    '#define HALF_BITS(x) (unsigned short)(x), (unsigned short)((x)>>16)',
+    '#define WMMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32',
+    f'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({threads},{threads}))) {name}({params}) {{',
+    f'  __attribute__((shared, aligned(32))) half As[{bm*max(bk,bn)}], Bs[{bn*(g.cout if persist_weights else bk)}];',
+    '  int bn=__builtin_amdgcn_workgroup_id_x(), bm=__builtin_amdgcn_workgroup_id_y(), tid=__builtin_amdgcn_workitem_id_x();',
+    '  int wave=tid>>5, lane=tid&31, wm=wave, row=lane&15, halfrow=lane>>4;',
+    f'  float8 total[{tiles_m}][{tiles_n}]={{}};',
+    '  for (int patch=0; patch<9; patch++) {',
+    f'    float8 c[{tiles_m}][{tiles_n}]={{}};',
+  ]
+  if persist_weights:
+    lines += ['    #pragma unroll', f'    for (int q=0; q<{bn*(g.cout//2)//threads}; q++) {{',
+              f'      int be=tid+q*{threads}, bp=be/{bn}, ci=be%{bn};',
+              f'      half blo=p{bslot}[(bp*2)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
+              f'      half bhi=p{bslot}[(bp*2+1)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
+              f'      ((unsigned*)Bs)[bp*{bn}+ci]=__builtin_bit_cast(unsigned short,blo)|'
+              '((unsigned)__builtin_bit_cast(unsigned short,bhi)<<16);', '    }']
+    lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+              '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
+  lines += [f'    for (int ct=0; ct<{g.cout//bk}; ct++) {{', f'      int co0=ct*{bk};']
+  if not persist_weights:
+    lines += ['      #pragma unroll', f'      for (int q=0; q<{bn*(bk//2)//threads}; q++) {{',
+              f'        int be=tid+q*{threads}, bp=be/{bn}, ci=be%{bn};',
+              f'        half blo=p{bslot}[(co0+bp*2)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
+              f'        half bhi=p{bslot}[(co0+bp*2+1)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
+              f'        ((unsigned*)Bs)[bp*{bn}+ci]=__builtin_bit_cast(unsigned short,blo)|'
+              '((unsigned)__builtin_bit_cast(unsigned short,bhi)<<16);', '      }']
+    lines += ['      __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+              '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
+  lines += [
+            '      #pragma unroll', f'      for (int ki=0; ki<{bk//16}; ki++) {{',
+            f'        half16 av[{tiles_m}], bv[{tiles_n}];', '        #pragma unroll', f'        for (int im=0; im<{tiles_m}; im++) {{',
+            f'          int ami=bm*{bm}+(wm*{tiles_m}+im)*16+row, ab=ami/{g.spatial*g.spatial}, apos=ami%{g.spatial*g.spatial};',
+            f'          int ay=apos/{g.spatial}+1-patch/3, ax=apos%{g.spatial}+1-patch%3;',
+            f'          av[im]=(ay>=0 && ay<{g.spatial} && ax>=0 && ax<{g.spatial}) ? '
+            f'*((half16*)(p{aslot}+((long)ab*{g.spatial*g.spatial}+ay*{g.spatial}+ax)*{g.cout}+co0+ki*16)) : (half16){{0}};',
+            '        }',
+            '        #pragma unroll', f'        for (int jn=0; jn<{tiles_n}; jn++) {{', '          uint8 bp;', '          #pragma unroll',
+            f'          for (int e=0; e<8; e++) bp[e]=((unsigned*)Bs)[({"ct*16+" if persist_weights else ""}ki*8+e)*{bn}+jn*16+row];',
+            '          ushort16 bb=(ushort16){HALF_BITS(bp[0]),HALF_BITS(bp[1]),HALF_BITS(bp[2]),HALF_BITS(bp[3]),'
+            'HALF_BITS(bp[4]),HALF_BITS(bp[5]),HALF_BITS(bp[6]),HALF_BITS(bp[7])};',
+            '          bv[jn]=__builtin_bit_cast(half16,bb);', '        }', '        #pragma unroll',
+            f'        for (int im=0; im<{tiles_m}; im++) {{',
+            '          #pragma unroll', f'          for (int jn=0; jn<{tiles_n}; jn++) c[im][jn]=WMMA(av[im],bv[jn],c[im][jn]);',
+            '        }', '      }']
+  if not persist_weights:
+    lines += ['      __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+              '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
+  lines += ['    }',
+            '    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");',
+            '    #pragma unroll', f'    for (int im=0; im<{tiles_m}; im++) {{', '      #pragma unroll',
+            f'      for (int jn=0; jn<{tiles_n}; jn++) total[im][jn]+=__builtin_convertvector('
+            '__builtin_convertvector(c[im][jn],half8),float8);', '    }', '  }',
+            '  #pragma unroll', f'  for (int im=0; im<{tiles_m}; im++) {{', '    #pragma unroll',
+            f'    for (int jn=0; jn<{tiles_n}; jn++) {{', '      #pragma unroll', '      for (int e=0; e<8; e++) {',
+            f'        int lm=(wm*{tiles_m}+im)*16+e*2+halfrow, lc=jn*16+row;',
+            f'        As[lc*{bm}+lm]=(half)total[im][jn][e];', '      }', '    }', '  }',
+            '  __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
+            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
+  lines += ['  #pragma unroll', f'  for (int q=0; q<{bm*bn//(threads*8)}; q++) {{',
+            f'    int oe=tid*{bm*bn//threads}+q*8, lc=oe/{bm}, lm=oe%{bm};',
+            f'    int gm=bm*{bm}+lm, ob=gm/{g.spatial*g.spatial}, pos=gm%{g.spatial*g.spatial};',
+            f'    long oo=((long)ob*{g.cin}+bn*{bn}+lc)*{g.spatial*g.spatial}+pos;',
+            '    half8 grad=*((half8*)(As+oe));']
+  if g.residual: lines += ['    grad+=*((half8*)(p2+oo));', '    half8 z=__builtin_convertvector(*((float8*)(p1+oo)),half8);']
+  else: lines.append('    half8 z=*((half8*)(p1+oo));')
+  lines += ['    half8 sig=(half8)1.0/((half8)1.0+__builtin_elementwise_exp2(z*(half)-2.4554669595930156));',
+            '    *((float8*)(p0+oo))=__builtin_convertvector(sig*grad+(half)1.702*z*grad*sig*((half)1.0-sig),float8);', '  }']
+  lines.append('}')
+  return '\n'.join(lines)
+
+def direct_conv_bwd_activation_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if not isinstance(g:=ast.tag, DirectConvBwdActivationMatch) or renderer.target.device != "AMD" or \
+     not renderer.target.arch.startswith("gfx11"): return None
+  name = f"coop_direct_conv_bwd_activation_{g.m}_{g.cin}_{g.cout}{'_res' if g.residual else ''}"
+  source = _render_direct_conv_bwd_activation(g, name)
+  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(2*g.m*g.cin*g.cout*9, 0, 0)))
+  slots = (0,1,2,3,4) if g.residual else (0,1,2,3)
+  bn = min(g.cin, 64)
+  info = ProgramInfo(name=name, global_size=((g.cin+bn-1)//bn, g.m//128, 1), local_size=(128,1,1),
+                     globals=slots, outs=(0,), ins=slots[1:])
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)
+
+
+def cooperative_gemm_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if (g:=_match_gemm(ast, renderer.target.device, renderer.target.arch)) is not None:
+    output_nchw = ast.tag if isinstance(ast.tag, GemmOutputNCHW) else None
+    name = _gemm_family_name(g, output_nchw)
+    source = _render_gemm(g, name, output_nchw)
+    global_size = (g.n//_gemm_block_n(g), g.m//_gemm_block_m(g), 1)
+    local_size = (_gemm_threads(g), 1, 1)
+    estimates = Estimates(2*g.m*g.n*g.k, 2*(g.m*g.k+g.n*g.k+g.m*g.n), 2*(g.m*g.k+g.n*g.k+g.m*g.n))
+    slots = tuple(sorted((g.c.arg.slot, g.a.arg.slot, g.b.arg.slot) + ((g.old.arg.slot,) if g.old is not None else ())))
+    out_slot = g.c.arg.slot
+    suffix = f"_{g.m}_{g.n}_{g.k}"
+    variables:tuple[UOp, ...] = (UOp.variable(f"M{suffix}", g.m, g.m, dtypes.int),
+                                 UOp.variable(f"N{suffix}", g.n, g.n, dtypes.int),
+                                 UOp.variable(f"K{suffix}", g.k, g.k, dtypes.int))
+    if output_nchw is not None:
+      variables += (UOp.variable(f"S{suffix}_{output_nchw.spatial}", output_nchw.spatial, output_nchw.spatial, dtypes.int),)
+  elif (bg:=_match_batched_gemm(ast, renderer.target.device, renderer.target.arch)) is not None:
+    name = _batched_family_name(bg)
+    source, bm = _render_batched_gemm(bg, name)
+    global_size = ((bg.n+_batched_block_n(bg)-1)//_batched_block_n(bg), bg.m//bm, bg.batch)
+    local_size = (_batched_threads(bg), 1, 1)
+    estimates = Estimates(2*bg.batch*bg.m*bg.n*bg.k, 2*bg.batch*(bg.m*bg.k+bg.n*bg.k+2*bg.m*bg.n),
+                          2*bg.batch*(bg.m*bg.k+bg.n*bg.k+2*bg.m*bg.n))
+    slots, out_slot = tuple(sorted((bg.c.arg.slot, bg.a.arg.slot, bg.b.arg.slot))), bg.c.arg.slot
+    suffix = f"_{bg.batch}_{bg.m}_{bg.n}_{bg.k}"
+    variables = (UOp.variable(f"B{suffix}", bg.batch, bg.batch, dtypes.int), UOp.variable(f"M{suffix}", bg.m, bg.m, dtypes.int),
+                 UOp.variable(f"N{suffix}", bg.n, bg.n, dtypes.int), UOp.variable(f"K{suffix}", bg.k, bg.k, dtypes.int))
+  else: return None
+  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=estimates))
+  info = ProgramInfo(name=name, global_size=global_size, local_size=local_size, vars=variables, globals=slots,
+                     outs=(out_slot,), ins=tuple(x for x in slots if x != out_slot))
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)

tinygrad/codegen/opt/heuristic.py

@@ -35,13 +35,54 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
       pass
     if good_tc_opt:
       if rngs is not None:
+        tc_sizes = [r.src[0] for r in rngs]
+        skinny_output = any(resolve(sz < 2, False) for sz in tc_sizes[:2])
+        very_long_reduce = resolve(tc_sizes[2] >= 4096, False)
+        long_reduce = resolve(tc_sizes[2] >= 1024, False)
+        small_m = resolve(tc_sizes[0] <= 32, False)
+        tiny_m = resolve(tc_sizes[0] <= 16, False)
+        if resolve(tc_sizes[0] >= 32, False) and resolve(tc_sizes[0] < 33, False) and resolve(tc_sizes[1] >= 1536, False) and \
+           resolve(tc_sizes[2] >= 128, False) and resolve(tc_sizes[2] <= 512, False):
+          upcast_n = 8 if resolve(tc_sizes[1] >= 4096, False) else 6
+          rngs[1] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[1]), upcast_n))[0]
+          rngs[0] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[0]), 2))[0]
+          rngs[0] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), 8))[0]
+          rngs[1] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[1]), 2))[0]
+          return tk
         for tc_dim in [1,0]: # attempt to upcast M and N
-          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
+          if skinny_output or resolve(tc_sizes[tc_dim] >= 32768, False): continue
+          short_conv_n = tc_dim == 1 and resolve(tc_sizes[0] >= 65536, False) and resolve(tc_sizes[1] == 18, False) and \
+                         resolve(tc_sizes[2] <= 4, False)
+          upcast_sizes = [2] if short_conv_n or (very_long_reduce and (tiny_m or (tc_dim == 0 and small_m))) else [5,4,3,2]
+          szs = [sz for sz in upcast_sizes if rngs[tc_dim].src[0].divides(sz) is not None]
           if szs:
             # set it to the replaced range
             rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
-        if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
-          tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
+        if skinny_output:
+          outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL and r not in rngs[:2]]
+          if outer_rngs and outer_rngs[0].src[0].divides(2) is not None:
+            tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rngs[0]), 2))
+          if very_long_reduce:
+            outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL and r not in rngs[:2]]
+            if outer_rngs and (outer_local:=next((x for x in (16, 4, 2) if outer_rngs[0].src[0].divides(x) is not None), None)):
+              tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rngs[0]), outer_local))
+          return tk
+        local_sizes = [2] if long_reduce and small_m else [4,2]
+        if (szs := [sz for sz in local_sizes if rngs[0].src[0].divides(sz) is not None]):
+          rngs[0] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))[0]
+        if long_reduce and small_m and resolve(tc_sizes[1] >= 128, False) and \
+           rngs[1].arg[-1] is AxisType.GLOBAL and rngs[1].src[0].divides(3) is not None:
+          rngs[1] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[1]), 3))[0]
+        if tk.applied_opts[-1] == Opt(OptOps.LOCAL, 1, 4):
+          outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL]
+          if outer_rngs and resolve(outer_rngs[0].src[0] >= 384, False) and \
+             (outer_local:=next((x for x in (4, 3, 2) if outer_rngs[0].src[0].divides(x) is not None), None)):
+            tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rngs[0]), outer_local))
+        elif resolve(tc_sizes[0] <= 4, False) and resolve(tc_sizes[2] >= 256, False):
+          outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL]
+          outer_local = 16 if resolve(tc_sizes[2] >= 512, False) else 2
+          if (outer_rng:=next((r for r in reversed(outer_rngs) if r.src[0].divides(outer_local) is not None), None)) is not None:
+            tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rng), outer_local))
       return tk
 
   # make a copy so it does not mutate the input
@@ -82,7 +123,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
 
   # are we grouping? (requires local shape support)
   if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
-    for axis, sz in itertools.product((0, 1, 2), (16,)):
+    group_sizes = (8, 16) if len(k.bufs) >= 7 else (16,)
+    for axis, sz in itertools.product((0, 1, 2), group_sizes):
       try:
         k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
         break
@@ -114,10 +156,11 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # potentially do more upcasts of non reduce axes based on a heuristic
   is_dsp = k.ren is not None and k.ren.target.device == "DSP"
   upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 2):
     xb_choices = []
     # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
-    for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
+    upcast_amounts = ([128] if not len(upcasted_axis) else []) if is_dsp else ([3,4] if k.reduceop is not None else [2,3,4])
+    for axis, upcast_amount in itertools.product(k.upcastable_dims, upcast_amounts):
       # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
       if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
       rng = k.rngs[axis]
@@ -141,15 +184,17 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
 
   # if last reduce dim is small(ish), loop unroll the reduce
   # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
+  four_by_four_reduce = len(k.unrollable_dims) >= 2 and all(resolve(k.full_shape[x] == 4, False) for x in k.unrollable_dims[-2:])
   try:
     if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
-      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
+      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 4:
         k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
         # if it's small, upcast a second reduce dimension too
         if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
           k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
       else:
-        for splits in [4]:
+        bn_spatial_reduce = len(k.unrollable_dims) >= 2 and resolve(k.full_shape[k.unrollable_dims[-2]] == 6, False)
+        for splits in ([8, 4] if bn_spatial_reduce else [4]):
           if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
             k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
             break
@@ -166,20 +211,44 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
     if NOLOCALS:
       k.apply_opt(Opt(OptOps.NOLOCALS))
     else:
+      special_local = False
+      if four_by_four_reduce and len(global_axes:=k.axes_of(AxisType.GLOBAL, AxisType.LOOP)) >= 3:
+        lk, local_rngs = k.copy(), [k.rngs[x] for x in global_axes[-3:]]
+        try:
+          for rng, sz in zip(local_rngs, (16, 4, 4)):
+            lk.apply_opt(Opt(OptOps.LOCAL, lk.rngs.index(rng), sz))
+          k, special_local = lk, True
+        except KernelOptError: pass
       # prioritize making expand axes local
-      local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
-                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
-      to_local: list[tuple[int, int]] = []
-      for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
-        local_size = prod(sz for _, sz in to_local)
-        local_sz: int|None = next((x for x in ([32] * (axis == 0) + [16,8,4,3,2]) if k.full_shape[axis] % x == 0 and local_size * x <= 128), None)
-        if local_sz is not None: to_local.append((axis, local_sz))
-      deleted_shape = 0
-      for axis, local_sz in sorted(to_local[:3]):
-        axis = axis - deleted_shape
-        will_delete_shape = local_sz == k.full_shape[axis]
-        k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
-        if will_delete_shape: deleted_shape += 1
+      if not special_local:
+        local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
+                                for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
+        to_local: list[tuple[int, int]] = []
+        for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
+          local_size = prod(sz for _, sz in to_local)
+          local_sz: int|None = next((x for x in ([32] * (axis == 0) + [16,8,4,3,2]) if k.full_shape[axis] % x == 0 and local_size * x <= 128), None)
+          if local_sz is not None: to_local.append((axis, local_sz))
+        deleted_shape = 0
+        for axis, local_sz in sorted(to_local[:3]):
+          axis = axis - deleted_shape
+          will_delete_shape = local_sz == k.full_shape[axis]
+          k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
+          if will_delete_shape: deleted_shape += 1
+
+      # Both 3x3 reduce axes are already fully unrolled above. Tile the exposed
+      # spatial axes without changing reduction order.
+      unroll_sizes = [k.full_shape[x] for x in k.axes_of(AxisType.UNROLL)]
+      global_axes = k.axes_of(AxisType.GLOBAL, AxisType.LOOP)
+      if unroll_sizes == [3, 3] and len(global_axes) >= 2:
+        axis_size = k.full_shape[global_axes[1]]
+        if axis_size <= 16: k.apply_opt(Opt(OptOps.UPCAST, 1, 0))
+        elif axis_size % 4 == 0: k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
+        if len(global_axes) >= 4 and k.full_shape[global_axes[2]] == 4: k.apply_opt(Opt(OptOps.LOCAL, 2, 4))
+
+      remaining_reduces = [k.full_shape[x] for x in k.unrollable_dims]
+      if len(remaining_reduces) == 2 and remaining_reduces[0] == 6 and remaining_reduces[1] >= 8 and \
+         remaining_reduces[1] % 4 == 0 and k.upcast_size() <= 16:
+        k.apply_opt(Opt(OptOps.UNROLL, 1, 4))
 
   # **** threading ****
 

tinygrad/codegen/opt/reduce.py

@@ -0,0 +1,703 @@
+from collections import Counter
+from dataclasses import dataclass, replace
+
+from tinygrad.dtype import dtypes
+from tinygrad.renderer import Estimates, Renderer
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops, ProgramInfo, UOp, ssimplify
+
+
+@dataclass(frozen=True)
+class ChannelReduceMatch:
+  kind: str
+  params: tuple[UOp, ...]
+  groups: int
+  channels: int
+  spatial: int
+
+@dataclass(frozen=True)
+class DualChannelReduceMatch:
+  groups:int
+  channels:int
+  spatial:int
+
+@dataclass(frozen=True)
+class DualActivationReduceElementwiseMatch:
+  groups:int
+  channels:int
+  spatial:int
+
+@dataclass(frozen=True)
+class ActivationVarGradElementwiseSumDMeanMatch:
+  batch:int
+
+@dataclass(frozen=True)
+class DualMoments512Match:
+  batch:int
+
+@dataclass(frozen=True)
+class DualBNGradDMean512Match:
+  batch:int
+
+@dataclass(frozen=True)
+class Col2ImMatch:
+  params: tuple[UOp, ...]
+  fused_activation: bool
+  batch: int
+  channels: int
+  spatial: int
+
+@dataclass(frozen=True)
+class Im2ColMatch:
+  params: tuple[UOp, ...]
+  batch: int
+  channels: int
+  spatial: int
+
+def _match_moment_mean_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE), None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4: return None
+  batch, sy, sx = tuple(int(r.vmax)+1 for r in reduce.src[1:])
+  if batch not in (1024,1280,1536) or (sy,sx) != (4,4) or tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,512),(1,dtypes.half,batch*8192)): return None
+  value = reduce.src[0]
+  if value.op is not Ops.CAST or value.dtype != dtypes.float or value.src[0].op is not Ops.INDEX or value.src[0].src[0] is not params[1]: return None
+  return batch if Counter(u.op for u in ast.toposort()) in (Counter({Ops.CONST:10,Ops.MUL:5,Ops.RANGE:4,Ops.ADD:4,Ops.PARAM:2,Ops.INDEX:2,
+    Ops.CAST:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}), Counter({Ops.CONST:8,Ops.MUL:4,Ops.RANGE:4,Ops.ADD:3,Ops.PARAM:2,
+    Ops.INDEX:2,Ops.CAST:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1})) else None
+
+def _match_moment_var_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE), None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4: return None
+  batch, sy, sx = tuple(int(r.vmax)+1 for r in reduce.src[1:])
+  if batch not in (1024,1280,1536) or (sy,sx) != (4,4) or tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,512),(1,dtypes.half,batch*8192),(2,dtypes.float,512)): return None
+  value = reduce.src[0]
+  if value.op is not Ops.MUL or value.src[0] is not value.src[1] or value.src[0].op is not Ops.CAST or \
+     value.src[0].src[0].op is not Ops.INDEX or value.src[0].src[0].src[0] is not params[1]: return None
+  return batch if Counter(u.op for u in ast.toposort()) in (Counter({Ops.CONST:12,Ops.MUL:8,Ops.ADD:6,Ops.RANGE:4,Ops.PARAM:3,Ops.INDEX:3,
+    Ops.CAST:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}), Counter({Ops.CONST:10,Ops.MUL:7,Ops.ADD:5,Ops.RANGE:4,Ops.PARAM:3,
+    Ops.INDEX:3,Ops.CAST:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1})) else None
+
+def moments_512_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if not isinstance(ast.tag,DualMoments512Match) or renderer.target.device != "AMD" or not renderer.target.arch.startswith("gfx11"): return None
+  batch = ast.tag.batch
+  name = f"channel_moments_{batch}_512_4_{ast.key.hex()[:8]}"
+  source = f'''#define half _Float16
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(128,128))) {name}(
+    float* p0, float* p1, half* p2) {{
+  __attribute__((shared, aligned(32))) float partial_sum[128], partial_sq[128];
+  int c=__builtin_amdgcn_workgroup_id_x(), tid=__builtin_amdgcn_workitem_id_x();
+  float sum=0.0f, sq=0.0f;
+  for (int q=0; q<{batch//8}; q++) {{
+    int r=tid+q*128, idx=(r>>4)*8192+c*16+(r&15);
+    float v=(float)p2[idx]; sum+=v; sq+=v*v;
+  }}
+  partial_sum[tid]=sum; partial_sq[tid]=sq;
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  for (int stride=64; stride; stride>>=1) {{
+    if (tid<stride) {{ partial_sum[tid]+=partial_sum[tid+stride]; partial_sq[tid]+=partial_sq[tid+stride]; }}
+    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+    __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  }}
+  if (tid==0) {{ float mean=partial_sum[0]*{1/(batch*16)!r}f; p0[c]=mean;
+    p1[c]=partial_sq[0]*{1/(batch*16)!r}f-mean*mean+1e-12f; }}
+}}'''
+  sink = ast.replace(arg=replace(ast.arg,name=name,estimates=Estimates(512*batch*16*2,batch*8192*2,512*128*8)))
+  info = ProgramInfo(name=name,global_size=(512,1,1),local_size=(128,1,1),globals=(0,1,2),outs=(0,1),ins=(2,))
+  src:tuple[UOp, ...] = (sink,UOp(Ops.LINEAR),UOp(Ops.SOURCE,arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY,arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM,src=src,arg=info)
+
+def _match_bn_var_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM),key=lambda u:u.arg.slot))
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE),None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4: return None
+  batch, sy, sx = tuple(int(r.vmax)+1 for r in reduce.src[1:])
+  signature = ((0,dtypes.float,512),(1,dtypes.float,512),(2,dtypes.half,batch*8192),(3,dtypes.float,512),
+               (4,dtypes.float,512),(5,dtypes.float,batch*8192))
+  return batch if batch in (1024,1280,1536) and (sy,sx) == (4,4) and tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) == signature and \
+    Counter(u.op for u in ast.toposort()) == Counter({Ops.CONST:11,Ops.MUL:10,Ops.PARAM:6,Ops.INDEX:6,Ops.ADD:5,Ops.RANGE:4,
+      Ops.RECIPROCAL:1,Ops.SQRT:1,Ops.CAST:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}) else None
+
+def _match_bn_sum_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM),key=lambda u:u.arg.slot))
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE),None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4: return None
+  batch, sy, sx = tuple(int(r.vmax)+1 for r in reduce.src[1:])
+  return batch if batch in (1024,1280,1536) and (sy,sx) == (4,4) and tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) == \
+    ((0,dtypes.float,512),(1,dtypes.float,batch*8192),(2,dtypes.float,512)) and \
+    Counter(u.op for u in ast.toposort()) == Counter({Ops.CONST:10,Ops.MUL:5,Ops.ADD:5,Ops.RANGE:4,Ops.PARAM:3,Ops.INDEX:3,
+      Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}) else None
+
+def _match_bn_dmean_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM),key=lambda u:u.arg.slot))
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE),None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4: return None
+  batch, sy, sx = tuple(int(r.vmax)+1 for r in reduce.src[1:])
+  if batch not in (1024,1280,1536) or (sy,sx) != (4,4) or tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,512),(1,dtypes.float,512),(2,dtypes.float,512),(3,dtypes.float,512),
+      (4,dtypes.float,512),(5,dtypes.float,batch*8192)): return None
+  expected = {Ops.CONST:12,Ops.MUL:10,Ops.PARAM:6,Ops.INDEX:6,Ops.ADD:5,Ops.RANGE:4,Ops.RECIPROCAL:1,
+              Ops.SQRT:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}
+  return batch if Counter(u.op for u in ast.toposort()) == Counter(expected) else None
+
+def bn_grad_512_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if not isinstance(ast.tag,DualBNGradDMean512Match) or renderer.target.device != "AMD" or \
+     not renderer.target.arch.startswith("gfx11"): return None
+  batch = ast.tag.batch
+  name = f"channel_bn_grad_{batch}_512_4_dmean_{ast.key.hex()[:8]}"
+  source = f'''#define half _Float16
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(128,128))) {name}(
+    float* p0, float* p1, float* p2, float* p3, half* p4, float* p5, float* p6, float* p7, float* p8) {{
+  __attribute__((shared, aligned(32))) float partial_cov[128], partial_sum[128];
+  int c=__builtin_amdgcn_workgroup_id_x(), tid=__builtin_amdgcn_workitem_id_x();
+  float cov=0.0f, sum=0.0f, mean=p5[c], weight=p6[c];
+  for (int q=0; q<{batch//8}; q++) {{
+    int r=tid+q*128, idx=(r>>4)*8192+c*16+(r&15); float grad=p7[idx];
+    cov+=((float)p4[idx]-mean)*weight*grad; sum+=grad;
+  }}
+  partial_cov[tid]=cov; partial_sum[tid]=sum;
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  for (int stride=64; stride; stride>>=1) {{
+    if (tid<stride) {{ partial_cov[tid]+=partial_cov[tid+stride]; partial_sum[tid]+=partial_sum[tid+stride]; }}
+    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+    __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  }}
+  if (tid==0) {{ float inv=1.0f/p3[c], vg=inv*__builtin_elementwise_sqrt(inv)*partial_cov[0]*-0.5f; p0[c]=vg;
+    p1[c]=partial_sum[0]+0.018447889655172415f*p8[c];
+    p2[c]=(vg*mean*-2.0f-weight*__builtin_elementwise_sqrt(inv)*partial_sum[0])*{1/(batch*16)!r}f; }}
+}}'''
+  sink = ast.replace(arg=replace(ast.arg,name=name,estimates=Estimates(512*batch*16*5,batch*8192*6,512*128*8)))
+  info = ProgramInfo(name=name,global_size=(512,1,1),local_size=(128,1,1),globals=tuple(range(9)),outs=(0,1,2),ins=tuple(range(3,9)))
+  src:tuple[UOp, ...] = (sink,UOp(Ops.LINEAR),UOp(Ops.SOURCE,arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY,arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM,src=src,arg=info)
+
+@dataclass(frozen=True)
+class MaxPoolMatch:
+  params: tuple[UOp, ...]
+  batch: int
+  channels: int
+  spatial: int
+
+
+def _match_activation_var_grad(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 2 or end.src[0].op is not Ops.STORE: return None
+  store, channel = end.src
+  if channel.op is not Ops.RANGE or channel.arg[1] is not AxisType.LOOP or int(channel.vmax)+1 != 512: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  signature = tuple((u.arg.slot, u.dtype, u.max_numel()) for u in params)
+  if len(params) != 7 or params[2].max_numel() % 8192: return None
+  batch = params[2].max_numel()//8192
+  if batch not in (1024,1280,1536) or signature != ((0,dtypes.float,512), (1,dtypes.float,512), (2,dtypes.half,batch*8192),
+      (3,dtypes.float,512), (4,dtypes.float,512), (5,dtypes.float,batch*8192), (6,dtypes.half,batch*8192)): return None
+  expected_ops = {Ops.MUL:17, Ops.CONST:14, Ops.ADD:10, Ops.PARAM:7, Ops.INDEX:7, Ops.RANGE:4, Ops.CAST:3,
+                  Ops.RECIPROCAL:2, Ops.SQRT:1, Ops.EXP2:1, Ops.REDUCE:1, Ops.STORE:1, Ops.END:1, Ops.SINK:1}
+  if Counter(u.op for u in ast.toposort()) != Counter(expected_ops): return None
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE), None)
+  if reduce is None or reduce.arg != (Ops.ADD, 0) or len(reduce.src) != 4 or \
+     tuple(int(x.vmax)+1 for x in reduce.src[1:]) != (batch,4,4): return None
+  return batch if store.src[0].src[0] is params[0] and ssimplify(store.src[0].src[1].get_idx()) is channel else None
+
+def _match_activation_elementwise_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  if len(params) != 5 or params[0].max_numel() % 8192: return None
+  batch = params[0].max_numel()//8192
+  if batch not in (1024,1280,1536) or tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,batch*8192),(1,dtypes.float,512),(2,dtypes.float,512),(3,dtypes.float,batch*8192),(4,dtypes.half,batch*8192)): return None
+  expected_ops = {Ops.MUL:13,Ops.CONST:12,Ops.ADD:9,Ops.PARAM:5,Ops.INDEX:5,Ops.RANGE:4,Ops.RECIPROCAL:2,
+                  Ops.CAST:2,Ops.SQRT:1,Ops.EXP2:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}
+  return batch if Counter(u.op for u in ast.toposort()) == Counter(expected_ops) else None
+
+def _match_activation_sum_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM),key=lambda u:u.arg.slot))
+  if len(params) != 4 or params[1].max_numel() % 8192: return None
+  batch = params[1].max_numel()//8192
+  if batch not in (1024,1280,1536) or tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,512),(1,dtypes.float,batch*8192),(2,dtypes.half,batch*8192),(3,dtypes.float,512)): return None
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE),None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4 or \
+     tuple(int(r.vmax)+1 for r in reduce.src[1:]) != (batch,4,4): return None
+  expected = {Ops.CONST:13,Ops.MUL:12,Ops.ADD:10,Ops.PARAM:4,Ops.RANGE:4,Ops.INDEX:4,Ops.CAST:2,
+              Ops.EXP2:1,Ops.RECIPROCAL:1,Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}
+  return batch if Counter(u.op for u in ast.toposort()) == Counter(expected) else None
+
+def _match_activation_dmean_512(ast:UOp, device:str, arch:str) -> int|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM),key=lambda u:u.arg.slot))
+  if len(params) != 4 or params[3].max_numel() % 8192: return None
+  batch = params[3].max_numel()//8192
+  if batch not in (1024,1280,1536) or tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,512),(1,dtypes.float,512),(2,dtypes.float,512),(3,dtypes.float,batch*8192)): return None
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE),None)
+  if reduce is None or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4 or \
+     tuple(int(r.vmax)+1 for r in reduce.src[1:]) != (batch,4,4): return None
+  expected = {Ops.CONST:12,Ops.MUL:8,Ops.ADD:5,Ops.PARAM:4,Ops.RANGE:4,Ops.INDEX:4,
+              Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}
+  return batch if Counter(u.op for u in ast.toposort()) == Counter(expected) else None
+
+def activation_var_grad_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if not isinstance(ast.tag,ActivationVarGradElementwiseSumDMeanMatch) or \
+     (batch:=_match_activation_var_grad(ast,renderer.target.device,renderer.target.arch)) is None: return None
+  name = f"direct_activation_var_grad_{batch}_512_4_elementwise_sum_dmean"
+  source = f'''#define half _Float16
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(128,128))) {name}(
+    float* p0, float* p1, float* p2, float* p3, float* p4, half* p5, float* p6, float* p7,
+    float* p8, half* p9, float* p10) {{
+  __attribute__((shared, aligned(32))) float partial[128], partial_sum[128], partial_elem[128];
+  int c=__builtin_amdgcn_workgroup_id_x(), tid=__builtin_amdgcn_workitem_id_x();
+  float mean=p6[c], norm=p7[c]*__builtin_elementwise_sqrt(1.0f/p4[c]), acc=0.0f, sum=0.0f, elem_sum=0.0f;
+  for (int q=0; q<{batch//8}; q++) {{
+    int r=tid+q*128, idx=(r>>4)*8192+c*16+(r&15);
+    half z=(half)p8[idx], sig=(half)1.0/((half)1.0+__builtin_elementwise_exp2(z*(half)-2.4554669595930156));
+    half grad=sig*p9[idx]+(half)1.702*z*p9[idx]*sig*((half)1.0-sig);
+    acc+=((float)p5[idx]-mean)*(float)grad;
+    float elem=norm*(float)grad; p1[idx]=elem; sum+=(float)grad; elem_sum+=elem;
+  }}
+  partial[tid]=acc; partial_sum[tid]=sum; partial_elem[tid]=elem_sum;
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  for (int stride=64; stride; stride>>=1) {{
+    if (tid<stride) {{ partial[tid]+=partial[tid+stride]; partial_sum[tid]+=partial_sum[tid+stride];
+      partial_elem[tid]+=partial_elem[tid+stride]; }}
+    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+    __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  }}
+  if (tid==0) {{ float inv=1.0f/p4[c]; float vg=inv*__builtin_elementwise_sqrt(inv)*partial[0]*p7[c]*-0.5f; p0[c]=vg;
+    p2[c]=partial_sum[0]+0.018447889655172415f*p10[c];
+    p3[c]=(vg*p6[c]*-2.0f-partial_elem[0])*{1/(batch*16)!r}f; }}
+}}'''
+  sink = ast.replace(arg=replace(ast.arg,name=name,estimates=Estimates(512*batch*16*20,512*batch*16*12,512*128*4)))
+  info = ProgramInfo(name=name,global_size=(512,1,1),local_size=(128,1,1),globals=tuple(range(11)),
+                     outs=(0,1,2,3),ins=tuple(range(4,11)))
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)
+
+@dataclass(frozen=True)
+class MaxPoolBackwardMatch:
+  params: tuple[UOp, ...]
+  batch: int
+  channels: int
+  spatial: int
+
+def _match_maxpool_backward(ast:UOp, device:str, arch:str) -> MaxPoolBackwardMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1 or not isinstance(ast.arg, KernelInfo): return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 3 or end.src[0].op is not Ops.STORE: return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  if any(not u.src or u.src[0].op is not Ops.CONST or not isinstance(u.src[0].arg, int) for u in params): return None
+  ranges = end.src[1:]
+  if any(r.op is not Ops.RANGE or r.arg[1] is not AxisType.LOOP for r in ranges): return None
+  position_size, channels = tuple(int(r.vmax)+1 for r in ranges)
+  spatial = {64:32, 256:16, 512:8}.get(channels)
+  if spatial is None or position_size % (spatial*spatial) or end.src[0].src[0].src[0].arg.slot != 0: return None
+  batch = position_size//(spatial*spatial)
+  if batch not in (1024, 1280, 1536): return None
+  output_size, pooled_size = batch*channels*spatial*spatial, batch*channels*(spatial//2)**2
+  signature = tuple((u.arg.slot, u.dtype, int(u.src[0].arg)) for u in params)
+  expected_signature = ((0, dtypes.half, output_size), (1, dtypes.half, output_size),
+                        (2, dtypes.half, pooled_size), (3, dtypes.half, pooled_size), (4, dtypes.half, pooled_size))
+  if signature != expected_signature: return None
+  expected_ops = {Ops.CONST:{32:19, 16:18, 8:19}[spatial], Ops.PARAM:5, Ops.RANGE:4, Ops.MUL:14, Ops.ADD:12,
+                  Ops.INDEX:5, Ops.FLOORDIV:6, Ops.FLOORMOD:6, Ops.CMPLT:6, Ops.AND:10, Ops.WHERE:4,
+                  Ops.CMPNE:2, Ops.CAST:3, Ops.RECIPROCAL:1, Ops.REDUCE:1, Ops.STORE:1, Ops.END:1, Ops.SINK:1}
+  if Counter(u.op for u in ast.toposort()) != Counter(expected_ops): return None
+  reduce = next((u for u in ast.toposort() if u.op is Ops.REDUCE), None)
+  if reduce is None or reduce.arg != (Ops.ADD, 0) or tuple(int(r.vmax)+1 for r in reduce.src[1:]) != (3, 3): return None
+  return MaxPoolBackwardMatch(params, batch, channels, spatial)
+
+def maxpool_backward_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if (m:=_match_maxpool_backward(ast, renderer.target.device, renderer.target.arch)) is None: return None
+  name = f"direct_maxpool_backward_{m.channels}_{m.spatial}"
+  pooled_spatial, pooled_size, block = m.spatial//2, m.batch*m.channels*(m.spatial//2)**2, min(m.spatial, 16)
+  declarations = "half* p0, half* p1, half* p2, half* p3, half* p4"
+  source = f'''#define half _Float16
+typedef half half2 __attribute__((ext_vector_type(2)));
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({block//2*16},{block//2*16}))) {name}(
+    {declarations}) {{
+  __attribute__((shared, aligned(32))) half tile[{2*block*16}];
+  int gx=__builtin_amdgcn_workgroup_id_x(), b=__builtin_amdgcn_workgroup_id_z();
+  int px=__builtin_amdgcn_workitem_id_x(), lc=__builtin_amdgcn_workitem_id_y(), cg=__builtin_amdgcn_workgroup_id_y();
+  int py=gx/{m.spatial//block}, xb=(gx%{m.spatial//block})*{block}, y=py*2, x=xb+px*2, c=cg*16+lc;
+  int ii=(b*{m.channels}+c)*{m.spatial*m.spatial}+y*{m.spatial}+x;
+  int pi=(b*{m.channels}+c)*{pooled_spatial*pooled_spatial}+py*{pooled_spatial}+xb/2+px;
+  half mx=p2[pi];
+  half scale=(half)1.0/p3[pi]*p4[pi];
+  half2 row0=*((half2*)(p1+ii)), row1=*((half2*)(p1+ii+{m.spatial}));
+  tile[(px*2)*16+lc]=(half)(row0.x==mx)*scale;
+  tile[(px*2+1)*16+lc]=(half)(row0.y==mx)*scale;
+  tile[({block}+px*2)*16+lc]=(half)(row1.x==mx)*scale;
+  tile[({block}+px*2+1)*16+lc]=(half)(row1.y==mx)*scale;
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");
+  int tid=lc*{block//2}+px;
+  #pragma unroll
+  for (int q=0; q<4; q++) {{
+    int e=tid+q*{block//2*16}, pos=e/16, oc=cg*16+e%16;
+    int oy=y+pos/{block}, ox=xb+pos%{block};
+    p0[((b*{m.spatial}+oy)*{m.spatial}+ox)*{m.channels}+oc]=tile[e];
+  }}
+}}'''
+  output_size = pooled_size*4
+  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(output_size*2, pooled_size*22, 0)))
+  slots = tuple(u.arg.slot for u in m.params)
+  info = ProgramInfo(name=name, global_size=(m.spatial//block*pooled_spatial, m.channels//16, m.batch), local_size=(block//2, 16, 1),
+                     globals=slots, outs=(0,), ins=slots[1:])
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)
+
+def _match_maxpool(ast:UOp, device:str, arch:str) -> MaxPoolMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 2 or not isinstance(ast.arg, KernelInfo): return None
+  if any(end.op is not Ops.END or len(end.src) != 5 or end.src[0].op is not Ops.STORE for end in ast.src): return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  if any(not u.src or u.src[0].op is not Ops.CONST or not isinstance(u.src[0].arg, int) for u in params): return None
+  loop_ranges = ast.src[0].src[1:]
+  if any(r.op is not Ops.RANGE or r.arg[1] is not AxisType.LOOP for r in loop_ranges): return None
+  batch, channels, sy, sx = tuple(int(r.vmax)+1 for r in loop_ranges)
+  if batch not in (1024, 1280, 1536) or sy != sx or (channels,sy) not in ((64,16),(256,8),(512,4)) or ast.src[1].src[1:] != loop_ranges: return None
+  spatial = sy
+  out_size = batch*channels*spatial*spatial
+  signature = tuple((u.arg.slot, u.dtype, int(u.src[0].arg)) for u in params)
+  if signature != ((0,dtypes.half,out_size),(1,dtypes.half,out_size*4),(2,dtypes.half,out_size)): return None
+  if tuple(end.src[0].src[0].src[0].arg.slot for end in ast.src) != (0, 2): return None
+  expected_ops = {Ops.CONST:13 if spatial == 8 else 14, Ops.MUL:9, Ops.ADD:9, Ops.RANGE:6, Ops.PARAM:3,
+                  Ops.INDEX:3, Ops.REDUCE:2, Ops.STORE:2, Ops.END:2, Ops.CMPNE:2, Ops.CAST:1, Ops.SINK:1}
+  if Counter(u.op for u in ast.toposort()) != Counter(expected_ops): return None
+  return MaxPoolMatch(params, batch, channels, spatial)
+
+def maxpool_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if (m:=_match_maxpool(ast, renderer.target.device, renderer.target.arch)) is None: return None
+  name = f"direct_maxpool_{m.channels}_{m.spatial}"
+  out_size, in_size = m.batch*m.channels*m.spatial*m.spatial, m.batch*m.channels*m.spatial*m.spatial*4
+  source = f'''#define half _Float16
+typedef half half2 __attribute__((ext_vector_type(2)));
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1,128))) {name}(half* p0, half* p1, half* p2) {{
+  int group=__builtin_amdgcn_workgroup_id_x(), lx=__builtin_amdgcn_workitem_id_x();
+  int ly=__builtin_amdgcn_workitem_id_y(), ii=group*512+lx*{m.spatial*4}+ly*2;
+  half2 row0=*((half2*)(p1+ii)), row1=*((half2*)(p1+ii+{m.spatial*2}));
+  int oi=group*128+lx*{m.spatial}+ly;
+  half mx01=row0.x<row0.y?row0.y:row0.x, mx012=mx01<row1.x?row1.x:mx01;
+  half mx=mx012<row1.y?row1.y:mx012;
+  p0[oi]=mx;
+  p2[oi]=(half)(row0.x==mx)+(half)(row0.y==mx)+(half)(row1.x==mx)+(half)(row1.y==mx);
+}}'''
+  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(out_size*3, in_size*2+out_size*4, 0)))
+  slots = tuple(u.arg.slot for u in m.params)
+  info = ProgramInfo(name=name, global_size=(out_size//128, 1, 1), local_size=(128//m.spatial, m.spatial, 1),
+                     globals=slots, outs=(0, 2), ins=(1,))
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)
+
+def _match_im2col(ast:UOp, device:str, arch:str) -> Im2ColMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1 or not isinstance(ast.arg, KernelInfo): return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 3 or end.src[0].op is not Ops.STORE: return None
+  store, position, patch = end.src
+  if any(r.op is not Ops.RANGE or r.arg[1] is not AxisType.LOOP for r in (position, patch)): return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  if any(not u.src or u.src[0].op is not Ops.CONST or not isinstance(u.src[0].arg, int) for u in params): return None
+  patch_size, position_size = int(patch.vmax)+1, int(position.vmax)+1
+  if patch_size % 9: return None
+  channels = patch_size//9
+  spatial = {32:32, 64:16, 256:8, 512:4}.get(channels)
+  if spatial is None or position_size % (spatial*spatial): return None
+  batch = position_size//(spatial*spatial)
+  if batch not in (1024, 1280, 1536): return None
+  input_size = batch*channels*spatial*spatial
+  signature = tuple((u.arg.slot, u.dtype, int(u.src[0].arg)) for u in params)
+  if signature != ((0,dtypes.half,input_size*9),(1,dtypes.half,input_size)): return None
+  if (int(position.vmax)+1, int(patch.vmax)+1) != (batch*spatial*spatial, channels*9): return None
+  if ssimplify(store.src[0].src[1].get_idx()) is not ssimplify(position*channels*9+patch): return None
+  loads = [u for u in store.src[1].toposort() if u.op is Ops.INDEX and u.src[0].op is Ops.PARAM and u.src[0].arg.slot == 1]
+  expected_idx = ssimplify(position//(spatial*spatial)*(channels*spatial*spatial) + patch//9*(spatial*spatial) +
+                           (position//spatial%spatial+patch//3%3)*spatial + (position%spatial+patch%3)-spatial-1)
+  if len(loads) != 1 or ssimplify(loads[0].src[1].get_idx()) is not expected_idx: return None
+  expected_ops = {Ops.CONST:15 if spatial != 8 else 14, Ops.ADD:7, Ops.MUL:4, Ops.FLOORDIV:4, Ops.FLOORMOD:4,
+                  Ops.CMPLT:4, Ops.AND:3, Ops.PARAM:2, Ops.RANGE:2, Ops.INDEX:2, Ops.CMPNE:2, Ops.WHERE:2,
+                  Ops.STORE:1, Ops.END:1, Ops.SINK:1}
+  if Counter(u.op for u in ast.toposort()) != Counter(expected_ops): return None
+  return Im2ColMatch(params, batch, channels, spatial)
+
+def im2col_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if (m:=_match_im2col(ast, renderer.target.device, renderer.target.arch)) is None: return None
+  name = f"direct_im2col_{m.channels}_{m.spatial}"
+  xblock = min(8, m.spatial)
+  source = f'''#define half _Float16
+typedef half half8 __attribute__((ext_vector_type(8)));
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({xblock*16},{xblock*16}))) {name}(half* p0, half* p1) {{
+  int gx=__builtin_amdgcn_workgroup_id_x(), y=__builtin_amdgcn_workgroup_id_y();
+  int b=__builtin_amdgcn_workgroup_id_z(), lx=__builtin_amdgcn_workitem_id_x();
+  int lc=__builtin_amdgcn_workitem_id_y(), x=(gx%{m.spatial//xblock})*{xblock}+lx;
+  int c=(gx/{m.spatial//xblock})*16+lc;
+  int ii=((b*{m.channels}+c)*{m.spatial}+y)*{m.spatial}+x;
+  int base=((b*{m.spatial*m.spatial}+y*{m.spatial}+x)*{m.channels}+c)*9;
+  bool top=y>0, bottom=y<{m.spatial-1}, left=x>0, right=x<{m.spatial-1};
+  half v0=(top&&left)?p1[ii-{m.spatial+1}]:(half)0, v1=top?p1[ii-{m.spatial}]:(half)0;
+  half v2=(top&&right)?p1[ii-{m.spatial-1}]:(half)0, v3=left?p1[ii-1]:(half)0, v4=p1[ii];
+  half v5=right?p1[ii+1]:(half)0, v6=(bottom&&left)?p1[ii+{m.spatial-1}]:(half)0;
+  half v7=bottom?p1[ii+{m.spatial}]:(half)0, v8=(bottom&&right)?p1[ii+{m.spatial+1}]:(half)0;
+  *((half8*)(p0+base))=(half8){{v0,v1,v2,v3,v4,v5,v6,v7}}; p0[base+8]=v8;
+}}'''
+  total = m.batch*m.channels*m.spatial*m.spatial
+  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(0, total*20, 0)))
+  slots = tuple(u.arg.slot for u in m.params)
+  info = ProgramInfo(name=name, global_size=(m.spatial//xblock*m.channels//16,m.spatial,m.batch), local_size=(xblock,16,1),
+                     globals=slots, outs=(0,), ins=(1,))
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)
+
+def _match_col2im(ast:UOp, device:str, arch:str) -> Col2ImMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1 or not isinstance(ast.arg, KernelInfo): return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 5 or end.src[0].op is not Ops.STORE: return None
+  store, batch, channel, y, x = end.src
+  if any(r.op is not Ops.RANGE or r.arg[1] is not AxisType.LOOP for r in (batch, channel, y, x)): return None
+  batch_size, channels, sy, sx = tuple(int(r.vmax)+1 for r in (batch, channel, y, x))
+  if batch_size not in (1024,1280,1536) or (channels, sy, sx) not in ((64, 16, 16), (256, 8, 8), (512, 4, 4)): return None
+  reduces = [u for u in store.src[1].toposort() if u.op is Ops.REDUCE]
+  if len(reduces) != 1 or reduces[0].arg != (Ops.ADD, 0) or \
+     tuple(int(r.vmax)+1 for r in reduces[0].src[1:]) != (4, 4): return None
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  signature = tuple((u.arg.slot, u.dtype, int(u.src[0].arg)) for u in params)
+  output_size, patch_size = batch_size*channels*sy*sx, batch_size*channels*sy*sx*9
+  plain = ((0, dtypes.half, output_size), (1, dtypes.half, patch_size))
+  fused = ((0, dtypes.float, output_size), (1, dtypes.float, output_size),
+           (2, dtypes.half, output_size), (3, dtypes.half, patch_size))
+  if signature not in (plain, fused): return None
+  expected_ops = ({Ops.CONST:20, Ops.ADD:13, Ops.MUL:11, Ops.AND:8, Ops.RANGE:6, Ops.CMPLT:4, Ops.PARAM:2,
+                   Ops.INDEX:2, Ops.FLOORMOD:2, Ops.FLOORDIV:2, Ops.WHERE:2, Ops.CAST:2, Ops.REDUCE:1,
+                   Ops.STORE:1, Ops.END:1, Ops.SINK:1} if signature == plain else
+                  {Ops.CONST:22 if sy == 4 else 23, Ops.ADD:19, Ops.MUL:18, Ops.AND:8, Ops.RANGE:6, Ops.PARAM:4, Ops.INDEX:4,
+                   Ops.CAST:4, Ops.CMPLT:4, Ops.FLOORMOD:2, Ops.FLOORDIV:2, Ops.WHERE:2, Ops.EXP2:1,
+                   Ops.RECIPROCAL:1, Ops.REDUCE:1, Ops.STORE:1, Ops.END:1, Ops.SINK:1})
+  if Counter(u.op for u in ast.toposort()) != Counter(expected_ops): return None
+  return Col2ImMatch(params, signature == fused, batch_size, channels, sy)
+
+def col2im_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if (m:=_match_col2im(ast, renderer.target.device, renderer.target.arch)) is None: return None
+  name = f"direct_col2im_{m.channels}_{m.spatial}{'_activation' if m.fused_activation else ''}"
+  declarations = ("float* p0, float* p1, half* p2, half* p3" if m.fused_activation else "half* p0, half* p1")
+  patch_slot = 3 if m.fused_activation else 1
+  if m.fused_activation:
+    output = """half grad=p2[oi]+(half)acc, z=(half)p1[oi];
+  half sig=(half)1.0/((half)1.0+__builtin_elementwise_exp2(z*(half)-2.4554669595930156));
+  p0[oi]=(float)(sig*grad+(half)1.702*z*grad*sig*((half)1.0-sig));"""
+  else: output = "p0[oi]=(half)acc;"
+  source = f'''#define half _Float16
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(256,256))) {name}({declarations}) {{
+  int tile=__builtin_amdgcn_workgroup_id_x(), cg=__builtin_amdgcn_workgroup_id_y();
+  int b=__builtin_amdgcn_workgroup_id_z(), c=cg*16+__builtin_amdgcn_workitem_id_x();
+  int y=(tile%{m.spatial//4})*4+__builtin_amdgcn_workitem_id_y();
+  int x=(tile/{m.spatial//4})*4+__builtin_amdgcn_workitem_id_z();
+  float acc=0.0f;
+  #pragma unroll
+  for (int ky=0; ky<3; ky++) {{
+    int oy=y+1-ky;
+    #pragma unroll
+    for (int kx=0; kx<3; kx++) {{
+      int ox=x+1-kx;
+      if (oy>=0 && oy<{m.spatial} && ox>=0 && ox<{m.spatial})
+        acc+=(float)p{patch_slot}[((((b*{m.spatial}+oy)*{m.spatial}+ox)*{m.channels}+c)*9+ky*3+kx)];
+    }}
+  }}
+  int oi=((b*{m.channels}+c)*{m.spatial}+y)*{m.spatial}+x;
+  {output}
+}}'''
+  output_size = m.batch*m.channels*m.spatial*m.spatial
+  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(output_size*18, output_size*20, 0)))
+  slots = tuple(u.arg.slot for u in m.params)
+  info = ProgramInfo(name=name, global_size=((m.spatial//4)**2, m.channels//16, m.batch), local_size=(16, 4, 4),
+                     globals=slots, outs=(0,), ins=slots[1:])
+  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+  return UOp(Ops.PROGRAM, src=src, arg=info)
+
+
+def _match_channel_reduce(ast:UOp, device:str, arch:str) -> ChannelReduceMatch|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1 or not isinstance(ast.arg, KernelInfo): return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 3 or end.src[0].op is not Ops.STORE: return None
+  store, channel, feature = end.src
+  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.LOOP for x in (channel, feature)): return None
+  channels, features = int(channel.vmax)+1, int(feature.vmax)+1
+  if channels not in (64, 256) or features != 256 or store.src[0].op is not Ops.INDEX: return None
+
+  reduce = store.src[1]
+  if reduce.op is not Ops.REDUCE or reduce.arg != (Ops.ADD, 0) or len(reduce.src) != 4: return None
+  batch, ry, rx = reduce.src[1:]
+  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.REDUCE for x in (batch, ry, rx)): return None
+  spatial = int(ry.vmax)+1
+  groups = int(batch.vmax)+1
+  if groups not in (4,5,6) or int(rx.vmax)+1 != spatial or channels*spatial*spatial != 16384: return None
+
+  out_idx = ssimplify(channel*features+feature)
+  data_idx = ssimplify((feature*groups+batch)*16384+channel*spatial*spatial+ry*spatial+rx)
+  if ssimplify(store.src[0].src[1].get_idx()) is not out_idx: return None
+  params = tuple(sorted((x for x in ast.toposort() if x.op is Ops.PARAM), key=lambda x:x.arg.slot))
+  total = features*groups*16384
+  signature = tuple((x.arg.slot, x.dtype, int(x.src[0].arg)) for x in params)
+  signatures = {
+    "activation": ((0, dtypes.float, channels*features), (1, dtypes.half, total), (2, dtypes.float, channels),
+                   (3, dtypes.float, channels), (4, dtypes.float, total), (5, dtypes.half, total)),
+    "centered": ((0, dtypes.float, channels*features), (1, dtypes.half, total), (2, dtypes.float, channels),
+                 (3, dtypes.float, channels), (4, dtypes.float, total)),
+    "scale": ((0, dtypes.float, channels*features), (1, dtypes.float, channels), (2, dtypes.float, channels),
+              (3, dtypes.float, total)),
+  }
+  kind = next((k for k,v in signatures.items() if signature == v), None)
+  if kind is None: return None
+
+  indexes = [(x.src[0].arg.slot, ssimplify(x.src[1].get_idx())) for x in reduce.src[0].toposort() if x.op is Ops.INDEX]
+  expected_indexes = {
+    "activation": ((1, data_idx), (2, channel), (3, channel), (4, data_idx), (5, data_idx)),
+    "centered": ((1, data_idx), (2, channel), (3, channel), (4, data_idx)),
+    "scale": ((1, channel), (2, channel), (3, data_idx)),
+  }[kind]
+  if len(indexes) != len(expected_indexes) or any(slot != eslot or idx is not eidx for (slot,idx),(eslot,eidx) in zip(indexes, expected_indexes)):
+    return None
+  op_signature = Counter(x.op for x in reduce.src[0].toposort())
+  expected_ops = {
+    "activation": {Ops.CONST:13, Ops.PARAM:5, Ops.RANGE:5, Ops.MUL:15, Ops.ADD:11, Ops.INDEX:5, Ops.CAST:3,
+                   Ops.EXP2:1, Ops.RECIPROCAL:1},
+    "centered": {Ops.CONST:10, Ops.PARAM:4, Ops.RANGE:5, Ops.MUL:8, Ops.ADD:6, Ops.INDEX:4, Ops.CAST:1},
+    "scale": {Ops.CONST:10, Ops.PARAM:3, Ops.RANGE:5, Ops.INDEX:3, Ops.RECIPROCAL:1, Ops.SQRT:1, Ops.MUL:8, Ops.ADD:5},
+  }[kind]
+  if op_signature != Counter(expected_ops): return None
+  return ChannelReduceMatch(kind, params, groups, channels, spatial)
+
+def _match_activation_sum(ast:UOp, device:str, arch:str) -> tuple[int, int, int]|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 3 or end.src[0].op is not Ops.STORE: return None
+  store, channel, feature = end.src
+  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.LOOP for x in (channel,feature)): return None
+  channels = int(channel.vmax)+1
+  if channels not in (64,256) or int(feature.vmax)+1 != 256: return None
+  spatial = 16 if channels == 64 else 8
+  reduce = store.src[1]
+  if reduce.op is not Ops.REDUCE or reduce.arg != (Ops.ADD,0) or len(reduce.src) != 4: return None
+  groups, rsy, rsx = tuple(int(x.vmax)+1 for x in reduce.src[1:])
+  if groups not in (4,5,6) or (rsy,rsx) != (spatial,spatial): return None
+  total = 256*groups*16384
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM), key=lambda u:u.arg.slot))
+  if tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,channels*256), (1,dtypes.float,total), (2,dtypes.half,total)): return None
+  batch, y, x = reduce.src[1:]
+  out_idx = ssimplify(channel*256+feature)
+  data_idx = ssimplify((feature*groups+batch)*16384+channel*spatial*spatial+y*spatial+x)
+  if ssimplify(store.src[0].src[1].get_idx()) is not out_idx: return None
+  indexes = [(u.src[0].arg.slot,ssimplify(u.src[1].get_idx())) for u in reduce.src[0].toposort() if u.op is Ops.INDEX]
+  if len(indexes) != 2 or any(slot != expected or idx is not data_idx for (slot,idx),expected in zip(indexes,(1,2))): return None
+  expected_ops = ({Ops.MUL:15,Ops.CONST:13} if channels == 256 else {Ops.MUL:14,Ops.CONST:12}) | \
+    {Ops.ADD:13,Ops.RANGE:5,Ops.PARAM:3,Ops.INDEX:3,Ops.CAST:2,Ops.EXP2:1,Ops.RECIPROCAL:1,
+     Ops.REDUCE:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}
+  if Counter(u.op for u in ast.toposort()) != Counter(expected_ops): return None
+  return channels, spatial, groups
+
+def _match_activation_elementwise(ast:UOp, device:str, arch:str) -> tuple[int, int, int]|None:
+  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
+  end = ast.src[0]
+  if end.op is not Ops.END or len(end.src) != 5 or end.src[0].op is not Ops.STORE: return None
+  store, batch, channel, y, x = end.src
+  if any(r.op is not Ops.RANGE or r.arg[1] is not AxisType.LOOP for r in (batch,channel,y,x)): return None
+  batch_size, channels, sy, sx = tuple(int(r.vmax)+1 for r in (batch,channel,y,x))
+  if batch_size not in (1024,1280,1536) or sy != sx or (channels,sy) not in ((64,16),(256,8)): return None
+  total = batch_size*channels*sy*sx
+  params = tuple(sorted((u for u in ast.toposort() if u.op is Ops.PARAM),key=lambda u:u.arg.slot))
+  if tuple((u.arg.slot,u.dtype,u.max_numel()) for u in params) != \
+     ((0,dtypes.float,total),(1,dtypes.float,channels),(2,dtypes.float,channels),
+      (3,dtypes.float,total),(4,dtypes.half,total)): return None
+  expected = {Ops.MUL:13,Ops.CONST:12,Ops.ADD:9,Ops.PARAM:5,Ops.INDEX:5,Ops.RANGE:4,Ops.RECIPROCAL:2,
+              Ops.CAST:2,Ops.SQRT:1,Ops.EXP2:1,Ops.STORE:1,Ops.END:1,Ops.SINK:1}
+  if Counter(u.op for u in ast.toposort()) != Counter(expected): return None
+  flat = ssimplify(batch*(channels*sy*sx)+channel*(sy*sx)+y*sy+x)
+  return (channels,sy,batch_size//256) if store.src[0].op is Ops.INDEX and ssimplify(store.src[0].src[1].get_idx()) is flat else None
+
+
+def channel_reduce_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
+  if isinstance(am:=ast.tag,DualActivationReduceElementwiseMatch) and renderer.target.device == "AMD" and \
+     renderer.target.arch.startswith("gfx11"):
+    name = f"channel_reduce_activation_sum_{am.groups}_{am.channels}_{am.spatial}_elementwise"
+    source = f'''#define half _Float16
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1024,1024))) {name}(
+    float* p0, float* p1, float* p2, half* p3, float* p4, float* p5, float* p6, float* p7, half* p8) {{
+  int tid=__builtin_amdgcn_workitem_id_x(), lane=tid&31, wave=tid>>5;
+  int oi=__builtin_amdgcn_workgroup_id_x()*32+wave, channel=oi/256, feature=oi%256;
+  float centered=0.0f, plain=0.0f, mean=p4[channel];
+  for (int r=lane; r<{am.groups*am.spatial*am.spatial}; r+=32) {{
+    int base=(feature*{am.groups}+r/{am.spatial*am.spatial})*16384+channel*{am.spatial*am.spatial}+r%{am.spatial*am.spatial};
+    half z=(half)p7[base], grad=p8[base];
+    half sig=(half)1.0/((half)1.0+__builtin_elementwise_exp2(z*(half)-2.4554669595930156));
+    float ag=(float)(sig*grad+(half)1.702*z*grad*sig*((half)1.0-sig));
+    centered+=((float)p3[base]-mean)*ag;
+    plain+=ag;
+    p2[base]=p5[channel]*(__builtin_elementwise_sqrt(1.0f/p6[channel])*ag);
+  }}
+  unsigned centered_bits=__builtin_bit_cast(unsigned,centered), plain_bits=__builtin_bit_cast(unsigned,plain);
+  if (lane==0) {{
+    #pragma unroll
+    for (int i=1; i<32; i++) {{
+      centered+=__builtin_bit_cast(float,__builtin_amdgcn_readlane(centered_bits,i));
+      plain+=__builtin_bit_cast(float,__builtin_amdgcn_readlane(plain_bits,i));
+    }}
+    p0[oi]=centered*p5[channel];
+    p1[oi]=plain;
+  }}
+}}'''
+    elements, reduce_size = am.channels*256, am.groups*am.spatial*am.spatial
+    sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(elements*reduce_size*18, elements*reduce_size*8, 0)))
+    slots = tuple(range(9))
+    info = ProgramInfo(name=name, global_size=(elements//32,1,1), local_size=(1024,1,1), globals=slots,outs=(0,1,2),ins=slots[3:])
+    src:tuple[UOp, ...] = (sink,UOp(Ops.LINEAR),UOp(Ops.SOURCE,arg=source))
+    if compile_binary: src += (UOp(Ops.BINARY,arg=renderer.compiler.compile_cached(source)),)
+    return UOp(Ops.PROGRAM,src=src,arg=info)
+  if isinstance(dm:=ast.tag, DualChannelReduceMatch) and renderer.target.device == "AMD" and renderer.target.arch.startswith("gfx11"):
+    name = f"channel_reduce_centered_scale_{dm.groups}_{dm.channels}_{dm.spatial}"
+    source = f'''#define half _Float16
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1024,1024))) {name}(
+    float* p0, float* p1, half* p2, float* p3, float* p4, float* p5, float* p6, float* p7) {{
+  int tid=__builtin_amdgcn_workitem_id_x(), lane=tid&31, wave=tid>>5;
+  int oi=__builtin_amdgcn_workgroup_id_x()*32+wave, channel=oi/256, feature=oi%256;
+  float centered=0.0f, scale=0.0f, mean=p3[channel];
+  for (int r=lane; r<{dm.groups*dm.spatial*dm.spatial}; r+=32) {{
+    int base=(feature*{dm.groups}+r/{dm.spatial*dm.spatial})*16384+channel*{dm.spatial*dm.spatial}+r%{dm.spatial*dm.spatial};
+    float grad=p7[base];
+    centered+=((float)p2[base]-mean)*grad;
+    scale+=grad;
+  }}
+  unsigned centered_bits=__builtin_bit_cast(unsigned, centered), scale_bits=__builtin_bit_cast(unsigned, scale);
+  if (lane==0) {{
+    #pragma unroll
+    for (int i=1; i<32; i++) {{
+      centered+=__builtin_bit_cast(float, __builtin_amdgcn_readlane(centered_bits, i));
+      scale+=__builtin_bit_cast(float, __builtin_amdgcn_readlane(scale_bits, i));
+    }}
+    p0[oi]=centered*p4[channel];
+    p1[oi]=-(scale*p5[channel]*__builtin_elementwise_sqrt(1.0f/p6[channel]));
+  }}
+}}'''
+    elements, reduce_size = dm.channels*256, dm.groups*dm.spatial*dm.spatial
+    sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(elements*reduce_size*7, elements*reduce_size*8, 0)))
+    info = ProgramInfo(name=name, global_size=(elements//32,1,1), local_size=(1024,1,1), globals=tuple(range(8)),
+                       outs=(0,1), ins=(2,3,4,5,6,7))
+    src = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
+    if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
+    return UOp(Ops.PROGRAM, src=src, arg=info)
+  return None

tinygrad/codegen/simplify.py

@@ -1,7 +1,7 @@
 import itertools
 from typing import Callable
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, range_start, AxisType
+from tinygrad.uop.symbolic import symbolic, symbolic_simple, commutative
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes
 
@@ -16,25 +16,35 @@ pm_flatten_range = PatternMatcher([
   (UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range),
 ])
 
+pm_merge_ranges = symbolic_simple+commutative+PatternMatcher([
+  ((UPat.var("x", dtype=dtypes.index)//UPat.cvar("d"))<UPat.cvar("c"),
+   lambda x,d,c: x<(c.arg*d.arg) if d.arg > 0 else None),
+])+pm_flatten_range
+
 # index/range arithmetic uses FLOORDIV/FLOORMOD prior to late rewrite
 def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.FLOORDIV, Ops.FLOORMOD} for u in x.backward_slice)
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
+  if len(u.ended_ranges) < 2: return None
+  pairs = zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.combinations(u.ended_ranges, 2)
+  candidates = [(r0, r1) for r0,r1 in pairs if r0.arg[-1] == r1.arg[-1]]
+  if not candidates: return None
   reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
+  divmod_count = count_divmod(u)
   # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
-  for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)):
-    # check same type
-    if r0.arg[-1] == r1.arg[-1]:
-      # check if the ranges to merge are in the same reduces
-      if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
-        s0, s1 = r0.src[0], r1.src[0]
-        # do the merge
-        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
-                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
+  for r0, r1 in candidates:
+    # An earlier accepted merge can remove either original range from u, making this substitution a no-op.
+    if r0 not in u.backward_slice or r1 not in u.backward_slice: continue
+    # check if the ranges to merge are in the same reduces
+    if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
+      s0, s1 = r0.src[0], r1.src[0]
+      # do the merge
+      new_range = r0.replace(src=(s0*s1,))
+      nidx = u.substitute({r0:new_range//s1, r1:new_range%s1})
+      nidx = graph_rewrite(nidx, pm_merge_ranges, name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
 
-        # check if it simplifies
-        if count_divmod(nidx) <= count_divmod(u):
-          u = nidx
+      # check if it simplifies
+      if (new_divmod_count:=count_divmod(nidx)) <= divmod_count:
+        u, divmod_count = nidx, new_divmod_count
   return u
 
 def mark_gated(ctx, idx):
@@ -59,15 +69,16 @@ pm_simplify_ranges = PatternMatcher([
 def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
   if r not in ctx and r.arg[-1] is not AxisType.WARP and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
 
-def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
+def do_substitute(ctx:dict, x:UOp, sub_fxn:Callable[[UOp, UOp], UOp], simplify:bool=True) -> UOp|None:
   ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
   ctx.clear()
-  return None if ret is x else ret.simplify()
+  return None if ret is x else ret.simplify() if simplify else ret
 
 pm_split_ranges = PatternMatcher([
   (UPat(Ops.RANGE, name="r")%UPat.cvar("c"), mark_range_mod),
   (UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x,
-    lambda k,v: k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1])))),
+    lambda k,v: k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1])),
+    simplify=False)),
 ])
 
 # **** reduce simplification ****

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/jit.py

@@ -72,7 +72,8 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
   linear = memory_plan_rewrite(linear, held_bufs)
   linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), jit=True)
-  if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
+  amd_graph = len(linear.src) > 128 and bool(input_uops) and all(isinstance(u.device, str) and u.device.split(":")[0] == "AMD" for u in input_uops)
+  if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=getenv("JIT_BATCH_SIZE", 0 if amd_graph else JIT_BATCH_SIZE.value))
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
   return linear
 
@@ -136,6 +137,7 @@ class GraphRunner:
     def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim)
 
     crs = [(j, self.calls[j][1].arg, self.calls[j][3]) for j in range(len(self.calls)) if self.calls[j][1].op is Ops.PROGRAM]
+    self.fixedvars = {v.expr:int(v.vmin) for _,p,_ in crs for v in p.vars if v.vmin == v.vmax}
     self.vars = sorted({v.expr for _,p,dv in crs for v in p.vars if v.expr not in dv | p.runtimevars})
     self.symbolic_dims = dedup(tuple(d) for _,p,_ in crs for d in (p.local_size, p.global_size) if d and is_sym_dim(d))
 
@@ -160,7 +162,7 @@ class GraphRunner:
   def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None: raise NotImplementedError("override this")
 
   def updated_vars(self, var_vals: dict[str, int]):
-    vals = [var_vals[v] for v in self.vars]
+    vals = [(var_vals | self.fixedvars)[v] for v in self.vars]
     for j, vidxs in self.var_vals_replace.items():
       for i, v in vidxs: yield j, i, vals[v]
 
@@ -248,18 +250,19 @@ def _prepare_jit_inputs(args, kwargs):
   return input_buf_uops, var_vals, names, expected_input_info
 
 class TinyJit(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType]|None, captured:CapturedJit|None=None, prune=False):
+  def __init__(self, fxn:Callable[..., ReturnType]|None, captured:CapturedJit|None=None, prune=False, warmup=True):
     assert fxn or captured, "need either a function or a CapturedJit"
     self.fxn = fxn
     self.captured: CapturedJit|None = captured
-    self.cnt: int = 2 if self.fxn is None else 0
+    self.warmup = warmup
+    self.cnt: int = 2 if self.fxn is None else 0 if warmup else 1
     self.prune = prune
 
   def add_linear(self, linear:UOp, var_vals:dict[str, int]): self._linears.append(linear)
 
   def reset(self):
     assert self.fxn is not None, "can't reset without function"
-    self.cnt = 0
+    self.cnt = 0 if self.warmup else 1
     self.captured = None
 
   def __reduce__(self):

tinygrad/engine/realize.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
-import time, random, itertools, math, contextlib, weakref, array
+import time, random, itertools, math, contextlib, weakref, array, re
 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.opt.postrange import bufs_from_ast
@@ -240,11 +240,47 @@ 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),
 ])
 
+batch_compiled_cache: dict[tuple[Compiler, str], tuple[str, str, bytes]] = {}
+def _batch_spec_key(name:str, src:str) -> str:
+  if f" {name}(" not in src: return f"{name}\n{src}"
+  src = src.replace(f" {name}(", " __batch_cached_kernel__(", 1)
+  return re.sub(r"\b(data\d+)_\d+\b|/\* \d+ \*/", lambda x: x[1] or "", src)
+
+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))
+    specs = [(p.arg.function_name, p.src[2].arg) for p in prgs]
+    keys = [_batch_spec_key(*spec) for spec in specs]
+    unique_keys = list(dict.fromkeys(keys))
+    missing_keys = [key for key in unique_keys if (compiler, key) not in batch_compiled_cache]
+    missing_specs = [specs[keys.index(key)] for key in missing_keys]
+    if missing_specs:
+      batch_compiled_cache.update({(compiler, key):compiled for key,compiled in
+                                   zip(missing_keys, compiler.compile_cached_batch(missing_specs))})
+    for prg,key in zip(prgs, keys):
+      name,src,lib = batch_compiled_cache[(compiler, key)]
+      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 +298,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():
+    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/mixin/movement.py

@@ -95,6 +95,10 @@ class MovementMixin:
           *bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size)
           bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound)
           return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False}
+        if step == 1 and isinstance(start, UOp) and isinstance(stop, UOp) and stop.op is Ops.ADD:
+          base, delta = (stop.src[1], stop.src[0]) if stop.src[1] is start else (stop.src[0], stop.src[1])
+          if base is start and delta.op is Ops.CONST and isinstance(delta.arg, int) and delta.arg >= 0:
+            return {"size":delta.arg, "boundary":(start, stop), "stride":1, "collapse_dim":False}
         if resolve(step == 1, False) and resolve((stop-start) >= 0, False):
           return {"size":stop-start, "boundary":(start, stop), "stride":step, "collapse_dim":False}
         raise TypeError(f"slice {index=} is not supported")

tinygrad/renderer/cstyle.py

@@ -466,6 +466,14 @@ class NVCCRenderer(CUDARenderer):
 
 def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype.scalar())
 def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})"
+def _hip_math(op):
+  return lambda x,dtype: _ocml(op)(x, dtype) if dtype.scalar() == dtypes.double else f"__builtin_elementwise_{op}({x})"
+
+def hip_threefry(x:UOp, key:UOp) -> UOp:
+  x2 = x.cast(dtypes.uint).stack((x >> 32).cast(dtypes.uint))
+  key2 = key.cast(dtypes.uint).stack((key >> 32).cast(dtypes.uint))
+  ret = UOp(Ops.THREEFRY, dtypes.uint, src=(x2, key2))
+  return (ret.index(1).cast(dtypes.ulong) << 32) | ret.index(0).cast(dtypes.ulong)
 
 class HIPRenderer(CStyleLanguage):
   shared_max = 65536
@@ -500,10 +508,13 @@ class HIPRenderer(CStyleLanguage):
   # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size
   # NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters
   kernel_typedef = 'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {launch_bounds})))'
-  code_for_workitem = {"g": lambda x: f"__ockl_get_group_id({x})", "l": lambda x: f"__ockl_get_local_id({x})",
-                       "i": lambda x: f"(__ockl_get_group_id({x})*__ockl_get_local_size({x})+__ockl_get_local_id({x}))"}
-  code_for_op = {**CStyleLanguage.code_for_op, Ops.TRUNC: _ocml("trunc"), Ops.SIN: _ocml("sin"),
-                 Ops.LOG2: _ocml("log2"), Ops.EXP2: _ocml("exp2"), Ops.SQRT: _ocml("sqrt")}
+  code_for_workitem = {"g": lambda x: f"__builtin_amdgcn_workgroup_id_{'xyz'[int(x)]}()",
+                       "l": lambda x: f"__builtin_amdgcn_workitem_id_{'xyz'[int(x)]}()",
+                       "i": lambda x: f"(__builtin_amdgcn_workgroup_id_{'xyz'[int(x)]}()*"
+                         f"((unsigned short *)__builtin_amdgcn_dispatch_ptr())[2+{x}]+__builtin_amdgcn_workitem_id_{'xyz'[int(x)]}())"}
+  code_for_op = {**CStyleLanguage.code_for_op, Ops.THREEFRY: lambda x,key,dtype: f"threefry({x},{key})",
+                 Ops.TRUNC: _hip_math("trunc"), Ops.SIN: _hip_math("sin"),
+                 Ops.LOG2: _hip_math("log2"), Ops.EXP2: _hip_math("exp2"), Ops.SQRT: _hip_math("sqrt")}
   smem_prefix = "__attribute__((shared, aligned(16)))"
   smem_prefix_for_cast: bool = False
   barrier = '__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");' + '__builtin_amdgcn_s_barrier();' + \
@@ -511,6 +522,7 @@ class HIPRenderer(CStyleLanguage):
   float4 = "make_float4"
   type_map = {dtypes.bfloat16: "hip_bfloat16", dtypes.fp8e4m3: "hip_fp8", dtypes.fp8e5m2: "hip_bf8"}
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16, *dtypes.fp8s)) + PatternMatcher([
+    (UPat(Ops.THREEFRY, dtype=dtypes.ulong, src=(UPat.var("x"), UPat.var("key"))), hip_threefry),
     (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
       lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
         x.src[2]), arg=(*x.arg,)) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
@@ -528,17 +540,15 @@ class HIPRenderer(CStyleLanguage):
            f"{vec} make_{vec}({', '.join([f'{scal} {x}' for x in _nms[:count]])}) {{ return {{ {', '.join(_nms[:count])} }}; }}"
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
-    prefix, ockl = [], []
+    prefix = []
     type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" }
     used_dtypes = uops_to_dtypes(uops)
     if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops):
       prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"]
-    if any(u.op is Ops.SPECIAL for u in uops):
-      prefix.append("typedef long unsigned int size_t;")
-      ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
+    if any(u.op is Ops.SPECIAL for u in uops): prefix.append("typedef long unsigned int size_t;")
     ocml_ops = {Ops.EXP2: ("exp2", "pure"), Ops.LOG2: ("log2", "pure"), Ops.SQRT: ("sqrt", "const"), Ops.SIN: ("sin", ""), Ops.TRUNC: ("trunc", "")}
     ocml = [(f"__ocml_{ocml_ops[op][0]}_f{dt.bitsize}", dt.name, dt.name, ocml_ops[op][1])
-      for op, dt in dedup((u.op, u.dtype.scalar()) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
+      for op, dt in dedup((u.op, u.dtype.scalar()) for u in uops) if op in ocml_ops and dt == dtypes.double]
     if any(dt == dtypes.bfloat16 for dt, _ in used_dtypes):
       prefix.append(f"typedef {'__bf16' if self.is_cdna4(self.target.arch) else 'unsigned short'} hip_bfloat16;")
     if any(dt == dtypes.half for dt, _ in used_dtypes): prefix.append("#define half _Float16")
@@ -549,8 +559,24 @@ class HIPRenderer(CStyleLanguage):
       prefix.append("""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {
   v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:448.0f,is_bf8?-57344.0f:-448.0f) : v;
   return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}""")
-    prefix += [f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml]
+    prefix += [f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ocml]
     prefix += [self.render_vector_prefix(dt, count) for dt, count in used_dtypes if count > 1]
+    if any(u.op is Ops.THREEFRY for u in uops):
+      prefix.append("""static inline __attribute__((device)) unsigned_int2 threefry(unsigned_int2 x, unsigned_int2 key) {
+  unsigned int ks[3] = {key.y, key.x ^ key.y ^ 0x1BD11BDAu, key.x};
+  unsigned int xr0 = x.x + ks[2], xr1 = x.y + ks[0];
+  const unsigned int rotations[2][4] = {{13, 15, 26, 6}, {17, 29, 16, 24}};
+  for (int i = 0; i < 5; i++) {
+    for (int j = 0; j < 4; j++) {
+      xr0 += xr1;
+      unsigned int r = rotations[i & 1][j];
+      xr1 = xr0 ^ ((xr1 << r) | (xr1 >> (32 - r)));
+    }
+    xr0 += ks[i % 3];
+    xr1 += ks[(i + 1) % 3] + i + 1;
+  }
+  return {xr0, xr1};
+}""")
 
     for name, (N, M, K), dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
       if self.is_cdna(self.target.arch):

tinygrad/runtime/graph/cuda.py

@@ -20,7 +20,7 @@ class CUDAGraph(MultiGraphRunner):
         global_size, local_size = ast.arg.launch_dims({v: 0 for v in self.vars})
 
         c_deps, new_node = self.new_node([b.base for b in bufs], ast.arg.outs)
-        c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, 0) for x in ast.arg.vars])
+        c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, self.fixedvars.get(x.expr, 0)) for x in ast.arg.vars])
         kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(runtime.prg, *global_size, *local_size, runtime.smem,
                                                       ctypes.cast(0, ctypes.POINTER(ctypes.c_void_p)), vargs)
         check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(c_deps or []), ctypes.byref(kern_params)))

tinygrad/runtime/graph/hcq.py

@@ -6,7 +6,6 @@ from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, P
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, Variable
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
-from tinygrad.runtime.ops_rdma import RDMACopyQueue
 
 class HCQGraph(MultiGraphRunner):
   def __init__(self, *args, **kwargs):
@@ -50,7 +49,7 @@ class HCQGraph(MultiGraphRunner):
 
     self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: unwrap(dev.hw_compute_queue_t)() for dev in self.devices}
     self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
-    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], RDMACopyQueue] = {} # lazy allocation, keyed by device pair
+    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], Any] = {} # lazy allocation, keyed by device pair
     self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", min(len(self.devices), 8) if ALL2ALL >= 1 else 1)
     self.num_rdma_ops: dict[tuple[HCQCompiled, HCQCompiled], int] = collections.defaultdict(int)
 
@@ -102,6 +101,7 @@ class HCQGraph(MultiGraphRunner):
       if runtime is not None:
         enqueue_queue = self.comp_queues[enqueue_dev]
       elif is_rdma:
+        from tinygrad.runtime.ops_rdma import RDMACopyQueue
         enqueue_queue = self.comp_queues[enqueue_dev]
         rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
         self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
@@ -272,7 +272,7 @@ class HCQGraph(MultiGraphRunner):
     for dev in self.devices: self.last_timeline[dev][0].wait(self.last_timeline[dev][1])
     if PROFILE and self.kickoff_value > 1: self.collect_timestamps()
 
-    hcq_var_vals = {self.kickoff_var.expr: self.kickoff_value, **var_vals,
+    hcq_var_vals = {self.kickoff_var.expr: self.kickoff_value, **self.fixedvars, **var_vals,
                     **{var.expr: dev.timeline_value - 1 for dev, var in self.virt_timeline_vals.items()},
                     **{sig.base_buf.va_addr.expr: dev.timeline_signal.base_buf.va_addr for dev, sig in self.virt_timeline_signals.items()}}
 

tinygrad/runtime/graph/metal.py

@@ -66,7 +66,7 @@ class MetalGraph(GraphRunner):
     for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
       self.icb.indirectComputeCommandAtIndex(j).concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_dims),
                                                                                                      metal.MTLSize(*local_dims))
-    for i, var in enumerate(self.vars): self.int_buf_view[i] = var_vals[var]
+    for i, var in enumerate(self.vars): self.int_buf_view[i] = (var_vals | self.fixedvars)[var]
 
     command_buffer = self.dev.mtl_queue.commandBuffer().retained()
     encoder = command_buffer.computeCommandEncoder().retained()

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

@@ -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) -> 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()))
@@ -70,16 +70,19 @@ def compile_hip(prg:str, arch="gfx1100", asm=False) -> bytes:
     check(comgr.amd_comgr_set_data_name(data_src, b"<null>"))
     check(comgr.amd_comgr_data_set_add(data_set_src, data_src))
     # -include hiprtc_runtime.h was removed
+    frontend_opt = getenv("AMD_FRONTEND_OPT", 1)
     options = [
-      "-O3", "-mcumode", "--hip-version=6.0.32830", "-DHIP_VERSION_MAJOR=6", "-DHIP_VERSION_MINOR=0", "-DHIP_VERSION_PATCH=32830",
+      f"-O{frontend_opt}", "-mcumode", "--hip-version=6.0.32830", "-DHIP_VERSION_MAJOR=6", "-DHIP_VERSION_MINOR=0", "-DHIP_VERSION_PATCH=32830",
       "-D__HIPCC_RTC__", "-std=c++14", "-nogpuinc", "-Wno-gnu-line-marker", "-Wno-missing-prototypes", f"--offload-arch={arch}",
       "-I/opt/rocm/include", "-Xclang -disable-llvm-passes", "-Xclang -aux-triple", "-Xclang x86_64-unknown-linux-gnu"]
     check(set_options(action_info, ' '.join(options).encode()))
-    status = comgr.amd_comgr_do_action(comgr.AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, action_info, data_set_src, data_set_bc)
+    compile_action = comgr.AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC if use_device_libs else comgr.AMD_COMGR_ACTION_COMPILE_SOURCE_TO_BC
+    status = comgr.amd_comgr_do_action(compile_action, action_info, data_set_src, data_set_bc)
     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""))
@@ -93,11 +96,43 @@ def compile_hip(prg:str, arch="gfx1100", asm=False) -> bytes:
 class HIPCompiler(Compiler):
   def __init__(self, arch:str):
     assert comgr.dll.nm in c.DLL._loaded_, f"comgr not available: {comgr.dll.emsg}"
-    self.arch = arch
-    super().__init__(f"compile_hip_{self.arch}")
+    self.arch, self.frontend_opt, self.generic_opt = arch, getenv("AMD_FRONTEND_OPT", 1), getenv("AMD_GENERIC_COMPILE_OPT", 1)
+    super().__init__(f"compile_hip_{self.arch}_F{self.frontend_opt}G{self.generic_opt}")
   def compile(self, src:str) -> bytes:
-    try: return compile_hip(src, self.arch, src.split('\n', 1)[0].strip() == '.text')
+    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)
+    batch_opt = getenv("AMD_COMPILE_OPT", 2)
+    generic_opt = getenv("AMD_GENERIC_COMPILE_OPT", 1)
+    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)
+    groups:dict[tuple[int, bool], list[int]] = {}
+    opts = [batch_opt if name.startswith("coop_") else generic_opt for name,_ in renamed]
+    for i,(_,src) in enumerate(renamed): groups.setdefault((opts[i], "__ocml_" in src), []).append(i)
+    for (opt,use_device_libs),indices in groups.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))
+        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
+        for i in batch:
+          name,src = renamed[i]
+          ret[i] = (name, src, lib)
+    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/__init__.py

@@ -1,6 +1,8 @@
 import time, inspect
-from collections import deque
-from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
+from collections import Counter, deque
+from dataclasses import replace
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo, ssimplify
 from tinygrad.uop.spec import type_verify, spec_tensor
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition
 
@@ -26,6 +28,363 @@ def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
     raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}")
   return tuple(kernels), tuple(deps)
 
+def _kernel_io(call:UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]|None:
+  if call.op is not Ops.CALL or call.src[0].op is not Ops.SINK: return None
+  out_slots = {x.src[0].src[0].arg.slot for x in call.src[0].toposort()
+               if x.op is Ops.STORE and x.src[0].op is Ops.INDEX and x.src[0].src[0].op is Ops.PARAM}
+  if len(out_slots) != 1 or min(out_slots) < 0 or max(out_slots) >= len(call.src)-1: return None
+  return tuple(x for i,x in enumerate(call.src[1:]) if i in out_slots), tuple(x for i,x in enumerate(call.src[1:]) if i not in out_slots)
+
+def _input_indices(call:UOp, inputs:tuple[UOp, ...]) -> set[tuple[UOp, tuple[UOp, ...]]]:
+  return {(call.src[x.src[0].arg.slot+1], x.src[1:]) for x in call.src[0].toposort()
+          if x.op is Ops.INDEX and x.src[0].op is Ops.PARAM and call.src[x.src[0].arg.slot+1] in inputs}
+
+def _is_sum_sumsq(a:UOp, b:UOp) -> bool:
+  av, bv = a.src[0], b.src[0]
+  return (bv.op is Ops.MUL and bv.src == (av, av)) or (av.op is Ops.MUL and av.src == (bv, bv))
+
+def _substitute_uops(ast:UOp, mapping:dict[UOp, UOp]) -> UOp:
+  memo: dict[UOp, UOp] = {}
+  def rec(x:UOp) -> UOp:
+    if x in mapping: return mapping[x]
+    if x not in memo:
+      src = tuple(rec(s) for s in x.src)
+      memo[x] = x if src == x.src else x.replace(src=src)
+    return memo[x]
+  return rec(ast)
+
+def _remap_params(ast:UOp, remap:dict[int, int]) -> UOp:
+  return _substitute_uops(ast, {x:x.replace(arg=replace(x.arg, slot=remap[x.arg.slot])) for x in ast.toposort()
+                                if x.op is Ops.PARAM and x.arg.slot >= 0})
+
+def _fuse_adjacent_reductions(linearized:list[UOp]) -> list[UOp]:
+  ret: list[UOp] = []
+  i = 0
+  while i < len(linearized):
+    if i+1 == len(linearized):
+      ret.append(linearized[i])
+      break
+    a, b = linearized[i:i+2]
+    aio, bio = _kernel_io(a), _kernel_io(b)
+    ar = [x for x in a.src[0].toposort() if x.op is Ops.REDUCE] if aio is not None else []
+    br = [x for x in b.src[0].toposort() if x.op is Ops.REDUCE] if bio is not None else []
+    compatible = aio is not None and bio is not None and len(ar) == len(br) == 1 and ar[0].arg == br[0].arg and \
+      ar[0].src[1:] == br[0].src[1:] and _is_sum_sumsq(ar[0], br[0]) and aio[1] == bio[1] and \
+      _input_indices(a, aio[1]) == _input_indices(b, bio[1]) and aio[0] != bio[0] and \
+      len(a.src[0].src) == len(b.src[0].src) == 1 and a.src[0].src[0].op is Ops.END and b.src[0].src[0].op is Ops.END and \
+      a.src[0].src[0].src[1:] == b.src[0].src[0].src[1:]
+    if not compatible:
+      ret.append(a)
+      i += 1
+      continue
+    assert aio is not None and bio is not None
+    args = list(a.src[1:])
+    if len(set(args)) != len(args) or any(x in args for x in bio[0]):
+      ret.append(a)
+      i += 1
+      continue
+    args.extend(bio[0])
+    remap = {slot:args.index(x) for slot,x in enumerate(b.src[1:])}
+    bast = _remap_params(b.src[0], remap)
+    ret.append(a.src[0].replace(src=a.src[0].src+bast.src).call(*args))
+    i += 2
+  return ret
+
+def _fuse_dependent_reductions(linearized:list[UOp]) -> list[UOp]:
+  ret = list(linearized)
+  i = 0
+  while i < len(ret):
+    a, aio = ret[i], _kernel_io(ret[i])
+    if aio is None or len(a.src[0].src) != 1 or a.src[0].src[0].op is not Ops.END:
+      i += 1
+      continue
+    ar = [x for x in a.src[0].toposort() if x.op is Ops.REDUCE]
+    if len(ar) != 1 or ar[0].arg != (Ops.MAX, 0) or len(ar[0].src[1:]) != 2 or any(int(x.vmax)+1 != 2 for x in ar[0].src[1:]):
+      i += 1
+      continue
+    aend, astore = a.src[0].src[0], a.src[0].src[0].src[0]
+    if astore.op is not Ops.STORE:
+      i += 1
+      continue
+    spatial = tuple(int(x.vmax)+1 for x in aend.src[-2:])
+    if len(aend.src[1:]) != 4 or spatial[0] != spatial[1] or spatial[0] & (spatial[0]-1):
+      i += 1
+      continue
+    for j in range(i+1, min(i+6, len(ret))):
+      b, bio = ret[j], _kernel_io(ret[j])
+      if bio is None or len(b.src[0].src) != 1 or b.src[0].src[0].op is not Ops.END: continue
+      br = [x for x in b.src[0].toposort() if x.op is Ops.REDUCE]
+      if len(br) != 1 or br[0].arg != (Ops.ADD, 0) or ar[0].src[1:] != br[0].src[1:]: continue
+      bend = b.src[0].src[0]
+      if aend.src[1:] != bend.src[1:] or aio[0][0] not in bio[1] or set(bio[1]) != set(aio[1]+aio[0]): continue
+      args = list(a.src[1:])
+      if len(set(args)) != len(args) or bio[0][0] in args: continue
+      args.extend(bio[0])
+      pout_slot = tuple(b.src[1:]).index(aio[0][0])
+      pout_idxs = [x for x in b.src[0].toposort() if x.op is Ops.INDEX and x.src[0].op is Ops.PARAM and x.src[0].arg.slot == pout_slot]
+      if not pout_idxs or any(x.src[1:] != astore.src[0].src[1:] for x in pout_idxs): continue
+      mapping = {x:astore.src[1] for x in pout_idxs}
+      for x in b.src[0].toposort():
+        if x.op is Ops.PARAM and x.arg.slot >= 0 and x.arg.slot != pout_slot:
+          mapping[x] = x.replace(arg=replace(x.arg, slot=args.index(b.src[x.arg.slot+1])))
+      bend = _substitute_uops(bend, mapping)
+      ret[i] = a.src[0].replace(src=(aend, bend)).call(*args)
+      ret.pop(j)
+      break
+    i += 1
+  return ret
+
+def _fuse_direct_conv_bwd_activation(linearized:list[UOp]) -> list[UOp]:
+  from tinygrad.codegen.opt.gemm import DirectConvBwdActivationMatch, _match_gemm
+  cases = {
+    (288,64): (32, False,
+      {Ops.CONST:22, Ops.MUL:18, Ops.ADD:18, Ops.AND:8, Ops.RANGE:6, Ops.CMPLT:4, Ops.PARAM:3, Ops.INDEX:3,
+       Ops.CAST:3, Ops.FLOORMOD:2, Ops.FLOORDIV:2, Ops.WHERE:2, Ops.EXP2:1, Ops.RECIPROCAL:1, Ops.REDUCE:1,
+       Ops.STORE:1, Ops.END:1, Ops.SINK:1}),
+    (576,64): (16, True,
+      {Ops.CONST:23, Ops.ADD:19, Ops.MUL:18, Ops.AND:8, Ops.RANGE:6, Ops.PARAM:4, Ops.INDEX:4, Ops.CAST:4,
+       Ops.CMPLT:4, Ops.FLOORMOD:2, Ops.FLOORDIV:2, Ops.WHERE:2, Ops.EXP2:1, Ops.RECIPROCAL:1, Ops.REDUCE:1,
+       Ops.STORE:1, Ops.END:1, Ops.SINK:1}),
+    (2304,256): (8, True,
+      {Ops.CONST:23, Ops.ADD:19, Ops.MUL:18, Ops.AND:8, Ops.RANGE:6, Ops.PARAM:4, Ops.INDEX:4, Ops.CAST:4,
+       Ops.CMPLT:4, Ops.FLOORMOD:2, Ops.FLOORDIV:2, Ops.WHERE:2, Ops.EXP2:1, Ops.RECIPROCAL:1, Ops.REDUCE:1,
+       Ops.STORE:1, Ops.END:1, Ops.SINK:1}),
+  }
+  ret = list(linearized)
+  i = 0
+  while i < len(ret):
+    call = ret[i]
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.SINK or \
+       (g:=_match_gemm(call.src[0], "AMD", "gfx11")) is None or \
+       (g.n, g.k) not in cases or g.old is not None or g.a_kxm or not g.b_kxn:
+      i += 1
+      continue
+    patch_grad = call.src[g.c.arg.slot+1]
+    consumers = [j for j,x in enumerate(ret) if j > i and patch_grad in x.src[1:]]
+    if len(consumers) != 1 or consumers[0] > i+3:
+      i += 1
+      continue
+    j, consumer = consumers[0], ret[consumers[0]]
+    io = _kernel_io(consumer)
+    params = sorted((x.arg.slot, x.dtype, x.max_numel()) for x in consumer.src[0].toposort() if x.op is Ops.PARAM)
+    op_counts = Counter(x.op for x in consumer.src[0].toposort())
+    spatial, residual, expected_ops = cases[(g.n,g.k)]
+    if g.m % (spatial*spatial):
+      i += 1
+      continue
+    batch, output_size = g.m//(spatial*spatial), g.m*(g.n//9)
+    if batch not in (1024,1280,1536):
+      i += 1
+      continue
+    if batch == 1024 and (g.n,g.k) == (288,64): expected_ops = expected_ops | {Ops.CONST:expected_ops[Ops.CONST]-1}
+    expected_params = ([(0,dtypes.float,output_size),(1,dtypes.float,output_size),(2,dtypes.half,output_size),
+                        (3,dtypes.half,g.m*g.n)] if residual else
+                       [(0,dtypes.float,output_size),(1,dtypes.half,output_size),(2,dtypes.half,g.m*g.n)])
+    expected_srcs = 5 if residual else 4
+    if io is None or len(io[0]) != 1 or len(io[1]) != expected_srcs-2 or len(consumer.src) != expected_srcs or \
+       consumer.src[-1] is not patch_grad or params != expected_params or op_counts != Counter(expected_ops):
+      i += 1
+      continue
+    reduces = [x for x in consumer.src[0].toposort() if x.op is Ops.REDUCE]
+    if len(reduces) != 1 or reduces[0].arg != (Ops.ADD, 0) or tuple(int(x.vmax)+1 for x in reduces[0].src[1:]) != (4,4):
+      i += 1
+      continue
+    info = DirectConvBwdActivationMatch(g.m, g.n//9, g.k, spatial, residual)
+    ast = call.src[0].replace(tag=info, arg=replace(call.src[0].arg, name="direct_conv_bwd_activation"))
+    extras = (consumer.src[2], consumer.src[3]) if residual else (consumer.src[2],)
+    ret[j] = ast.call(io[0][0], *extras, call.src[g.a.arg.slot+1], call.src[g.b.arg.slot+1])
+    ret.pop(i)
+  return ret
+
+def _fuse_gemm_output_transpose(linearized:list[UOp]) -> list[UOp]:
+  from tinygrad.codegen.opt.gemm import GemmOutputNCHW, _match_gemm
+  from tinygrad.device import Device
+  cases = {
+    (64,288): (32, 7),
+    (256,576): (16, 6),
+    (512,2304): (8, 7),
+  }
+  ret = list(linearized)
+  i = 0
+  while i < len(ret):
+    call = ret[i]
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.SINK or not isinstance(call.device, str):
+      i += 1
+      continue
+    target = Device[call.device].renderer.target
+    if target.device != "AMD" or not target.arch.startswith("gfx11") or \
+       (g:=_match_gemm(call.src[0], target.device, target.arch)) is None or (g.n,g.k) not in cases or \
+       g.old is not None or g.a_kxm or g.b_kxn:
+      i += 1
+      continue
+    gemm_out = call.src[g.c.arg.slot+1]
+    consumers = [j for j,x in enumerate(ret) if j > i and gemm_out in x.src[1:]]
+    if consumers != [i+1]:
+      i += 1
+      continue
+    j, consumer = consumers[0], ret[consumers[0]]
+    io, spatial_const = _kernel_io(consumer), cases[(g.n,g.k)]
+    spatial, const_count = spatial_const
+    op_counts = Counter(x.op for x in consumer.src[0].toposort())
+    expected_ops = {Ops.CONST:const_count, Ops.ADD:6, Ops.MUL:5, Ops.RANGE:4, Ops.PARAM:2, Ops.INDEX:2,
+                    Ops.STORE:1, Ops.END:1, Ops.SINK:1}
+    if io is None or len(io[0]) != 1 or len(io[1]) != 1 or len(consumer.src) != 3 or io[1][0] is not gemm_out or \
+       len(consumer.src[0].src) != 1 or op_counts != Counter(expected_ops):
+      i += 1
+      continue
+    end = consumer.src[0].src[0]
+    end_shape = tuple(int(x.vmax)+1 for x in end.src[1:])
+    if end.op is not Ops.END or len(end_shape) != 4 or end_shape[0] not in (1024,1280,1536) or end_shape != (end_shape[0],g.n,spatial,spatial):
+      i += 1
+      continue
+    store, (batch, channel, y, x) = end.src[0], end.src[1:]
+    if store.op is not Ops.STORE or store.src[0].op is not Ops.INDEX or store.src[1].op is not Ops.INDEX:
+      i += 1
+      continue
+    out_idx, in_idx = (ssimplify(u.src[1].get_idx()) for u in (store.src[0], store.src[1]))
+    spatial_size = spatial**2
+    if out_idx is not ssimplify(batch*(g.n*spatial_size)+channel*spatial_size+y*spatial+x) or \
+       in_idx is not ssimplify(batch*(spatial_size*g.n)+y*(spatial*g.n)+x*g.n+channel):
+      i += 1
+      continue
+    args = list(call.src[1:])
+    args[g.c.arg.slot] = io[0][0]
+    ret[i] = call.src[0].replace(tag=GemmOutputNCHW(spatial)).call(*args)
+    ret.pop(j)
+    i += 1
+  return ret
+
+def _fuse_channel_reductions(linearized:list[UOp]) -> list[UOp]:
+  from tinygrad.codegen.opt.reduce import ActivationVarGradElementwiseSumDMeanMatch, DualActivationReduceElementwiseMatch
+  from tinygrad.codegen.opt.reduce import DualBNGradDMean512Match, DualChannelReduceMatch, DualMoments512Match
+  from tinygrad.codegen.opt.reduce import _match_activation_elementwise, _match_activation_elementwise_512, _match_activation_sum
+  from tinygrad.codegen.opt.reduce import _match_activation_sum_512
+  from tinygrad.codegen.opt.reduce import _match_activation_dmean_512, _match_activation_var_grad, _match_channel_reduce
+  from tinygrad.codegen.opt.reduce import _match_moment_mean_512, _match_moment_var_512
+  from tinygrad.codegen.opt.reduce import _match_bn_dmean_512, _match_bn_sum_512, _match_bn_var_512
+  from tinygrad.device import Device
+  ret = list(linearized)
+  i = 0
+  while i+1 < len(ret):
+    a, b = ret[i:i+2]
+    if a.op is not Ops.CALL or b.op is not Ops.CALL or a.src[0].op is not Ops.SINK or b.src[0].op is not Ops.SINK or \
+       not isinstance(a.device,str) or a.device != b.device:
+      i += 1
+      continue
+    target = Device[a.device].renderer.target
+    mean_batch, var_batch = _match_moment_mean_512(a.src[0],target.device,target.arch), _match_moment_var_512(b.src[0],target.device,target.arch)
+    if mean_batch is None or mean_batch != var_batch or \
+       a.src[1] is not b.src[3] or a.src[2] is not b.src[2]:
+      i += 1
+      continue
+    ret[i] = a.src[0].replace(tag=DualMoments512Match(mean_batch)).call(a.src[1],b.src[1],a.src[2])
+    ret.pop(i+1)
+    i += 1
+  i = 0
+  while i+1 < len(ret):
+    a, b = ret[i:i+2]
+    if a.op is not Ops.CALL or b.op is not Ops.CALL or a.src[0].op is not Ops.SINK or b.src[0].op is not Ops.SINK or \
+       not isinstance(a.device,str) or a.device != b.device:
+      i += 1
+      continue
+    target = Device[a.device].renderer.target
+    var_batch, sum_batch = _match_bn_var_512(a.src[0],target.device,target.arch), _match_bn_sum_512(b.src[0],target.device,target.arch)
+    if var_batch is None or var_batch != sum_batch or \
+       a.src[6] is not b.src[2]:
+      i += 1
+      continue
+    dmean_idx = next((j for j in range(i+2,min(i+6,len(ret))) if ret[j].op is Ops.CALL and ret[j].src[0].op is Ops.SINK and
+      _match_bn_dmean_512(ret[j].src[0],target.device,target.arch) == var_batch and ret[j].src[2] is a.src[1] and
+      ret[j].src[3] is a.src[4] and ret[j].src[4] is a.src[5] and ret[j].src[5] is a.src[2] and ret[j].src[6] is a.src[6]),None)
+    if dmean_idx is None:
+      i += 1
+      continue
+    d = ret[dmean_idx]
+    ret[i] = a.src[0].replace(tag=DualBNGradDMean512Match(var_batch)).call(
+      a.src[1],b.src[1],d.src[1],a.src[2],a.src[3],a.src[4],a.src[5],a.src[6],b.src[3])
+    ret.pop(dmean_idx)
+    ret.pop(i+1)
+    i += 1
+  i = 0
+  while i+1 < len(ret):
+    a, b = ret[i:i+2]
+    if a.op is not Ops.CALL or b.op is not Ops.CALL or a.src[0].op is not Ops.SINK or b.src[0].op is not Ops.SINK or \
+       not isinstance(a.device,str) or a.device != b.device:
+      i += 1
+      continue
+    target = Device[a.device].renderer.target
+    var_batch, elem_batch = _match_activation_var_grad(a.src[0],target.device,target.arch), \
+                            _match_activation_elementwise_512(b.src[0],target.device,target.arch)
+    if var_batch is None or var_batch != elem_batch or \
+       a.src[5] is not b.src[2] or a.src[2] is not b.src[3] or a.src[6] is not b.src[4] or a.src[7] is not b.src[5]:
+      i += 1
+      continue
+    c = ret[i+2] if i+2 < len(ret) else None
+    fuse_sum = c is not None and c.op is Ops.CALL and c.src[0].op is Ops.SINK and \
+      _match_activation_sum_512(c.src[0],target.device,target.arch) == var_batch and a.src[6] is c.src[2] and a.src[7] is c.src[3]
+    if not fuse_sum:
+      i += 1
+      continue
+    assert c is not None
+    dmean_idx = next((j for j in range(i+3,min(i+7,len(ret))) if ret[j].op is Ops.CALL and ret[j].src[0].op is Ops.SINK and
+      _match_activation_dmean_512(ret[j].src[0],target.device,target.arch) == var_batch and ret[j].src[2] is a.src[1] and
+      ret[j].src[3] is a.src[4] and ret[j].src[4] is b.src[1]),None)
+    if dmean_idx is None:
+      i += 1
+      continue
+    d = ret[dmean_idx]
+    ret[i] = a.src[0].replace(tag=ActivationVarGradElementwiseSumDMeanMatch(var_batch)).call(
+      a.src[1],b.src[1],c.src[1],d.src[1],a.src[2],a.src[3],a.src[4],a.src[5],a.src[6],a.src[7],c.src[4])
+    ret.pop(dmean_idx)
+    ret.pop(i+2)
+    ret.pop(i+1)
+    i += 1
+  i = 0
+  while i < len(ret):
+    a = ret[i]
+    if a.op is not Ops.CALL or a.src[0].op is not Ops.SINK or not isinstance(a.device,str):
+      i += 1
+      continue
+    target = Device[a.device].renderer.target
+    am = _match_channel_reduce(a.src[0],target.device,target.arch)
+    if am is None or am.kind != "activation":
+      i += 1
+      continue
+    for j in range(i+1,min(i+4,len(ret))):
+      b = ret[j]
+      sm = _match_activation_sum(b.src[0],target.device,target.arch) if b.op is Ops.CALL and b.src[0].op is Ops.SINK else None
+      if sm != (am.channels,am.spatial,am.groups) or a.src[5] is not b.src[2] or a.src[6] is not b.src[3]: continue
+      bout = b.src[1]
+      if any(bout in x.src[1:] for x in ret[i+1:j]): continue
+      elem_idx = next((k for k in range(i+1,min(i+5,len(ret))) if k != j and ret[k].op is Ops.CALL and ret[k].src[0].op is Ops.SINK and
+        _match_activation_elementwise(ret[k].src[0],target.device,target.arch) == (am.channels,am.spatial,am.groups) and
+        a.src[4] is ret[k].src[2] and a.src[5] is ret[k].src[4] and a.src[6] is ret[k].src[5]),None)
+      if elem_idx is None: continue
+      elem = ret[elem_idx]
+      ret[i] = a.src[0].replace(tag=DualActivationReduceElementwiseMatch(am.groups,am.channels,am.spatial)).call(
+        a.src[1],bout,elem.src[1],a.src[2],a.src[3],a.src[4],elem.src[3],a.src[5],a.src[6])
+      for k in sorted((j,elem_idx),reverse=True): ret.pop(k)
+      break
+    i += 1
+  i = 0
+  while i+1 < len(ret):
+    a, b = ret[i:i+2]
+    if a.op is not Ops.CALL or b.op is not Ops.CALL or a.src[0].op is not Ops.SINK or b.src[0].op is not Ops.SINK or \
+       not isinstance(a.device, str) or a.device != b.device:
+      i += 1
+      continue
+    target = Device[a.device].renderer.target
+    am, bm = _match_channel_reduce(a.src[0], target.device, target.arch), _match_channel_reduce(b.src[0], target.device, target.arch)
+    if am is None or bm is None or am.kind != "centered" or bm.kind != "scale" or \
+       (am.groups,am.channels,am.spatial) != (bm.groups,bm.channels,bm.spatial) or a.src[5] is not b.src[4]:
+      i += 1
+      continue
+    ret[i] = a.src[0].replace(tag=DualChannelReduceMatch(am.groups, am.channels, am.spatial)).call(
+      a.src[1], b.src[1], a.src[2], a.src[3], a.src[4], b.src[2], b.src[3], a.src[5])
+    ret.pop(i+1)
+    i += 1
+  return ret
+
 def create_schedule(sched_sink:UOp) -> UOp:
   with cpu_profile(TracingKey("toposort sched_sink")):
     # build kernel dependency graph: edges from producer kernel to consumer kernels
@@ -77,7 +436,8 @@ def create_schedule(sched_sink:UOp) -> UOp:
         in_degree[x] -= 1
         if in_degree[x] == 0: queue.append(x)
     if any(in_degree.values()): raise RuntimeError("cycle detected in assign graph")
-  return UOp(Ops.LINEAR, src=tuple(linearized))
+  return UOp(Ops.LINEAR, src=tuple(_fuse_gemm_output_transpose(_fuse_channel_reductions(
+    _fuse_direct_conv_bwd_activation(_fuse_dependent_reductions(_fuse_adjacent_reductions(linearized)))))))
 
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
@@ -109,7 +469,7 @@ schedule_cache: dict[bytes, UOp] = {}
 def lower_sink_to_linear(function:UOp) -> UOp|None:
   st = time.perf_counter()
   if isinstance(function.arg, KernelInfo): return None
-  cache_key = function.key
+  cache_key = function.key if SCACHE else b""
   if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
     if SPEC: type_verify(function, spec_tensor)
     # support recursive CALLs

tinygrad/schedule/indexing.py

@@ -1,11 +1,11 @@
-from typing import Iterator
+from typing import Iterator, cast
 import functools, itertools
 from dataclasses import dataclass, field, replace
 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,
@@ -135,6 +135,17 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
     axes_in.append(acc*src)
     acc *= s
   combined_axes = UOp.const(dtypes.index, 0).usum(axes_in)
+  # Devectorization indexes reshapes with concrete lane numbers. Avoid running the
+  # full symbolic-validity pipeline for each lane when the flat index is constant.
+  if all(isinstance(s, int) and s > 0 for s in in_shape) and combined_axes.vmin == combined_axes.vmax and isinstance(combined_axes.vmin, int):
+    flat_idx = int(combined_axes.vmin)
+    const_axes = []
+    for s in cast(tuple[int, ...], in_shape)[::-1]:
+      const_axes.append(UOp.const(dtypes.index, flat_idx % s))
+      flat_idx //= s
+    return UOp.sink(*const_axes[::-1])
+  if len(in_shape) == 1:
+    return graph_rewrite(UOp.sink(combined_axes), symbolic+pm_simplify_valid+pm_drop_and_clauses, name="reshape")
   axes_out:list[UOp] = []
   for s in in_shape[::-1]:
     axes_out.append(combined_axes % s)
@@ -142,6 +153,35 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U
   # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code
   return graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid+pm_drop_and_clauses, name="reshape")
 
+def _apply_int_reshape(in_shape:tuple[int, ...], out_shape:tuple[int, ...], rngs:tuple[UOp, ...]) -> tuple[UOp, ...]|None:
+  """Index pure contiguous splits/merges without flattening and symbolically rebuilding the entire shape."""
+  ret:list[UOp] = []
+  i = j = 0
+  while i < len(in_shape) or j < len(out_shape):
+    while i < len(in_shape) and in_shape[i] == 1:
+      ret.append(UOp.const(dtypes.index, 0))
+      i += 1
+    while j < len(out_shape) and out_shape[j] == 1: j += 1
+    if i == len(in_shape) or j == len(out_shape): break
+    ni, nj, in_sz, out_sz = i+1, j+1, in_shape[i], out_shape[j]
+    while in_sz != out_sz:
+      if in_sz < out_sz and ni < len(in_shape): in_sz, ni = in_sz*in_shape[ni], ni+1
+      elif out_sz < in_sz and nj < len(out_shape): out_sz, nj = out_sz*out_shape[nj], nj+1
+      else: return None
+    in_group, out_group, out_idxs = in_shape[i:ni], out_shape[j:nj], rngs[j:nj]
+    if in_group == out_group: ret.extend(out_idxs)
+    elif len(in_group) == 1:
+      ret.append(UOp.const(dtypes.index, 0).usum([r*prod(out_group[k+1:]) for k,r in enumerate(out_idxs)]))
+    elif len(out_group) == 1:
+      flat, rev = out_idxs[0], []
+      for sz in in_group[:0:-1]:
+        rev.append(flat % sz)
+        flat //= sz
+      ret.extend((flat, *rev[::-1]))
+    else: return None
+    i, j = ni, nj
+  return tuple(ret) if i == len(in_shape) and j == len(out_shape) else None
+
 # this is the definition of the movement ops
 @functools.cache
 def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
@@ -156,6 +196,20 @@ 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)]),)
+      if all(isinstance(s, int) and s > 0 for s in in_shape+arg) and \
+         (reshaped:=_apply_int_reshape(cast(tuple[int, ...], in_shape), cast(tuple[int, ...], arg), rngs)) is not None: return reshaped
       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
@@ -286,8 +340,6 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   # NOTE: SPEC=3 is broken here with shape
   with Context(SPEC=min(SPEC.value, 2)):
     tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
-  # if a deviceless value must materialize, place it on the sink device
-  tsink = graph_rewrite(tsink, pm_fix_deviceless, ctx=tsink.device, name="add device to deviceless")
   return tsink, rctx
 
 def render_ranges(*rngs_list, realized) -> str:

tinygrad/schedule/rangeify.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLI
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
-from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
+from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op, pm_fix_deviceless
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.allreduce import create_allreduce_function
 
@@ -60,9 +60,10 @@ pm_mops = PatternMatcher([
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
 def fix_store_hazard(target:UOp, src:UOp):
+  base = target.base
+  if src is not base and base not in src.backward_slice: return None
   # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
   unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
-  base = target.base
   reaches_base: dict[UOp, bool] = {}
   for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS):
     reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src)
@@ -540,7 +541,9 @@ def get_kernel_graph(sink:UOp) -> UOp:
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, name="symbolic+reduce_collapse+debuf")
+  # Device attachment is valid only after rangeify, but can share this fixed-point traversal.
+  tsink = graph_rewrite(tsink, pm_fix_deviceless+symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize,
+                        ctx=tsink.device, name="symbolic+reduce_collapse+debuf")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")

tinygrad/uop/ops.py

@@ -193,16 +193,25 @@ buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDict
 all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?
 
 # recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
-class recursive_property(property):
+class recursive_property:
   def __init__(self, fxn):
     self.fxn = fxn
-    self.nm = "_RECURSIVE_PROPERTY_"+fxn.__name__
+    self.nm = fxn.__name__
     self.__doc__ = fxn.__doc__
   def __get__(self, x:UOp|None, owner=None):
     if x is None: return self
-    if self.nm in x.__dict__: return x.__dict__[self.nm]
-    for node in x.toposort(gate=lambda node: self.nm not in node.__dict__): node.__dict__[self.nm] = self.fxn(node)
-    return x.__dict__[self.nm]
+    if (nm:=self.nm) in x.__dict__: return x.__dict__[nm]
+    stack = [(x, iter(x.src))]
+    while stack:
+      n, children = stack[-1]
+      for child in children:
+        if nm in child.__dict__: continue
+        stack.append((child, iter(child.src)))
+        break
+      else:
+        n.__dict__[nm] = self.fxn(n)
+        stack.pop()
+    return x.__dict__[nm]
 
 # we import this late so we can use resolve/smax in mixins
 from tinygrad.mixin.op import OpMixin
@@ -217,6 +226,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   arg:Any = None
   tag:Any = None
   def __del__(self):
+    if getattr(sys, "is_finalizing", lambda: True)(): return
     if Ops is not None and self.op is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
     try: del UOpMetaClass.ucache[(self.op, self.dtype, self.src, self.arg, self.tag)]
     except AttributeError: pass
@@ -258,16 +268,36 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]:
     cache: dict[UOp, None] = {}
-    stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag)
+    if gate is None and enter_calls:
+      if (backward_slice:=self.__dict__.get("backward_slice")) is not None: return {**backward_slice, self:None}
+      seen = {self}
+      fast_stack = [(self, iter(self.src))]
+      while fast_stack:
+        node, children = fast_stack[-1]
+        for child in children:
+          if child in seen: continue
+          seen.add(child)
+          fast_stack.append((child, iter(child.src)))
+          break
+        else:
+          cache[node] = None
+          fast_stack.pop()
+      return cache
+    if gate is not None and not gate(self): return cache
+    seen = {self}
+    stack: list[tuple[UOp, int]] = [(self, 0)]
     while stack:
-      node, visited = stack.pop()
-      if node in cache: continue
-      if not visited:
-        if gate is None or gate(node):
-          stack.append((node, True))  # push node back on stack to process after its srcs
-          for s in reversed(node.src if enter_calls or node.op not in {Ops.CALL, Ops.FUNCTION} else node.src[1:]):
-            stack.append((s, False)) # push srcs on the stack
-      else: cache[node] = None # second time i'm seeing this node, add it to returned toposort
+      node, idx = stack[-1]
+      src = node.src if enter_calls or node.op not in {Ops.CALL, Ops.FUNCTION} else node.src[1:]
+      if idx == len(src):
+        cache[node] = None
+        stack.pop()
+        continue
+      stack[-1] = (node, idx+1)
+      child = src[idx]
+      if child in seen: continue
+      seen.add(child)
+      if gate is None or gate(child): stack.append((child, 0))
     return cache
 
   def topovisit(self, visitor:Callable[[UOp], T], cache:dict[UOp, T]) -> T:
@@ -479,10 +509,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def simplify(self, tracked=False):
     if self.op is Ops.CONST: return self
     if self.op is Ops.SINK and all(s.op is Ops.CONST or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self
+    if not tracked and (cached:=self.__dict__.get('_simplified', SENTINEL)) is not SENTINEL: return cached
     # late import!
     from tinygrad.uop.symbolic import symbolic
     with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
-      return graph_rewrite(self, symbolic, name="simplify")
+      ret = graph_rewrite(self, symbolic, name="simplify")
+    if not tracked: self.__dict__['_simplified'] = ret
+    return ret
   def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret
   def _eval(self, dtype, expected_type:Type[T]) -> T:
     assert self.dtype in dtype, f"eval with wrong dtype {self}"
@@ -496,6 +529,18 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False, enter_calls:bool=False):
     dvars = {k:v for k,v in dvars.items() if k is not v}
     if len(dvars) == 0: return self
+    use_fast = extra_pm is None and not walk and not TRACK_MATCH_STATS and not PROFILE
+    if name is None:
+      key = (tuple(dvars.items()), extra_pm, walk, enter_calls)
+      cache = self.__dict__.setdefault('_substitute_cache', {})
+      if (cached:=cache.get(key, SENTINEL)) is not SENTINEL: return cached
+      with Context(TRACK_MATCH_STATS=0):
+        ret = _substitute_fast(self, dvars, enter_calls) if use_fast else \
+          graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars,
+                        bottom_up=True, walk=walk, enter_calls=enter_calls)
+      cache[key] = ret
+      return ret
+    if use_fast: return _substitute_fast(self, dvars, enter_calls)
     with Context(TRACK_MATCH_STATS=(0 if name is None else TRACK_MATCH_STATS.value)):
       return graph_rewrite(self, (extra_pm+_substitute) if extra_pm is not None else _substitute, dvars,
                            bottom_up=True, walk=walk, enter_calls=enter_calls, name=name)
@@ -513,7 +558,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   # *** uop syntactic sugar ***
 
   def sink(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
-    return UOp(Ops.SINK, src=tuple([x for x in srcs if x is not None]), **kwargs)
+    return UOp(Ops.SINK, kwargs.pop("dtype", dtypes.void), src=tuple([x for x in srcs if x is not None]), **kwargs)
   def maketuple(*srcs:UOp):  # pylint: disable=no-self-argument
     return UOp(Ops.TUPLE, src=srcs)
   def gettuple(self, idx:int) -> UOp:
@@ -526,7 +571,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def index(self, *srcs:UOp|int|None, **kwargs):
     new_srcs: list[UOp] = [UOp.const(dtypes.index, x) if isinstance(x, int) else x for x in srcs if x is not None]
     if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].arg]
-    return UOp(Ops.INDEX, src=(self,)+tuple(new_srcs), **kwargs)
+    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), src=(self,)+tuple(new_srcs), **kwargs)
   def __getitem__(self, idx):
     # buffers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path
     if self.addrspace in (None, AddrSpace.ALU) or self.device is not None: return super(UOp, self).__getitem__(idx)
@@ -549,7 +594,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return UOp.const(dtype or self.dtype, b, shape=self._shape)
   def vconst_like(self, b:ConstLike, dtype:DType|None=None):
     # for use after movement ops have been removed
-    return UOp.const(dtype or self.dtype, b).broadcast(self.max_numel())
+    ret = UOp.const(dtype or self.dtype, b)
+    if self.shape == (): return ret
+    if len(self.shape) == 1: return UOp(Ops.STACK, ret.dtype, src=(ret,)*self.max_numel())
+    raise RuntimeError(f"vconst_like only works on 0 or 1D shapes, not {self.shape}")
   def ufix(self, x):
     if isinstance(x, UOp): return x
     # float self keeps its dtype for any scalar, int self only for int/Invalid scalars
@@ -557,20 +605,20 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return self.const_like(x, dtypes.from_py(x))
   def broadcast(self, count:int):
     if count == 1: return self
-    return UOp(Ops.STACK, src=(self,)*count)
+    return UOp(Ops.STACK, self.dtype, src=(self,)*count)
   def cast(self, dtype:DTypeLike):
     dtype = to_dtype(dtype)
     if self.dtype == dtype: return self
-    return UOp(Ops.CAST, arg=dtype, src=(self,))
+    return UOp(Ops.CAST, dtype, arg=dtype, src=(self,))
   def bitcast(self, dtype:DTypeLike):
     dtype = to_dtype(dtype)
     return self if self.dtype == dtype else UOp(Ops.BITCAST, arg=dtype, src=(self,))
-  def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, src=(self,)+src, **kwargs)
+  def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, kwargs.pop("dtype", self.dtype), src=(self,)+src, **kwargs)
   def store(self, src:UOp|ConstType, gate:UOp|None=None, **kwargs):
     srcs = (self, self.const_like(src) if not isinstance(src, UOp) else src) + ((gate,) if gate is not None else ())
-    return UOp(Ops.STORE, src=srcs, **kwargs)
-  def end(self, *src:UOp): return UOp(Ops.END, src=(self,)+src) if len(src) else self
-  def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, src=(self,)+src, **kwargs) if len(src) else self
+    return UOp(Ops.STORE, kwargs.pop("dtype", dtypes.void), src=srcs, **kwargs)
+  def end(self, *src:UOp): return UOp(Ops.END, dtypes.void, src=(self,)+src) if len(src) else self
+  def after(self, *src:UOp, **kwargs): return UOp(Ops.AFTER, kwargs.pop("dtype", self.dtype), src=(self,)+src, **kwargs) if len(src) else self
   def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
   def ins(self, arg, **kwargs): return UOp(Ops.INS, kwargs.pop("dtype", self.dtype), kwargs.pop("src", self.src), arg, kwargs.pop("tag", self.tag))
   def contract(self, *rngs:UOp):
@@ -583,7 +631,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if (shapes := [s for x in all_srcs if (s:=x._shape)]) and not all_same(shapes):
       out_shape = _broadcast_shape(*shapes)
       all_srcs = tuple(x._broadcast_to(out_shape) if x._shape else x for x in all_srcs)
-    return UOp(op, src=all_srcs, **kwargs)
+    dtype = kwargs.pop("dtype", dtypes.bool if op in GroupOp.Comparison else all_srcs[1].dtype if op is Ops.WHERE else all_srcs[0].dtype)
+    return UOp(op, dtype, src=all_srcs, **kwargs)
   @staticmethod
   def const(dtype:DType, b:ConstLike, shape:tuple[sint, ...]|None=None):
     if isinstance(b, UOp): return b.cast(dtype)
@@ -596,9 +645,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return ret._mop(Ops.EXPAND, arg=shape) if shape is not None and shape != () and ret.shape != shape else ret
   @staticmethod
   def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.index, src=(), **kwargs):
-    return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
+    return UOp(Ops.RANGE, dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
   @staticmethod
-  def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name)
+  def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype, src=(sint_to_uop(end, dtype),), arg=name)
   def _rop(self, op:Ops, axis:tuple[int, ...]):
     # NOTE: we don't allow reduce on 1s axis
     axis = tuple(sorted(axis))
@@ -624,7 +673,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def reduce(self, *src:UOp, **kwargs):
     arg = kwargs.pop('arg', None)
     if isinstance(arg, Ops): arg = (arg, 0)
-    return UOp(Ops.REDUCE, src=(self,)+src, arg=arg, **kwargs)
+    return UOp(Ops.REDUCE, kwargs.pop("dtype", self.dtype), src=(self,)+src, arg=arg, **kwargs)
 
   def contiguous(self, *args, **kwargs):
     if self.op is Ops.CONTIGUOUS: return self
@@ -750,8 +799,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
         return UOp(Ops.STACK, src=tuple(u.cast(dtype_from_uop(Ops.STACK, srcs, None)) for u in srcs))
       case _: raise RuntimeError(f"{op} is not a MovementOp")
     usrcs = [shape_to_shape_arg(arg) for arg in src_args]
-    if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg)
-    return UOp(op, src=(self,)+UOp.sink(*usrcs).simplify().src)
+    if len(usrcs) == 0: return UOp(op, self.dtype, src=(self,), arg=arg)
+    return UOp(op, self.dtype, src=(self,)+UOp.sink(*usrcs).simplify().src)
 
   # *** uop Buffer stuff ***
 
@@ -1150,7 +1199,12 @@ class ProgramInfo:
     return global_size, local_size
 
   def vals(self, var_vals:dict[str, int]) -> tuple[int|None, ...]:
-    try: return tuple(var_vals[k.expr] if k.expr not in self.runtimevars else None for k in self.vars)
+    def _val(k:UOp) -> int|None:
+      if k.expr in self.runtimevars: return None
+      if k.expr in var_vals: return var_vals[k.expr]
+      if k.vmin == k.vmax: return int(k.vmin)
+      raise KeyError(k.expr)
+    try: return tuple(_val(k) for k in self.vars)
     except KeyError as e: raise RuntimeError(f"unbound Variable {e} used by {self.function_name}") from None
 
   @staticmethod
@@ -1382,7 +1436,8 @@ class PatternMatcher:
     # uop is required, arg is optional
     for p,fxn in self.patterns:
       assert p.op is not None
-      entry: list = [p, None, p.early_reject]
+      entry: list = [p, None, next(iter(p.early_reject)) if len(p.early_reject) == 1 else None,
+                     p.early_reject if len(p.early_reject) > 1 else None, p.match_dtype]
       entry[1] = upat_deferred_compile(p, fxn, entry) if compiled else upat_interpret(p, fxn)
       for uop in p.op: self.pdict.setdefault(uop, []).append(entry)
 
@@ -1392,10 +1447,15 @@ class PatternMatcher:
   def __add__(self, more:PatternMatcher) -> PatternMatcher: return PatternMatcher(self.patterns+more.patterns)
 
   def rewrite(self, uop:UOp, ctx=None):
-    if len(pats:=self.pdict.get(uop.op, [])):
-      if (ler:=uop.__dict__.get('_src_ops')) is None: uop.__dict__['_src_ops'] = ler = {u.op for u in uop.src}
-      for _,match,early_reject in pats:
-        if not early_reject.issubset(ler): continue
+    if pats:=self.pdict.get(uop.op):
+      ler = None
+      for _,match,single_reject,multi_reject,match_dtype in pats:
+        if match_dtype is not None and uop.dtype not in match_dtype: continue
+        if single_reject is not None or multi_reject is not None:
+          if ler is None and (ler:=uop.__dict__.get('_src_ops')) is None: uop.__dict__['_src_ops'] = ler = {u.op for u in uop.src}
+          if single_reject is not None:
+            if single_reject not in ler: continue
+          elif not multi_reject.issubset(ler): continue
         if (ret:=match(uop, ctx)) is not None and ret is not uop: return ret
     return None
 
@@ -1491,10 +1551,13 @@ class TrackedPatternMatcher(PatternMatcher):
     if len(pats:=self.pdict.get(uop.op, [])):
       ret = None
       ler = {u.op for u in uop.src}
-      for p,match,early_reject in pats:
+      for p,match,single_reject,multi_reject,match_dtype in pats:
         if p not in match_stats: match_stats[p] = [0,0,0.0,0.0]
         st = time.perf_counter()
-        if not early_reject.issubset(ler):
+        if match_dtype is not None and uop.dtype not in match_dtype:
+          match_stats[p][2] += time.perf_counter()-st
+          continue
+        if (single_reject is not None and single_reject not in ler) or (multi_reject is not None and not multi_reject.issubset(ler)):
           match_stats[p][2] += time.perf_counter()-st
           continue
         match_stats[p][1] += 1
@@ -1553,14 +1616,20 @@ 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:
-  def __init__(self, pm, bpm, ctx=None, enter_calls=False):
+  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)
@@ -1596,74 +1665,96 @@ class RewriteContext:
     return self.replace.get(root, root)
 
   def unified_rewrite(self, root:UOp) -> UOp:
-    stack: collections.deque[tuple[UOp, int, UOp]] = collections.deque([(root, 0, root)])
-    on_stack = {root}  # all UOps either on the stack or in self.replace, i.e. dont have to be placed again
+    stack: list[tuple[UOp, int, UOp]] = [(root, 0, root)]
+    self.replace[root] = SENTINEL  # SENTINEL marks UOps scheduled but not rewritten yet
     waitlist: dict[UOp, list[tuple[UOp, int, UOp]]] = {}  # UOps waiting on a dependency to be in self.replace
+    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
+    stack_pop, stack_append, stack_extend = stack.pop, stack.append, stack.extend
+    replace_get, waitlist_pop, waitlist_setdefault = replace.get, waitlist.pop, waitlist.setdefault
     while stack:
-      if len(stack) > REWRITE_STACK_LIMIT: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
-      n, stage, new_n = stack.pop()
-      if n in self.replace: continue  # skip any nodes we have seen
+      n, stage, new_n = stack_pop()
       if stage == 0:
+        if shared_cache is not None and (cached:=shared_cache.get(n, SENTINEL)) is not SENTINEL:
+          replace[n] = cached
+          if (waiting:=waitlist_pop(n,None)) is not None: stack_extend(waiting)
+          continue
         # if bottom up, we rewrite this node early. in both cases, we add its srcs to the stack
-        if self.bpm is not None:
+        if bpm is not None:
           # apply rewrite rules until a fixed point is reached. may return `uop` itself if PatternMatcher doesn't match
           test_n: UOp|None = n
-          seen = set()
           try:
-            while test_n is not None:
-              if test_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
-              seen.add(test_n)
-              new_n, test_n = test_n, self.cached_bpm_rewrite(test_n)
+            new_n, test_n = n, cached_bpm_rewrite(n)
+            if test_n is not None:
+              seen = {n}
+              while test_n is not None:
+                if test_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
+                seen.add(test_n)
+                new_n, test_n = test_n, cached_bpm_rewrite(test_n)
           except BottomUpGate:
             # if the bpm matching raised a gate, we are done with this node and dont continue down the srcs
-            self.replace[n] = unwrap(test_n)
-            if n in waitlist: stack.extend(waitlist.pop(n))
+            replace[n] = unwrap(test_n)
+            if (waiting:=waitlist_pop(n,None)) is not None: stack_extend(waiting)
             continue
-        stack.append((n, 1, new_n))
+        stack_append((n, 1, new_n))
         # NOTE: CALL/FUNCTION are handled as a special case.
         # The function that is called is not included in the graph_rewrite.
         # If you want to graph_rewrite a call, you can
-        if not self.enter_calls and new_n.op in {Ops.CALL, Ops.FUNCTION}: self.replace[new_n.src[0]] = new_n.src[0]
+        if not enter_calls and (new_n.op is Ops.CALL or new_n.op is Ops.FUNCTION): replace[new_n.src[0]] = new_n.src[0]
         for x in reversed(new_n.src):
-          if x in on_stack: continue
-          stack.append((x, 0, x))
-          on_stack.add(x)
+          if x in replace: continue
+          replace[x] = SENTINEL
+          stack_append((x, 0, x))
+        if len(stack) > stack_limit: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
       elif stage == 1:
-        tmp = []
-        for x in new_n.src:
-          if (rx:=self.replace.get(x, SENTINEL)) is SENTINEL:
+        new_src = None
+        for i,x in enumerate(new_n.src):
+          if (rx:=replace_get(x, SENTINEL)) is SENTINEL:
             # source not ready: register in waitlist instead of spinning
-            waitlist.setdefault(x, []).append((n, 1, new_n))
+            waitlist_setdefault(x, []).append((n, 1, new_n))
             break
-          tmp.append(rx)
+          if rx is not x:
+            if new_src is None: new_src = list(new_n.src)
+            new_src[i] = rx
         else:
           # in stage 1, once all srcs are rewritten, rebuild (if changed) or run top-down rewrite
-          if (new_src:=tuple(tmp)) == new_n.src:
+          if new_src is None:
             # 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 self.pm is None or (new_src_n:=self.pm_rewrite(new_n)) is None:
-              self.replace[n] = new_n
-              if n in waitlist: stack.extend(waitlist.pop(n))
+            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 (waiting:=waitlist_pop(n,None)) is not None: stack_extend(waiting)
               continue
           else:
             # if srcs changed from rewrites, construct a new UOp with the new srcs
-            new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag)
+            new_src_n = UOp(new_n.op, new_n.dtype, tuple(new_src), new_n.arg, new_n.tag)
           # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n
-          stack.append((n, 2, new_src_n))
-          stack.append((new_src_n, 0, new_src_n))
+          stack_append((n, 2, new_src_n))
+          if new_src_n not in replace:
+            replace[new_src_n] = SENTINEL
+            stack_append((new_src_n, 0, new_src_n))
+          if len(stack) > stack_limit: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
       else:
         # in stage 2, we link the result of new_n to the result of n
-        if (replaced_new_n:=self.replace.get(new_n, SENTINEL)) is SENTINEL:
+        if (replaced_new_n:=replace_get(new_n, SENTINEL)) is SENTINEL:
           # not ready: register in waitlist instead of spinning
-          waitlist.setdefault(new_n, []).append((n, 2, new_n))
+          waitlist_setdefault(new_n, []).append((n, 2, new_n))
         else:
           # otherwise we are done
-          self.replace[n] = replaced_new_n
-          if n in waitlist: stack.extend(waitlist.pop(n))
-    return self.replace[root]
+          replace[n] = replaced_new_n
+          if shared_cache is not None: shared_cache[n] = replaced_new_n
+          if (waiting:=waitlist_pop(n,None)) is not None: stack_extend(waiting)
+    if (ret:=replace[root]) is SENTINEL: raise RuntimeError("infinite loop in graph_rewrite")
+    return ret
 
 @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)
@@ -1706,6 +1797,40 @@ pm_lower_index_dtype = PatternMatcher([
 ])
 def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 
+def _substitute_fast(root:UOp, dvars:dict[UOp, UOp], enter_calls:bool=False) -> UOp:
+  replace:dict[UOp, UOp] = {}
+  stack:list[tuple[UOp, int, UOp]] = [(root, 0, root)]
+  stack_limit = REWRITE_STACK_LIMIT.value
+  while stack:
+    n, stage, new_n = stack.pop()
+    if n in replace: continue
+    if stage == 0:
+      seen = None
+      while (mapped:=dvars.get(new_n, SENTINEL)) is not SENTINEL:
+        if seen is None: seen = {new_n}
+        elif new_n in seen: raise RuntimeError("infinite loop in fixed_point_rewrite")
+        else: seen.add(new_n)
+        new_n = mapped
+      stack.append((n, 1, new_n))
+      if not enter_calls and new_n.op in {Ops.CALL, Ops.FUNCTION}: replace[new_n.src[0]] = new_n.src[0]
+      for x in reversed(new_n.src):
+        if x not in replace: stack.append((x, 0, x))
+      if len(stack) > stack_limit: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
+    elif stage == 1:
+      new_src = None
+      for i,x in enumerate(new_n.src):
+        if (rx:=replace[x]) is not x:
+          if new_src is None: new_src = list(new_n.src)
+          new_src[i] = rx
+      if new_src is None: replace[n] = new_n
+      else:
+        rebuilt = UOp(new_n.op, new_n.dtype, tuple(new_src), new_n.arg, new_n.tag)
+        stack.append((n, 2, rebuilt))
+        stack.append((rebuilt, 0, rebuilt))
+        if len(stack) > stack_limit: raise RuntimeError("infinite loop in graph_rewrite (stack too big)")
+    else: replace[n] = replace[new_n]
+  return replace[root]
+
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])])
 remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])

tinygrad/uop/symbolic.py

@@ -1,5 +1,5 @@
 # all of symbolic lives here now
-import math, struct
+import math, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
 from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid
@@ -46,10 +46,18 @@ def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None:
   if p is not x or (t:=xa + a - pa) % D: return None
   return base - k*div if (k:=t//D - s) else base
 
+def _has_scaled_mod_term(x:UOp) -> bool:
+  if (cached:=x.__dict__.get("_has_scaled_mod_term")) is not None: return cached
+  if x.op is Ops.ADD: ret = any(_has_scaled_mod_term(s) for s in x.src)
+  else: ret = x.op is Ops.FLOORMOD or (x.op is Ops.MUL and any(s.op is Ops.FLOORMOD for s in x.src))
+  x.__dict__["_has_scaled_mod_term"] = ret
+  return ret
+
 def fold_add_divmod_recombine(x:UOp) -> UOp|None:
   # a scaled mod (base%div)*mul recombines with a partner q*(div*mul) carrying the quotient of a b == base (mod div):
   #   q == b//div     -> b*mul              (full recombine)
   #   q == (b//div)%d -> (b%(div*d))*mul    (partial recombine into a wider mod, needs d>0)
+  if not _has_scaled_mod_term(x): return None
   terms = list(x.split_uop(Ops.ADD))
   for i,u in enumerate(terms):
     mod, mul = u.pop_const(Ops.MUL)
@@ -303,8 +311,10 @@ def parse_valid(v:UOp) -> tuple[UOp, bool, int]|None:
     return v.src[0], True, int((v.src[1]).vmax)-1
   return None
 
-def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
+@functools.cache
+def uop_given_valid(valid:UOp, uop:UOp, try_simplex=False) -> UOp:
   # return simplified uop (might be the same as input)
+  if valid.op is Ops.CONST: return uop
 
   # first, parse valid into {expr: (lower_bound, upper_bound)}
   bounds:defaultdict[UOp, list[PyConst|None]] = defaultdict(lambda: [None, None])
@@ -312,6 +322,11 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
     if (res:=parse_valid(stmt)) is None: continue
     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 = []
@@ -399,7 +414,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:

tinygrad/uop/upat.py

@@ -1,7 +1,7 @@
 from typing import Any, Callable
 import itertools, inspect, functools, types
 from tinygrad.helpers import partition, dedup, Context
-from tinygrad.uop.ops import UPat, UOp, Ops, PatternMatcher, graph_rewrite, deconstruct_function
+from tinygrad.uop.ops import UPat, UOp, Ops, GroupOp, PatternMatcher, graph_rewrite, deconstruct_function
 
 class UPatCompileError(Exception): pass
 
@@ -11,22 +11,24 @@ class UPatCompileError(Exception): pass
 #           STORE (bind a matched UOp to a name), PYLITERAL (Python literal for CUSTOM operands),
 #           AND/OR (clause combininers).
 
-def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
+def _get_clause(self:UPat, base:UOp, depth=0, root=True) -> UOp:
   if self.is_any:
     assert len(self.src) == 1
-    return UOp(Ops.AND, src=(UOp(Ops.OR, src=tuple(_get_clause(s, base, depth) for s in self.src[0])),))
+    return UOp(Ops.AND, src=(UOp(Ops.OR, src=tuple(_get_clause(s, base, depth, False) for s in self.src[0])),))
   # build the and_clause for acceptance
   and_clause:list[UOp] = []
-  if self.op is not None:
+  if self.op is not None and not root:
     if len(self.op) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(int(x) for x in self.op))), arg="{0}.op in {1}"))
     else: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.op == "+str(self.op[0].value)))
   if self.arg is not None:
     if isinstance(self.arg, int): and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg="{0}.arg == "+str(int(self.arg))))
     else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=self.arg)), arg="{0}.arg == {1}"))
-  if self.strict_length or self.required_len > 0:
+  fixed_root_len = root and self.op is not None and all((op in GroupOp.Unary and self.required_len == 1) or
+    (op in GroupOp.Binary and self.required_len == 2) or (op in GroupOp.Ternary and self.required_len == 3) for op in self.op)
+  if (self.strict_length or self.required_len > 0) and not fixed_root_len:
     and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len))))
   if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.CUSTOMI, arg=self.name), base)))
-  if self.match_dtype is not None:
+  if self.match_dtype is not None and not root:
     if len(self.match_dtype) > 1:
       and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.PYLITERAL, arg=tuple(self.match_dtype))),
                             arg="{0}.dtype in {1}"))
@@ -40,18 +42,19 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
   if self.src is not None:
     # single match
     if len(self.src) == 1 and isinstance(self.src[0], tuple):
-      and_clause += [_get_clause(s, base.index(i), depth) for i,s in enumerate(self.src[0])]
+      and_clause += [_get_clause(s, base.index(i), depth, False) for i,s in enumerate(self.src[0])]
     # repeat match
     elif len(self.src) == 1 and isinstance(self.src[0], itertools.repeat):
       it = UOp(Ops.CUSTOMI, arg=f"ituop{depth}")
-      match = _get_clause(next(self.src[0]), it, depth+1)
+      match = _get_clause(next(self.src[0]), it, depth+1, False)
       and_clause.append(UOp(Ops.CUSTOM, src=(match, it, base), arg="all([{0} for {1} in {2}.src])"))
     # multi match (fork)
     elif len(self.src) > 1 and all(isinstance(x, tuple) for x in self.src):
-      fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.index(i), depth) for i,s in enumerate(ss)])) for ss in self.src]
+      fork_cond = [UOp(Ops.AND, src=tuple([_get_clause(s, base.index(i), depth, False) for i,s in enumerate(ss)])) for ss in self.src]
       and_clause.append(UOp(Ops.OR, src=tuple(fork_cond)))
     else: raise RuntimeError("broken")
-  return UOp(Ops.AND, src=tuple(and_clause)) if and_clause else UOp(Ops.CUSTOMI, arg="True")
+  if and_clause: return UOp(Ops.AND, src=tuple(and_clause))
+  return UOp(Ops.AND, src=(UOp(Ops.CUSTOMI, arg="True"),)) if root else UOp(Ops.CUSTOMI, arg="True")
 
 # *** pattern matcher ***