diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index ea54cfe44ce4..9acd96421de6 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -2056,8 +2056,34 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi #if defined __AVX2__ + UNUSED(kmask1); + UNUSED(kmask2); + UNUSED(kmask3); + UNUSED(utmp); + const __m256i m4 = _mm256_set1_epi8(0xF); + // Unpack the 12 packed 6-bit scales/mins of a superblock with SIMD instead + // of the scalar utmp[] shuffle. The scalar version needs ~25 GP + // instructions plus a store/reload of utmp[] to get the 16 bytes back into + // a vector register; the sequence below does it with two byte shuffles and + // a handful of masked shifts, entirely in registers. + // + // Input bytes u[0..11] (u[0..3] = utmp[0], u[4..7] = utmp[1], u[8..11] = utmp[2]) + // produce the 16 bytes { scales[0..7], mins[0..7] }: + // out[ 0.. 3] = u[0..3] & 0x3f + // out[ 4.. 7] = (u[8..11] & 0x0f) | ((u[0..3] & 0xc0) >> 2) + // out[ 8..11] = u[4..7] & 0x3f + // out[12..15] = (u[8..11] >> 4) | ((u[4..7] & 0xc0) >> 2) + const __m128i sc_shuf_a = _mm_setr_epi8( 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 8, 9, 10, 11); + const __m128i sc_shuf_b = _mm_setr_epi8(-1, -1, -1, -1, 0, 1, 2, 3, -1, -1, -1, -1, 4, 5, 6, 7); + const __m128i sc_mask_a = _mm_setr_epi8(0x3f, 0x3f, 0x3f, 0x3f, 0x0f, 0x0f, 0x0f, 0x0f, + 0x3f, 0x3f, 0x3f, 0x3f, 0x00, 0x00, 0x00, 0x00); + const __m128i sc_mask_b = _mm_setr_epi8(0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x0f, 0x0f, 0x0f, 0x0f); + const __m128i sc_mask_c = _mm_setr_epi8(0x00, 0x00, 0x00, 0x00, 0x30, 0x30, 0x30, 0x30, + 0x00, 0x00, 0x00, 0x00, 0x30, 0x30, 0x30, 0x30); + __m256 acc = _mm256_setzero_ps(); __m128 acc_m = _mm_setzero_ps(); @@ -2066,17 +2092,20 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d); const float dmin = -y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); - memcpy(utmp, x[i].scales, 12); - utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); - const uint32_t uaux = utmp[1] & kmask1; - utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); - utmp[2] = uaux; - utmp[0] &= kmask1; + // reads 16 bytes starting at scales[0]; the trailing 4 bytes belong to + // qs[] of the same (144 byte) superblock, so this never reads out of bounds + const __m128i sc_raw = _mm_loadu_si128((const __m128i *)x[i].scales); + const __m128i sc_a = _mm_shuffle_epi8(sc_raw, sc_shuf_a); + const __m128i sc_b = _mm_shuffle_epi8(sc_raw, sc_shuf_b); + const __m128i sc_mins = _mm_or_si128( + _mm_or_si128(_mm_and_si128(sc_a, sc_mask_a), + _mm_and_si128(_mm_srli_epi16(sc_a, 4), sc_mask_b)), + _mm_and_si128(_mm_srli_epi16(sc_b, 2), sc_mask_c)); const uint8_t * GGML_RESTRICT q4 = x[i].qs; const int8_t * GGML_RESTRICT q8 = y[i].qs; - const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); + const __m256i mins_and_scales = _mm256_cvtepu8_epi16(sc_mins); const __m256i q8sums = _mm256_loadu_si256((const __m256i*)y[i].bsums); const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));