tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 15:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

put that back

Browse Source
This commit is contained in:
geohot
2025-03-26 15:34:57 +08:00
parent 2e4cae342b
commit 880b4a5e47
2 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/replay_pkl.py
+8 -1
View File
@@ -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
tinygrad/runtime/ops_dsp.py
+5 -5
View File
@@ -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") + \