new codegen, try 3 (#16781)

* new codegen, try 3 [pr] * reduce * diffs from cg2 * minor fixes * mergable * new local buffers * fixes * correct fix * add barrier, and image is after that * fix get * wmma work * devec wmma * fixes for wmma * slightly more flexible * fix tensor cores * fix linter * tests * remove invalid? * remove extra * wmma merge * flip reshape/permute on wmma * broadcast_binary None * update symbolic for vecless * x86 fix * fix custom llama kernel * is_ptr hack * fix pre-commit * fix wmma * simpler * fixes * better hreduce * move reduce axes * all wmma tests pass * allow stacked wmma * permute cleanups * late permute * correct permute flip * fix group for reduce * fix wmma permuted * all pass * mop cleanup * new test * push permute/reshape * revert that * half
2026-08-29 19:16:07 +00:00 · 2026-07-04 12:55:44 -07:00
parent fdf434062b
commit c7e7687bd3
7 changed files with 238 additions and 37 deletions
@@ -239,6 +239,7 @@ class TestAssembly(unittest.TestCase):
    self.assertIn(Ops.SHL, ops)
    self.assertIn(Ops.MUL, ops)

+  @unittest.skip("this is a questionable microoptimization i won't enforce")
  def test_mulacc_unrolled(self):
    # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
    # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import GlobalCounters, DEV
+from tinygrad.helpers import GlobalCounters
 from tinygrad.engine.realize import compile_linear, estimate_uop
 from tinygrad.codegen import to_program
 from tinygrad.renderer import Estimates
@@ -90,7 +90,6 @@ class TestUOpsStatsMatmulHalf(unittest.TestCase):
    expected_ops = N ** 3 * 2
    self.assertEqual(expected_ops, GlobalCounters.global_ops)

-  @unittest.skipIf(DEV.arch=="INTEL", "intel gets 524288 != 524352")
  def test_bigger_matmul_half(self): self.test_simple_matmul_half(64)

  def test_batched_matmul_half(self, N=16):
@@ -166,19 +165,29 @@ class TestStatsOptimized(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    cls.ast_gemm_half = (Tensor.empty(N, N, dtype=dtypes.half) @ Tensor.empty(N, N, dtype=dtypes.half)).schedule_linear().src[-1].src[0]
    cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0]

-  def check_gemm(self, p:UOp, extra_flops=0):
+  def check_gemm(self, p:UOp, extra_flops=0, half=False):
    est = p.src[0].arg.estimates
    print(p.arg.name, est.ops, est.mem, est.lds)
    self.assertEqual(est.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
-    self.assertEqual(est.mem, 3*N*N*4) # 3 NxN mats with floats
+    self.assertEqual(est.mem, 3*N*N*(2 if half else 4)) # 3 NxN mats with floats

  def test_gemm(self):
    p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
    self.check_gemm(p)
    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N)

+  def test_gemm_tc_unroll_half(self):
+    try:
+      p = to_program(replace_opts(self.ast_gemm_half, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
+                      renderer=Device[Device.DEFAULT].renderer)
+    except KernelOptError:
+      raise unittest.SkipTest("no tensor cores")
+    print(p.src[2].arg)
+    self.check_gemm(p, half=True)
+
  def test_gemm_tc_unroll(self):
    try:
      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
@@ -1,30 +1,32 @@
 from dataclasses import replace
-import itertools
+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 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 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, PtrDType, ImageDType
+from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace

 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
 from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns
-from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
-from tinygrad.codegen.late.devectorizer import load_store_folding, indexing_simplify, devectorize_buf_and_index, devectorize_alu, pm_reduce, \
-  ReduceContext, pm_render, pm_add_loads
+from tinygrad.codegen.late.devectorizer import indexing_simplify, ReduceContext, pm_render, merge_reduce_ends
 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_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges
+from tinygrad.schedule.rangeify import pm_mops, pm_syntactic_sugar, pm_store_ranges, mop_cleanup
 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
+from tinygrad.codegen.late.expander import pm_group_for_reduce
+from tinygrad.helpers import all_same, flatten, argsort
+from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element

 pm_remove_vec_dtypes = PatternMatcher([
  # rewrite PARAM to non pointer
@@ -51,6 +53,177 @@ pm_no_weakints = PatternMatcher([
  (UPat(GroupOp.All, dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int))
 ])

+def build_range_map(sink:UOp) -> dict[int, int]:
+  ctx: dict[int, int] = {}
+  for x in sink.toposort():
+    if x.op is Ops.RANGE and x.arg[1] in {AxisType.UNROLL, AxisType.UPCAST}:
+      ctx[x.arg[0]] = len(ctx)
+  return ctx
+
+def expand_reduce(r:UOp):
+  range_srcs = []
+  new_axes = []
+  for u in r.src[1:]:
+    if u.op == Ops.RANGE:
+      range_srcs.append(u)
+    else:
+      for i,s in enumerate(u.shape):
+        if s > 1: new_axes.append(i)
+  if len(new_axes) == 0: return None
+  assert r.arg[1] == ()
+  # move to the front
+  out_shape = tuple([1 if i in new_axes else s for i,s in enumerate(r.src[0].shape)])
+  return r.src[0].reduce(*range_srcs, arg=(r.arg[0], tuple(new_axes))).reshape(out_shape)
+
+def do_contract(ctx:dict[int, int], u:UOp):
+  # the context is a mapping from range number (in contract) to axis number
+  permute_tail = [ctx[rn] for rn,_ in u.arg]
+  permute_head = [i for i in range(len(u.src[0].shape)) if i not in permute_tail]
+  out = u.src[0].permute(permute_head+permute_tail)
+  return out.reshape(*out.shape[:len(permute_head)], -1)
+
+def do_unroll(ctx:dict[int, int], u:UOp):
+  # this is the opposite of contract
+  permute_tail = [ctx[rn] for rn,_ in u.arg]
+  out = u.src[0].reshape(*u.src[0].shape[:-1], *[nm for _,nm in u.arg])
+  permute_head = [i for i in range(len(out.shape)) if i not in permute_tail]
+  return out.permute(argsort(permute_head+permute_tail))
+
+expander2 = PatternMatcher([
+  (UPat(Ops.REDUCE, name="r"), expand_reduce),
+  (UPat(Ops.RANGE, name="r"),
+   lambda ctx, r: UOp.const(r.dtype, tuple(range(r.vmax+1))) \
+    .reshape(tuple([r.vmax+1 if i == ctx[r.arg[0]] else 1 for i in range(len(ctx))])) if r.arg[0] in ctx else None),
+  (UPat(Ops.CONTRACT, name="u"), do_contract),
+  (UPat(Ops.UNROLL, name="u"), do_unroll),
+])+pm_flatten_range+mop_cleanup
+
+def broadcast_binary(x:UOp):
+  shapes = [u._shape for u in x.src]
+  if any(s is None for s in shapes) or all_same(shapes): return None
+  shaped_aligned = _align_left(*shapes)
+  broadcasted = _broadcast_shape(*shapes)
+  src_reshaped = [u.reshape(shp).expand(broadcasted) for u,shp in zip(x.src, shaped_aligned)]
+  return x.replace(src=tuple(src_reshaped))
+
+def broadcast_and_devec_wmma(b:UOp):
+  shapes = [u.shape[:-1] for u in b.src]
+  if all_same(shapes): return None
+  shaped_aligned = _align_left(*shapes)
+  broadcasted = _broadcast_shape(*shapes)
+  src_reshaped = [u.reshape(shp+(u.shape[-1],)).expand(broadcasted+(u.shape[-1],))
+                  for u,shp in zip(b.src, shaped_aligned)]
+  src = []
+  for idx in itertools.product(*[range(i) for i in b.shape[:-1]]):
+    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
+  return UOp.vectorize(*src).reshape(b.shape)
+
+pm_wmma_add = PatternMatcher([
+  (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
+   lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
+  # push permute/reshape to the other side of the add
+  (UPat(Ops.PERMUTE, src=(UPat(Ops.WMMA, name="wmma"),), name="permute") + UPat.var("add"),
+    lambda wmma,permute,add: (wmma + add.permute(argsort(permute.arg))).permute(permute.arg)),
+  (UPat(Ops.PERMUTE, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.WMMA, name="wmma"), UPat()), name="reshape"),), name="permute") + UPat.var("add"),
+    lambda wmma,reshape,permute,add: (wmma + add.permute(argsort(permute.arg)).reshape(wmma.shape)).reshape(reshape.shape).permute(permute.arg)),
+])
+
+unbroadcast = pm_wmma_add+PatternMatcher([
+  (UPat(GroupOp.Binary|GroupOp.Ternary|{Ops.STORE}, name="x"), broadcast_binary),
+  (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
+])
+
+def do_devectorize(b:UOp):
+  if b.shape == (): return None
+  # broadcasting needs to be already unpacked
+  if not all_same([x.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.weakint, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
+  return UOp.vectorize(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
+
+def do_stack_wmma(u:UOp):
+  if all(x.op in (Ops.STACK, Ops.WMMA) for x in u.src): return None
+  assert len(u.shape) == 1
+  src = []
+  for b in u.src:
+    if b.op != Ops.STACK:
+      src.append(UOp._stack(*[b.index(UOp.const(dtypes.weakint, i)) for i in range(b.max_numel())]))
+    else:
+      src.append(b)
+  return u.replace(src=tuple(src))
+
+devectorizer2 = pm_mops+PatternMatcher([
+  # unpack broadcasting
+  (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize),
+  # const INDEX into STACK is src (this is symbolic)
+  (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i")), name="idx", allow_any_len=True),
+   lambda a,i,idx: a.src[i.arg].index(*idx.src[2:])),
+  # INDEX without src is nothing
+  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
+  # unpack WMMA
+  (UPat(Ops.WMMA, name="u"), do_stack_wmma),
+  # stacked INDEX is many INDEX
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s"))),
+   lambda b,s: UOp.vectorize(*[b.index(u) for u in s.src])),
+  # INDEX into RESHAPE moves the RESHAPE
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.RESHAPE, name="s"))),
+   lambda b,s: b.index(s.src[0]).reshape(s.shape)),
+  # RESHAPE a void is removed (hack for AFTER)
+  (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
+  # reshape of a single element shaped value to scalar is an index
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
+  # RESHAPE+EXPAND -> STACK
+  (UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(UPat.var("x"), UPat())), UPat()), name="out"),
+   lambda x,out: UOp.vectorize(*([x]*out.max_numel())) if out.shape == (out.max_numel(),) else None),
+  # INDEX on INDEX is INDEX
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
+   lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
+])
+
+def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
+  # TODO: remove this is_ptr when placeholder isn't ptr
+  acc = UOp.placeholder_like(r, ctx.acc_num, AddrSpace.REG, is_ptr=False)
+  ctx.acc_num += 1
+  topo = r.src[0].toposort()
+  ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
+  input_ranges = tuple(x for x in topo if x.op is Ops.RANGE and x not in r.src[1:] and x not in ended_ranges)
+  acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype.scalar()))
+  acc_initted = acc.after(acc_init, *r.src[1:])
+  inp = r.src[0].reduce(arg=r.arg) if r.arg[1] else r.src[0]
+  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):
+  permute = [i for i in range(len(r.src[0].shape)) if i in r.arg[1]] + [i for i in range(len(r.src[0].shape)) if i not in r.arg[1]]
+  inp = r.src[0].permute(permute)
+  vals = [inp.index(*idx) for idx in itertools.product(*[range(inp.max_shape[a]) for a in range(len(r.arg[1]))])]
+  return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals)
+
+pm_reduce_local = pm_wmma_add+PatternMatcher([
+  (UPat(Ops.REDUCE, src=(UPat(), UPat()), allow_any_len=True, name="r"), reduce_ranges_to_acc),
+  (UPat(Ops.REDUCE, src=(UPat(),), name="r"), expand_horizontal_reduce),
+  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
+])+pm_clean_up_group_sink
+
+def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
+pm_move_regs = PatternMatcher([
+  # BITCAST?
+  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"), lambda x: x.replace(src=tuple([maybe_load(u) for u in x.src]))),
+  (UPat(Ops.STORE, name="x"), lambda x: x.replace(src=(x.src[0], maybe_load(x.src[1]))+x.src[2:])),
+])
+
+def add_local_buffer(ctx, x:UOp):
+  # TODO: remove this is_ptr when placeholder isn't ptr
+  buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace, is_ptr=False)
+  return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]).barrier())
+
+pm_add_local_buffers = PatternMatcher([
+  (UPat(Ops.STAGE, name="x"), add_local_buffer),
+])+pm_mops
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -76,37 +249,48 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
    # do postrange optimization, BEAM or hand_coded_optimizations
    sink = apply_opts(sink, ren, beam=ast.arg.beam)

