From b110945afcf1c6d0b2eff3c5158a69470451489d Mon Sep 17 00:00:00 2001 From: Millaguie Date: Mon, 10 Aug 2026 18:28:15 +0200 Subject: [PATCH] ggml : add AVX-512 DT3 vec_dot Decode both planes of a block with VBMI byte permutes: the *3 multiply chain (wrapping, so it commutes with the permutation) is computed once on the whole 56-byte block, and masked vpermb picks each element's byte from the chain vector of its base-3 digit. The qh lanes never see 3^4, which would read the padding 5th digit of the qh bytes. The trits reach the integer product as xi in {0, 1, 2} via the same avg trick as tq1_0, with VNNI dpbusd against the q8_0 bytes and sum(y) subtracted. The per-q8_0-block sums and the float accumulation keep the exact operation order of the generic implementation, so the result is bit-identical to it (checked by test-dt3). 2.2x over the (autovectorized) generic on a Ryzen AI 9 HX 370. --- ggml/src/ggml-cpu/arch-fallback.h | 1 - ggml/src/ggml-cpu/arch/x86/quants.c | 173 ++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index 7e0ec4ced..8209016cc 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/ggml/src/ggml-cpu/arch-fallback.h @@ -87,7 +87,6 @@ #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) // quants.c #define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 -#define ggml_vec_dot_dt3_q8_0_generic ggml_vec_dot_dt3_q8_0 // repack.cpp #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c index ea54cfe44..cddff1af8 100644 --- a/ggml/src/ggml-cpu/arch/x86/quants.c +++ b/ggml/src/ggml-cpu/arch/x86/quants.c @@ -1571,6 +1571,179 @@ void ggml_vec_dot_tq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo #endif } +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) && defined(__AVX512VNNI__) +// multiply bytes by 3 with wraparound (there is no 8-bit SIMD multiply) +static inline __m512i dt3_mul3_epi8(const __m512i v) { + return _mm512_add_epi8(v, _mm512_add_epi8(v, v)); +} + +// bring the top base-3 digit of each byte down to xi = ((uint8_t) q * 3) >> 8, +// in {0, 1, 2}, with the same avg trick as ggml_vec_dot_tq1_0_q8_K +static inline __m512i dt3_decode_epi8(__m512i q) { + // cancel the +1 from avg so that it behaves like a halving add + q = _mm512_subs_epu8(q, _mm512_set1_epi8(1)); + // multiply by 3 and get the top 2 bits + q = _mm512_avg_epu8(q, _mm512_avg_epu8(q, _mm512_setzero_si512())); + return _mm512_and_si512(_mm512_srli_epi16(q, 6), _mm512_set1_epi8(3)); +} +#endif + +void ggml_vec_dot_dt3_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + assert(n % QK_DT3 == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_dt3 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + const int nb = n / QK_DT3; + +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) && defined(__AVX512VNNI__) + // Source byte of each of the 128 elements of a plane, in element order, as + // offsets into the 56-byte block for plane 0 (see unpack_plane_dt3): + // element m + n*16 (0..79) is digit n of qs[m], m in [0,16) + // element 80 + m + n*8 (80..119) is digit n of qs[16 + m], m in [0,8) + // element 120 + j + n*2 (120..127) is digit n of qh[j], j in [0,2) + // The low vector covers elements 0..63 and the high vector 64..127, so + // that each aligns with two full q8_0 blocks of the other operand. + static const uint8_t kidx_lo[64] = { + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 0 of qs[0..15] + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 1 + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 2 + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 3 + }; + static const uint8_t kidx_hi[64] = { + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 4 of qs[0..15] + 16, 17, 18, 19, 20, 21, 22, 23, // digit 0 of qs[16..23] + 16, 17, 18, 19, 20, 21, 22, 23, // digit 1 + 16, 17, 18, 19, 20, 21, 22, 23, // digit 2 + 16, 17, 18, 19, 20, 21, 22, 23, // digit 3 + 16, 17, 18, 19, 20, 21, 22, 23, // digit 4 + 48, 49, 48, 49, 48, 49, 48, 49, // digits 0,0,1,1,2,2,3,3 of qh[0],qh[1] + }; + // plane 1 offsets: qs starts 24 bytes later, qh starts 2 bytes later + static const uint8_t koff_hi[64] = { + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, + 24, 24, 24, 24, 24, 24, 24, 24, 2, 2, 2, 2, 2, 2, 2, 2, + }; + + // digit -> multiplier blend masks, one bit per byte lane (digit n of a + // byte is extracted by multiplying by 3^n and reading the top 2 bits; + // the qh lanes of the high vector never see 3^4, which would read the + // padding 5th digit of the qh bytes) + const __mmask64 mask_lo_3 = (__mmask64) 0x00000000FFFF0000; // lanes 16..31 + const __mmask64 mask_lo_9 = (__mmask64) 0x0000FFFF00000000; // lanes 32..47 + const __mmask64 mask_lo_27 = (__mmask64) 0xFFFF000000000000; // lanes 48..63 + const __mmask64 mask_hi_1 = (__mmask64) 0x0300000000FF0000; // lanes 16..23, 56, 57 + const __mmask64 mask_hi_3 = (__mmask64) 0x0C000000FF000000; // lanes 24..31, 58, 59 + const __mmask64 mask_hi_9 = (__mmask64) 0x300000FF00000000; // lanes 32..39, 60, 61 + const __mmask64 mask_hi_27 = (__mmask64) 0xC000FF0000000000; // lanes 40..47, 62, 63 + + const __m512i idx_lo0 = _mm512_loadu_si512(kidx_lo); + const __m512i idx_hi0 = _mm512_loadu_si512(kidx_hi); + const __m512i idx_lo1 = _mm512_add_epi8(idx_lo0, _mm512_set1_epi8(24)); + const __m512i idx_hi1 = _mm512_add_epi8(idx_hi0, _mm512_loadu_si512(koff_hi)); + + const __m512i zero = _mm512_setzero_si512(); + const __m512i ones = _mm512_set1_epi8(1); + + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + // the block is 56 bytes; the masked load must not read past the end + const __m512i blk = _mm512_maskz_loadu_epi8((__mmask64) ((UINT64_C(1) << sizeof(block_dt3)) - 1), &x[i]); + + // multiplying by 3^n (with wraparound) commutes with the byte + // permutation, so the multiply chain is computed once on the whole + // block and shared by both planes; masked vpermb then picks each + // lane's byte from the chain vector of its digit + const __m512i v3 = dt3_mul3_epi8(blk); + const __m512i v9 = dt3_mul3_epi8(v3); + const __m512i v27 = dt3_mul3_epi8(v9); + const __m512i v81 = dt3_mul3_epi8(v27); + + __m512i xi_lo[2]; + __m512i xi_hi[2]; + + for (int p = 0; p < 2; p++) { + const __m512i idx_lo = p == 0 ? idx_lo0 : idx_lo1; + const __m512i idx_hi = p == 0 ? idx_hi0 : idx_hi1; + + __m512i q_lo = _mm512_permutexvar_epi8(idx_lo, blk); + q_lo = _mm512_mask_permutexvar_epi8(q_lo, mask_lo_3, idx_lo, v3); + q_lo = _mm512_mask_permutexvar_epi8(q_lo, mask_lo_9, idx_lo, v9); + q_lo = _mm512_mask_permutexvar_epi8(q_lo, mask_lo_27, idx_lo, v27); + xi_lo[p] = dt3_decode_epi8(q_lo); + + __m512i q_hi = _mm512_permutexvar_epi8(idx_hi, v81); + q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_1, idx_hi, blk); + q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_3, idx_hi, v3); + q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_9, idx_hi, v9); + q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_27, idx_hi, v27); + xi_hi[p] = dt3_decode_epi8(q_hi); + } + + // one DT3 block (128 weights) maps to four q8_0 blocks (4 * 32 = 128) + const __m512i y_lo = _mm512_inserti64x4(_mm512_castsi256_si512( + _mm256_loadu_si256((const __m256i *) y[4*i + 0].qs)), + _mm256_loadu_si256((const __m256i *) y[4*i + 1].qs), 1); + const __m512i y_hi = _mm512_inserti64x4(_mm512_castsi256_si512( + _mm256_loadu_si256((const __m256i *) y[4*i + 2].qs)), + _mm256_loadu_si256((const __m256i *) y[4*i + 3].qs), 1); + + // t = xi - 1, so t.y = xi.y - sum(y) + const __m512i sy_lo = _mm512_dpbusd_epi32(zero, ones, y_lo); + const __m512i sy_hi = _mm512_dpbusd_epi32(zero, ones, y_hi); + + const __m512i t1_lo = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_lo[0], y_lo), sy_lo); + const __m512i t1_hi = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_hi[0], y_hi), sy_hi); + const __m512i t2_lo = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_lo[1], y_lo), sy_lo); + const __m512i t2_hi = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_hi[1], y_hi), sy_hi); + + // reduce each q8_0 block (8 consecutive int32 lanes) to its sum with + // an hadd tree; integer addition order does not affect the result: + // ab = [A01 A23 B01 B23 | A45 A67 B45 B67] + // abcd = [A0123 B0123 C0123 D0123 | A4567 B4567 C4567 D4567] + // sv = [sum(A) sum(B) sum(C) sum(D)] + const __m256i ab1 = _mm256_hadd_epi32(_mm512_castsi512_si256(t1_lo), _mm512_extracti64x4_epi64(t1_lo, 1)); + const __m256i cd1 = _mm256_hadd_epi32(_mm512_castsi512_si256(t1_hi), _mm512_extracti64x4_epi64(t1_hi, 1)); + const __m256i ab2 = _mm256_hadd_epi32(_mm512_castsi512_si256(t2_lo), _mm512_extracti64x4_epi64(t2_lo, 1)); + const __m256i cd2 = _mm256_hadd_epi32(_mm512_castsi512_si256(t2_hi), _mm512_extracti64x4_epi64(t2_hi, 1)); + const __m256i abcd1 = _mm256_hadd_epi32(ab1, cd1); + const __m256i abcd2 = _mm256_hadd_epi32(ab2, cd2); + const __m128i sv1 = _mm_add_epi32(_mm256_castsi256_si128(abcd1), _mm256_extracti128_si256(abcd1, 1)); + const __m128i sv2 = _mm_add_epi32(_mm256_castsi256_si128(abcd2), _mm256_extracti128_si256(abcd2, 1)); + + int sumi1[4]; + int sumi2[4]; + _mm_storeu_si128((__m128i *) sumi1, sv1); + _mm_storeu_si128((__m128i *) sumi2, sv2); + + const float d1 = GGML_CPU_FP16_TO_FP32(x[i].d[0]); + const float d2 = GGML_CPU_FP16_TO_FP32(x[i].d[1]); + + // same accumulation order as the generic implementation + for (int k = 0; k < 4; k++) { + const float dy = GGML_CPU_FP16_TO_FP32(y[4*i + k].d); + sumf += dy * (d1*sumi1[k] + d2*sumi2[k]); + } + } + + *s = sumf; + +#else + UNUSED(x); + UNUSED(y); + UNUSED(nb); + ggml_vec_dot_dt3_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); +#endif +} + void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { assert(nrc == 1); UNUSED(nrc);