From e4615e0cd99794aefafeb5216fa055b637221ea3 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 17 Mar 2025 16:20:46 +0800 Subject: [PATCH] dsp work try 3 --- tinygrad/codegen/devectorizer.py | 1 + tinygrad/runtime/ops_dsp.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 430cb29aaa..87120089ab 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -12,6 +12,7 @@ from tinygrad.renderer import Renderer # ***** load/store grouping ***** def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None): + if getenv("UNSAFE_DISABLE_MASK", 0): mask = None # first, extract all the relevant offsets offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict) for i in range(vec.dtype.count): diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index be76d29747..d27c50f808 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -16,7 +16,7 @@ dsp_pm = PatternMatcher([ lambda x: UOp(Ops.CUSTOM, dtypes.uchar.vec(128), src=tuple(x.gep(tuple(range(i, i+32))) for i in range(0, 128, 32)), arg="__builtin_HEXAGON_V6_vpackhub_sat_128B(__builtin_HEXAGON_V6_vpackwh_sat_128B({3}, {2}), __builtin_HEXAGON_V6_vpackwh_sat_128B({1}, {0}))")), (UPat(Ops.GEP, name="x"), lambda x: UOp(Ops.CUSTOM, x.dtype, x.src+x.src, - "__builtin_shufflevector({0}, {1}, "+','.join([str(y) for y in x.arg])+")") if len(x.arg) > 1 else None), + "__builtin_shufflevector({0}, {1}, "+','.join([str(y) for y in x.arg])+")") if len(x.arg) > 1 and x.src[0].dtype.count > 1 else None), ]) dsp_pm_late = PatternMatcher([