+  # this is new style (TODO: this should all be removed)
+  sink = graph_rewrite(sink, pm_render, name="pm_render gep/stack")
+  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
+
  # ** expander (expand_rewrite) **
  sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic")

  # expand
-  sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+  #sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
+  sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander")
+  sink = graph_rewrite(sink, pm_group_for_reduce, name="group for reduce")

  # add locals
-  sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
+  sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
+  #sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")

  # ** devectorizer (full_graph_rewrite) **
  # remove reduce
-  sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")
+  #sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")
+  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(), name="remove_reduce")

  # 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")

  # **** optimizations are done, now we lower to actual code ****

+  sink = graph_rewrite(sink, symbolic_simple+unbroadcast, name="*** unbroadcast")
+
  # add loads and remove invalids
-  sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
+  #sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
+  sink = graph_rewrite(sink, pm_move_regs, name="** add loads")

  # devectorize
-  sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding, ctx=ren, name="devectorize")
-
-  # this is new style (TODO: this should all be removed)
-  sink = graph_rewrite(sink, pm_render, name="pm_render gep/stack")
-  sink = graph_rewrite(sink, pm_remove_vec_dtypes, name="transform to new style")
+  #sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding, ctx=ren, name="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")

+  # 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, pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
@@ -290,6 +290,12 @@ class Scheduler:
            # axes to range number (was done in lowerer)
            tc_upcast_axes = tuple([tuple([(self.rngs[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
            tc_reduce_axes = tuple([self.rngs[a].arg[0] for a in tc_reduce_axes])
+            def with_missing_tc_axes(arg):
+              ret = list(arg)
+              for rn,_ in tc_upcast_axes[0]+tc_upcast_axes[1]:
+                if rn not in [x[0] for x in ret]: ret.append((rn, 1))
+              return tuple(ret)
+            tc_upcast_axes = tuple(with_missing_tc_axes(v) for v in tc_upcast_axes)

            # construct the op
            # TODO: remove tc_upcast_axes from the arg
@@ -269,7 +269,7 @@ SCACHE = ContextVar("SCACHE", 1)
 # allow use of atomics for embedding backward
 USE_ATOMICS = ContextVar("USE_ATOMICS", 0)
 # don't allow broadcast
-DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 1)
+DISALLOW_BROADCAST = ContextVar("DISALLOW_BROADCAST", 0)

@dataclass(frozen=True)
 class Metadata:
@@ -85,7 +85,7 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str:
 def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp:
  if len(arg) == 0: return UOp(Ops.STACK)
  elif len(arg) == 1: return UOp.const(dtypes.weakint, arg[0])
-  else: return UOp(Ops.STACK, dtypes.weakint.vec(len(arg)), tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg))
+  else: return UOp(Ops.STACK, dtypes.weakint, tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg))

 def consumer_map_from_toposort(lst:Iterable[UOp]):
  ret: dict[UOp, dict[UOp, None]] = {}
@@ -289,16 +289,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
        return self.src[0].as_shape if len(self.src) >= 1 else None

      # wmma output shape = accumulator shape (src[2])
-      case Ops.WMMA | Ops.SHAPED_WMMA: return self.src[2]._shape
+      case Ops.WMMA:
+        in0, in1, out0 = self.arg[6]
+        wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
+        return wmma_b + (prod([x for _,x in out0]),)
+      case Ops.SHAPED_WMMA: return self.src[2]._shape

      # 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:
        return self.src[0]._shape
-      # REDUCE with empty axis is passthrough (lowered form)
-      case Ops.REDUCE if len(self.arg[1]) == 0:
-        # these can mismatch if there's a horizonal reduce
-        return (self.dtype.count,) if self.dtype.count > 1 else ()

      # TODO: disallow shape changing bitcast
      case Ops.BITCAST:
@@ -465,8 +465,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  def _stack(self, *srcs):
    # TODO: this should become the real stack
    return UOp(Ops.STACK, self.dtype, (self,)+srcs)
-  def vectorize(self, *srcs):
-    return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs)
+  def vectorize(self, *srcs): return self._stack(*srcs)
  def index(self, *srcs:UOp|int|None, ptr=False, **kwargs):
    new_srcs: list[UOp] = [UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in srcs if x is not None]
    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple(new_srcs), **kwargs)
@@ -1027,18 +1026,20 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  # *** uop high level syntactic sugar ***

  @staticmethod
-  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL):
+  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, is_ptr=True):
    if addrspace is AddrSpace.GLOBAL:
