From 880b4a5e47b4783a542a71accabe2d288a88b55b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 26 Mar 2025 15:34:57 +0800 Subject: [PATCH] put that back --- extra/replay_pkl.py | 9 ++++++++- tinygrad/runtime/ops_dsp.py | 10 +++++----- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py index e44bf9c316..1598b8cd3e 100644 --- a/extra/replay_pkl.py +++ b/extra/replay_pkl.py @@ -106,9 +106,16 @@ if __name__ == "__main__": #k.apply_opt(Opt(OptOps.UPCAST, 0, 2)) """ # blocked - if len(k.full_shape) == 6 and k.full_shape[-4:] == (32,3,3,3): + if knum == 1: k.apply_opt(Opt(OptOps.UPCAST, 2, 32)) k.apply_opt(Opt(OptOps.UPCAST, 1, 4)) + elif knum == 4: + k.apply_opt(Opt(OptOps.UNROLL, 0, 8)) + k.apply_opt(Opt(OptOps.UPCAST, 2, 32)) + k.apply_opt(Opt(OptOps.UPCAST, 1, 4)) + elif knum == 66: + k.apply_opt(Opt(OptOps.UNROLL, 0, 4)) + k.apply_opt(Opt(OptOps.UPCAST, 0, 8)) elif k.full_shape[-3:] == (32,3,3): #if k.full_shape[-4]%4 != 0: k.apply_opt(Opt(OptOps.PADTO, len(k.full_shape)-4, 4)) # 3x3 dwconv diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 7a654069e4..ef01aa68f9 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -138,11 +138,11 @@ dsp_pm = PatternMatcher([ # lambda g,cast: g.src[0].cast(cast.dtype.scalar()).broadcast(len(g.arg)) if all(x==0 for x in g.arg) else None), # REDUCE int4 -> 2xint2, int8 -> 4xint2 - #(UPat(Ops.REDUCE, dtype=dtypes.int.vec(4), name="r"), - # lambda r: UOp(Ops.CAT, r.dtype, (gep_on_reduce(r.gep((0,1)), r), gep_on_reduce(r.gep((2,3)), r)))), - #(UPat(Ops.REDUCE, dtype=dtypes.int.vec(8), name="r"), - # lambda r: UOp(Ops.CAT, r.dtype, (gep_on_reduce(r.gep((0,1)), r), gep_on_reduce(r.gep((2,3)), r), - # gep_on_reduce(r.gep((4,5)), r), gep_on_reduce(r.gep((6,7)), r)))), + (UPat(Ops.REDUCE, dtype=dtypes.int.vec(4), name="r"), + lambda r: UOp(Ops.CAT, r.dtype, (gep_on_reduce(r.gep((0,1)), r), gep_on_reduce(r.gep((2,3)), r)))), + (UPat(Ops.REDUCE, dtype=dtypes.int.vec(8), name="r"), + lambda r: UOp(Ops.CAT, r.dtype, (gep_on_reduce(r.gep((0,1)), r), gep_on_reduce(r.gep((2,3)), r), + gep_on_reduce(r.gep((4,5)), r), gep_on_reduce(r.gep((6,7)), r)))), # __builtin_HEXAGON_V6_vrmpybus (UPat(dtype=dtypes.int.vec(32), name="a0")*UPat(name="a1") + UPat(name="b0")*UPat(name="b1") + \