From e1d2bec4a4889a0e91c5a0bb45b2634d8868d6b8 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Sat, 22 Mar 2025 18:52:56 +0800
Subject: [PATCH] opt

---
 extra/replay_pkl.py        | 14 ++++++++++++++
 tinygrad/codegen/kernel.py |  2 +-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py
index a235558bb8..da192585e8 100644
--- a/extra/replay_pkl.py
+++ b/extra/replay_pkl.py
@@ -116,6 +116,11 @@ if __name__ == "__main__":
             k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
             k.apply_opt(Opt(OptOps.UPCAST, 1, 24))
             k.apply_opt(Opt(OptOps.UPCAST, 0, 16))
+          elif knum == 9:
+            # same as 6
+            k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
+            k.apply_opt(Opt(OptOps.UPCAST, 1, 24))
+            k.apply_opt(Opt(OptOps.UPCAST, 0, 16))
           elif knum in [7,11]:
             k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
             k.apply_opt(Opt(OptOps.UPCAST, 1, 144))
@@ -129,6 +134,15 @@ if __name__ == "__main__":
             k.apply_opt(Opt(OptOps.UPCAST, 1, 64))
             #k.apply_opt(Opt(OptOps.UPCAST, 0, 2))
             pass
+          elif knum == 1:
+            k.apply_opt(Opt(OptOps.UPCAST, 2, 32))
+            k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
+            k.apply_opt(Opt(OptOps.UPCAST, 0, 2))
+          elif knum == 26:
+            #k.apply_opt(Opt(OptOps.UNROLL, 1, 0))
+            #k.apply_opt(Opt(OptOps.UNROLL, 0, 0))
+            #k.apply_opt(Opt(OptOps.UPCAST, 2, 128))
+            pass
           #elif knum == 18:
           #  k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
           #  k.apply_opt(Opt(OptOps.UPCAST, 1, 192))
diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py
index df44027fd9..8fcc027599 100644
--- a/tinygrad/codegen/kernel.py
+++ b/tinygrad/codegen/kernel.py
@@ -666,7 +666,7 @@ class Kernel:
       print(self.name)
       if getenv("RAWAST"): print(self.ast)
       for i,(buf,st) in enumerate([(buf,st) for buf,st in zip(self.bufs, self.sts) if buf.op not in {Ops.CONST, Ops.VALID}]):
-        print(f"{i:2d}: {str(st.shape):25s} {str(buf.src[0].dtype).replace('dtypes.',''):20s} {str(st.real_strides()):20s}",
+        print(f"{i:2d}: {str(st.shape):25s} {str(buf.src[0].dtype).replace('dtypes.',''):20s} {str(st.real_strides()):30s}",
               str(st) if DEBUG >= 4 else "")
       print(self.applied_opts)
       if DEBUG >= 5: print(modified_ast)