no junk ops

test/unit/test_function.py

@@ -557,9 +557,9 @@ class TestFunctionTuple(unittest.TestCase):
     def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
     with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
 
-  def test_vector_load_is_program_input(self):
+  def test_shrink_load_is_program_input(self):
     out, inp = UOp.param(0, dtypes.float, (1,)), UOp.param(1, dtypes.float, (8,))
-    values = UOp(Ops.VLOAD, dtypes.float, (inp[0],), arg=8)
+    values = UOp(Ops.SHRINK, src=(inp, UOp.const(dtypes.weakint, 0), UOp.const(dtypes.weakint, 8))).load()
     sink = out[0].store(values.index(0)).sink(arg=KernelInfo(name="vector_load"))
     info = ProgramInfo.from_sink(sink)
     self.assertEqual(info.outs, (0,))

tinygrad/codegen/__init__.py

@@ -120,8 +120,17 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([
   (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
 ])
 
-def do_devectorize(b:UOp):
-  if b.shape == () or b.tag == "vectorized": return None
+@functools.cache
+def _uses_shrink_memory(x:UOp) -> bool:
+  if x.op in (Ops.LOAD, Ops.STORE): return x.src[0].op is Ops.SHRINK
+  if x.op is Ops.AFTER: return _uses_shrink_memory(x.src[0])
+  if x.op in GroupOp.Elementwise or x.op in (Ops.STACK, Ops.RESHAPE, Ops.PERMUTE):
+    return any(_uses_shrink_memory(y) for y in x.src)
+  return False
+
+def do_devectorize(ctx:Renderer|tuple[set[UOp], Renderer], b:UOp):
+  preserved = ctx[0] if isinstance(ctx, tuple) else set()
+  if b.shape == () or b in preserved or _uses_shrink_memory(b): return None
   # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
   if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None
   src = []
@@ -187,6 +196,7 @@ def fix_group_for_reduce(x:UOp):
 @dataclass
 class ReduceContext:
   acc_num: int = 0
+  renderer: Renderer|None = None
 
 def merge_reduce_ends(sink:UOp):
   # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
@@ -220,7 +230,8 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
   acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")
   return acc.after(acc_out)
 
-def expand_horizontal_reduce(r:UOp):
+def expand_horizontal_reduce(ctx:ReduceContext, r:UOp):
+  if ctx.renderer is not None and ctx.renderer.has_native_reduce(r): return None
   inp = r.src[0]
   vals = [inp.index(*idx) for idx in itertools.product(*[range(inp.max_shape[a]) for a in range(r.arg[1])])]
   return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals)
@@ -317,7 +328,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander")
 
   # remove reduce
-  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(), name="remove reduces")
+  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(renderer=ren), name="remove reduces")
 
   # add locals
   sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
@@ -330,7 +341,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+  sink = graph_rewrite(sink, symbolic_simple, name="pre-devectorize symbolic")
+  native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  sink = graph_rewrite(sink, devectorizer2, ctx=(native_reduce_uops, ren), name="devectorize2")
 
   # simplify indexing
   sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
@@ -340,7 +353,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # do memory coalescing (late)
   sink = memory_coalescing(sink, ren)
-  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  sink = graph_rewrite(sink, symbolic_simple, name="pre-image symbolic", bottom_up=True)
+  native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  sink = graph_rewrite(sink, ew_devectorizer+pm_simplify_add_image, name="add images", ctx=(native_reduce_uops, ren), bottom_up=True)
 
   # extra symbolic before decomp. crashes without this?
   sink = graph_rewrite(sink, sym, name="extra symbolic")

tinygrad/codegen/late/coalesce.py

@@ -106,7 +106,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     # TODO: this should handle images too, it's just memory coalescing
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
-      if u.tag == "vectorized": continue
+      if u.src[0].op is Ops.SHRINK: continue
       assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue

tinygrad/llm/cpu.py

@@ -5,7 +5,7 @@ from tinygrad import Tensor, nn, UOp, getenv, dtypes
 from tinygrad.dtype import DType, AddrSpace
 from tinygrad.llm.gguf import _GGML_QUANT
 from tinygrad.renderer import Estimates
-from tinygrad.uop.ops import Ops, GroupOp, KernelInfo, AxisType
+from tinygrad.uop.ops import Ops, KernelInfo, AxisType
 if TYPE_CHECKING:
   from tinygrad.llm.model import ExpertWeights, FFNBlock, Linear
 
@@ -17,14 +17,33 @@ def _concrete_int(value:int|UOp) -> int:
 
 def _dot_byte_parts(a:tuple[UOp, ...], b:tuple[UOp, ...]) -> UOp:
   av, bv = UOp.stack(*a), UOp.stack(*b)
