From ffffe1bdb593553baf6aac484d917e42f66ea030 Mon Sep 17 00:00:00 2001 From: Wencey Wang Date: Sun, 17 May 2026 20:42:36 +0200 Subject: [PATCH] Add intrinsics for LSX and LASX This adds actual intrinsics of SIMD on loongarch64, with considerable performance improvements. Benchmark on `Loongson-3A6000-HV Not Specified CPU @ 2.5GHz` Baseline GCC 14 (debian trixie): ``` =========================== Total time (ms) : 17816 Nodes searched : 2336177 Nodes/second : 131128 ``` With this PR: GCC 16 (debian sid): ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : g++ (GNUC) 16.1.0 on Linux Compilation architecture : loongarch64-lasx Compilation settings : 64bit LASX LSX Compiler __VERSION__ macro : 16.1.0 =========================== Total time (ms) : 3518 Nodes searched : 2336177 Nodes/second : 664063 ``` ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : g++ (GNUC) 16.1.0 on Linux Compilation architecture : loongarch64-lsx Compilation settings : 64bit LSX Compiler __VERSION__ macro : 16.1.0 =========================== Total time (ms) : 4944 Nodes searched : 2336177 Nodes/second : 472527 ``` clang 22 (debian sid): ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : clang++ 22.1.5 on Linux Compilation architecture : loongarch64-lasx Compilation settings : 64bit LASX LSX Compiler __VERSION__ macro : Debian Clang 22.1.5 (1) =========================== Total time (ms) : 3401 Nodes searched : 2336177 Nodes/second : 686908 ``` ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : clang++ 22.1.5 on Linux Compilation architecture : loongarch64-lsx Compilation settings : 64bit LSX Compiler __VERSION__ macro : Debian Clang 22.1.5 (1) =========================== Total time (ms) : 5501 Nodes searched : 2336177 Nodes/second : 424682 ``` GCC 14 (debian trixie): ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : g++ (GNUC) 14.2.0 on Linux Compilation architecture : loongarch64-lasx Compilation settings : 64bit LASX LSX Compiler __VERSION__ macro : 14.2.0 =========================== Total time (ms) : 3559 Nodes searched : 2336177 Nodes/second : 656413 ``` ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : g++ (GNUC) 14.2.0 on Linux Compilation architecture : loongarch64-lsx Compilation settings : 64bit LSX Compiler __VERSION__ macro : 14.2.0 =========================== Total time (ms) : 5212 Nodes searched : 2336177 Nodes/second : 448230 ``` clang 19 (debian trixie): ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : clang++ 19.1.7 on Linux Compilation architecture : loongarch64-lasx Compilation settings : 64bit LASX LSX Compiler __VERSION__ macro : Debian Clang 19.1.7 (3+b1) =========================== Total time (ms) : 4830 Nodes searched : 2336177 Nodes/second : 483680 ``` ``` ./stockfish compiler Stockfish dev-20260516-52e8d9ef by the Stockfish developers (see AUTHORS file) Compiled by : clang++ 19.1.7 on Linux Compilation architecture : loongarch64-lsx Compilation settings : 64bit LSX Compiler __VERSION__ macro : Debian Clang 19.1.7 (3+b1) =========================== Total time (ms) : 5568 Nodes searched : 2336177 Nodes/second : 419572 ``` closes https://github.com/official-stockfish/Stockfish/pull/6815 No functional change --- .github/workflows/tests.yml | 31 +++ src/Makefile | 2 + src/misc.cpp | 6 + src/nnue/layers/affine_transform.h | 22 ++- .../layers/affine_transform_sparse_input.h | 42 ++-- src/nnue/layers/clipped_relu.h | 31 +++ src/nnue/layers/sqr_clipped_relu.h | 30 +++ src/nnue/nnue_accumulator.cpp | 35 ++++ src/nnue/nnue_common.h | 13 ++ src/nnue/nnue_feature_transformer.h | 4 +- src/nnue/simd.h | 187 ++++++++++++++++++ 11 files changed, 385 insertions(+), 18 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 03fe1b7d1..ed1640e76 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -56,6 +56,14 @@ jobs: base_image: "ppc64le/alpine:latest" platform: linux/ppc64le shell: bash + - name: Linux GCC loongarch64 + os: ubuntu-22.04 + compiler: g++ + comp: gcc + run_loongarch64_tests: true + base_image: "loongarch64/alpine:3.21" + platform: linux/loong64 + shell: bash - name: macOS 15 Apple Clang os: macos-15-intel compiler: clang++ @@ -375,6 +383,29 @@ jobs: docker run --rm --platform ${{ matrix.config.platform }} -v ${{ github.workspace }}:/app sf_builder ../tests/signature.sh $benchref + # loongarch64 tests + + - name: Test loongarch64 build + if: matrix.config.run_loongarch64_tests + run: | + echo "cd src && export LDFLAGS='-static' && make clean && make -j4 ARCH=loongarch64 build" > script.sh + docker run --rm --platform ${{ matrix.config.platform }} -v ${{ github.workspace }}:/app sf_builder + ../tests/signature.sh $benchref + + - name: Test loongarch64-lsx build + if: matrix.config.run_loongarch64_tests + run: | + echo "cd src && export LDFLAGS='-static' && make clean && make -j4 ARCH=loongarch64-lsx build" > script.sh + docker run --rm --platform ${{ matrix.config.platform }} -v ${{ github.workspace }}:/app sf_builder + ../tests/signature.sh $benchref + + - name: Test loongarch64-lasx build + if: matrix.config.run_loongarch64_tests + run: | + echo "cd src && export LDFLAGS='-static' && make clean && make -j4 ARCH=loongarch64-lasx build" > script.sh + docker run --rm --platform ${{ matrix.config.platform }} -v ${{ github.workspace }}:/app sf_builder + ../tests/signature.sh $benchref + # Other tests - name: Check perft and search reproducibility diff --git a/src/Makefile b/src/Makefile index 9b9aa9ffd..1794d9ba9 100644 --- a/src/Makefile +++ b/src/Makefile @@ -823,12 +823,14 @@ ifeq ($(dotprod),yes) endif ifeq ($(lasx),yes) + CXXFLAGS += -DUSE_LASX ifeq ($(comp),$(filter $(comp),gcc clang mingw icx)) CXXFLAGS += -mlasx endif endif ifeq ($(lsx),yes) + CXXFLAGS += -DUSE_LSX ifeq ($(comp),$(filter $(comp),gcc clang mingw icx)) CXXFLAGS += -mlsx endif diff --git a/src/misc.cpp b/src/misc.cpp index 2cb3ce5d4..d144a26fc 100644 --- a/src/misc.cpp +++ b/src/misc.cpp @@ -263,6 +263,12 @@ std::string compiler_info() { compiler += " NEON_DOTPROD"; #elif defined(USE_NEON) compiler += " NEON"; +#endif +#if defined(USE_LASX) + compiler += " LASX"; +#endif +#if defined(USE_LSX) + compiler += " LSX"; #endif compiler += (HasPopCnt ? " POPCNT" : ""); diff --git a/src/nnue/layers/affine_transform.h b/src/nnue/layers/affine_transform.h index 9c7699ac3..70310f2ed 100644 --- a/src/nnue/layers/affine_transform.h +++ b/src/nnue/layers/affine_transform.h @@ -40,7 +40,7 @@ namespace Stockfish::Eval::NNUE::Layers { -#if defined(USE_SSSE3) || defined(USE_NEON_DOTPROD) +#if defined(USE_SSSE3) || defined(USE_NEON_DOTPROD) || defined(USE_LSX) || defined(USE_LASX) #define ENABLE_SEQ_OPT #endif @@ -218,6 +218,16 @@ class AffineTransform { #define vec_add_dpbusd_32(acc, a, b) \ SIMD::dotprod_m128_add_dpbusd_epi32(acc, vreinterpretq_s8_s32(a), \ vreinterpretq_s8_s32(b)) + #elif defined(USE_LASX) + using vec_t = __m256i; + #define vec_set_32 __lasx_xvreplgr2vr_w + #define vec_add_32 __lasx_xvadd_w + #define vec_add_dpbusd_32 SIMD::lasx_m256_add_dpbusd_epi32 + #elif defined(USE_LSX) + using vec_t = __m128i; + #define vec_set_32 __lsx_vreplgr2vr_w + #define vec_add_32 __lsx_vadd_w + #define vec_add_dpbusd_32 SIMD::lsx_m128_add_dpbusd_epi32 #endif static constexpr IndexType OutputSimdWidth = sizeof(vec_t) / sizeof(OutputType); @@ -302,6 +312,16 @@ class AffineTransform { SIMD::dotprod_m128_add_dpbusd_epi32(acc, vreinterpretq_s8_s32(a), \ vreinterpretq_s8_s32(b)) #define vec_hadd SIMD::neon_m128_hadd + #elif defined(USE_LASX) + using vec_t = __m256i; + #define vec_setzero() __lasx_xvldi(0) + #define vec_add_dpbusd_32 SIMD::lasx_m256_add_dpbusd_epi32 + #define vec_hadd SIMD::lasx_m256_hadd + #elif defined(USE_LSX) + using vec_t = __m128i; + #define vec_setzero() __lsx_vldi(0) + #define vec_add_dpbusd_32 SIMD::lsx_m128_add_dpbusd_epi32 + #define vec_hadd SIMD::lsx_m128_hadd #endif const auto inputVector = reinterpret_cast(input); diff --git a/src/nnue/layers/affine_transform_sparse_input.h b/src/nnue/layers/affine_transform_sparse_input.h index ab77dc4bd..b0651c99f 100644 --- a/src/nnue/layers/affine_transform_sparse_input.h +++ b/src/nnue/layers/affine_transform_sparse_input.h @@ -58,7 +58,7 @@ class AffineTransformSparseInput { static constexpr IndexType PaddedOutputDimensions = ceil_to_multiple(OutputDimensions, MaxSimdWidth); -#if (USE_SSSE3 | (USE_NEON >= 8)) +#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) static constexpr IndexType ChunkSize = 4; #else static constexpr IndexType ChunkSize = 1; @@ -70,7 +70,7 @@ class AffineTransformSparseInput { using NNZOutputType = std::uint16_t; #endif -#if (USE_SSSE3 | (USE_NEON >= 8)) +#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) alignas(CacheLineSize) static constexpr struct OffsetIndices { NNZOutputType offset_indices[256][8]; @@ -110,7 +110,7 @@ class AffineTransformSparseInput { } static constexpr IndexType get_weight_index(IndexType i) { -#if (USE_SSSE3 | (USE_NEON >= 8)) +#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) return get_weight_index_scrambled(i); #else return i; @@ -147,7 +147,7 @@ class AffineTransformSparseInput { // Forward propagation void propagate(const InputType* input, OutputType* output) const { -#if (USE_SSSE3 | (USE_NEON >= 8)) +#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) #if defined(USE_AVX512) using invec_t = __m512i; using outvec_t = __m512i; @@ -175,6 +175,18 @@ class AffineTransformSparseInput { using outvec_t = int32x4_t; #define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a)) #define vec_add_dpbusd_32 SIMD::neon_m128_add_dpbusd_epi32 + #elif defined(USE_LASX) + using invec_t = __m256i; + using outvec_t = __m256i; + #define vec_add_32 __lasx_xvadd_w + #define vec_set_32 __lasx_xvreplgr2vr_w + #define vec_add_dpbusd_32 SIMD::lasx_m256_add_dpbusd_epi32 + #elif defined(USE_LSX) + using invec_t = __m128i; + using outvec_t = __m128i; + #define vec_add_32 __lsx_vadd_w + #define vec_set_32 __lsx_vreplgr2vr_w + #define vec_add_dpbusd_32 SIMD::lsx_m128_add_dpbusd_epi32 #endif constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType); constexpr IndexType NumChunks = ceil_to_multiple(InputDimensions, 8) / ChunkSize; @@ -182,7 +194,7 @@ class AffineTransformSparseInput { // If we're using high-latency dot product instructions, split the accumulators // to create 3 separate dependency chains and merge at the end constexpr IndexType NumRegs = - #if defined(USE_VNNI) || defined(USE_NEON_DOTPROD) + #if defined(USE_VNNI) || defined(USE_NEON_DOTPROD) || defined(USE_LASX) 3 * NumAccums; #else NumAccums; @@ -203,13 +215,13 @@ class AffineTransformSparseInput { // convince GCC to not do weird pointer arithmetic in the following loop const std::int8_t* weights_cp = weights; - #if defined(USE_VNNI) || defined(USE_NEON_DOTPROD) - #if defined(USE_VNNI) - for (IndexType k = NumAccums; k < NumRegs; ++k) - acc[k] = vec_zero(); - #else + #if defined(USE_VNNI) || defined(USE_NEON_DOTPROD) || defined(USE_LASX) + #if defined(USE_NEON_DOTPROD) for (IndexType k = NumAccums; k < NumRegs; ++k) acc[k] = vdupq_n_s32(0); + #else + for (IndexType k = NumAccums; k < NumRegs; ++k) + acc[k] = vec_zero(); #endif while (start < end - 2) @@ -236,12 +248,12 @@ class AffineTransformSparseInput { vec_add_dpbusd_32(acc[k + 2 * NumAccums], in2, col2[k]); } } - #if defined(USE_VNNI) - for (IndexType k = 0; k < NumAccums; ++k) - acc[k] = vec_add_32(vec_add_32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]); - #else + #if defined(USE_NEON_DOTPROD) for (IndexType k = 0; k < NumAccums; ++k) acc[k] = vaddq_s32(vaddq_s32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]); + #else + for (IndexType k = 0; k < NumAccums; ++k) + acc[k] = vec_add_32(vec_add_32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]); #endif #endif while (start < end) @@ -271,7 +283,7 @@ class AffineTransformSparseInput { } private: -#if (USE_SSSE3 | (USE_NEON >= 8)) +#if (defined(USE_SSSE3) || defined(USE_LSX) || defined(USE_LASX) || (USE_NEON >= 8)) #if defined(__GNUC__) || defined(__clang__) #define RESTRICT __restrict__ diff --git a/src/nnue/layers/clipped_relu.h b/src/nnue/layers/clipped_relu.h index 9ce85d3fa..009e2e5ce 100644 --- a/src/nnue/layers/clipped_relu.h +++ b/src/nnue/layers/clipped_relu.h @@ -154,6 +154,37 @@ class ClippedReLU { out[i] = vmax_s8(vqmovn_s16(shifted), Zero); } constexpr IndexType Start = NumChunks * (SimdWidth / 2); + +#elif defined(USE_LASX) + constexpr IndexType NumChunks = InputDimensions / 32; + const auto in = reinterpret_cast(input); + const auto out = reinterpret_cast<__m256i*>(output); + for (IndexType i = 0; i < NumChunks; ++i) + { + const __m256i packed0 = SIMD::lasx_packus_32(in[i * 4 + 0], in[i * 4 + 1]); + const __m256i packed1 = SIMD::lasx_packus_32(in[i * 4 + 2], in[i * 4 + 3]); + const __m256i words0 = __lasx_xvsrli_h(packed0, WeightScaleBits); + const __m256i words1 = __lasx_xvsrli_h(packed1, WeightScaleBits); + const __m256i packed = __lasx_xvssrani_b_h(words1, words0, 0); + const __m256i swaped = __lasx_xvpermi_d(packed, 0xD8); + __lasx_xvst(__lasx_xvshuf4i_w(swaped, 0xD8), out + i, 0); + } + constexpr IndexType Start = NumChunks * 32; + +#elif defined(USE_LSX) + constexpr IndexType NumChunks = InputDimensions / 16; + const auto in = reinterpret_cast(input); + const auto out = reinterpret_cast<__m128i*>(output); + for (IndexType i = 0; i < NumChunks; ++i) + { + const __m128i packed0 = SIMD::lsx_packus_32(in[i * 4 + 0], in[i * 4 + 1]); + const __m128i packed1 = SIMD::lsx_packus_32(in[i * 4 + 2], in[i * 4 + 3]); + const __m128i words0 = __lsx_vsrli_h(packed0, WeightScaleBits); + const __m128i words1 = __lsx_vsrli_h(packed1, WeightScaleBits); + out[i] = __lsx_vssrani_b_h(words1, words0, 0); + } + constexpr IndexType Start = NumChunks * 16; + #else constexpr IndexType Start = 0; #endif diff --git a/src/nnue/layers/sqr_clipped_relu.h b/src/nnue/layers/sqr_clipped_relu.h index 53412d014..e45f6f0f1 100644 --- a/src/nnue/layers/sqr_clipped_relu.h +++ b/src/nnue/layers/sqr_clipped_relu.h @@ -90,6 +90,36 @@ class SqrClippedReLU { } constexpr IndexType Start = NumChunks * 16; +#elif defined(USE_LASX) + constexpr IndexType NumChunks = InputDimensions / 32; + const auto in = reinterpret_cast(input); + const auto out = reinterpret_cast<__m256i*>(output); + for (IndexType i = 0; i < NumChunks; ++i) + { + const __m256i words0 = __lasx_xvssrani_h_w(in[i * 4 + 1], in[i * 4 + 0], 0); + const __m256i words1 = __lasx_xvssrani_h_w(in[i * 4 + 3], in[i * 4 + 2], 0); + const __m256i sqr0 = __lasx_xvmuh_h(words0, words0); + const __m256i sqr1 = __lasx_xvmuh_h(words1, words1); + const __m256i packed = __lasx_xvssrlni_b_h(sqr1, sqr0, 3); + const __m256i permed = __lasx_xvpermi_d(packed, 0xD8); + __lasx_xvst(__lasx_xvshuf4i_w(permed, 0xD8), out + i, 0); + } + constexpr IndexType Start = NumChunks * 32; + +#elif defined(USE_LSX) + constexpr IndexType NumChunks = InputDimensions / 16; + const auto in = reinterpret_cast(input); + const auto out = reinterpret_cast<__m128i*>(output); + for (IndexType i = 0; i < NumChunks; ++i) + { + const __m128i words0 = __lsx_vssrani_h_w(in[i * 4 + 1], in[i * 4 + 0], 0); + const __m128i words1 = __lsx_vssrani_h_w(in[i * 4 + 3], in[i * 4 + 2], 0); + const __m128i sqr0 = __lsx_vmuh_h(words0, words0); + const __m128i sqr1 = __lsx_vmuh_h(words1, words1); + out[i] = __lsx_vssrlni_b_h(sqr1, sqr0, 3); + } + constexpr IndexType Start = NumChunks * 16; + #else constexpr IndexType Start = 0; #endif diff --git a/src/nnue/nnue_accumulator.cpp b/src/nnue/nnue_accumulator.cpp index dc6813416..a1480fd89 100644 --- a/src/nnue/nnue_accumulator.cpp +++ b/src/nnue/nnue_accumulator.cpp @@ -516,6 +516,41 @@ Bitboard get_changed_pieces(const std::array& oldPieces, sameBB |= static_cast(equalMask) << i; } return ~sameBB; +#elif defined(USE_LASX) + static_assert(sizeof(Piece) == 1); + + Bitboard changed = 0; + + for (int i = 0; i < 64; i += 32) + { + const __m256i old_v = __lasx_xvld(reinterpret_cast(&oldPieces[i]), 0); + const __m256i new_v = __lasx_xvld(reinterpret_cast(&newPieces[i]), 0); + const __m256i diff = __lasx_xvxor_v(old_v, new_v); + const __m256i mask = __lasx_xvmsknz_b(diff); + const auto lo = static_cast(__lasx_xvpickve2gr_d(mask, 0)); + const auto hi = static_cast(__lasx_xvpickve2gr_d(mask, 2)); + + changed |= (static_cast(lo) | (static_cast(hi) << 16)) << i; + } + + return changed; +#elif defined(USE_LSX) + static_assert(sizeof(Piece) == 1); + + Bitboard changed = 0; + + for (int i = 0; i < 64; i += 16) + { + const __m128i old_v = __lsx_vld(reinterpret_cast(&oldPieces[i]), 0); + const __m128i new_v = __lsx_vld(reinterpret_cast(&newPieces[i]), 0); + const __m128i diff = __lsx_vxor_v(old_v, new_v); + const __m128i mask = __lsx_vmsknz_b(diff); + + changed |= static_cast(static_cast(__lsx_vpickve2gr_d(mask, 0))) + << i; + } + + return changed; #elif defined(USE_NEON) uint8x16x4_t old_v = vld4q_u8(reinterpret_cast(oldPieces.data())); uint8x16x4_t new_v = vld4q_u8(reinterpret_cast(newPieces.data())); diff --git a/src/nnue/nnue_common.h b/src/nnue/nnue_common.h index bd546d277..7bed4872b 100644 --- a/src/nnue/nnue_common.h +++ b/src/nnue/nnue_common.h @@ -42,6 +42,13 @@ #elif defined(USE_SSE2) #include +#elif defined(USE_LASX) + #include + #include + +#elif defined(USE_LSX) + #include + #elif defined(USE_NEON) #include #endif @@ -71,11 +78,17 @@ constexpr const std::size_t Leb128MagicStringSize = sizeof(Leb128MagicString) - #if defined(USE_AVX2) constexpr std::size_t SimdWidth = 32; +#elif defined(USE_LASX) +constexpr std::size_t SimdWidth = 32; + #elif defined(USE_SSE2) constexpr std::size_t SimdWidth = 16; #elif defined(USE_NEON) constexpr std::size_t SimdWidth = 16; + +#elif defined(USE_LSX) +constexpr std::size_t SimdWidth = 16; #endif constexpr std::size_t MaxSimdWidth = 32; diff --git a/src/nnue/nnue_feature_transformer.h b/src/nnue/nnue_feature_transformer.h index 6b913e80f..f74ecae7f 100644 --- a/src/nnue/nnue_feature_transformer.h +++ b/src/nnue/nnue_feature_transformer.h @@ -104,7 +104,7 @@ class FeatureTransformer { // | 1 | 3 | 5 | 7 | // Vector 1 // | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | // Packed Result return {0, 2, 4, 6, 1, 3, 5, 7}; -#elif defined(USE_AVX2) +#elif defined(USE_AVX2) || defined(USE_LASX) // _mm256_packus_epi16 after permutation: // | 0 | 2 | | 4 | 6 | // Vector 0, 2 // | 1 | 3 | | 5 | 7 | // Vector 1, 3 @@ -300,7 +300,7 @@ class FeatureTransformer { // the multiplication. constexpr int shift = - #if defined(USE_SSE2) + #if defined(USE_SSE2) || defined(USE_LASX) || defined(USE_LSX) 7; #else 6; diff --git a/src/nnue/simd.h b/src/nnue/simd.h index 6a549eb3f..44cafe569 100644 --- a/src/nnue/simd.h +++ b/src/nnue/simd.h @@ -33,6 +33,13 @@ #elif defined(USE_NEON) #include + +#elif defined(USE_LASX) + #include + #include + +#elif defined(USE_LSX) + #include #endif #include "../types.h" @@ -226,6 +233,150 @@ inline int16x8_t vaddw_high_s8(int16x8_t a, int8x16_t b) { return vaddw_s8(a, vg inline int16x8_t vsubw_high_s8(int16x8_t a, int8x16_t b) { return vsubw_s8(a, vget_high_s8(b)); } #endif +#elif USE_LASX +using vec_t = __m256i; +using vec_i8_t = __m128i; +using vec128_t = __m128i; +using psqt_vec_t = __m256i; +using vec_uint_t = __m256i; + +inline __m256i lasx_load256(const __m256i* a) { + return __lasx_xvld(reinterpret_cast(a), 0); +} + +inline void lasx_store256(__m256i* a, __m256i b) { __lasx_xvst(b, reinterpret_cast(a), 0); } + +inline __m256i lasx_packus_16(__m256i a, __m256i b) { + #if defined(__clang__) && defined(__has_builtin) && __has_builtin(__builtin_lasx_xvssrani_bu_h) + return (__m256i) __builtin_lasx_xvssrani_bu_h((v32i8) b, (v32i8) a, 0); + #else + return __lasx_xvssrani_bu_h(b, a, 0); + #endif +} + +inline __m256i lasx_packus_32(__m256i a, __m256i b) { + #if defined(__clang__) && defined(__has_builtin) && __has_builtin(__builtin_lasx_xvssrani_hu_w) + return (__m256i) __builtin_lasx_xvssrani_hu_w((v16i16) b, (v16i16) a, 0); + #else + return __lasx_xvssrani_hu_w(b, a, 0); + #endif +} + + #define vec_load(a) lasx_load256(a) + #define vec_store(a, b) lasx_store256(a, b) + #define vec_add_16(a, b) __lasx_xvadd_h(a, b) + #define vec_sub_16(a, b) __lasx_xvsub_h(a, b) + #define vec_mulhi_16(a, b) __lasx_xvmuh_h(a, b) + #define vec_zero() __lasx_xvldi(0) + #define vec_set_16(a) __lasx_xvreplgr2vr_h(a) + #define vec_max_16(a, b) __lasx_xvmax_h(a, b) + #define vec_min_16(a, b) __lasx_xvmin_h(a, b) + #define vec_slli_16(a, b) __lasx_xvslli_h(a, b) + // Inverse permuted at load time + #define vec_packus_16(a, b) lasx_packus_16(a, b) + #define vec_load_psqt(a) lasx_load256(a) + #define vec_store_psqt(a, b) lasx_store256(a, b) + #define vec_add_psqt_32(a, b) __lasx_xvadd_w(a, b) + #define vec_sub_psqt_32(a, b) __lasx_xvsub_w(a, b) + #define vec_zero_psqt() __lasx_xvldi(0) + #define vec_nnz(a) lasx_vec_nnz(a) + #define vec_convert_8_16(a) lasx_cvtepi8_epi16(a) + + #define vec128_zero __lsx_vldi(0) + #define vec128_set_16(a) __lsx_vreplgr2vr_h(a) + #define vec128_load(a) (*(a)) + #define vec128_storeu(a, b) *(a) = (b) + #define vec128_add(a, b) __lsx_vadd_h(a, b) + + #define NumRegistersSIMD 24 + #define MaxChunkSize 32 + +inline __m256i lasx_cvtepi8_epi16(__m128i a) { + #if defined(__has_builtin) && __has_builtin(__builtin_lasx_cast_128) + return __lasx_vext2xv_h_b(__lasx_cast_128(a)); + #elif defined(__GNUC__) && !defined(__clang__) + __m256i out; + __asm__("vext2xv.h.b %u0, %u1" : "=f"(out) : "f"(a)); + return out; + #else + int64_t lo = (int64_t) __lsx_vpickve2gr_d(a, 0); + int64_t hi = (int64_t) __lsx_vpickve2gr_d(a, 1); + __m256i v = __lasx_xvldi(0); + v = __lasx_xvinsgr2vr_d(v, lo, 0); + v = __lasx_xvinsgr2vr_d(v, hi, 2); + return __lasx_xvsllwil_h_b(v, 0); + #endif +} + +inline int lasx_vec_nnz(__m256i a) { + const __m256i cmp = __lasx_xvslt_w(__lasx_xvldi(0), a); + const __m256i msk = __lasx_xvmskltz_w(cmp); + return ((int) __lasx_xvpickve2gr_w(msk, 0) & 0xF) + | (((int) __lasx_xvpickve2gr_w(msk, 4) & 0xF) << 4); +} + +#elif USE_LSX +using vec_t = __m128i; +using vec_i8_t = std::uint64_t; +using vec128_t = __m128i; +using psqt_vec_t = __m128i; +using vec_uint_t = __m128i; + +inline __m128i lsx_packus_16(__m128i a, __m128i b) { + #if defined(__clang__) && defined(__has_builtin) && __has_builtin(__builtin_lsx_vssrani_bu_h) + return (__m128i) __builtin_lsx_vssrani_bu_h((v16i8) b, (v16i8) a, 0); + #else + return __lsx_vssrani_bu_h(b, a, 0); + #endif +} + +inline __m128i lsx_packus_32(__m128i a, __m128i b) { + #if defined(__clang__) && defined(__has_builtin) && __has_builtin(__builtin_lsx_vssrani_hu_w) + return (__m128i) __builtin_lsx_vssrani_hu_w((v8i16) b, (v8i16) a, 0); + #else + return __lsx_vssrani_hu_w(b, a, 0); + #endif +} + + #define vec_load(a) (*(a)) + #define vec_store(a, b) *(a) = (b) + #define vec_add_16(a, b) __lsx_vadd_h(a, b) + #define vec_sub_16(a, b) __lsx_vsub_h(a, b) + #define vec_mulhi_16(a, b) __lsx_vmuh_h(a, b) + #define vec_zero() __lsx_vldi(0) + #define vec_set_16(a) __lsx_vreplgr2vr_h(a) + #define vec_max_16(a, b) __lsx_vmax_h(a, b) + #define vec_min_16(a, b) __lsx_vmin_h(a, b) + #define vec_slli_16(a, b) __lsx_vslli_h(a, b) + // Inverse permuted at load time + #define vec_packus_16(a, b) lsx_packus_16(a, b) + #define vec_load_psqt(a) (*(a)) + #define vec_store_psqt(a, b) *(a) = (b) + #define vec_add_psqt_32(a, b) __lsx_vadd_w(a, b) + #define vec_sub_psqt_32(a, b) __lsx_vsub_w(a, b) + #define vec_zero_psqt() __lsx_vldi(0) + +inline int lsx_vec_nnz(__m128i a) { + const __m128i cmp = __lsx_vslt_w(__lsx_vldi(0), a); + const __m128i msk = __lsx_vmskltz_w(cmp); + return ((int) __lsx_vpickve2gr_w(msk, 0) & 0xF); +} + #define vec_nnz(a) lsx_vec_nnz(a) + +inline __m128i vec_convert_8_16(std::uint64_t x) { + __m128i v = __lsx_vldrepl_d(reinterpret_cast(&x), 0); + return __lsx_vsllwil_h_b(v, 0); +} + + #define vec128_zero __lsx_vldi(0) + #define vec128_set_16(a) __lsx_vreplgr2vr_h(a) + #define vec128_load(a) (*(a)) + #define vec128_storeu(a, b) *(a) = (b) + #define vec128_add(a, b) __lsx_vadd_h(a, b) + + #define NumRegistersSIMD 24 + #define MaxChunkSize 16 + #else #undef VECTOR @@ -382,6 +533,42 @@ dotprod_m128_add_dpbusd_epi32(int32x4_t& acc, int8x16_t a, int8x16_t b) { } #endif +#if defined(USE_LASX) + +[[maybe_unused]] static int lasx_m256_hadd(__m256i sum, int bias) { + __m256i v = sum; + v = __lasx_xvadd_w(v, __lasx_xvshuf4i_w(v, 0x4E)); // [C,D,A,B] per lane + v = __lasx_xvadd_w(v, __lasx_xvshuf4i_w(v, 0xB1)); // [B,A,D,C] per lane + int lo_sum = (int) __lasx_xvpickve2gr_w(v, 0); + int hi_sum = (int) __lasx_xvpickve2gr_w(v, 4); + return lo_sum + hi_sum + bias; +} + +[[maybe_unused]] static void lasx_m256_add_dpbusd_epi32(__m256i& acc, __m256i a, __m256i b) { + __m256i tmp = __lasx_xvmulwev_h_bu_b(a, b); + tmp = __lasx_xvmaddwod_h_bu_b(tmp, a, b); + acc = __lasx_xvadd_w(acc, __lasx_xvhaddw_w_h(tmp, tmp)); +} + +#endif // USE_LASX + +#if defined(USE_LSX) + +[[maybe_unused]] static int lsx_m128_hadd(__m128i sum, int bias) { + sum = __lsx_vadd_w(sum, __lsx_vshuf4i_w(sum, 0x4E)); // [C,D,A,B] + sum = __lsx_vadd_w(sum, __lsx_vshuf4i_w(sum, 0xB1)); // [B,A,D,C] + return __lsx_vpickve2gr_w(sum, 0) + bias; +} + +[[maybe_unused]] static void lsx_m128_add_dpbusd_epi32(__m128i& acc, __m128i a, __m128i b) { + // tmp[i] = a[2i]*b[2i] + a[2i+1]*b[2i+1] + __m128i tmp = __lsx_vmulwev_h_bu_b(a, b); + tmp = __lsx_vmaddwod_h_bu_b(tmp, a, b); + acc = __lsx_vadd_w(acc, __lsx_vhaddw_w_h(tmp, tmp)); +} + +#endif // USE_LSX + // Compute optimal SIMD register count for feature transformer accumulation. template