Repository: tinygrad/tinygrad
Pull request: Make vectorization of CONST explicit
Resolves #5313

Problem statement:
Use UOps.VECTORIZE for consts
const renderers have an extra call for rendering a vectorized const. This can simplify to:
```
  # returns a str expression of the const with the given type
  def render_const(self, x:ConstType, dtype:DType) -> str:
    if math.isnan(x): return self.render_cast("NAN", dtype)
    elif math.isinf(x): return self.render_cast(("-" if x < 0 else "") + "INFINITY", dtype)
    elif dtype == dtypes.bool: return "1" if x else "0"
    elif dtype == dtypes.float: return f"{x}f"
    return str(x)
```

Current UOp:
`UOp(UOps.CONST, dtypes.float.vec(4), (), 0.0)`
Desired:
`UOp(UOps.VECTORIZE, dtypes.float.vec(4), (UOp(UOps.CONST, dtypes.float, (), 0.0), ))`

Should also fix DEFINE_ACC's extra type_verify assert.

Hints from issue discussion:
new to project, can I take this one?

Commit message: remove dead code in PTX CONST render

Patch to review:
diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py
index 2f598de887fe5..1f8f597b14ad8 100644
--- a/test/test_uop_graph.py
+++ b/test/test_uop_graph.py
@@ -119,15 +119,6 @@ def test_const_cast(self):
     self.assertEqual(out.op, UOps.CONST)
     self.assertEqual(out.arg, 0)
 
-  def test_const_vectorize_fold(self):
-    c0 = UOp(UOps.CONST, dtypes.half, arg=0.0)
-    out = UOp(UOps.VECTORIZE, dtypes.half.vec(2), (c0, c0))
-    g = UOpGraph([out])
-    self.assertEqual(len(g.uops), 1)
-    out = g.uops[-1]
-    self.assertEqual(out.op, UOps.CONST)
-    self.assertEqual(out.arg, 0.0)
-
   def test_noop_vectorize_fold(self):
     d0 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), arg=(0, True))
     idx = UOp.const(dtypes.int, 0)
@@ -171,6 +162,178 @@ def _test_vec(geps):
     xy2 = tuple(UOp(UOps.GEP, dtypes.float, (val2, ), i) for i in range(2))
     self.assertIs(_test_vec(xy1+xy2).op, UOps.VECTORIZE)
 
+  def test_gep_vec_const_fold(self):
+    c0 = UOp.const(dtypes.float, 0.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    vec2 = UOp(UOps.VECTORIZE, dtypes.float.vec(2), (c0, c1))
+    gep0 = UOp(UOps.GEP, dtypes.float, (vec2,), 0)
+    gep1 = UOp(UOps.GEP, dtypes.float, (vec2,), 1)
+    g = UOpGraph([gep0, gep1])
+    self.assert_equiv_uops(g.uops[0], c0)
+    self.assert_equiv_uops(g.uops[1], c1)
+
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    vec4 = UOp(UOps.VECTORIZE, dtypes.float.vec(4), (c0, c1, c2, c3))
+    gep2 = UOp(UOps.GEP, dtypes.float, (vec4,), 2)
+    gep3 = UOp(UOps.GEP, dtypes.float, (vec4,), 3)
+    g = UOpGraph([gep0, gep1, gep2, gep3])
+    self.assert_equiv_uops(g.uops[0], c0)
+    self.assert_equiv_uops(g.uops[1], c1)
+    self.assert_equiv_uops(g.uops[2], c2)
+    self.assert_equiv_uops(g.uops[3], c3)
+
+    c4 = UOp.const(dtypes.float, 4.0)
+    c5 = UOp.const(dtypes.float, 5.0)
+    c6 = UOp.const(dtypes.float, 6.0)
+    c7 = UOp.const(dtypes.float, 7.0)
+    vec8 = UOp(UOps.VECTORIZE, dtypes.float.vec(8), (c0, c1, c2, c3, c4, c5, c6, c7))
+    gep4 = UOp(UOps.GEP, dtypes.float, (vec8,), 4)
+    gep5 = UOp(UOps.GEP, dtypes.float, (vec8,), 5)
+    gep6 = UOp(UOps.GEP, dtypes.float, (vec8,), 6)
+    gep7 = UOp(UOps.GEP, dtypes.float, (vec8,), 7)
+    g = UOpGraph([gep0, gep1, gep2, gep3, gep4, gep5, gep6, gep7])
+    self.assert_equiv_uops(g.uops[0], c0)
+    self.assert_equiv_uops(g.uops[1], c1)
+    self.assert_equiv_uops(g.uops[2], c2)
+    self.assert_equiv_uops(g.uops[3], c3)
+    self.assert_equiv_uops(g.uops[4], c4)
+    self.assert_equiv_uops(g.uops[5], c5)
+    self.assert_equiv_uops(g.uops[6], c6)
+    self.assert_equiv_uops(g.uops[7], c7)
+
+  def test_wmma_vectorize_fold(self):
+    for i in [2, 4, 8]:
+      vec = UOp(UOps.VECTORIZE, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
+      var = UOp.var(dtype=dtypes.half.vec(i))
+      acc = UOp.var('acc', dtype=dtypes.half.vec(i))
+      wmma = UOp(UOps.WMMA, dtypes.half.vec(i), (vec, var, acc))
+      g = UOpGraph([wmma])
+      self.assert_equiv_uops(g.uops[0], acc)
+      self.assertEqual(len(g.uops), 1)
+
+    for i in [2, 4, 8]:
+      var = UOp.var(dtype=dtypes.half.vec(i))
+      vec = UOp(UOps.VECTORIZE, dtypes.half.vec(i), tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
+      acc = UOp.var('acc', dtype=dtypes.half.vec(i))
+      wmma = UOp(UOps.WMMA, dtypes.half.vec(i), (var, vec, acc))
+      g = UOpGraph([wmma])
+      self.assert_equiv_uops(g.uops[0], acc)
+      self.assertEqual(len(g.uops), 1)
+
+  def test_wmma_vectorize_no_fold(self):
+    for i in [4, 8]:
+      vec = UOp(UOps.VECTORIZE, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) + tuple(UOp.var(f"tmp{j}", dtype=dtypes.half) for j in range(i//2)))
+      var = UOp.var(f"tmp{i}", dtype=dtypes.half.vec(i))
+      acc = UOp.var('acc', dtype=dtypes.half.vec(i))
+      wmma = UOp(UOps.WMMA, dtypes.half.vec(i), (vec, var, acc))
+      g = UOpGraph([wmma])
+      self.assert_equiv_uops(g.uops[-1], wmma)
+
+    for i in [4, 8]:
+      var = UOp.var(f"tmp{i}", dtype=dtypes.half.vec(i))
+      vec = UOp(UOps.VECTORIZE, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) + tuple(UOp.var(f"tmp{j}", dtype=dtypes.half) for j in range(i//2)))
+      acc = UOp.var('acc', dtype=dtypes.half.vec(i))
+      wmma = UOp(UOps.WMMA, dtypes.half.vec(i), (var, vec, acc))
+      g = UOpGraph([wmma])
+      self.assert_equiv_uops(g.uops[-1], wmma)
+
+    for i in [2, 4, 8]:
+      vec = UOp(UOps.VECTORIZE, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
+      var = UOp.var(f"tmp{i}", dtype=dtypes.half.vec(i))
+      acc = UOp.var('acc', dtype=dtypes.half.vec(i))
+      wmma = UOp(UOps.WMMA, dtypes.half.vec(i), (vec, var, acc))
+      g = UOpGraph([wmma])
+      self.assert_equiv_uops(g.uops[-1], wmma)
+
+    for i in [2, 4, 8]:
+      var = UOp.var(f"tmp{i}", dtype=dtypes.half.vec(i))
+      vec = UOp(UOps.VECTORIZE, dtypes.half.vec(i),
+                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
+      acc = UOp.var('acc', dtype=dtypes.half.vec(i))
+      wmma = UOp(UOps.WMMA, dtypes.half.vec(i), (var, vec, acc))
+      g = UOpGraph([wmma])
+      self.assert_equiv_uops(g.uops[-1], wmma)
+
+  def test_phi_vec_const_fold(self):
+    c0 = UOp.const(dtypes.float, 0.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    vec2 = UOp(UOps.VECTORIZE, dtypes.float.vec(2), (c0, c1))
+    x2 = UOp(UOps.DEFINE_VAR, dtypes.float.vec(2))
+    phi2 = UOp(UOps.PHI, dtypes.float.vec(2), (vec2, x2))
+    g = UOpGraph([phi2])
+    self.assert_equiv_uops(g.uops[0], x2)
+
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    vec4 = UOp(UOps.VECTORIZE, dtypes.float.vec(4), (c0, c1, c2, c3))
+    x4 = UOp(UOps.DEFINE_VAR, dtypes.float.vec(4))
+    phi4 = UOp(UOps.PHI, dtypes.float.vec(4), (vec4, x4))
+    g = UOpGraph([phi4])
+    self.assert_equiv_uops(g.uops[0], x4)
+
+    c4 = UOp.const(dtypes.float, 4.0)
+    c5 = UOp.const(dtypes.float, 5.0)
+    c6 = UOp.const(dtypes.float, 6.0)
+    c7 = UOp.const(dtypes.float, 7.0)
+    vec8 = UOp(UOps.VECTORIZE, dtypes.float.vec(8), (c0, c1, c2, c3, c4, c5, c6, c7))
+    x8 = UOp(UOps.DEFINE_VAR, dtypes.float.vec(8))
+    phi8 = UOp(UOps.PHI, dtypes.float.vec(8), (vec8, x8))
+    g = UOpGraph([phi8])
+    self.assert_equiv_uops(g.uops[0], x8)
+
+  def test_gep_acc_vec_fold(self):
+    c0 = UOp.const(dtypes.float, 0.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    vec2 = UOp(UOps.VECTORIZE, dtypes.float.vec(2), (c0, c1))
+    acc2 = UOp(UOps.DEFINE_ACC, dtypes.float.vec(2), (vec2,))
+    gep0 = UOp(UOps.GEP, dtypes.float, (acc2,), 0)
+    gep1 = UOp(UOps.GEP, dtypes.float, (acc2,), 1)
+    g = UOpGraph([gep0, gep1])
+    self.assert_equiv_uops(g.uops[0], c0)
+    self.assert_equiv_uops(g.uops[1], c1)
+
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    vec4 = UOp(UOps.VECTORIZE, dtypes.float.vec(4), (c0, c1, c2, c3))
+    acc4 = UOp(UOps.DEFINE_ACC, dtypes.float.vec(4), (vec4,))
+    gep0 = UOp(UOps.GEP, dtypes.float, (acc4,), 0)
+    gep1 = UOp(UOps.GEP, dtypes.float, (acc4,), 1)
+    gep2 = UOp(UOps.GEP, dtypes.float, (acc4,), 2)
+    gep3 = UOp(UOps.GEP, dtypes.float, (acc4,), 3)
+    g = UOpGraph([gep0, gep1, gep2, gep3])
+    self.assert_equiv_uops(g.uops[0], c0)
+    self.assert_equiv_uops(g.uops[1], c1)
+    self.assert_equiv_uops(g.uops[2], c2)
+    self.assert_equiv_uops(g.uops[3], c3)
+
+    c4 = UOp.const(dtypes.float, 4.0)
+    c5 = UOp.const(dtypes.float, 5.0)
+    c6 = UOp.const(dtypes.float, 6.0)
+    c7 = UOp.const(dtypes.float, 7.0)
+    vec8 = UOp(UOps.VECTORIZE, dtypes.float.vec(8), (c0, c1, c2, c3, c4, c5, c6, c7))
+    acc8 = UOp(UOps.DEFINE_ACC, dtypes.float.vec(8), (vec8,))
+    gep0 = UOp(UOps.GEP, dtypes.float, (acc8,), 0)
+    gep1 = UOp(UOps.GEP, dtypes.float, (acc8,), 1)
+    gep2 = UOp(UOps.GEP, dtypes.float, (acc8,), 2)
+    gep3 = UOp(UOps.GEP, dtypes.float, (acc8,), 3)
+    gep4 = UOp(UOps.GEP, dtypes.float, (acc8,), 4)
+    gep5 = UOp(UOps.GEP, dtypes.float, (acc8,), 5)
+    gep6 = UOp(UOps.GEP, dtypes.float, (acc8,), 6)
+    gep7 = UOp(UOps.GEP, dtypes.float, (acc8,), 7)
+    g = UOpGraph([gep0, gep1, gep2, gep3, gep4, gep5, gep6, gep7])
+    self.assert_equiv_uops(g.uops[0], c0)
+    self.assert_equiv_uops(g.uops[1], c1)
+    self.assert_equiv_uops(g.uops[2], c2)
+    self.assert_equiv_uops(g.uops[3], c3)
+    self.assert_equiv_uops(g.uops[4], c4)
+    self.assert_equiv_uops(g.uops[5], c5)
+    self.assert_equiv_uops(g.uops[6], c6)
+    self.assert_equiv_uops(g.uops[7], c7)
+
   def test_cast_alu_fold(self):
     d0 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.bool), arg=(0, True))
     d1 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg=(1, False))
diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py
index ae4653bccf6e8..c84913e2b467f 100644
--- a/tinygrad/codegen/uopgraph.py
+++ b/tinygrad/codegen/uopgraph.py
@@ -157,8 +157,10 @@ def index_collapse(idx,rng,buf,add,mul,ld,reduce_allow_any_len):
   *[(UOp(UOps.VECTORIZE, dtypes.float.vec(i), tuple(UOp(UOps.GEP, dtypes.float,
                          src=(UOp.var('x', dtype=dtypes.float.vec(i)),), arg=j) for j in range(i))), lambda x: x) for i in [2, 4, 8]],
   # tensor core with a 0 input is acc
-  (UOp(UOps.WMMA, src=(UOp.const(None, 0.0), UOp.var(), UOp.var('acc'))), lambda acc: acc),
-  (UOp(UOps.WMMA, src=(UOp.var(), UOp.const(None, 0.0), UOp.var('acc'))), lambda acc: acc),
+  *[(UOp(UOps.WMMA, src=(UOp(UOps.VECTORIZE, src=tuple(UOp.const(None, 0.0) for _ in range(i))), UOp.var(), UOp.var('acc'))),
+     lambda acc: acc) for i in [2, 4, 8]],
+  *[(UOp(UOps.WMMA, src=(UOp.var(), UOp(UOps.VECTORIZE, src=tuple(UOp.const(None, 0.0) for _ in range(i))), UOp.var('acc'))),
+     lambda acc: acc) for i in [2, 4, 8]],
   # tensor core cleanups
   *[(UOp(UOps.REDUCE, src=(UOp(UOps.EXPAND, src=tuple(UOp(UOps.GEP, dtypes.float, src=(UOp.var('x'),), arg=i) for i in range(j))).name("expand"),))
     .name("reduce_allow_any_len"), reduce_before_expand) for j in [2,4,8]],
@@ -167,9 +169,9 @@ def index_collapse(idx,rng,buf,add,mul,ld,reduce_allow_any_len):
   # threefry
   (UOp(UOps.ALU, dtype=dtypes.uint64, src=(UOp.var("x"), UOp.var("seed")), arg=BinaryOps.THREEFRY), threefry2x32),
   # sum collapse to mul (with possible GEP)
-  (UPat(UOps.PHI, src=(UPat(UOps.DEFINE_ACC, name="phi_input", src=[UPat(UOps.CONST), UPat(UOps.RANGE, name="loop")]),
+  (UPat(UOps.PHI, src=(UPat(UOps.DEFINE_ACC, name="phi_input", src=[UPat({UOps.VECTORIZE, UOps.CONST}), UPat(UOps.RANGE, name="loop")]),
                        UPat(UOps.ALU, BinaryOps.ADD, src=(UPat(name="val1"), UPat(name="val2"))))), sum_collapse),
-  (UPat(UOps.PHI, src=(UPat(UOps.GEP, name="phi_input", src=(UPat(UOps.DEFINE_ACC, src=[UPat(UOps.CONST), UPat(UOps.RANGE, name="loop")]),)),
+  (UPat(UOps.PHI, src=(UPat(UOps.GEP, name="phi_input", src=(UPat(UOps.DEFINE_ACC, src=[UPat(UOps.VECTORIZE), UPat(UOps.RANGE, name="loop")]),)),
                        UPat(UOps.ALU, BinaryOps.ADD, src=(UPat(name="val1"), UPat(name="val2"))))), sum_collapse),
   # extra arange loop folding because we don't fold adds. TODO: fold adds
   (UOp(UOps.REDUCE, src=((UOp.var("idx") + UOp.cvar("mval") * UOp(UOps.RANGE, src=(UOp.var("loop_start"), UOp.var("loop_end"))).name("rng") +
@@ -203,14 +205,18 @@ def index_collapse(idx,rng,buf,add,mul,ld,reduce_allow_any_len):
   (UOp.max(UOp.cvar('c'), -(UOp(UOps.RANGE).name('s')+UOp.cvar('c2'))), lambda c,s,c2: -(s+c2) if -(s.src[1].arg-1+c2.arg) >= c.arg else None),
   # const rules
   (UOp(UOps.GEP, src=(UOp.cvar("c"),)).name("root"), lambda root, c: root.const(c.arg)),
+  *[(UPat(UOps.GEP, name="root", src=(UPat(UOps.VECTORIZE, name="v", src=[UPat(UOps.CONST) for _ in range(i)]))),
+     lambda root, v: root.const(v.src[root.arg].arg)) for i in [2, 4, 8]],
   (UPat(UOps.CAST, name="root", src=UPat(UOps.CONST, name="c")), lambda root, c: root.const(c.arg)),
-  (UPat(UOps.VECTORIZE, name="root", src=UPat(UOps.CONST, name="c")), lambda root, c: root.const(c.arg)),
-  # a phi on a DEFINE_ACC without loops or a CONST is a noop. this is for correctness, not just speed
+  # a phi on a DEFINE_ACC without loops, CONST, or a vectorized CONST is a noop. this is for correctness, not just speed
   (UOp(UOps.PHI, src=(UOp(UOps.DEFINE_ACC).name("acc"), UOp.var("acc"))), lambda acc: UOp.cast(acc.src[0], acc.dtype)),
   (UOp(UOps.PHI, src=(UOp(UOps.DEFINE_ACC, src=(UOp.cvar(),)), UOp.var("x"))), lambda x: x),
   (UOp(UOps.PHI, src=(UOp.cvar(), UOp.var("x"))), lambda x: x),
+  *[(UPat(UOps.PHI, src=(UPat(UOps.VECTORIZE, src=[UPat(UOps.CONST) for _ in range(i)]), UPat(name="x"))), lambda x: x) for i in [2, 4, 8]],
   # a DEFINE_ACC without inputs is a const + GEP on a const is the const
-  (UOp(UOps.DEFINE_ACC, src=(UOp.cvar(),)).name("root"), lambda root: UOp.cast(root.src[0], root.dtype)),
+  (UOp(UOps.DEFINE_ACC, src=(UOp.var(),)).name("root"), lambda root: UOp.cast(root.src[0], root.dtype)),
+  *[(UPat(UOps.GEP, name="root", src=(UPat(UOps.DEFINE_ACC, src=(UPat(UOps.VECTORIZE, name="v", src=[UPat(UOps.CONST) for _ in range(i)]))))),
+     lambda root, v: root.const(v.src[root.arg].arg)) for i in [2, 4, 8]],
   (UOp(UOps.GEP, src=(UOp.cvar("x"),)).name("root"), lambda root,x: root.const(x.arg)),
   # a conditional with the same results either way is a noop, also fold const conditionals
   (UOp.var().where(UOp.var("val"), UOp.var("val")), lambda val: val),
@@ -307,11 +313,11 @@ def index_collapse(idx,rng,buf,add,mul,ld,reduce_allow_any_len):
   (UOp(UOps.CAST).name("root"), lambda root: root.src[0] if str(root.dtype) == str(root.src[0].dtype) else None),
   (UOp(UOps.VECTORIZE).name("root"), lambda root: root.src[0] if str(root.dtype) == str(root.src[0].dtype) else None),
   # fold gated LOAD/STORE
-  (UOp.load(UOp.var("buf"), UOp.var("idx"), UOp.const(dtypes.bool, True), UOp.cvar("var")), lambda buf,idx,var: UOp.load(buf, idx, dtype=var.dtype)),
-  (UOp.load(UOp.var("buf"), UOp.var("idx"), UOp.const(dtypes.bool, True), UOp.cvar("var"), UOp.var("barrier")),
+  (UOp.load(UOp.var("buf"), UOp.var("idx"), UOp.const(dtypes.bool, True), UOp.var("var")), lambda buf,idx,var: UOp.load(buf, idx, dtype=var.dtype)),
+  (UOp.load(UOp.var("buf"), UOp.var("idx"), UOp.const(dtypes.bool, True), UOp.var("var"), UOp.var("barrier")),
    lambda buf,idx,var,barrier: UOp.load(buf, idx, barrier, dtype=var.dtype)),
-  (UOp.load(UOp.var(), UOp.var(), UOp.const(dtypes.bool, False), UOp.cvar("var")), lambda var: var),
-  (UOp.load(UOp.var(), UOp.var(), UOp.const(dtypes.bool, False), UOp.cvar("var"), UOp.var()), lambda var: var),
+  (UOp.load(UOp.var(), UOp.var(), UOp.const(dtypes.bool, False), UOp.var("var")), lambda var: var),
+  (UOp.load(UOp.var(), UOp.var(), UOp.const(dtypes.bool, False), UOp.var("var"), UOp.var()), lambda var: var),
   (UOp.store(UOp.var("buf"), UOp.var("idx"), UOp.var("val"), UOp.const(dtypes.bool, True)), UOp.store),
   (UOp.store(UOp.var(), UOp.var(), UOp.var(), UOp.const(dtypes.bool, False)), lambda: UOp(UOps.NOOP)),
   # remove NOOPs from SINK
@@ -374,7 +380,7 @@ def do_reduce_with_expand(root):
   expands = [x for x in root.src[1:] if x.op is UOps.EXPAND]
   expands_reduce = [x for x in expands if root.src[0].op is UOps.EXPAND and all(y in root.src[0].arg for y in x.arg)]
   expands_non_reduce = [x for x in expands if x not in expands_reduce]
-  const = UOp.const(root.dtype.scalar(), dtypes.as_const(0, root.dtype.scalar()) if root.arg is ReduceOps.SUM else dtypes.min(root.dtype.scalar()))
+  const = UOp.const(root.dtype, dtypes.as_const(0, root.dtype.scalar()) if root.arg is ReduceOps.SUM else dtypes.min(root.dtype.scalar()))
   ret = acc = UOp(UOps.DEFINE_ACC, root.dtype, (const,) + tuple(x for x in root.src[1:] if x.op is not UOps.EXPAND), (acc_number,))
   acc_number += 1
   alu_op = {ReduceOps.SUM:BinaryOps.ADD, ReduceOps.MAX:BinaryOps.MAX}[cast(ReduceOps, root.arg)]
diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py
index 74989703f8983..57fcecd20433d 100644
--- a/tinygrad/codegen/uops.py
+++ b/tinygrad/codegen/uops.py
@@ -75,6 +75,9 @@ def const(self:Union[UOp, DType, None], b:ConstType|Variable): return UOp._const
   def _const(dtype:Optional[DType], b:ConstType|Variable):
     # TODO: fix dtype of b.max after Variable is just an UOp
     if isinstance(b, Variable): return UOp(UOps.DEFINE_VAR, dtype, (UOp.const(dtypes.int, b.min), UOp.const(dtypes.int, cast(int,b.max))), b)
+    if dtype is not None and dtype != dtype.scalar():
+      return UOp(UOps.VECTORIZE, dtype, src=tuple(
+        UOp(UOps.CONST, dtype.scalar(), arg=dtypes.as_const(b, dtype.scalar())) for _ in range(dtype.count)))
     return UOp(UOps.CONST, dtype, arg=dtypes.as_const(b, dtype) if dtype is not None else b)
   def alu(self, arg, *src:UOp):
     return UOp(UOps.ALU, dtypes.bool if arg in {BinaryOps.CMPLT, BinaryOps.CMPNE} else (self, *src)[-1].dtype, (self,)+src, arg)
@@ -99,9 +102,9 @@ def divides(self, v):
       if self.arg is BinaryOps.MUL: return any(x.divides(v) for x in self.src)
     return False # generic false if we aren't sure
   @functools.cached_property
-  def vmin(self) -> UOp: return x if (x:=self._min_max[0]) is not None else self.const(dtypes.min(cast(DType, self.dtype)))
+  def vmin(self) -> UOp: return x if (x:=self._min_max[0]) is not None else self.const(dtypes.min(cast(DType, self.dtype))).vmin
   @functools.cached_property
-  def vmax(self) -> UOp: return x if (x:=self._min_max[1]) is not None else self.const(dtypes.max(cast(DType, self.dtype)))
+  def vmax(self) -> UOp: return x if (x:=self._min_max[1]) is not None else self.const(dtypes.max(cast(DType, self.dtype))).vmax
   @functools.cached_property
   def _min_max(self) -> Tuple[Optional[UOp], Optional[UOp]]:
     # NOTE: returned UOp is assumed to be CONST
@@ -109,11 +112,13 @@ def _min_max(self) -> Tuple[Optional[UOp], Optional[UOp]]:
     # TODO: UOps.SPECIAL is UOps.DEFINE_VAR
     if self.op is UOps.SPECIAL: return self.const(0), self.const(self.arg[1]-1) if isinstance(self.arg[1], int) else None
     if self.op is UOps.CONST: return self, self
+    if self.op is UOps.VECTORIZE and all(x.op == UOps.CONST for x in self.src):
+      return min(self.src, key=lambda x: x.arg), max(self.src, key=lambda x: x.arg)
     if self.op is UOps.ALU:
       if self.arg is UnaryOps.NEG and self.dtype != dtypes.bool and not dtypes.is_unsigned(cast(DType, self.dtype)):
-        return self.const(-self.src[0].vmax.arg), self.const(-self.src[0].vmin.arg)
+        return self.const(-self.src[0].vmax.arg).vmin, self.const(-self.src[0].vmin.arg).vmax
       if self.arg is BinaryOps.ADD:
-        return self.const(self.src[0].vmin.arg+self.src[1].vmin.arg), self.const(self.src[0].vmax.arg+self.src[1].vmax.arg)
+        return self.const(self.src[0].vmin.arg+self.src[1].vmin.arg).vmin, self.const(self.src[0].vmax.arg+self.src[1].vmax.arg).vmin
     return None, None
 
 class UPat:
@@ -183,10 +188,11 @@ def type_verify(uops):
   for u in uops:
     uop, arg, src, dtype = u.op, u.arg, u.src, u.dtype
     if uop in {UOps.CONST, UOps.DEFINE_ACC}:
+      if uop is UOps.CONST:
+        assert dtype is not None and dtype == dtype.scalar(), f"consts should be scalar, got {dtype}"
+        assert type(arg) is type(dtypes.as_const(arg, dtype)), f"type of {arg=} does not match {dtype}"
       if uop is UOps.DEFINE_ACC:
-        assert dtype is not None and src[0].dtype == dtype.scalar(), f"type of {src[0].dtype=} must be a scalar {dtype.scalar()}"
-        arg = src[0].arg
-      assert dtype is not None and type(arg) is type(dtypes.as_const(arg, dtype)), f"type of {arg=} does not match {dtype}"
+        assert dtype is not None and src[0].dtype == dtype, f"dtype mismatch {src[0].dtype=} != {dtype=}"
     if uop in {UOps.CAST, UOps.BITCAST, UOps.VECTORIZE}: assert arg is None and dtype is not None # type is the output type, not an arg
     if uop is UOps.CAST: assert dtype.count == 1 and len(src) == 1
     if uop is UOps.VECTORIZE:
diff --git a/tinygrad/renderer/assembly.py b/tinygrad/renderer/assembly.py
index c5b25ccdfd10c..e2c3728e293e8 100644
--- a/tinygrad/renderer/assembly.py
+++ b/tinygrad/renderer/assembly.py
@@ -159,7 +159,7 @@ def _cast(a, dtype:DType, atype:DType, bitcast=False, u=None, pred=False):
         elif uop is UOps.DEFINE_ACC:
           if dtype.count > 1:
             r[u] = [ssa('acc', dtype=self.types[dtype.scalar()]) for _ in range(dtype.count)]
-            for uu in r[u]: kk(f"mov.b{self.types[dtype.scalar()][1:]} {uu}, {const(src[0].arg, dtype.scalar())};")
+            for uu in r[u]: kk(f"mov.b{self.types[dtype.scalar()][1:]} {uu}, {const(src[0].src[0].arg, dtype.scalar())};")
           else: kk(f"mov.{f'b{self.types[dtype][1:]}' if dtype != dtypes.bool else 'pred'} {ssa('acc', u)}, {const(src[0].arg, dtype)};")
         elif uop is UOps.SPECIAL:
           assert args[0][0] != "i", "idx not supported"
@@ -170,9 +170,7 @@ def _cast(a, dtype:DType, atype:DType, bitcast=False, u=None, pred=False):
           bufs.append((args.expr, dtype))
           r[u] = f"%{args.expr}"
           kk(*self.render_load(args.expr, ssa('dat', u, self.types[dtype]), dtype, ss=".param"))
-        elif uop is UOps.CONST:
-          if dtype.count > 1: r[u] = [const(args, dtype.scalar(), mov=True) for _ in range(dtype.count)]
-          else: r[u] = const(args, dtype, mov=True)
+        elif uop is UOps.CONST: r[u] = const(args, dtype, mov=True)
         elif uop is UOps.GEP: r[u] = r[src[0]][u.arg]
         elif uop is UOps.LOAD:
           assert src[0].dtype == dtypes.int64, "load isn't int64"
diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py
index 5dd07ea7cd652..a3e85b30c3c90 100644
--- a/tinygrad/renderer/cstyle.py
+++ b/tinygrad/renderer/cstyle.py
@@ -171,7 +171,7 @@ def ssa(prefix:str, u:Optional[UOp]=None):
           bufs.append((nm:=f"data{args[0]}", (dtype,args[1])))
           r[u] = nm
         elif uop is UOps.WMMA: kk(f"{self.render_dtype(dtype)} {ssa('wmma',u)} = __{args[0]}({r[src[0]]}, {r[src[1]]}, {r[src[2]]});")
-        elif uop is UOps.DEFINE_ACC: kk(f"{self.render_dtype(dtype)} {ssa('acc',u)} = {self.render_const(src[0].arg, dtype)};")
+        elif uop is UOps.DEFINE_ACC: kk(f"{self.render_dtype(dtype)} {ssa('acc',u)} = {r[src[0]]};")
         elif uop is UOps.CONST: r[u] = self.render_const(args, dtype) if args >= 0 else f"({self.render_const(args, dtype)})"
         elif uop is UOps.GEP:
           assert src[0].dtype is not None
diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py
index f523a3ae6c33b..aca011ab0d7de 100644
--- a/tinygrad/runtime/ops_python.py
+++ b/tinygrad/runtime/ops_python.py
@@ -86,10 +86,9 @@ def __call__(self, *bufs, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tup
         elif uop is UOps.SPECIAL:
           if arg[0][0] == 'g': ul[i] = [idxs[2-int(arg[0][-1])]] * warp_size
           elif arg[0][0] == 'l': ul[i] = [x[2-int(arg[0][-1])] for x in warp]
-        elif uop is UOps.CONST:
-          ul[i] = [[arg] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [arg] * warp_size
+        elif uop is UOps.CONST: ul[i] = [arg] * warp_size
         elif uop is UOps.DEFINE_ACC:
-          ul[i] = [[inp[0][0]] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [inp[0][0]] * warp_size
+          ul[i] = [[inp[0][0][0]] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [inp[0][0]] * warp_size
         elif uop is UOps.RANGE:
           if i not in ul: ul[i] = [inp[0][0]] * warp_size
           else:


Code review comments:
diff_hunk: @@ -99,21 +102,23 @@ def divides(self, v):
       if self.arg is BinaryOps.MUL: return any(x.divides(v) for x in self.src)
     return False # generic false if we aren't sure
   @functools.cached_property
-  def vmin(self) -> UOp: return x if (x:=self._min_max[0]) is not None else self.const(dtypes.min(cast(DType, self.dtype)))
+  def vmin(self) -> UOp: return x if (x:=self._min_max[0]) is not None else self.const(dtypes.min(cast(DType, self.dtype))).vmin
   @functools.cached_property
-  def vmax(self) -> UOp: return x if (x:=self._min_max[1]) is not None else self.const(dtypes.max(cast(DType, self.dtype)))
+  def vmax(self) -> UOp: return x if (x:=self._min_max[1]) is not None else self.const(dtypes.max(cast(DType, self.dtype))).vmax
   @functools.cached_property
   def _min_max(self) -> Tuple[Optional[UOp], Optional[UOp]]:
     # NOTE: returned UOp is assumed to be CONST
     if self.op in (UOps.DEFINE_VAR, UOps.RANGE): return self.src[0], self.src[1] if isinstance(self.src[1].arg, int) else None
     # TODO: UOps.SPECIAL is UOps.DEFINE_VAR
     if self.op is UOps.SPECIAL: return self.const(0), self.const(self.arg[1]-1) if isinstance(self.arg[1], int) else None
     if self.op is UOps.CONST: return self, self
+    if self.op is UOps.VECTORIZE and all(x.op == UOps.CONST for x in self.src):
+      return min(self.src, key=lambda x: x.arg), max(self.src, key=lambda x: x.arg)
path: tinygrad/codegen/uops.py
text: @user1:
hmm, I think this is extra complexity. you always have the same const across a vectorize.
Can it just be self.const(dtypes.max(self.dtype.scalar())) in the above?

@author:
> you always have the same const across a vectorize.

Not true. I've seen different value consts being vectorized in test_real_world.py. I was surprised by this and I recorded it in my notes while I was hacking at the PTX renderer (see the unhighlighted white text above the diff):
![Pasted Graphic 3](https://github.com/user-attachments/assets/58100559-cb53-47a7-a5fa-883ba1e06e8e)

I agree that it is extra complexity though -- I definitely need to cleanup the vmin & vmax logic. Will try to have a fix out shortly, but I need to understand the upstream behavior a little better. chenyu's generic max folding was a sore point for my rebasing haha

@user1:
ohh right you can have a const folded valid mixed with invalid.

@author:
I think I've got the vmin/vmax stuff reasonably clean, if the extra `sconst` def is acceptable. The `sconst` def's purpose is to save lines and satisfy mypy, otherwise we'd have casts everywhere.

I checked the process_replay and it's consistent with what we were seeing before the cleanup.