-  return UOp(Ops.DOT, dtypes.int32, (av, bv), arg=4)
+  return _dot_byte_vectors(av, bv)
+
+def _dot_byte_vectors(a:UOp, b:UOp, scale:UOp|None=None) -> UOp:
+  assert a.shape == b.shape and len(a.shape) == 1 and a.shape[0] % 4 == 0
+  product = a.cast(dtypes.int32) * b.cast(dtypes.int32)
+  if scale is not None: product = product * scale
+  return product.reshape(a.shape[0]//4, 4)._rop(Ops.ADD, (1,))
+
+def _unpack_nibbles(packed:UOp, values:tuple[int, ...]) -> UOp:
+  assert packed.shape == (16,) and len(values) == 16
+  def lookup(index:UOp) -> UOp:
+    value = UOp.stack(*(UOp.const(dtypes.int8, values[-1]) for _ in range(packed.shape[0])))
+    for i in range(14, -1, -1): value = index.eq(i).where(UOp.const(dtypes.int8, values[i]), value)
+    return value
+  lowv, highv = lookup(packed & 15), lookup((packed >> 4) & 15)
+  return UOp.stack(lowv, highv).reshape(32)
 
 def _dot_bytes(a:tuple[UOp, ...], b:tuple[UOp, ...]) -> UOp:
   parts = _dot_byte_parts(a, b)
   return sum((parts.index(i) for i in range(len(a)//4)), UOp.const(dtypes.int32, 0))
 
 def _contiguous_vector_load(ptr:UOp, lanes:int, dtype:DType|None=None) -> UOp:
-  return UOp(Ops.VLOAD, dtype or ptr.dtype, (ptr,), arg=lanes)
+  assert ptr.op is Ops.INDEX
+  buf, coords = ptr.src[0], ptr.src[1:]
+  index = sum((coord * math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(dtypes.weakint, 0))
+  address = UOp(Ops.SHRINK, src=(buf.flatten(), index, UOp.const(dtypes.weakint, lanes)))
+  return address.load(dtype=dtype or ptr.dtype)
 
 def _contiguous_vector_ptr(buf:UOp, index:UOp, lanes:int) -> UOp:
   return UOp(Ops.SHRINK, src=(buf, index, UOp.const(dtypes.weakint, lanes)))
@@ -32,14 +51,14 @@ def _contiguous_vector_ptr(buf:UOp, index:UOp, lanes:int) -> UOp:
 def _dot_bytes_ptr(a:UOp, b:UOp) -> UOp:
   av = _contiguous_vector_load(a, 32, dtypes.int8)
   bv = _contiguous_vector_load(b, 32, dtypes.int8)
-  return UOp(Ops.DOT, dtypes.int32, (av, bv), arg=4)
+  return _dot_byte_vectors(av, bv)
 
 def _dot_nibbles_ptr(packed:UOp, x:UOp, values:tuple[int, ...]) -> UOp:
   assert len(values) == 16
   pv = _contiguous_vector_load(packed, 16)
-  qvalues = UOp(Ops.UNPACK_LUT, dtypes.int8, (pv,), arg=values)
+  qvalues = _unpack_nibbles(pv, values)
   xv = _contiguous_vector_load(x, 32, dtypes.int8)
-  return UOp(Ops.DOT, dtypes.int32, (qvalues, xv), arg=4)
+  return _dot_byte_vectors(qvalues, xv)
 
 def _dot_nibbles_pair_ptr(packed:UOp, x:UOp, values:tuple[int, ...]) -> tuple[UOp, UOp]:
   assert len(values) == 16
@@ -56,7 +75,7 @@ def _dot_q6_ptr(block:UOp, x:UOp, subgroup:int) -> UOp:
   hi = ((_contiguous_vector_load(hi_ptr, 32) >> high_shift) & 3) << 4
   qvalues = ((lo | hi) - 32).bitcast(dtypes.int8)
   xv = _contiguous_vector_load(x, 32, dtypes.int8)
-  return UOp(Ops.DOT, dtypes.int32, (qvalues, xv), arg=4)
+  return _dot_byte_vectors(qvalues, xv)
 
 def _rms_f16_product_ptr(x:UOp, norm:UOp, weight:UOp, scale:UOp) -> UOp:
   return _contiguous_vector_load(x, 8).cast(dtypes.float32) * _contiguous_vector_load(norm, 8).cast(dtypes.float32) * \
@@ -250,23 +269,18 @@ def _load_f16(raw:UOp, offset:UOp) -> UOp:
   return bits.bitcast(dtypes.float16).cast(dtypes.float32)
 
 def _load_f16x8_ptr(raw:UOp) -> UOp:
-  return UOp(Ops.VLOAD, dtypes.float16, (raw,), arg=8).cast(dtypes.float32).rtag("vectorized")
+  return _contiguous_vector_load(raw, 8, dtypes.float16).cast(dtypes.float32)
 
 def _vector_reg(reg:UOp, *deps:UOp) -> UOp:
   reg = reg.after(*deps)
   return UOp(Ops.SHRINK, src=(reg, UOp.const(dtypes.weakint, 0), UOp.const(dtypes.weakint, reg.max_numel())))
 
 def _vector_acc_init(reg:UOp, *deps:UOp) -> UOp:
-  return reg.after(_vector_reg(reg, *deps).store(reg.const_like(0), tag="vectorized"))
-
-def _vectorized_tree(x:UOp) -> UOp:
-  if x.max_numel() == 1: return x
-  src = tuple(_vectorized_tree(y) if y.max_numel() > 1 else y for y in x.src)
-  return x.replace(src=src, tag="vectorized") if x.op in GroupOp.Elementwise else x.replace(src=src)
+  return reg.after(_vector_reg(reg, *deps).store(reg.const_like(0)))
 
 def _vector_acc_update(reg:UOp, value:UOp, *deps:UOp) -> UOp:
-  previous = _vector_reg(reg, *deps).load(tag="vectorized")
-  return _vector_reg(reg, *deps).store((previous + _vectorized_tree(value)).rtag("vectorized"), tag="vectorized")
+  previous = _vector_reg(reg, *deps).load()
+  return _vector_reg(reg, *deps).store(previous + value)
 
 def _finite_exp2(x:UOp) -> UOp:
   # The causal-convolution activation is finite. Bounding the exponent preserves sigmoid saturation while avoiding
@@ -512,7 +526,7 @@ def _cpu_expert_weighted_grouped_uop(out:UOp, raw:UOp, probs:UOp, head:UOp, next
   qvalues, scales = [], []
   for subgroup in range(8):
     packed = _contiguous_vector_load(raw[base + 8 + subgroup * 16], 16)
-    qvalues.append(UOp(Ops.UNPACK_LUT, dtypes.int8, (packed,), arg=values))
+    qvalues.append(_unpack_nibbles(packed, values))
     low_byte = raw[base + 4 + subgroup // 2].load()
     low = (low_byte >> (4 * (subgroup % 2))) & 15
     high = (high_word >> (2 * subgroup)) & 3
@@ -526,10 +540,10 @@ def _cpu_expert_weighted_grouped_uop(out:UOp, raw:UOp, probs:UOp, head:UOp, next
   route = matched_routes[route_idx].load().cast(dtypes.weakint)
   input_idx = route // routes_per_input
   block_acc = UOp.placeholder((8,), dtypes.int32, slot=4, addrspace=AddrSpace.REG)
-  stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0), tag="vectorized")
+  stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0))
   for subgroup in range(8):
     xv = _contiguous_vector_load(xq[route, block, subgroup * 32], 32, dtypes.int8)
-    stage = _vector_acc_update(block_acc, UOp(Ops.DOT, dtypes.int32, (qvalues[subgroup], xv), arg=4) * scales[subgroup], stage)
+    stage = _vector_acc_update(block_acc, _dot_byte_vectors(qvalues[subgroup], xv, scales[subgroup]), stage)
   dot = sum((block_acc.after(stage).index(i) for i in range(8)), UOp.const(dtypes.int32, 0)).cast(dtypes.float32)
   contribution = dot * xd[route, block] * block_scale * matched_probs[route_idx].load()
   updated = totals[output_lane, input_idx].store(totals.after(route_loop)[output_lane, input_idx].load() + contribution)
@@ -573,7 +587,7 @@ def _cpu_expert_silu_uop(out:UOp, raw0:UOp, raw1:UOp, sel:UOp, xq:UOp, xd:UOp, l
   vectorized = ggml_type in (14, 23) or (ggml_type == 21 and repacked)
   acc0 = UOp.placeholder((8 if vectorized else 1,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
   acc1 = UOp.placeholder((8 if vectorized else 1,), dtypes.float32, slot=1, addrspace=AddrSpace.REG)
-  initialized = UOp.group(*((_vector_reg(acc, job).store(acc.const_like(0), tag="vectorized") if vectorized else
+  initialized = UOp.group(*((_vector_reg(acc, job).store(acc.const_like(0)) if vectorized else
                              acc.after(job).store(acc.const_like(0))) for acc in (acc0, acc1)))
   if ggml_type == 23:
     block = UOp.range(in_features // 256, 110)
@@ -684,9 +698,9 @@ def _cpu_expert_silu_grouped_uop(out:UOp, raw0:UOp, raw1:UOp, head:UOp, next_rou
     qvalues:list[UOp] = []
     scales:list[UOp] = []
     for subgroup in range(0, 8, 2):
-      packed = _contiguous_vector_load(raw[row_base + meta_size + block * 128 + subgroup * 16], 32)
-      unpacked = UOp(Ops.UNPACK_LUT, dtypes.int8, (packed,), arg=values)
-      qvalues.extend(UOp.stack(*(unpacked.index(i) for i in range(off, off + 32))) for off in (0, 32))
+      for offset in (0, 16):
+        packed = _contiguous_vector_load(raw[row_base + meta_size + block * 128 + subgroup * 16 + offset], 16)
+        qvalues.append(_unpack_nibbles(packed, values))
     for subgroup in range(8):
       scale_byte = raw[meta + 2 + subgroup // 2].load()
       scales.append(1 + 2 * ((scale_byte >> (4 * (subgroup % 2))) & 15).cast(dtypes.int32))
@@ -698,11 +712,11 @@ def _cpu_expert_silu_grouped_uop(out:UOp, raw0:UOp, raw1:UOp, head:UOp, next_rou
     route = matched[route_idx].load().cast(dtypes.weakint)
     xidx = (route.cast(dtypes.uint32) // routes_per_input).cast(dtypes.weakint)
     block_acc = UOp.placeholder((8,), dtypes.int32, slot=6 + projection, addrspace=AddrSpace.REG)
-    stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0), tag="vectorized")
+    stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0))
     for subgroup in range(8):
       xv = _contiguous_vector_load(xq[xidx, block, subgroup * 32], 32, dtypes.int8)
-      parts = UOp(Ops.DOT, dtypes.int32, (qvalues[subgroup], xv), arg=4)
-      stage = _vector_acc_update(block_acc, parts * scales[subgroup], stage)
+      parts = _dot_byte_vectors(qvalues[subgroup], xv, scales[subgroup])
+      stage = _vector_acc_update(block_acc, parts, stage)
     block_sum = block_acc.after(stage)
     value = sum((block_sum.index(i) for i in range(8)), UOp.const(dtypes.int32, 0)).cast(dtypes.float32) * \
       xd[xidx, block] * _load_f16(raw, meta)
@@ -771,7 +785,7 @@ def _moe_stage1_uop(rhidden:UOp, shidden:UOp, rgate:UOp, rup:UOp, sgate:UOp, sup
   expert = sel[route].load().cast(dtypes.weakint)
   racc0 = UOp.placeholder((8 if expert_repacked else 1,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
   racc1 = UOp.placeholder((8 if expert_repacked else 1,), dtypes.float32, slot=1, addrspace=AddrSpace.REG)
-  rinit = UOp.group(*((_vector_reg(acc, routed_job).store(acc.const_like(0), tag="vectorized") if expert_repacked else
+  rinit = UOp.group(*((_vector_reg(acc, routed_job).store(acc.const_like(0)) if expert_repacked else
                        acc.after(routed_job).store(acc.const_like(0))) for acc in (racc0, racc1)))
   if expert_repacked:
     rblock = UOp.range(dim // 256, 100)
@@ -802,7 +816,7 @@ def _moe_stage1_uop(rhidden:UOp, shidden:UOp, rgate:UOp, rup:UOp, sgate:UOp, sup
   shared_output, groups = shared_begin + shared_job, dim // 32
   sacc0 = UOp.placeholder((8,), dtypes.float32, slot=2, addrspace=AddrSpace.REG)
   sacc1 = UOp.placeholder((8,), dtypes.float32, slot=3, addrspace=AddrSpace.REG)
-  sinit = UOp.group(*(_vector_reg(acc, shared_job).store(acc.const_like(0), tag="vectorized") for acc in (sacc0, sacc1)))
+  sinit = UOp.group(*(_vector_reg(acc, shared_job).store(acc.const_like(0)) for acc in (sacc0, sacc1)))
   if shared_repacked and groups % 8 == 0:
     sblock = UOp.range(groups // 8, 101)
     svalues = []
@@ -1264,7 +1278,7 @@ def _attention_prefill_online_uop(out:UOp, q:UOp, cache:UOp, start_pos:UOp) -> U
   row_sums = tuple(UOp.placeholder((1,), dtypes.float32, slot=2*token_tile+i, addrspace=AddrSpace.REG) for i in range(token_tile))
   clear = UOp.range(chunks, 90)
   cleared = UOp.group(*(_contiguous_vector_ptr(numerator.after(job), clear * 8, 8).store(
-    UOp.stack(*(UOp.const(dtypes.float32, 0) for _ in range(8))), tag="vectorized") for numerator in numerators)).end(clear)
+    UOp.stack(*(UOp.const(dtypes.float32, 0) for _ in range(8)))) for numerator in numerators)).end(clear)
   initialized = UOp.group(cleared,
     *(row_max.after(job).store(-math.inf) for row_max in row_maxes),
     *(row_sum.after(job).store(0.0) for row_sum in row_sums))
@@ -1299,8 +1313,7 @@ def _attention_prefill_online_uop(out:UOp, q:UOp, cache:UOp, start_pos:UOp) -> U
     nbase = vchunk * 8
     previous = _contiguous_vector_load(numerator.after(initialized, position)[nbase], 8)
     updated = previous * old_scale + values * weight
-    numerator_updates.append(_contiguous_vector_ptr(numerator.after(qk_done), nbase, 8).store(
-      _vectorized_tree(updated), tag="vectorized"))
+    numerator_updates.append(_contiguous_vector_ptr(numerator.after(qk_done), nbase, 8).store(updated))
   values_done = UOp.group(*numerator_updates).end(vchunk)
   state_updates = [row_max[0].store(next_max) for row_max,next_max in zip(row_maxes, next_maxes)]
   state_updates += [row_sum[0].store(row_sum.after(initialized, position)[0].load() * old_scale + weight)
@@ -1312,7 +1325,7 @@ def _attention_prefill_online_uop(out:UOp, q:UOp, cache:UOp, start_pos:UOp) -> U
   for token_offset,(numerator,row_sum) in enumerate(zip(numerators, row_sums)):
     query = bh * tokens + token_base + token_offset
     value = _contiguous_vector_load(numerator.after(positions_done)[output * 8], 8) / row_sum.after(positions_done)[0].load()
-    stores.append(_contiguous_vector_ptr(outf, query * dim + output * 8, 8).store(_vectorized_tree(value), tag="vectorized"))
+    stores.append(_contiguous_vector_ptr(outf, query * dim + output * 8, 8).store(value))
   return UOp.group(*stores).end(output, job, core).sink(
     arg=KernelInfo(name=f"attention_prefill_online_uop_{batch}_{heads}_{tokens}_{kv_heads}_{dim}_{cache_len}",
                    opts_to_apply=())).rtag("cpu_parallel")
@@ -1364,8 +1377,7 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta
   current = UOp.placeholder((dim * dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
   init_chunk = UOp.range(dim * dim // 8, 90)
   initial_values = _contiguous_vector_load(statef[bh * dim * dim + init_chunk * 8], 8).cast(dtypes.float32)
-  initialized_state = _contiguous_vector_ptr(current, init_chunk * 8, 8).store(
-    _vectorized_tree(initial_values), tag="vectorized").end(init_chunk)
+  initialized_state = _contiguous_vector_ptr(current, init_chunk * 8, 8).store(initial_values).end(init_chunk)
 
   token = UOp.range(tokens, 91)
   token_base = (bh * tokens + token) * dim
@@ -1381,8 +1393,8 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta
   row = UOp.range(dim, 93)
   state_k = UOp.placeholder((8,), dtypes.float32, slot=3, addrspace=AddrSpace.REG)
   state_q = UOp.placeholder((8,), dtypes.float32, slot=4, addrspace=AddrSpace.REG)
-  dot_init = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0), tag="vectorized"),
-                       _vector_reg(state_q, row).store(state_q.const_like(0), tag="vectorized"))
+  dot_init = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0)),
+                       _vector_reg(state_q, row).store(state_q.const_like(0)))
   col = UOp.range(chunks, 94)
   state_vec = _contiguous_vector_load(current.after(initialized_state, token)[row * dim + col * 8], 8)
   q_vec = _contiguous_vector_load(qf[token_base + col * 8], 8)
@@ -1399,8 +1411,7 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta
   current_values = _contiguous_vector_load(current.after(initialized_state, token)[update_base], 8)
   key_values = _contiguous_vector_load(kf[token_base + update * 8], 8)
   next_values = current_values * av + delta * key_values
-  updated_state = _contiguous_vector_ptr(current.after(dots), update_base, 8).store(
-    _vectorized_tree(next_values), tag="vectorized").end(update)
+  updated_state = _contiguous_vector_ptr(current.after(dots), update_base, 8).store(next_values).end(update)
   rows_done = UOp.group(saved_core, updated_state).end(row)
 
   norm_acc = UOp.placeholder((8,), dtypes.float32, slot=5, addrspace=AddrSpace.REG)
@@ -1416,8 +1427,7 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta
 
   save_chunk = UOp.range(dim * dim // 8, 98)
   saved_values = _contiguous_vector_load(current.after(token_done)[save_chunk * 8], 8).cast(next_state.dtype)
-  saved_state = _contiguous_vector_ptr(nextf, bh * dim * dim + save_chunk * 8, 8).store(
-    _vectorized_tree(saved_values), tag="vectorized").end(save_chunk)
+  saved_state = _contiguous_vector_ptr(nextf, bh * dim * dim + save_chunk * 8, 8).store(saved_values).end(save_chunk)
   return saved_state.end(bh).sink(
     arg=KernelInfo(name=f"gated_delta_prefill_uop_{batch}_{heads}_{tokens}_{dim}_{state.dtype.name}",
                    opts_to_apply=())).rtag("cpu_parallel")
@@ -1446,8 +1456,8 @@ def _gated_delta_uop(core:UOp|None, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UO
   row = UOp.range(dim, 91)
   state_k = UOp.placeholder((8,), dtypes.float32, slot=2, addrspace=AddrSpace.REG)
   state_q = UOp.placeholder((8,), dtypes.float32, slot=3, addrspace=AddrSpace.REG)
-  initialized = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0), tag="vectorized"),
-                          _vector_reg(state_q, row).store(state_q.const_like(0), tag="vectorized"))
+  initialized = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0)),
+                          _vector_reg(state_q, row).store(state_q.const_like(0)))
   col = UOp.range(chunks, 92)
   state_vec = _contiguous_vector_load(statef[bh * dim * dim + row * dim + col * 8], 8).cast(dtypes.float32)
   q_vec = _contiguous_vector_load(qf[bh * dim + col * 8], 8)
@@ -1464,8 +1474,8 @@ def _gated_delta_uop(core:UOp|None, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UO
   state_base = bh * dim * dim + row * dim + update * 8
   state_values = _contiguous_vector_load(statef[state_base], 8).cast(dtypes.float32)
   key_values = _contiguous_vector_load(kf[bh * dim + update * 8], 8)
-  next_values = _vectorized_tree(state_values * av + delta * key_values).cast(next_state.dtype)
-  update_state = _contiguous_vector_ptr(nextf.after(dots), state_base, 8).store(next_values, tag="vectorized").end(update)
+  next_values = (state_values * av + delta * key_values).cast(next_state.dtype)
+  update_state = _contiguous_vector_ptr(nextf.after(dots), state_base, 8).store(next_values).end(update)
   rows_done = UOp.group(saved_core, update_state).end(row)
 
   if normalize:
@@ -1498,14 +1508,14 @@ def _gated_delta_uop(core:UOp|None, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UO
       norm_values = _contiguous_vector_load(norm_weight[quant_group * 32 + offset], 8).cast(dtypes.float32)
       gate_values = _contiguous_vector_load(gatef[base + offset], 8).cast(dtypes.float32)
       gate_sigmoid = UOp.stack(*(gate_values.index(i).sigmoid() for i in range(8)))
-      value_vecs.append(_vectorized_tree(source_values * scale * norm_values * gate_values * gate_sigmoid))
+      value_vecs.append(source_values * scale * norm_values * gate_values * gate_sigmoid)
     values = tuple(value_vecs[i // 8].index(i % 8) for i in range(32))
     amax = functools.reduce(lambda a,b:a.maximum(b), (value.abs() for value in values))
     d = (amax / 127).maximum(1e-8)
     quant_stores = []
     for chunk_idx,value_vec in enumerate(value_vecs):
       quant_values = UOp.stack(*((value_vec.index(i) / d).round().maximum(-127).minimum(127).cast(dtypes.int8) for i in range(8)))
-      quant_stores.append(_contiguous_vector_ptr(quantf, base + chunk_idx * 8, 8).store(quant_values, tag="vectorized"))
+      quant_stores.append(_contiguous_vector_ptr(quantf, base + chunk_idx * 8, 8).store(quant_values))
     stores = UOp.group(scalef[bh * (dim // 32) + quant_group].store(d), *quant_stores).end(quant_group)
     name = f"gated_delta_q8_uop_{batch}_{heads}_{dim}_{state.dtype.name}"
   return stores.end(bh).sink(arg=KernelInfo(name=name, opts_to_apply=())).rtag("cpu_parallel")
@@ -1611,8 +1621,8 @@ def _gdn_qkv_uop(q:UOp, k:UOp, v:UOp, conv:UOp, k_heads:int, v_heads:int, dim:in
 
   qsum = UOp.placeholder((8,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
   ksum = UOp.placeholder((8,), dtypes.float32, slot=1, addrspace=AddrSpace.REG)
-  initialized = UOp.group(_vector_reg(qsum, job).store(qsum.const_like(0), tag="vectorized"),
-                          _vector_reg(ksum, job).store(ksum.const_like(0), tag="vectorized"))
+  initialized = UOp.group(_vector_reg(qsum, job).store(qsum.const_like(0)),
+                          _vector_reg(ksum, job).store(ksum.const_like(0)))
   chunk = UOp.range(dim // 8, 100)
   qvalue = _contiguous_vector_load(conv[qbase + chunk * 8], 8).cast(dtypes.float32)
   kvalue = _contiguous_vector_load(conv[kbase + chunk * 8], 8).cast(dtypes.float32)
@@ -1627,7 +1637,7 @@ def _gdn_qkv_uop(q:UOp, k:UOp, v:UOp, conv:UOp, k_heads:int, v_heads:int, dim:in
   kvalues = _contiguous_vector_load(conv[kbase + out_chunk * 8], 8).cast(dtypes.float32) * kscale
   vvalues = _contiguous_vector_load(conv[vbase + out_chunk * 8], 8).cast(dtypes.float32)
   stores = UOp.group(*(
-    target[batch_idx, head, token, out_chunk * 8 + lane].store(_vectorized_tree(values).index(lane))
+    target[batch_idx, head, token, out_chunk * 8 + lane].store(values.index(lane))
     for target,values in ((q, qvalues), (k, kvalues), (v, vvalues)) for lane in range(8))).end(out_chunk)
   return stores.end(job, core).sink(
     arg=KernelInfo(name=f"gdn_qkv_uop_{batch}_{tokens}_{k_heads}_{v_heads}_{dim}", opts_to_apply=())).rtag("cpu_parallel")
@@ -1744,7 +1754,7 @@ def _f16_matvec_uop(out:UOp, x:UOp, weight:UOp) -> UOp:
   token_base = token_block * token_tile
   xf, wf, outf = x.flatten(), weight.flatten(), out.flatten()
   accs = tuple(UOp.placeholder((8,), dtypes.float32, slot=i, addrspace=AddrSpace.REG) for i in range(token_tile))
-  accs = tuple(acc.after(_vector_reg(acc, output_job, token_block).store(acc.const_like(0), tag="vectorized")) for acc in accs)
+  accs = tuple(acc.after(_vector_reg(acc, output_job, token_block).store(acc.const_like(0))) for acc in accs)
   chunk = UOp.range(in_features // 8, 92)
   weights = _load_f16x8_ptr(wf[output * in_features + chunk * 8])
   updates = []
@@ -1752,9 +1762,8 @@ def _f16_matvec_uop(out:UOp, x:UOp, weight:UOp) -> UOp:
     xbase = (token_base + token) * in_features + chunk * 8
     values = _load_f16x8_ptr(xf[xbase]) if x.dtype == dtypes.float16 else \
       UOp.stack(*(xf[xbase + lane].load() for lane in range(8)))
-    previous = _vector_reg(acc, chunk).load(tag="vectorized")
-    updates.append(_vector_reg(acc, chunk).store(
-      (previous + (values * weights).rtag("vectorized")).rtag("vectorized"), tag="vectorized"))
+    previous = _vector_reg(acc, chunk).load()
+    updates.append(_vector_reg(acc, chunk).store(previous + values * weights))
   done = UOp.group(*updates).end(chunk)
   stores = [outf[(token_base + token) * out_features + output].store(
     sum((acc.after(done).index(lane) for lane in range(8)), UOp.const(dtypes.float32, 0)).cast(out.dtype))

tinygrad/renderer/__init__.py

@@ -30,12 +30,12 @@ class Estimates:
         if u.op in {Ops.INDEX, Ops.SHRINK}:
           excluded = excluded.union(set(UOp.sink(*u.src[1:]).toposort(lambda x: x.op is not Ops.END)))
     for u in uops:
-      if u.op in {Ops.LOAD, Ops.VLOAD, Ops.STORE}:
+      if u.op in {Ops.LOAD, Ops.STORE}:
         buf = u
         while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
         if buf.op is Ops.PARAM:
           # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
-          elements = u.max_numel() if u.op is Ops.VLOAD else u.src[0].max_numel()
+          elements = u.src[0].max_numel()
           accessed = mem.get((buf, u.op), 0) + elements * u.src[0].dtype.scalar().itemsize * mults
           mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize)
       if u.op is Ops.RANGE:
@@ -47,7 +47,7 @@ class Estimates:
       elif u.op is Ops.END: mults = mult_stack.pop(-1)
       elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
       elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1
-      elif u.op in {Ops.LOAD, Ops.VLOAD} and u.src[0].addrspace != AddrSpace.REG:
+      elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG:
         lds += u.max_numel() * u.dtype.scalar().itemsize * mults
       elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG:
         lds += u.max_numel() * u.src[1].dtype.scalar().itemsize * mults
@@ -62,6 +62,7 @@ class Renderer:
   suffix: str = ""
   # TODO: make this generic with a list of supported types
   supports_float4: bool = True
+  def has_native_reduce(self, x:UOp) -> bool: return False
   has_local: bool = True
   has_threads: bool = False
   has_shared: bool = True

tinygrad/renderer/cstyle.py

@@ -22,7 +22,7 @@ base_rewrite = PatternMatcher([
 
   # casting
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \
-    if x.max_numel() > 1 and (x.addrspace is AddrSpace.REG or x.tag == "vectorized") else None),
+    if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"),
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: ctx[x.src[0]] if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"),
@@ -131,6 +131,7 @@ class CStyleLanguage(Renderer):
   float4: str|None = None
   float4_style: tuple[str, str] = ('(', ')')
   gep_arr_threshold: int = 4
+  def should_inline_where(self, x:UOp) -> bool: return False
   type_map: dict[DType, str] = {}
   infinity: str = "INFINITY"
   nan: str = "NAN"
@@ -239,7 +240,7 @@ class CStyleLanguage(Renderer):
       if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
         (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
-        (u.op in {Ops.STACK, Ops.DOT, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and
+        (u.op in {Ops.STACK, Ops.REDUCE, *GroupOp.ALU, Ops.CAST, Ops.BITCAST} and (u.op is not Ops.WHERE or self.should_inline_where(u)) and
          child_count[u] == 1 and not getenv("EXPAND_SSA"))):
         r[u] = l
       else:
@@ -260,6 +261,23 @@ class ClangRenderer(CStyleLanguage):
   gep_arr_threshold = 0
   has_local = False
   has_threads = bool(getenv("THREADS", 1))
+  def should_inline_where(self, x:UOp) -> bool: return self._render_lut_lookup(self, x) is not None
+  @staticmethod
+  def _byte_dot_reduce(x:UOp) -> tuple[UOp, UOp, UOp|None]|None:
+    if x.op is not Ops.REDUCE or x.arg != (Ops.ADD, 1) or len(x.src) != 1: return None
+    permute = x.src[0]
+    if permute.op is not Ops.PERMUTE or permute.arg != (1, 0) or permute.src[0].op is not Ops.RESHAPE: return None
+    product, scale = permute.src[0].src[0], None
+    if product.op is Ops.MUL:
+      for base, candidate_scale in (product.src, product.src[::-1]):
+        if base.op is Ops.MUL and all(s.op is Ops.CAST and s.dtype is dtypes.int32 and s.src[0].dtype is dtypes.int8 for s in base.src):
+          product, scale = base, candidate_scale
+          break
+    if product.op is not Ops.MUL or not all(s.op is Ops.CAST and s.dtype is dtypes.int32 and s.src[0].dtype is dtypes.int8 for s in product.src):
+      return None
+    a, b = (s.src[0] for s in product.src)
+    return (a, b, scale) if a.shape == b.shape and a.shape in ((16,), (32,)) else None
+  def has_native_reduce(self, x:UOp) -> bool: return self._byte_dot_reduce(x) is not None
   global_max = (NUM_CPU_THREADS.value, 0, 0)
   infinity = "__builtin_inff()"
   nan = '__builtin_nanf("")'
@@ -283,6 +301,32 @@ class ClangRenderer(CStyleLanguage):
            f"__builtin_ia32_pmaddubsw{suffix}(__builtin_elementwise_abs({ctx[a]}), __builtin_ia32_psignb{suffix}({ctx[b]}, {ctx[a]}))), " \
            f"(short __attribute__((ext_vector_type({lanes})))){{{','.join([multiplier] * lanes)}}})"
   @staticmethod
+  def _uniform_vector_const(x:UOp):
+    return x.src[0].arg if x.op is Ops.STACK and x.src and all(v.op is Ops.CONST and v.arg == x.src[0].arg for v in x.src) else None
+  @staticmethod
+  def _render_lut_lookup(ctx, x:UOp):
+    values:dict[int, int] = {}
+    index:UOp|None = None
+    node = x
+    while node.op is Ops.WHERE and node.src[0].op in (Ops.CMPEQ, Ops.CMPNE):
+      cond, true_value, false_value = node.src
+      left, right = cond.src
+      key = ClangRenderer._uniform_vector_const(right)
+      if key is None: left, right, key = right, left, ClangRenderer._uniform_vector_const(left)
+      if not isinstance(key, int) or not 0 <= key < 16 or (index is not None and left is not index): return None
+      index = left
+      selected, node = (true_value, false_value) if cond.op is Ops.CMPEQ else (false_value, true_value)
+      value = ClangRenderer._uniform_vector_const(selected)
+      if not isinstance(value, int): return None
+      values[key] = value
+    default = ClangRenderer._uniform_vector_const(node)
+    if index is None or not isinstance(default, int) or index.shape not in ((16,), (32,)): return None
+    lut_values = tuple(values.get(i, default) for i in range(16))
+    lanes, suffix = index.shape[0], "128" if index.shape == (16,) else "256"
+    charv = f"signed char __attribute__((ext_vector_type({lanes})))"
+    lut = f"({charv}){{{','.join(map(str, lut_values * (lanes // 16)))}}}"
+    return f"__builtin_ia32_pshufb{suffix}({lut}, __builtin_bit_cast({charv}, {ctx[index]}))"
+  @staticmethod
   def _render_contiguous_stack(ctx, x:UOp):
     if not x.src or not all(v.op is Ops.INDEX and len(v.src) == 2 and v.src[0] is x.src[0] and
                             v.src[1].op is Ops.CONST for v in x.src): return None
@@ -290,19 +334,25 @@ class ClangRenderer(CStyleLanguage):
     if indices != tuple(range(indices[0], indices[0] + len(indices))): return None
     return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[0]]}, {','.join(map(str, indices))})"
   @staticmethod
-  def _render_unpack_lut(ctx, x:UOp, packed:UOp):
-    if packed.shape not in ((16,), (32,)): return None
-    lanes, suffix = packed.shape[0], "128" if packed.shape[0] == 16 else "256"
-    charv = f"signed char __attribute__((ext_vector_type({lanes})))"
-    pv = f"__builtin_bit_cast({charv}, {ctx[packed]})"
-    lut_values = x.arg if lanes == 16 else x.arg * 2
-    lut = f"({charv}){{{','.join(map(str, lut_values))}}}"
-    mask = f"({charv}){{{','.join(['15'] * lanes)}}}"
-    lo = f"__builtin_ia32_pshufb{suffix}({lut}, {pv} & {mask})"
-    hi = f"__builtin_ia32_pshufb{suffix}({lut}, ({pv} >> 4) & {mask})"
-    if lanes == 16: indices = tuple(range(32))
-    else: indices = (*range(16), *range(32, 48), *range(16, 32), *range(48, 64))
-    return f"__builtin_shufflevector({lo}, {hi}, {','.join(map(str, indices))})"
+  def _render_concat_stack(ctx, x:UOp):
+    if len(x.src) != 2 or x.src[0].shape != x.src[1].shape or len(x.src[0].shape) != 1: return None
+    lanes = x.src[0].shape[0]
+    return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[1]]}, {','.join(map(str, range(lanes*2)))})"
+  @staticmethod
+  def _render_vector_permute(ctx, x:UOp):
+    old_shape, order = x.src[0].shape, x.arg
+    new_shape = tuple(old_shape[i] for i in order)
+    indices = []
+    for flat in range(x.max_numel()):
+      coord, rem = [], flat
+      for size in reversed(new_shape):
+        coord.append(rem % size)
+        rem //= size
+      new_coord = tuple(reversed(coord))
+      old_coord = tuple(new_coord[order.index(i)] for i in range(len(order)))
+      old_flat = sum(c * math.prod(old_shape[i+1:]) for i,c in enumerate(old_coord))
+      indices.append(old_flat)
+    return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[0]]}, {','.join(map(str, indices))})"
   @staticmethod
   def _render_vector_load(ctx, x:UOp, address:UOp):
     vec_type = ctx._render_dtype(x.dtype, x.max_numel(), AddrSpace.REG)
@@ -312,23 +362,25 @@ class ClangRenderer(CStyleLanguage):
   @staticmethod
   def _render_vector_store(ctx, address:UOp, value:UOp):
     vec_type = ctx._render_dtype(value.dtype, value.max_numel(), AddrSpace.REG)
-    # Like VLOAD, explicit vector stores can target unaligned packed data or register-backed arrays.
+    # Explicit vector stores can target unaligned packed data or register-backed arrays.
     return f"do {{ {vec_type} _v = {ctx[value]}; __builtin_memcpy({ctx[address]}, &_v, sizeof(_v)); }} while (0);"
   string_rewrite = PatternMatcher([
+    (UPat(Ops.WHERE, name="x"), lambda ctx,x: ClangRenderer._render_lut_lookup(ctx, x)),
+    (UPat(Ops.MUL, dtypes.int32, src=(UPat(Ops.REDUCE, name="dot"), UPat.var("scale"))),
+     lambda ctx,dot,scale: ClangRenderer._render_byte_dot(ctx, pair[0], pair[1], scale)
+     if (pair:=ClangRenderer._byte_dot_reduce(dot)) is not None and pair[2] is None and scale.vmin >= -32768 and scale.vmax <= 32767 else None),
+    (UPat(Ops.REDUCE, name="x"), lambda ctx,x: ClangRenderer._render_byte_dot(ctx, *pair)
+     if (pair:=ClangRenderer._byte_dot_reduce(x)) is not None else None),
+    (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})"
+     if x.max_numel() > 1 else None),
+    (UPat(Ops.RESHAPE, name="x"), lambda ctx,x: ctx[x.src[0]]),
+    (UPat(Ops.PERMUTE, name="x"), lambda ctx,x: ClangRenderer._render_vector_permute(ctx, x)),
+    (UPat(Ops.STACK, name="x"), lambda ctx,x: ClangRenderer._render_concat_stack(ctx, x)),
     (UPat(Ops.STACK, name="x"), lambda ctx,x: ClangRenderer._render_contiguous_stack(ctx, x)),
-    (UPat(Ops.MUL, dtypes.int32,
-          src=(UPat(Ops.DOT, dtypes.int32, src=(UPat.var("a", dtypes.int8), UPat.var("b", dtypes.int8)), arg=4),
-               UPat.var("scale", dtypes.int32))),
-     lambda ctx,a,b,scale: ClangRenderer._render_byte_dot(ctx, a, b, scale)
-     if scale.vmin >= -32768 and scale.vmax <= 32767 else None),
-    (UPat(Ops.DOT, dtypes.int32, src=(UPat.var("a", dtypes.int8), UPat.var("b", dtypes.int8)), arg=4),
-     lambda ctx,a,b: ClangRenderer._render_byte_dot(ctx, a, b)),
-    (UPat(Ops.UNPACK_LUT, dtypes.int8, src=(UPat.var("packed", dtypes.uint8),), name="x"),
-     lambda ctx,x,packed: ClangRenderer._render_unpack_lut(ctx, x, packed)),
     (UPat(Ops.STORE, src=(UPat(Ops.SHRINK, name="address"), UPat.var("value")), name="x"),
      lambda ctx,x,address,value: ClangRenderer._render_vector_store(ctx, address, value)
-     if x.tag == "vectorized" and value.max_numel() > 1 else None),
-    (UPat(Ops.VLOAD, src=(UPat.var("address"),), name="x"),
+     if value.max_numel() > 1 else None),
+    (UPat(Ops.LOAD, src=(UPat(Ops.SHRINK, name="address"),), name="x"),
      lambda ctx,x,address: ClangRenderer._render_vector_load(ctx, x, address)),
   ]) + base_rewrite
 

tinygrad/uop/__init__.py

@@ -52,16 +52,12 @@ class Ops(FastEnum):
   INDEX = auto(); SHRINK = auto()
 
   # load/store before math
-  LOAD = auto(); VLOAD = auto(); STORE = auto()
+  LOAD = auto(); STORE = auto()
 
   # ** 4 -- math **
 
   # tensor core math op, not elementwise
   WMMA = auto()
-  # grouped dot product. arg is the number of input lanes reduced into each output lane
-  DOT = auto()
-  # unpack low/high nibbles through a byte lookup table. arg is the tuple of table values
-  UNPACK_LUT = auto()
 
   # UnaryOps
   CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto()

tinygrad/uop/ops.py

@@ -150,8 +150,6 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
     case Ops.WMMA:
       # WMMA output dtype is the accumulator dtype (src[2])
       return src[2].dtype
-    case Ops.DOT | Ops.UNPACK_LUT | Ops.VLOAD:
-      return None
     case Ops.GETTUPLE:
       # GETTUPLE extracts from a TUPLE (possibly through a FUNCTION)
       in_tuple = src[0].src[0] if src[0].op is Ops.FUNCTION else src[0]
@@ -396,17 +394,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       case Ops.WMMA:
         wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
         return wmma_b + (self.src[2].shape[-1],)
-      case Ops.DOT:
-        assert len(self.src) == 2 and self.src[0].shape == self.src[1].shape and len(self.src[0].shape) and \
-               isinstance(self.arg, int) and self.arg > 0 and self.src[0].shape[-1] % self.arg == 0
-        return self.src[0].shape[:-1] + (self.src[0].shape[-1] // self.arg,)
-      case Ops.UNPACK_LUT:
-        assert len(self.src) == 1 and isinstance(self.arg, tuple)
-        return self.src[0].shape[:-1] + (self.src[0].shape[-1] * 2,)
-      case Ops.VLOAD:
-        assert len(self.src) == 1 and self.src[0].shape == () and isinstance(self.arg, int) and self.arg > 0
-        return (self.arg,)
-
       # passthrough ops
       case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \
            Ops.COPY | Ops.ALLREDUCE | Ops.STORE | Ops.END:
@@ -872,12 +859,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.PARAM: return self.arg.addrspace
     if self.op is Ops.BUFFER: return self.arg.addrspace
     if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU
-    if self.op in {Ops.LOAD, Ops.VLOAD}: return AddrSpace.ALU # LOAD brings things into the ALU
+    if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
     if self.op in {Ops.CUSTOM, Ops.CUSTOMI}: return AddrSpace.ALU
     if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END, Ops.UNSHARD}:
       return self.src[0].addrspace
     if self.op in GroupOp.Movement: return self.src[0].addrspace
-    if self.op in {Ops.STACK, Ops.WMMA, Ops.DOT, Ops.UNPACK_LUT, Ops.GROUP} or self.op in GroupOp.Elementwise:
+    if self.op in {Ops.STACK, Ops.WMMA, Ops.GROUP} or self.op in GroupOp.Elementwise:
       ad = [x.addrspace for x in self.src if x.addrspace is not None]
       if not len(ad) or not all_same(ad): return None
       return ad[0]
@@ -1245,7 +1232,7 @@ class ProgramInfo:
     for u in sink.toposort():
       if u.op is Ops.PARAM and u.addrspace == AddrSpace.ALU: _vars.append(u)
       if u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU: _globals.append(u.arg.slot)
-      if u.op in (Ops.STORE, Ops.LOAD, Ops.VLOAD):
+      if u.op in (Ops.STORE, Ops.LOAD):
         if (idx:=u.src[0]).op in (Ops.INDEX, Ops.SHRINK) or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
           if (buf:=idx.src[0].buf_uop).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg.slot)
       if u.op is Ops.SPECIAL:

tinygrad/uop/spec.py

@@ -120,22 +120,11 @@ spec_shared = PatternMatcher([
    lambda uidx,gate,alt,load: validate_index(uidx, gate) if matches_dtype(alt, load.dtype) else False),
   (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat()), validate_index),
   (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index),
-  # VLOAD reinterprets a scalar address as a contiguous vector of its output dtype
-  (UPat(Ops.VLOAD, src=(UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted(),), name="x"),
-   lambda x,uidx: isinstance(x.arg, int) and x.arg > 0 and validate_index(uidx)),
   # STORE in tensor graph: store a value into a target
   (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: True),
 
   # WMMA has a <a, b, acc>
   (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 5),
-  # DOT reduces equal vector inputs in fixed-size groups
-  (UPat(Ops.DOT, src=(UPat(), UPat()), name="x"), lambda x:
-   isinstance(x.arg, int) and x.arg > 0 and x.src[0].shape == x.src[1].shape and len(x.src[0].shape) > 0 and
-   isinstance(x.src[0].shape[-1], int) and x.src[0].shape[-1] % x.arg == 0),
-  # UNPACK_LUT maps the low nibbles followed by the high nibbles through an immutable byte table
-  (UPat(Ops.UNPACK_LUT, dtypes.int8, src=(UPat(dtype=dtypes.uint8),), name="x"), lambda x:
-   len(x.src[0].shape) > 0 and isinstance(x.src[0].shape[-1], int) and isinstance(x.arg, tuple) and len(x.arg) == 16 and
-   all(isinstance(v, int) and -128 <= v <= 127 for v in x.arg)),
 ])
 
 def is_device(d): return isinstance(d, str) or (isinstance(d, tuple) and all(isinstance(s, str) for s in d))
@@ -217,6 +206,11 @@ spec_program = PatternMatcher([
   # allow special SHRINK
   (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True),
 
+  # register-vector views can remain when the renderer consumes their shape (for example a native horizontal reduction)
+  (UPat((Ops.RESHAPE, Ops.PERMUTE), name="x"), lambda x: x.addrspace is AddrSpace.ALU),
+  (UPat(Ops.REDUCE, src=(UPat(),), name="x"),
+   lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and x.arg[0] in GroupOp.Reduce and isinstance(x.arg[1], int)),
+
   # movement ops are not allowed in programs
   (UPat(GroupOp.Movement), lambda: False),