From c3c85c64eedfb207d7aba8288c3aef0a63f5ce4a Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Fri, 21 Mar 2025 09:24:33 +0800
Subject: [PATCH] simpler

---
 extra/replay_pkl.py         | 2 +-
 tinygrad/runtime/ops_dsp.py | 4 ++--
 tinygrad/spec.py            | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py
index 73259f08ab..f49109e7d0 100644
--- a/extra/replay_pkl.py
+++ b/extra/replay_pkl.py
@@ -87,7 +87,7 @@ if __name__ == "__main__":
             pass
           elif knum == 20:
             # 784x192 * 192x32 -> 784x32
-            k.apply_opt(Opt(OptOps.UNROLL, 0, 8))
+            k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
             k.apply_opt(Opt(OptOps.UPCAST, 1, 32))
             k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
           elif knum == 35:
diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py
index fcff61895e..ef6f3aca4e 100644
--- a/tinygrad/runtime/ops_dsp.py
+++ b/tinygrad/runtime/ops_dsp.py
@@ -81,8 +81,8 @@ def add_to_mul(c:UOp, x:UOp):
 
 dsp_pm_late = PatternMatcher([
   # prefetch L1
-  (UPat(Ops.LOAD, dtype=dtypes.uchar.vec(8), name="ld"),
-   lambda ld: ld.replace(src=ld.src+(UOp(Ops.CUSTOM, dtypes.void, src=(ld.src[0].src[0].src[0].index(ld.src[0].src[0].src[1]+16),),
+  (UPat(Ops.LOAD, dtype=(dtypes.uchar.vec(4), dtypes.uchar.vec(8)), name="ld"),
+   lambda ld: ld.replace(src=ld.src+(UOp(Ops.CUSTOM, dtypes.void, src=(ld.src[0].src[0].index(UOp.const(dtypes.int, 16)),),
                                          arg="__builtin_HEXAGON_Y2_dcfetch({0});"),)) if ld.src[-1].op is not Ops.CUSTOM else None),
 
   (UPat(Ops.CUSTOMI, dtype=dtypes.int.vec(32), name="c")+UPat.var("x"), add_to_mul),
diff --git a/tinygrad/spec.py b/tinygrad/spec.py
index b66c3d9def..ec3d3a9225 100644
--- a/tinygrad/spec.py
+++ b/tinygrad/spec.py
@@ -77,7 +77,7 @@ spec = PatternMatcher([
   # **** new style load/store ****
 
   # INDEX is used in new style load/store
-  (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat())), lambda: True),
+  (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.INDEX)), UPat())), lambda: True),
 
   # LOAD takes a <bufidx, alt?, gate?, barrier?>
   (UPat(Ops.LOAD, src=(UPat((Ops.INDEX, Ops.CAST)),)), lambda: True),