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

Split accumulator 3-way for avxvnni

Browse Source 
This does the same thing for x86-64-avxvnni as #6336, #6339.

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

No functional change
This commit is contained in:
mstembera
2025-10-14 17:25:48 +02:00
committed by Joost VandeVondele
parent fc5d296f9d
commit c956df4cbb
1 changed files with 36 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/nnue/layers/affine_transform_sparse_input.h
+36 -29
View File
@@ -254,6 +254,7 @@ class AffineTransformSparseInput {
#elif defined(USE_AVX2) #elif defined(USE_AVX2)
using invec_t = __m256i; using invec_t = __m256i;
using outvec_t = __m256i; using outvec_t = __m256i;
#define vec_add_32 _mm256_add_epi32
#define vec_set_32 _mm256_set1_epi32 #define vec_set_32 _mm256_set1_epi32
#define vec_add_dpbusd_32 SIMD::m256_add_dpbusd_epi32 #define vec_add_dpbusd_32 SIMD::m256_add_dpbusd_epi32
#elif defined(USE_SSSE3) #elif defined(USE_SSSE3)
@@ -272,21 +273,19 @@ class AffineTransformSparseInput {
#define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a)) #define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a))
#define vec_add_dpbusd_32 SIMD::neon_m128_add_dpbusd_epi32 #define vec_add_dpbusd_32 SIMD::neon_m128_add_dpbusd_epi32
#endif #endif
static constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType); constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType);
constexpr IndexType NumChunks = ceil_to_multiple<IndexType>(InputDimensions, 8) / ChunkSize; constexpr IndexType NumChunks = ceil_to_multiple<IndexType>(InputDimensions, 8) / ChunkSize;
constexpr IndexType NumAccums = OutputDimensions / OutputSimdWidth; constexpr IndexType NumAccums = OutputDimensions / OutputSimdWidth;
// If there's only one accumulator and we're using high-latency dot product instructions, // If we're using high-latency dot product instructions, split the accumulators
// split it to create three separate dependency chains and merge at the end // to create 3 separate dependency chains and merge at the end
constexpr bool SplitAccums = constexpr IndexType NumRegs =
#if defined(USE_VNNI) #if defined(USE_VNNI)
NumAccums == 1; 3 * NumAccums;
#else #else
false; NumAccums;
#endif #endif
constexpr IndexType NumRegs = SplitAccums ? 3 * NumAccums : NumAccums; std::uint16_t nnz[NumChunks];
std::uint16_t nnz[NumChunks]; IndexType count;
IndexType count;
const auto input32 = reinterpret_cast<const std::int32_t*>(input); const auto input32 = reinterpret_cast<const std::int32_t*>(input);
@@ -303,30 +302,34 @@ class AffineTransformSparseInput {
// convince GCC to not do weird pointer arithmetic in the following loop // convince GCC to not do weird pointer arithmetic in the following loop
const std::int8_t* weights_cp = weights; const std::int8_t* weights_cp = weights;
#if defined(USE_VNNI)
for (IndexType k = NumAccums; k < NumRegs; ++k)
acc[k] = vec_zero();
if constexpr (SplitAccums) while (start < end - 2)
{ {
acc[1] = acc[2] = vec_set_32(0); const std::ptrdiff_t i0 = *start++;
while (start < end - 2) const std::ptrdiff_t i1 = *start++;
const std::ptrdiff_t i2 = *start++;
const invec_t in0 = vec_set_32(input32[i0]);
const invec_t in1 = vec_set_32(input32[i1]);
const invec_t in2 = vec_set_32(input32[i2]);
const auto col0 =
reinterpret_cast<const invec_t*>(&weights_cp[i0 * OutputDimensions * ChunkSize]);
const auto col1 =
reinterpret_cast<const invec_t*>(&weights_cp[i1 * OutputDimensions * ChunkSize]);
const auto col2 =
reinterpret_cast<const invec_t*>(&weights_cp[i2 * OutputDimensions * ChunkSize]);
for (IndexType k = 0; k < NumAccums; ++k)
{ {
const std::ptrdiff_t i0 = *start++; vec_add_dpbusd_32(acc[k], in0, col0[k]);
const std::ptrdiff_t i1 = *start++; vec_add_dpbusd_32(acc[k + NumAccums], in1, col1[k]);
const std::ptrdiff_t i2 = *start++; vec_add_dpbusd_32(acc[k + 2 * NumAccums], in2, col2[k]);
const invec_t in0 = vec_set_32(input32[i0]);
const invec_t in1 = vec_set_32(input32[i1]);
const invec_t in2 = vec_set_32(input32[i2]);
const auto col0 =
reinterpret_cast<const invec_t*>(&weights_cp[i0 * OutputDimensions * ChunkSize]);
const auto col1 =
reinterpret_cast<const invec_t*>(&weights_cp[i1 * OutputDimensions * ChunkSize]);
const auto col2 =
reinterpret_cast<const invec_t*>(&weights_cp[i2 * OutputDimensions * ChunkSize]);
vec_add_dpbusd_32(acc[0], in0, *col0);
vec_add_dpbusd_32(acc[1], in1, *col1);
vec_add_dpbusd_32(acc[2], in2, *col2);
} }
acc[0] = vec_add_32(vec_add_32(acc[0], acc[1]), acc[2]);
} }
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
while (start < end) while (start < end)
{ {
const std::ptrdiff_t i = *start++; const std::ptrdiff_t i = *start++;
@@ -340,8 +343,12 @@ class AffineTransformSparseInput {
outvec_t* outptr = reinterpret_cast<outvec_t*>(output); outvec_t* outptr = reinterpret_cast<outvec_t*>(output);
for (IndexType k = 0; k < NumAccums; ++k) for (IndexType k = 0; k < NumAccums; ++k)
outptr[k] = acc[k]; outptr[k] = acc[k];
#undef vec_set_32 #undef vec_set_32
#undef vec_add_dpbusd_32 #undef vec_add_dpbusd_32
#ifdef vec_add_32
#undef vec_add_32
#endif
#else #else
// Use dense implementation for the other architectures. // Use dense implementation for the other architectures.
affine_transform_non_ssse3<InputDimensions, PaddedInputDimensions, OutputDimensions>( affine_transform_non_ssse3<InputDimensions, PaddedInputDimensions, OutputDimensions>(