-      ret = UOp(Ops.PARAM, dtype.ptr(prod(shape), addrspace), arg=ParamArg(slot, addrspace=addrspace))
+      # TODO: this should have a shape
+      ret = UOp(Ops.PARAM, dtype.ptr(prod(shape), addrspace) if is_ptr else dtype, arg=ParamArg(slot, addrspace=addrspace))
    else:
      assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
      buf_shape = (prod(shape),) + ((dtype.count,) if dtype.count > 1 else ())
-      ret = UOp(Ops.BUFFER, dtype.ptr(prod(shape), addrspace), src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, addrspace=addrspace))
+      ret = UOp(Ops.BUFFER, dtype.ptr(prod(shape), addrspace) if is_ptr else dtype,
+                src=(shape_to_shape_arg(buf_shape),), arg=ParamArg(slot, addrspace=addrspace))
    if len(shape) > 1: ret = ret.reshape(shape + ((dtype.count,) if addrspace in (AddrSpace.LOCAL, AddrSpace.REG) and dtype.count > 1 else ()))
    return ret
-  def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL):
+  def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL, is_ptr=True):
    assert all_int(self.shape), "no placeholder-like on symbolic shape"
-    return UOp.placeholder(self.max_shard_shape, self.dtype, slot, addrspace)
+    return UOp.placeholder(self.max_shard_shape, self.dtype, slot, addrspace, is_ptr=is_ptr)

  # set is store+end+after
  def set(self:UOp, val:UOp|ConstType, end:UOp|tuple[UOp, ...]|list[UOp]=()) -> UOp:
