iq: gather the AVX2 codebook from the grid table (STRATA_IQ256_GATHER)

The two IQ3 gate/up formats decode each 32-value half with eight scalar loads
assembled by _mm256_set_epi32.  A scratch build that substitutes a constant for
the grid vector shows that assembly is ~half the single-core kernel (IQ3_XXS
248 -> 130 us, IQ3_S 357 -> 140 us), and the two formats are ~58% of the
gate/up lookup work on the Swift 1.5 IQ3_XXS pack.

Add an opt-in AVX2 gather for IQ3_XXS (iq3xxs_grid) and IQ3_S (iq3s_grid),
the mirror of the AVX-512 kernels' STRATA_IQ_GATHER.  On the i7-12850HX
(Alder Lake, no AVX-512) it is bit-identical - a direct scalar-vs-gather output
dump matches for all five gate/up formats, native_expert_parity stays green -
and 1.32-1.34x on the two formats isolated, +1.0% serve end to end (on >= off in
5 of 5 interleaved pairs).  The pool is flat in the deterministic --stats (the
15-worker pool is not decode-bound), so it is a per-core win that compounds.

Default off, like the AVX-512 one: AVX2 gather throughput is microarchitecture-
dependent.  IQ2_XXS / IQ2_XS / IQ2_S use 64-bit grids, where an AVX2 i64 gather
measured a wash, and are unchanged.  The AVX-512 path is untouched.
This commit is contained in:
Yannick Loth authored and Niko1221 committed 2026-10-03 23:34:26 +02:00
1 parent 8ca4e160c2
commit 0736f05d8a
1 file changed
+21 -3
+21 -3
View File
@@ -95,6 +95,11 @@ inline void rows_ahead(const uint8_t* p) {
_mm_prefetch((const char*) p + 64, _MM_HINT_T0);
}
// E-2 on the AVX-2 path (the AVX-512 kernels' STRATA_IQ_GATHER): the IQ3 grids by one AVX2 gather instead of eight
// scalar loads assembled with set_epi32. AVX2 gather is a different instruction with worse throughput on some cores
// (opt-in, as on AVX-512); on the i7-12850HX it measures ~1.3x on the two 32-bit-grid formats, bit-exact.
static const bool gather = std::getenv("STRATA_IQ256_GATHER") != nullptr;
// ---- per format: one 32-value half (values 64*j + 32*half .. +31) -> grid magnitudes, sign vector, scales
template <int TY> struct Fmt32;
@@ -152,8 +157,13 @@ template <> struct Fmt32<18> { // IQ3_XXS: d, qs[64] grid bytes, 8 x u32 (4 x
static constexpr float K = 0.25f;
static inline void decode(const uint8_t* b, int j, int half, __m256i& g, __m256i& sgn, __m256i& sc) {
const uint8_t* q = b + 2 + 16 * j + 8 * half;
g = _mm256_set_epi32((int) iq3xxs_grid[q[7]], (int) iq3xxs_grid[q[6]], (int) iq3xxs_grid[q[5]], (int) iq3xxs_grid[q[4]],
(int) iq3xxs_grid[q[3]], (int) iq3xxs_grid[q[2]], (int) iq3xxs_grid[q[1]], (int) iq3xxs_grid[q[0]]);
if (gather) {
const __m256i idx = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i*) q));
g = _mm256_i32gather_epi32((const int*) iq3xxs_grid, idx, 4);
} else {
g = _mm256_set_epi32((int) iq3xxs_grid[q[7]], (int) iq3xxs_grid[q[6]], (int) iq3xxs_grid[q[5]], (int) iq3xxs_grid[q[4]],
(int) iq3xxs_grid[q[3]], (int) iq3xxs_grid[q[2]], (int) iq3xxs_grid[q[1]], (int) iq3xxs_grid[q[0]]);
}
const uint32_t w = u32(b + 2 + 64 + 8 * j + 4 * half);
sgn = _mm256_set_epi64x((long long) even_signs.v[(w >> 21) & 127], (long long) even_signs.v[(w >> 14) & 127],
(long long) even_signs.v[(w >> 7) & 127], (long long) even_signs.v[w & 127]);
@@ -167,9 +177,17 @@ template <> struct Fmt32<21> { // IQ3_S: d, qs[64], qh[8], signs[32], scales[4
static inline void decode(const uint8_t* b, int j, int half, __m256i& g, __m256i& sgn, __m256i& sc) {
const uint8_t* q = b + 2 + 16 * j + 8 * half;
const uint32_t h = b[66 + 2 * j + half];
if (gather) {
const __m256i bits = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
__m256i idx = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i*) q));
idx = _mm256_add_epi32(idx, _mm256_slli_epi32(
_mm256_and_si256(_mm256_srlv_epi32(_mm256_set1_epi32((int) h), bits), _mm256_set1_epi32(1)), 8));
g = _mm256_i32gather_epi32((const int*) iq3s_grid, idx, 4);
} else {
#define G3(k) (int) iq3s_grid[q[k] | (((h >> k) & 1u) << 8)]
g = _mm256_set_epi32(G3(7), G3(6), G3(5), G3(4), G3(3), G3(2), G3(1), G3(0));
g = _mm256_set_epi32(G3(7), G3(6), G3(5), G3(4), G3(3), G3(2), G3(1), G3(0));
#undef G3
}
const uint64_t m = u64(b + 74 + 8 * j);
sgn = sgn_vec(half ? (uint32_t) (m >> 32) : (uint32_t) m);
const uint8_t s = b[106 + j];