jirayu_sri/stockfish
1
0
Fork 0
mirror of https://github.com/official-stockfish/Stockfish.git synced 2026-07-22 12:47:08 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
fff35786bf4a5941a017066842f6977398ddac7e
stockfish/src/nnue/simd.h
T
anematodeandJoost VandeVondele fff35786bf Use bitset representation for nnz and move computation into feature transformer 
Passed avx2/bmi2 STC:
https://tests.stockfishchess.org/tests/view/69ff99eb9392f0c317213eda
LLR: 2.93 (-2.94,2.94) <0.00,2.00>
Total: 20992 W: 5469 L: 5192 D: 10331
Ptnml(0-2): 38, 2197, 5751, 2470, 40

Passed avxvnni STC:
https://tests.stockfishchess.org/tests/view/6a00022b9392f0c317213f7c
LLR: 2.93 (-2.94,2.94) <0.00,2.00>
Total: 47328 W: 12328 L: 12009 D: 22991
Ptnml(0-2): 112, 5148, 12847, 5423, 134

Passed NEON STC:
https://tests.stockfishchess.org/tests/view/69ff99d69392f0c317213ed8
LLR: 2.96 (-2.94,2.94) <0.00,2.00>
Total: 29600 W: 7664 L: 7376 D: 14560
Ptnml(0-2): 48, 3074, 8277, 3344, 57

Passed avx512icl non-regression:
https://tests.stockfishchess.org/tests/view/6a002b699392f0c317213f91
LLR: 2.94 (-2.94,2.94) <-1.75,0.25>
Total: 90112 W: 23130 L: 22975 D: 44007
Ptnml(0-2): 192, 9939, 24633, 10106, 186

Measurements from vondele (Neoverse V2, armv8-dotprod):

==== master ====
==== Bench: 2344696 ====
1 Nodes/second : 280724660
2 Nodes/second : 280647282
3 Nodes/second : 282055192
Average (over 3):  281142378
==== 50a44640a3 ====
==== Bench: 2344696 ====
1 Nodes/second : 284271937
2 Nodes/second : 285638071
3 Nodes/second : 284349426
Average (over 3):  284753144

The patch's benefit is non-uniform and is ~0 for avx512icl unfortunately -- although I think we should be able to find something there....

## Background/explanation

The idea here is to move the non-zero block computation back to the previous layer, and overlap the work better. Then, to avoid trying to emulate compress instructions on targets not supporting them (i.e., everything except for AVX512), we use a `pop_lsb` loop on a bitset enumerating the non-zero blocks, rather than first writing them out as indices.

