+namespace Stockfish::Eval::NNUE {
+
+ using BiasType = std::int16_t;
+ using WeightType = std::int16_t;
+ using PSQTWeightType = std::int32_t;
+
+ // 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
+
+ static_assert(PSQTBuckets % 8 == 0,
+ "Per feature PSQT values cannot be processed at granularity lower than 8 at a time.");
+
+ #ifdef USE_AVX512
+ typedef __m512i vec_t;
+ typedef __m256i psqt_vec_t;
+ #define vec_load(a) _mm512_load_si512(a)
+ #define vec_store(a,b) _mm512_store_si512(a,b)
+ #define vec_add_16(a,b) _mm512_add_epi16(a,b)
+ #define vec_sub_16(a,b) _mm512_sub_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()
+ #define NumRegistersSIMD 32
+
+ #elif USE_AVX2
+ typedef __m256i vec_t;
+ typedef __m256i psqt_vec_t;
+ #define vec_load(a) _mm256_load_si256(a)
+ #define vec_store(a,b) _mm256_store_si256(a,b)
+ #define vec_add_16(a,b) _mm256_add_epi16(a,b)
+ #define vec_sub_16(a,b) _mm256_sub_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()
+ #define NumRegistersSIMD 16
+
+ #elif USE_SSE2
+ typedef __m128i vec_t;
+ typedef __m128i psqt_vec_t;
+ #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_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()
+ #define NumRegistersSIMD (Is64Bit ? 16 : 8)
+
+ #elif USE_MMX
+ typedef __m64 vec_t;
+ typedef __m64 psqt_vec_t;
+ #define vec_load(a) (*(a))
+ #define vec_store(a,b) *(a)=(b)
+ #define vec_add_16(a,b) _mm_add_pi16(a,b)
+ #define vec_sub_16(a,b) _mm_sub_pi16(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_pi32(a,b)
+ #define vec_sub_psqt_32(a,b) _mm_sub_pi32(a,b)
+ #define vec_zero_psqt() _mm_setzero_si64()
+ #define NumRegistersSIMD 8
+
+ #elif USE_NEON
+ typedef int16x8_t vec_t;
+ typedef int32x4_t psqt_vec_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_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}
+ #define NumRegistersSIMD 16
+
+ #else
+ #undef VECTOR
+
+ #endif
+
+
+ #ifdef VECTOR
+
+ // Compute optimal SIMD register count for feature transformer accumulation.
+
+ // 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.
+ #pragma GCC diagnostic push
+ #pragma GCC diagnostic ignored "-Wignored-attributes"
+
+ template <typename SIMDRegisterType,
+ typename LaneType,
+ int NumLanes,
+ int MaxRegisters>
+ static constexpr int BestRegisterCount()
+ {
+ #define RegisterSize sizeof(SIMDRegisterType)
+ #define 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;
+ }
+
+ static constexpr int NumRegs = BestRegisterCount<vec_t, WeightType, TransformedFeatureDimensions, NumRegistersSIMD>();
+ static constexpr int NumPsqtRegs = BestRegisterCount<psqt_vec_t, PSQTWeightType, PSQTBuckets, NumRegistersSIMD>();
+
+ #pragma GCC diagnostic pop
+
+ #endif
+
+