diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py index 7380ac3bb2..c057cdc16d 100644 --- a/extra/replay_pkl.py +++ b/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)) diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index c04ce37cc5..f6cf53da1f 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/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): diff --git a/tinygrad/spec.py b/tinygrad/spec.py index c94677abcc..20063acdb5 100644 --- a/tinygrad/spec.py +++ b/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