@@ -126,7 +126,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)),
  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x),
  # ** zero folding **
-  (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x < x -> False
+  (UPat.var("x") < UPat.var("x"), lambda x: x.vconst_like(False, dtypes.bool)), # x < x -> False
  (UPat.var("x") % UPat.var("x"), lambda x: x.const_like(0)), # x%x -> 0
  (UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
  (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
@@ -137,7 +137,7 @@ symbolic_simple = propagate_invalid + PatternMatcher([
  ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"),
   lambda x,mask,c: x // c.arg if c.arg > 0 and c.arg & (c.arg-1) == 0 and mask.arg | (c.arg-1) == -1 else None),
  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
-   lambda x: x.const_like(False).cast(dtypes.bool.vec(x.dtype.count))), # x != x -> False (only ints)
+   lambda x: x.vconst_like(False, dtypes.bool)), # x != x -> False (only ints)
  # ** constant folding **
  (UPat(GroupOp.Unary, src=(UPat((Ops.CONST, Ops.STACK)),), name="a"), fold_const_alu),
  # NOTE: THREEFRY(const,const) folds via its decomposition
@@ -233,7 +233,7 @@ gep_pushing = PatternMatcher([
  # GEP on void is skipped
  (UPat(Ops.GEP, src=(UPat(dtype=dtypes.void, name="x"),)), lambda x: x),
  # GEP in order is removed
-  (UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].dtype.count)) else None),
+  (UPat(Ops.GEP, name="g"), lambda g: g.src[0] if not isinstance(g.dtype, PtrDType) and g.arg == tuple(range(g.src[0].max_numel())) else None),
  # push all GEPs through ALUs for index (TODO: remove this)
  (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, dtype=dtypes.weakint, name='gep'),
   lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \