prettier rendering

extra/replay_pkl.py

@@ -50,6 +50,10 @@ if __name__ == "__main__":
           elif knum == 3:
             k.apply_opt(Opt(op=OptOps.UNROLL, axis=0, arg=4))
             k.apply_opt(Opt(OptOps.UPCAST, 1, 128))
+          elif knum == 29:
+            k.apply_opt(Opt(OptOps.UPCAST, 1, 2))
+            k.apply_opt(Opt(OptOps.PADTO, 1, 128))
+            k.apply_opt(Opt(OptOps.UPCAST, 1, 128))
           else:
             k.hand_coded_optimizations()
           #if knum in [5]: k.apply_opt(Opt(OptOps.UPCAST, 1, 2))

tinygrad/runtime/ops_dsp.py

@@ -26,6 +26,9 @@ dsp_pm_late = PatternMatcher([
   (UPat.var("x")//UPat(Ops.VECTORIZE,src=UPat.var("y")), lambda x,y: x//UOp(Ops.CUSTOMI,x.dtype,(y,),arg="{0}") if x.op is not Ops.CUSTOMI else None),
   (UPat(Ops.DEFINE_ACC, src=(UPat(Ops.VECTORIZE, src=UPat(Ops.CONST, arg=0)),), dtype=dtypes.uchar.vec(128), name="d", allow_any_len=True),
    lambda d: d.replace(src=(UOp(Ops.CUSTOMI, d.dtype, arg="__builtin_HEXAGON_V6_vd0_128B()"),)+d.src[1:])),
+  # makes rendering nicer
+  (UPat(Ops.VECTORIZE, src=UPat(Ops.CONST, dtype=(dtypes.uint8, dtypes.int8)), name="v"),
+   lambda v: UOp(Ops.VECTORIZE, v.dtype, src=tuple(UOp.const(dtypes.int, x.arg) for x in v.src))),
 ])
 
 class DSPRenderer(ClangRenderer):

tinygrad/spec.py

@@ -109,7 +109,9 @@ spec = PatternMatcher([
 
   (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),
-  (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.count and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
+  # vectorize parents have to match dtype (unless it's a const for prettier rendering)
+  (UPat(Ops.VECTORIZE, name="x"),
+   lambda x: len(x.src)>1 and len(x.src) == x.dtype.count and all(y.op is Ops.CONST or x.dtype == y.dtype.vec(len(x.src)) for y in x.src)),
   (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: x.arg is None),
   (UPat(Ops.BARRIER, dtypes.void, src=UPat(Ops.STORE, allow_any_len=True)), lambda: True), # NOTE: all pointers must be local
   (UPat(Ops.BARRIER, dtypes.void), lambda: True), # BARRIERs can also happen at the end of loops