cpu : reduce AVX2 tiled accumulators together (#30296)

Assisted-by: Codex
This commit is contained in:
Trivalik authored and GitHub committed 2026-10-16 09:48:23 +03:00
1 parent c804754529
commit 3dbcb7fca5
1 file changed
+13 -18
+13 -18
View File
@@ -244,16 +244,6 @@ static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_t
const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO;
const int bs_off = seff * TILED_MICRO;
// Horizontal sum helper: converts 8x32-bit int lane inside a YMM register to scalar int32
auto hsum256_epi32 = [](const __m256i v) -> int32_t {
__m128i low = _mm256_castsi256_si128(v);
__m128i high = _mm256_extracti128_si256(v, 1);
__m128i sum128 = _mm_add_epi32(low, high);
sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(2, 3, 0, 1)));
sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(1, 0, 3, 2)));
return _mm_cvtsi128_si32(sum128);
};
for (int i = 0; i < TILED_MICRO; i++) {
const int ar = i0 + i;
const int d_off = seff * TILED_MICRO + ar;
@@ -362,15 +352,20 @@ static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_t
}
}
// OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating
// Reduce accumulators into a 256-bit vector across 8 columns
alignas(32) int32_t s1_vals[GROUP];
#pragma GCC unroll 8
for (int t = 0; t < GROUP; t++) {
s1_vals[t] = hsum256_epi32(acc[t]);
}
__m256i s1_vec = _mm256_load_si256((const __m256i *) s1_vals);
// Reduce the eight accumulators to column sums in acc[0]..acc[7] order, with additions modulo 2^32.
const __m256i ab = _mm256_hadd_epi32(acc[0], acc[1]);
const __m256i cd = _mm256_hadd_epi32(acc[2], acc[3]);
const __m256i ef = _mm256_hadd_epi32(acc[4], acc[5]);
const __m256i gh = _mm256_hadd_epi32(acc[6], acc[7]);
// Each 128-bit half now holds four column sums over input lanes 0..3 or 4..7.
const __m256i abcd = _mm256_hadd_epi32(ab, cd);
const __m256i efgh = _mm256_hadd_epi32(ef, gh);
const __m256i lo = _mm256_permute2x128_si256(abcd, efgh, 0x20);
const __m256i hi = _mm256_permute2x128_si256(abcd, efgh, 0x31);
__m256i s1_vec = _mm256_add_epi32(lo, hi);
// OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating
if constexpr (BIAS != 0 && ACTBIAS) {
// act-bias correction: raw = desired + 128*sum_s scales[s]*w_bsum[s]; corr precomputed in repack_src0
s1_vec = _mm256_sub_epi32(s1_vec, _mm256_set1_epi32(mins_row[0]));