From 9e19cdfbbead4ee649c04dc3b68306d6544458a5 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 26 Mar 2025 16:49:23 +0800 Subject: [PATCH] e kernels --- extra/replay_pkl.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py index 1a250da5c9..1c667a531c 100644 --- a/extra/replay_pkl.py +++ b/extra/replay_pkl.py @@ -135,8 +135,11 @@ if __name__ == "__main__": k.apply_opt(Opt(OptOps.UNROLL, 0, 8)) k.apply_opt(Opt(OptOps.UPCAST, 2, 32)) if k.full_shape[1]%4 == 0: k.apply_opt(Opt(OptOps.UPCAST, 1, 4)) - elif len(k.full_shape) == 1 and k.full_shape[0]%128 == 0: - k.apply_opt(Opt(OptOps.UPCAST, 0, 128)) + elif len(k.full_shape) == 1: + for sz in [128,64,32]: + if k.full_shape[0]%sz == 0: + k.apply_opt(Opt(OptOps.UPCAST, 0, sz)) + break """ #if knum in [7, 11, 14, 18]: