nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

l1prefetch back

Browse Source
This commit is contained in:
geohot
2025-03-26 19:51:18 +08:00
parent c4f5db8467
commit e283bec62e
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/ops_dsp.py
+2 -3
View File
@@ -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
tinygrad/spec.py
+3
View File
@@ -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 <bufidx, alt?, barrier?>
(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),