From af94addb3a0cf3f342d26093cdb5386b8ec7eef8 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 21 Mar 2025 17:46:45 +0800 Subject: [PATCH] ish --- extra/replay_pkl.py | 2 +- tinygrad/runtime/ops_dsp.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py index abb8a1a33b..af829f49cf 100644 --- a/extra/replay_pkl.py +++ b/extra/replay_pkl.py @@ -133,7 +133,7 @@ if __name__ == "__main__": out_strides = k.sts[0].real_strides() if len(out_strides) == 3: if full_shape[1] < 128: - if full_shape[2] <= 32: k.apply_opt(Opt(OptOps.UNROLL, 0, 0)) + if full_shape[2] <= 16: k.apply_opt(Opt(OptOps.UNROLL, 0, 0)) else: k.apply_opt(Opt(OptOps.UNROLL, 0, 8)) k.apply_opt(Opt(OptOps.UPCAST, 1, full_shape[1])) if out_strides[0] < 128: diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 8ddc7bbba8..e98683a36f 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -20,8 +20,12 @@ def multi_mul(a0, a1, b0, b1, c0, c1, d0, d1, acc=None): swizzle.append(64+i) swizzle.append(96+i) swizzle = tuple(swizzle) - if a0.op is not Ops.CAST: return None - if a1.op is not Ops.CAST: return None + if a0.op is not Ops.CAST: + #print("rejected on a0") + return None + if a1.op is not Ops.CAST: + #print("rejected on a1") + return None assert a0.op is Ops.CAST assert b0.op is Ops.CAST assert c0.op is Ops.CAST