An early AVX2 implementation worked on fishtest (https://tests.stockfishchess.org/tests/view/69c3410534b6988b1e472db4) and linrock demonstrated that it also worked for NEON (https://tests.stockfishchess.org/tests/view/69ce0dc73ddc0eccd617188f). Thanks to him for encouraging me to push this idea over the finish line.

To abstract over the particular NNZ representation (bitset or index list), we use `NNZInfo` and `NNZCursor`. The cursor is used to write to one or the other perspective of the NNZ list. I'd appreciate ideas on making the code cleaner/more readable as imo it's still a bit ugly.

## Fixing GCC 15 regression

Separately, vondele noticed a regression in ARM performance from GCC 15 caused by suboptimal codegen. See https://discord.com/channels/435943710472011776/813919248455827515/1502837886381461605 for more info, but the easiest fix that I could come up with was inserting a couple `asm` optimization barriers.

Single threaded data:
```
average: stockfish.master.13.3  1041619
average: stockfish.master.15.2   999214
average: stockfish.patched.15.2  1031369
```

There's definitely regressions elsewhere as well, which we shld chase down, but at least this should unblock the arm64 universal binary work.

closes https://github.com/official-stockfish/Stockfish/pull/6814

No functional change
2026-05-19 18:36:11 +02:00

630 lines
22 KiB
C++
Raw Blame History

/*
Stockfish, a UCI chess playing engine derived from Glaurung 2.1
Copyright (C) 2004-2026 The Stockfish developers (see AUTHORS file)
Stockfish is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
Stockfish is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
#ifndef NNUE_SIMD_H_INCLUDED
#define NNUE_SIMD_H_INCLUDED
#if defined(USE_AVX2)
#include <immintrin.h>
#elif defined(USE_SSE41)
#include <smmintrin.h>
#elif defined(USE_SSSE3)
#include <tmmintrin.h>
#elif defined(USE_SSE2)
#include <emmintrin.h>
#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"
#include "nnue_common.h"
namespace Stockfish::Eval::NNUE::SIMD {
// If vector instructions are enabled, we update and refresh the
// accumulator tile by tile such that each tile fits in the CPU's
// vector registers.
#define VECTOR
#ifdef USE_AVX512
using vec_t = __m512i;
using vec_i8_t = __m256i;
using vec128_t = __m128i;
using psqt_vec_t = __m256i;
using vec_uint_t = __m512i;
#define vec_load(a) _mm512_load_si512(a)
#define vec_store(a, b) _mm512_store_si512(a, b)
#define vec_convert_8_16(a) _mm512_cvtepi8_epi16(a)
#define vec_add_16(a, b) _mm512_add_epi16(a, b)
#define vec_sub_16(a, b) _mm512_sub_epi16(a, b)
#define vec_mulhi_16(a, b) _mm512_mulhi_epi16(a, b)
#define vec_zero() _mm512_setzero_epi32()
#define vec_set_16(a) _mm512_set1_epi16(a)
#define vec_max_16(a, b) _mm512_max_epi16(a, b)
#define vec_min_16(a, b) _mm512_min_epi16(a, b)
#define vec_slli_16(a, b) _mm512_slli_epi16(a, b)
// Inverse permuted at load time
#define vec_packus_16(a, b) _mm512_packus_epi16(a, b)
#define vec_load_psqt(a) _mm256_load_si256(a)
#define vec_store_psqt(a, b) _mm256_store_si256(a, b)
#define vec_add_psqt_32(a, b) _mm256_add_epi32(a, b)
#define vec_sub_psqt_32(a, b) _mm256_sub_epi32(a, b)
#define vec_zero_psqt() _mm256_setzero_si256()
#ifdef USE_SSSE3
#define vec_nnz(a) _mm512_cmpgt_epi32_mask(a, _mm512_setzero_si512())
#endif
#define vec128_zero _mm_setzero_si128()
#define vec128_set_16(a) _mm_set1_epi16(a)
#define vec128_load(a) _mm_load_si128(a)
#define vec128_storeu(a, b) _mm_storeu_si128(a, b)
#define vec128_add(a, b) _mm_add_epi16(a, b)
#define NumRegistersSIMD 16
#define MaxChunkSize 64
#elif USE_AVX2
using vec_t = __m256i;
using vec_i8_t = __m128i;
using vec128_t = __m128i;
using psqt_vec_t = __m256i;
using vec_uint_t = __m256i;
#define vec_load(a) _mm256_load_si256(a)
#define vec_store(a, b) _mm256_store_si256(a, b)
#define vec_convert_8_16(a) _mm256_cvtepi8_epi16(a)
#define vec_add_16(a, b) _mm256_add_epi16(a, b)
#define vec_sub_16(a, b) _mm256_sub_epi16(a, b)
#define vec_mulhi_16(a, b) _mm256_mulhi_epi16(a, b)
#define vec_zero() _mm256_setzero_si256()
#define vec_set_16(a) _mm256_set1_epi16(a)
#define vec_max_16(a, b) _mm256_max_epi16(a, b)
#define vec_min_16(a, b) _mm256_min_epi16(a, b)
#define vec_slli_16(a, b) _mm256_slli_epi16(a, b)
// Inverse permuted at load time
#define vec_packus_16(a, b) _mm256_packus_epi16(a, b)
#define vec_load_psqt(a) _mm256_load_si256(a)
#define vec_store_psqt(a, b) _mm256_store_si256(a, b)
#define vec_add_psqt_32(a, b) _mm256_add_epi32(a, b)
#define vec_sub_psqt_32(a, b) _mm256_sub_epi32(a, b)
#define vec_zero_psqt() _mm256_setzero_si256()
#ifdef USE_SSSE3
#if defined(USE_VNNI) && !defined(USE_AVXVNNI)
#define vec_nnz(a) _mm256_cmpgt_epi32_mask(a, _mm256_setzero_si256())
#else
#define vec_nnz(a) \
_mm256_movemask_ps( \
_mm256_castsi256_ps(_mm256_cmpgt_epi32(a, _mm256_setzero_si256())))
#endif
#endif
#define vec128_zero _mm_setzero_si128()
#define vec128_set_16(a) _mm_set1_epi16(a)
#define vec128_load(a) _mm_load_si128(a)
#define vec128_storeu(a, b) _mm_storeu_si128(a, b)
#define vec128_add(a, b) _mm_add_epi16(a, b)
#define NumRegistersSIMD 12
#define MaxChunkSize 32
#elif USE_SSE2
using vec_t = __m128i;
using vec_i8_t = std::uint64_t; // for the correct size -- will be loaded into an xmm reg
using vec128_t = __m128i;
using psqt_vec_t = __m128i;
using vec_uint_t = __m128i;
#define vec_load(a) (*(a))
#define vec_store(a, b) *(a) = (b)
#define vec_add_16(a, b) _mm_add_epi16(a, b)
#define vec_sub_16(a, b) _mm_sub_epi16(a, b)
#define vec_mulhi_16(a, b) _mm_mulhi_epi16(a, b)
#define vec_zero() _mm_setzero_si128()
#define vec_set_16(a) _mm_set1_epi16(a)
#define vec_max_16(a, b) _mm_max_epi16(a, b)
#define vec_min_16(a, b) _mm_min_epi16(a, b)
#define vec_slli_16(a, b) _mm_slli_epi16(a, b)
#define vec_packus_16(a, b) _mm_packus_epi16(a, b)
#define vec_load_psqt(a) (*(a))
#define vec_store_psqt(a, b) *(a) = (b)
#define vec_add_psqt_32(a, b) _mm_add_epi32(a, b)
#define vec_sub_psqt_32(a, b) _mm_sub_epi32(a, b)
#define vec_zero_psqt() _mm_setzero_si128()
#ifdef USE_SSSE3
#define vec_nnz(a) \
_mm_movemask_ps(_mm_castsi128_ps(_mm_cmpgt_epi32(a, _mm_setzero_si128())))
#endif
#ifdef __i386__
inline __m128i _mm_cvtsi64_si128(int64_t val) {
return _mm_loadl_epi64(reinterpret_cast<const __m128i*>(&val));
}
#endif
#ifdef USE_SSE41
#define vec_convert_8_16(a) _mm_cvtepi8_epi16(_mm_cvtsi64_si128(static_cast<int64_t>(a)))
#else
// Credit: Yoshie2000
inline __m128i vec_convert_8_16(uint64_t x) {
__m128i v8 = _mm_cvtsi64_si128(static_cast<int64_t>(x));
__m128i sign = _mm_cmpgt_epi8(_mm_setzero_si128(), v8);
return _mm_unpacklo_epi8(v8, sign);
}
#endif
#define vec128_zero _mm_setzero_si128()
#define vec128_set_16(a) _mm_set1_epi16(a)
#define vec128_load(a) _mm_load_si128(a)
#define vec128_storeu(a, b) _mm_storeu_si128(a, b)
#define vec128_add(a, b) _mm_add_epi16(a, b)
#define NumRegistersSIMD (Is64Bit ? 12 : 6)
#define MaxChunkSize 16
#elif USE_NEON
using vec_i8x8_t __attribute__((may_alias)) = int8x8_t;
using vec_i16x8_t __attribute__((may_alias)) = int16x8_t;
using vec_i8x16_t __attribute__((may_alias)) = int8x16_t;
using vec_u16x8_t __attribute__((may_alias)) = uint16x8_t;
using vec_i32x4_t __attribute__((may_alias)) = int32x4_t;
using vec_t __attribute__((may_alias)) = int16x8_t;
using vec_i8_t __attribute__((may_alias)) = int8x16_t;
using psqt_vec_t __attribute__((may_alias)) = int32x4_t;
using vec128_t __attribute__((may_alias)) = uint16x8_t;
using vec_uint_t __attribute__((may_alias)) = uint32x4_t;
#define vec_load(a) (*(a))
#define vec_store(a, b) *(a) = (b)
#define vec_add_16(a, b) vaddq_s16(a, b)
#define vec_sub_16(a, b) vsubq_s16(a, b)
#define vec_mulhi_16(a, b) vqdmulhq_s16(a, b)
#define vec_zero() vec_t{0}
#define vec_set_16(a) vdupq_n_s16(a)
#define vec_max_16(a, b) vmaxq_s16(a, b)
#define vec_min_16(a, b) vminq_s16(a, b)
#define vec_slli_16(a, b) vshlq_s16(a, vec_set_16(b))
#define vec_packus_16(a, b) reinterpret_cast<vec_t>(vcombine_u8(vqmovun_s16(a), vqmovun_s16(b)))
#define vec_load_psqt(a) (*(a))
#define vec_store_psqt(a, b) *(a) = (b)
#define vec_add_psqt_32(a, b) vaddq_s32(a, b)
#define vec_sub_psqt_32(a, b) vsubq_s32(a, b)
#define vec_zero_psqt() psqt_vec_t{0}
static constexpr std::uint32_t Mask[4] = {1, 2, 4, 8};
#define vec_nnz(a) \
vaddvq_u32(vandq_u32(vtstq_u32((uint32x4_t) a, (uint32x4_t) a), vld1q_u32(Mask)))
#define vec128_zero vdupq_n_u16(0)
#define vec128_set_16(a) vdupq_n_u16(a)
#define vec128_load(a) vld1q_u16(reinterpret_cast<const std::uint16_t*>(a))
#define vec128_storeu(a, b) vst1q_u16(reinterpret_cast<std::uint16_t*>(a), b)
#define vec128_add(a, b) vaddq_u16(a, b)
#define NumRegistersSIMD 16
#define MaxChunkSize 16
#ifndef __aarch64__
// Instructions that don't exist on 32-bit ARM
inline int16x8_t vaddw_high_s8(int16x8_t a, int8x16_t b) { return vaddw_s8(a, vget_high_s8(b)); }
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
#endif
struct Vec16Wrapper {
#ifdef VECTOR
using type = vec_t;
static type add(const type& lhs, const type& rhs) { return vec_add_16(lhs, rhs); }
static type sub(const type& lhs, const type& rhs) { return vec_sub_16(lhs, rhs); }
#else
using type = BiasType;
static type add(const type& lhs, const type& rhs) { return lhs + rhs; }
static type sub(const type& lhs, const type& rhs) { return lhs - rhs; }
#endif
};
struct Vec32Wrapper {
#ifdef VECTOR
using type = psqt_vec_t;
static type add(const type& lhs, const type& rhs) { return vec_add_psqt_32(lhs, rhs); }
static type sub(const type& lhs, const type& rhs) { return vec_sub_psqt_32(lhs, rhs); }
#else
using type = PSQTWeightType;
static type add(const type& lhs, const type& rhs) { return lhs + rhs; }
static type sub(const type& lhs, const type& rhs) { return lhs - rhs; }
#endif
};
enum UpdateOperation {
Add,
Sub
};
template<typename VecWrapper,
UpdateOperation... ops,
std::enable_if_t<sizeof...(ops) == 0, bool> = true>
typename VecWrapper::type fused(const typename VecWrapper::type& in) {
return in;
}
template<typename VecWrapper,
UpdateOperation update_op,
UpdateOperation... ops,
typename T,
typename... Ts,
std::enable_if_t<is_all_same_v<typename VecWrapper::type, T, Ts...>, bool> = true,
std::enable_if_t<sizeof...(ops) == sizeof...(Ts), bool> = true>
typename VecWrapper::type
fused(const typename VecWrapper::type& in, const T& operand, const Ts&... operands) {
switch (update_op)
{
case Add :
return fused<VecWrapper, ops...>(VecWrapper::add(in, operand), operands...);
case Sub :
return fused<VecWrapper, ops...>(VecWrapper::sub(in, operand), operands...);
default :
static_assert(update_op == Add || update_op == Sub,
"Only Add and Sub are currently supported.");
return typename VecWrapper::type();
}
}
#if defined(USE_AVX512)
[[maybe_unused]] static int m512_hadd(__m512i sum, int bias) {
return _mm512_reduce_add_epi32(sum) + bias;
}
[[maybe_unused]] static void m512_add_dpbusd_epi32(__m512i& acc, __m512i a, __m512i b) {
#if defined(USE_VNNI)
acc = _mm512_dpbusd_epi32(acc, a, b);
#else
__m512i product0 = _mm512_maddubs_epi16(a, b);
product0 = _mm512_madd_epi16(product0, _mm512_set1_epi16(1));
acc = _mm512_add_epi32(acc, product0);
#endif
}
#endif
#if defined(USE_AVX2)
[[maybe_unused]] static int m256_hadd(__m256i sum, int bias) {
__m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum, 1));
sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_PERM_BADC));
sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_PERM_CDAB));
return _mm_cvtsi128_si32(sum128) + bias;
}
[[maybe_unused]] static void m256_add_dpbusd_epi32(__m256i& acc, __m256i a, __m256i b) {
#if defined(USE_VNNI)
acc = _mm256_dpbusd_epi32(acc, a, b);
#else
__m256i product0 = _mm256_maddubs_epi16(a, b);
product0 = _mm256_madd_epi16(product0, _mm256_set1_epi16(1));
acc = _mm256_add_epi32(acc, product0);
#endif
}
#endif
#if defined(USE_SSSE3)
[[maybe_unused]] static int m128_hadd(__m128i sum, int bias) {
sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, 0x4E)); //_MM_PERM_BADC
sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, 0xB1)); //_MM_PERM_CDAB
return _mm_cvtsi128_si32(sum) + bias;
}
[[maybe_unused]] static void m128_add_dpbusd_epi32(__m128i& acc, __m128i a, __m128i b) {
__m128i product0 = _mm_maddubs_epi16(a, b);
product0 = _mm_madd_epi16(product0, _mm_set1_epi16(1));
acc = _mm_add_epi32(acc, product0);
}
#endif
#if defined(USE_NEON_DOTPROD)
[[maybe_unused]] static void
dotprod_m128_add_dpbusd_epi32(int32x4_t& acc, int8x16_t a, int8x16_t b) {
acc = vdotq_s32(acc, a, b);
}
#endif
#if defined(USE_NEON)
[[maybe_unused]] static int neon_m128_reduce_add_epi32(int32x4_t s) {
#if USE_NEON >= 8
return vaddvq_s32(s);
#else
return s[0] + s[1] + s[2] + s[3];
#endif
}
[[maybe_unused]] static int neon_m128_hadd(int32x4_t sum, int bias) {
return neon_m128_reduce_add_epi32(sum) + bias;
}
#endif
#if USE_NEON >= 8
[[maybe_unused]] static void neon_m128_add_dpbusd_epi32(int32x4_t& acc, int8x16_t a, int8x16_t b) {
int16x8_t product0 = vmull_s8(vget_low_s8(a), vget_low_s8(b));
int16x8_t product1 = vmull_high_s8(a, b);
int16x8_t sum = vpaddq_s16(product0, product1);
acc = vpadalq_s16(acc, sum);
}
#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>
class SIMDTiling {
#ifdef VECTOR
// We use __m* types as template arguments, which causes GCC to emit warnings
// about losing some attribute information. This is irrelevant to us as we
// only take their size, so the following pragma are harmless.
#if defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wignored-attributes"
#endif
template<typename SIMDRegisterType, typename LaneType, int NumLanes, int MaxRegisters>
static constexpr int BestRegisterCount() {
constexpr std::size_t RegisterSize = sizeof(SIMDRegisterType);
constexpr std::size_t LaneSize = sizeof(LaneType);
static_assert(RegisterSize >= LaneSize);
static_assert(MaxRegisters <= NumRegistersSIMD);
static_assert(MaxRegisters > 0);
static_assert(NumRegistersSIMD > 0);
static_assert(RegisterSize % LaneSize == 0);
static_assert((NumLanes * LaneSize) % RegisterSize == 0);
const int ideal = (NumLanes * LaneSize) / RegisterSize;
if (ideal <= MaxRegisters)
return ideal;
// Look for the largest divisor of the ideal register count that is smaller than MaxRegisters
for (int divisor = MaxRegisters; divisor > 1; --divisor)
if (ideal % divisor == 0)
return divisor;
return 1;
}
#if defined(__GNUC__)
#pragma GCC diagnostic pop
#endif
public:
static constexpr int NumRegs =
BestRegisterCount<vec_t, WeightType, TransformedFeatureWidth, NumRegistersSIMD>();
static constexpr int NumPsqtRegs =
BestRegisterCount<psqt_vec_t, PSQTWeightType, PSQTBuckets, NumRegistersSIMD>();
static constexpr IndexType TileHeight = NumRegs * sizeof(vec_t) / 2;
static constexpr IndexType PsqtTileHeight = NumPsqtRegs * sizeof(psqt_vec_t) / 4;
static_assert(HalfDimensions % TileHeight == 0, "TileHeight must divide HalfDimensions");
static_assert(PSQTBuckets % PsqtTileHeight == 0, "PsqtTileHeight must divide PSQTBuckets");
#endif
};
}
#endif
Reference in New Issue View Git Blame Copy Permalink