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
+22 -15
View File
@@ -254,6 +254,7 @@ class AffineTransformSparseInput {
#elif defined(USE_AVX2)
using invec_t = __m256i;
using outvec_t = __m256i;
#define vec_add_32 _mm256_add_epi32
#define vec_set_32 _mm256_set1_epi32
#define vec_add_dpbusd_32 SIMD::m256_add_dpbusd_epi32
#elif defined(USE_SSSE3)
@@ -272,19 +273,17 @@ class AffineTransformSparseInput {
#define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a))
#define vec_add_dpbusd_32 SIMD::neon_m128_add_dpbusd_epi32
#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 NumAccums = OutputDimensions / OutputSimdWidth;
// If there's only one accumulator and we're using high-latency dot product instructions,
// split it to create three separate dependency chains and merge at the end
constexpr bool SplitAccums =
// 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)
NumAccums == 1;
3 * NumAccums;
#else
false;
NumAccums;
#endif
constexpr IndexType NumRegs = SplitAccums ? 3 * NumAccums : NumAccums;
std::uint16_t nnz[NumChunks];
IndexType count;
@@ -303,10 +302,10 @@ 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)
for (IndexType k = NumAccums; k < NumRegs; ++k)
acc[k] = vec_zero();
if constexpr (SplitAccums)
{
acc[1] = acc[2] = vec_set_32(0);
while (start < end - 2)
{
const std::ptrdiff_t i0 = *start++;
@@ -321,12 +320,16 @@ class AffineTransformSparseInput {
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);
for (IndexType k = 0; k < NumAccums; ++k)
{
vec_add_dpbusd_32(acc[k], in0, col0[k]);
vec_add_dpbusd_32(acc[k + NumAccums], in1, col1[k]);
vec_add_dpbusd_32(acc[k + 2 * NumAccums], in2, col2[k]);
}
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)
{
const std::ptrdiff_t i = *start++;
@@ -340,8 +343,12 @@ class AffineTransformSparseInput {
outvec_t* outptr = reinterpret_cast<outvec_t*>(output);
for (IndexType k = 0; k < NumAccums; ++k)
outptr[k] = acc[k];
#undef vec_set_32
#undef vec_add_dpbusd_32
#ifdef vec_add_32
#undef vec_add_32
#endif
#else
// Use dense implementation for the other architectures.
affine_transform_non_ssse3<InputDimensions, PaddedInputDimensions, OutputDimensions>(