Fix DSP transcendentals (#9542)

This commit is contained in:
quortus
2025-03-22 11:08:18 +08:00
committed by GitHub
parent eddafb84e5
commit bdd44d4255
4 changed files with 8 additions and 4 deletions
+1 -1
View File
@@ -215,7 +215,7 @@ runs:
wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-19 main" | sudo tee /etc/apt/sources.list.d/llvm.list
sudo apt update -y || true
sudo apt install -y --no-install-recommends libllvm19
sudo apt install -y --no-install-recommends libllvm19 clang-19 lld-19
- name: Install LLVM (macOS)
if: inputs.llvm == 'true' && runner.os == 'macOS'
+2
View File
@@ -498,6 +498,8 @@ jobs:
cache-to: type=gha,mode=min
- name: Run test_tiny on DSP
run: DEBUG=2 DSP=1 python test/test_tiny.py
- name: Test transcendentals
run: CC=clang-19 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
- name: Test quantize onnx
run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
- name: Test LLVM=1 DEVECTORIZE=0
+4
View File
@@ -148,6 +148,7 @@ class TestTranscendentalVectorized(unittest.TestCase):
def test_log2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)
@unittest.skipIf(getenv("DSP"), "requires int division")
def test_sin_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)
@@ -155,5 +156,8 @@ class TestTranscendentalVectorized(unittest.TestCase):
# np.pow returns nan for negative values raised to a non-integral power
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.pow, np.pow, (0.001, 200), vec_size, param_range=(-10, 10))
def test_sqrt_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sqrt, np.sqrt, (0, 100), vec_size)
if __name__ == '__main__':
unittest.main()
+1 -3
View File
@@ -41,9 +41,7 @@ class DSPRenderer(ClangRenderer):
extra_matcher = dsp_pm_late+ClangRenderer.extra_matcher
string_rewrite = dsp_string+ClangRenderer.string_rewrite
type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
code_for_op = {**ClangRenderer.code_for_op, Ops.SIN: lambda x,dtype: f"__builtin_sin({x})",
Ops.LOG2: lambda x,dtype: f"__builtin_log2l({x})" if dtype == dtypes.float64 else f"__builtin_log2f({x})",
Ops.EXP2: lambda x,dtype: f"__builtin_exp2l({x})" if dtype == dtypes.float64 else f"__builtin_exp2f({x})"}
code_for_op = {k:v for k,v in ClangRenderer.code_for_op.items() if k != Ops.SQRT}
def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[DType,bool]]], uops:list[UOp], prefix=None) -> str:
ret = super().render_kernel(function_name, kernel, bufs, uops, prefix)