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
This commit is contained in:
Wencey Wang
2026-05-17 20:42:36 +02:00
committed by Joost VandeVondele
parent 0366d092a8
commit ffffe1bdb5
11 changed files with 385 additions and 18 deletions
+31
View File
@@ -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
+2
View File
@@ -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
+6
View File
@@ -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" : "");
+21 -1
View File
@@ -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<const vec_t*>(input);
+27 -15
View File
@@ -58,7 +58,7 @@ class AffineTransformSparseInput {
static constexpr IndexType PaddedOutputDimensions =
ceil_to_multiple<IndexType>(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<IndexType>(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__
+31
View File
@@ -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<const __m256i*>(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<const __m128i*>(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
+30
View File
@@ -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<const __m256i*>(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<const __m128i*>(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
+35
View File
@@ -516,6 +516,41 @@ Bitboard get_changed_pieces(const std::array<Piece, SQUARE_NB>& oldPieces,
sameBB |= static_cast<Bitboard>(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<const void*>(&oldPieces[i]), 0);
const __m256i new_v = __lasx_xvld(reinterpret_cast<const void*>(&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<std::uint16_t>(__lasx_xvpickve2gr_d(mask, 0));
const auto hi = static_cast<std::uint16_t>(__lasx_xvpickve2gr_d(mask, 2));
changed |= (static_cast<Bitboard>(lo) | (static_cast<Bitboard>(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<const void*>(&oldPieces[i]), 0);
const __m128i new_v = __lsx_vld(reinterpret_cast<const void*>(&newPieces[i]), 0);
const __m128i diff = __lsx_vxor_v(old_v, new_v);
const __m128i mask = __lsx_vmsknz_b(diff);
changed |= static_cast<Bitboard>(static_cast<std::uint16_t>(__lsx_vpickve2gr_d(mask, 0)))
<< i;
}
return changed;
#elif defined(USE_NEON)
uint8x16x4_t old_v = vld4q_u8(reinterpret_cast<const uint8_t*>(oldPieces.data()));
uint8x16x4_t new_v = vld4q_u8(reinterpret_cast<const uint8_t*>(newPieces.data()));
+13
View File
@@ -42,6 +42,13 @@
#elif defined(USE_SSE2)
#include <emmintrin.h>
#elif defined(USE_LASX)
#include <lasxintrin.h>
#include <lsxintrin.h>
#elif defined(USE_LSX)
#include <lsxintrin.h>
#elif defined(USE_NEON)
#include <arm_neon.h>
#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;
+2 -2
View File
@@ -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;
+187
View File
@@ -33,6 +33,13 @@
#elif defined(USE_NEON)
#include <arm_neon.h>
#elif defined(USE_LASX)
#include <lasxintrin.h>
#include <lsxintrin.h>
#elif defined(USE_LSX)
#include <lsxintrin.h>
#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<const void*>(a), 0);
}
inline void lasx_store256(__m256i* a, __m256i b) { __lasx_xvst(b, reinterpret_cast<void*>(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<const void*>(&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<IndexType TransformedFeatureWidth, IndexType HalfDimensions, IndexType PSQTBuckets>