From e283bec62e166822342ba6e03bcf46d76640edee Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 26 Mar 2025 19:51:18 +0800 Subject: [PATCH] l1prefetch back --- tinygrad/runtime/ops_dsp.py | 5 ++--- tinygrad/spec.py | 3 +++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 765ab3dab2..e2eaccb847 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -180,10 +180,9 @@ def add_to_mul(c:UOp, x:UOp): def prefetch_l1(ld:UOp, idx:UOp): if ld.src[-1].op is Ops.CUSTOM: return None ranges = sorted([x for x in ld.src[0].src[0].toposort if x.op is Ops.RANGE], key=lambda x: x.arg) - #ld_buf = idx.src[0].cast(idx.src[0].dtype.base.ptr()) - #x1 = UOp(Ops.CUSTOM, dtypes.void, src=(idx.src[0].index(idx.src[1]+UOp.const(dtypes.int, ld.dtype.count*2)),), arg="__builtin_HEXAGON_Y2_dcfetch({0});") + x1 = UOp(Ops.CUSTOM, dtypes.void, src=(idx.index(UOp.const(dtypes.int, ld.dtype.count*2)),), arg="__builtin_HEXAGON_Y2_dcfetch({0});") x2 = UOp(Ops.CUSTOM, dtypes.void, src=(idx.substitute({ranges[-1]: ranges[-1].src[0]}),), arg="__builtin_HEXAGON_Y2_dcfetch({0});") - return ld.replace(src=ld.src+(x2,)) + return ld.replace(src=ld.src+(x1, x2)) def vectorize_shuffle(vec:UOp): if not all(s.op in {Ops.GEP, Ops.CONST} for s in vec.src): return None diff --git a/tinygrad/spec.py b/tinygrad/spec.py index 5365498857..9e405f7a81 100644 --- a/tinygrad/spec.py +++ b/tinygrad/spec.py @@ -93,6 +93,9 @@ spec = PatternMatcher([ (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat()), name="idx"), validate_index), (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL)), UPat(), UPat(dtype=dtypes.bool, name="mask")), name="idx"), validate_index), + # double INDEX is used for l1 prefetch in DSP + (UPat(Ops.INDEX, src=(UPat(Ops.INDEX), UPat())), lambda: True), + # LOAD takes a (UPat(Ops.LOAD, src=(UPat((Ops.INDEX, Ops.CAST)),)), lambda: True), (UPat(Ops.LOAD, src=(UPat((Ops.INDEX, Ops.CAST)), UPat((Ops.IF, Ops.BARRIER)))), lambda: True),