+#include "simd.h"
+
+/*
+ This file contains the definition for a fully connected layer (aka affine transform).
+
+ - expected use-case is for when PaddedInputDimensions == 32 and InputDimensions <= 32.
+ - that's why AVX512 is hard to implement
+ - expected use-case is small layers
+ - inputs are processed in chunks of 4, weights are respectively transposed
+ - accumulation happens directly to int32s
+*/
+
+namespace Stockfish::Eval::NNUE::Layers {
+
+#if defined(USE_SSSE3) || defined(USE_NEON_DOTPROD)
+ #define ENABLE_SEQ_OPT
+#endif
+
+// Fallback implementation for older/other architectures.
+// Requires the input to be padded to at least 16 values.
+#ifndef ENABLE_SEQ_OPT
+
+template<IndexType InputDimensions, IndexType PaddedInputDimensions, IndexType OutputDimensions>
+static void affine_transform_non_ssse3(std::int32_t* output,
+ const std::int8_t* weights,
+ const std::int32_t* biases,
+ const std::uint8_t* input) {
+ #if defined(USE_SSE2) || defined(USE_NEON)
+ #if defined(USE_SSE2)
+ // At least a multiple of 16, with SSE2.
+ constexpr IndexType NumChunks = ceil_to_multiple<IndexType>(InputDimensions, 16) / 16;
+ const __m128i Zeros = _mm_setzero_si128();
+ const auto inputVector = reinterpret_cast<const __m128i*>(input);
+
+ #elif defined(USE_NEON)
+ constexpr IndexType NumChunks = ceil_to_multiple<IndexType>(InputDimensions, 16) / 16;
+ const auto inputVector = reinterpret_cast<const int8x8_t*>(input);
+ #endif
+
+ for (IndexType i = 0; i < OutputDimensions; ++i)
+ {
+ const IndexType offset = i * PaddedInputDimensions;
+
+ #if defined(USE_SSE2)
+ __m128i sumLo = _mm_cvtsi32_si128(biases[i]);
+ __m128i sumHi = Zeros;
+ const auto row = reinterpret_cast<const __m128i*>(&weights[offset]);
+ for (IndexType j = 0; j < NumChunks; ++j)
+ {
+ __m128i row_j = _mm_load_si128(&row[j]);
+ __m128i input_j = _mm_load_si128(&inputVector[j]);
+ __m128i extendedRowLo = _mm_srai_epi16(_mm_unpacklo_epi8(row_j, row_j), 8);
+ __m128i extendedRowHi = _mm_srai_epi16(_mm_unpackhi_epi8(row_j, row_j), 8);
+ __m128i extendedInputLo = _mm_unpacklo_epi8(input_j, Zeros);
+ __m128i extendedInputHi = _mm_unpackhi_epi8(input_j, Zeros);
+ __m128i productLo = _mm_madd_epi16(extendedRowLo, extendedInputLo);
+ __m128i productHi = _mm_madd_epi16(extendedRowHi, extendedInputHi);
+ sumLo = _mm_add_epi32(sumLo, productLo);
+ sumHi = _mm_add_epi32(sumHi, productHi);
+ }
+ __m128i sum = _mm_add_epi32(sumLo, sumHi);
+ __m128i sumHigh_64 = _mm_shuffle_epi32(sum, _MM_SHUFFLE(1, 0, 3, 2));
+ sum = _mm_add_epi32(sum, sumHigh_64);
+ __m128i sum_second_32 = _mm_shufflelo_epi16(sum, _MM_SHUFFLE(1, 0, 3, 2));
+ sum = _mm_add_epi32(sum, sum_second_32);
+ output[i] = _mm_cvtsi128_si32(sum);
+
+ #elif defined(USE_NEON)
+
+ int32x4_t sum = {biases[i]};
+ const auto row = reinterpret_cast<const int8x8_t*>(&weights[offset]);
+ for (IndexType j = 0; j < NumChunks; ++j)
+ {
+ int16x8_t product = vmull_s8(inputVector[j * 2], row[j * 2]);
+ product = vmlal_s8(product, inputVector[j * 2 + 1], row[j * 2 + 1]);
+ sum = vpadalq_s16(sum, product);
+ }
+ output[i] = sum[0] + sum[1] + sum[2] + sum[3];