From 6153d0e7946e63226dcbdc8f2eb9d9ec12e54cba Mon Sep 17 00:00:00 2001 From: Iwan Kawrakow Date: Mon, 16 Jun 2025 18:21:33 +0300 Subject: [PATCH] iq3_kt: AVX2 GEMM --- ggml/src/ggml.c | 16 ++--- ggml/src/iqk/iqk_gemm_ktquants.cpp | 108 +++++++++++++++++++++++++---- 2 files changed, 103 insertions(+), 21 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index edf25e35..cc056f89 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1601,11 +1601,6 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = { #else .vec_dot_type = GGML_TYPE_Q8_0_X4, #endif -//#ifdef __ARM_NEON -// .vec_dot_type = GGML_TYPE_F16, -//#else -// .vec_dot_type = GGML_TYPE_F32, -//#endif .nrows = 1, .row_meta_size = 4, }, @@ -1618,11 +1613,16 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = { .from_float = quantize_row_iq3_kt, .from_float_ref = (ggml_from_float_t)quantize_row_iq3_kt_ref, .vec_dot = vec_dot_iq3_kt_q8_k, -#ifdef __ARM_NEON - .vec_dot_type = GGML_TYPE_F16, +#if defined __AVX2__ + .vec_dot_type = GGML_TYPE_Q8_2_X4, #else - .vec_dot_type = GGML_TYPE_F32, + .vec_dot_type = GGML_TYPE_Q8_0_X4, #endif +//#ifdef __ARM_NEON +// .vec_dot_type = GGML_TYPE_F16, +//#else +// .vec_dot_type = GGML_TYPE_F32, +//#endif .nrows = 1, .row_meta_size = 4, }, diff --git a/ggml/src/iqk/iqk_gemm_ktquants.cpp b/ggml/src/iqk/iqk_gemm_ktquants.cpp index fc954b54..630c2e27 100644 --- a/ggml/src/iqk/iqk_gemm_ktquants.cpp +++ b/ggml/src/iqk/iqk_gemm_ktquants.cpp @@ -98,7 +98,7 @@ struct Trellis2 { }; -template +template struct Trellis3 { constexpr static uint32_t ka = 0xCBAC1FED; constexpr static uint32_t ka1 = ka*ka; @@ -127,7 +127,11 @@ struct Trellis3 { auto dot = _mm256_maddubs_epi16(v8, _mm256_set1_epi32(0x01010101)); auto i8 = _mm256_add_epi32(_mm256_set1_epi32(-126), _mm256_madd_epi16(dot, _mm256_set1_epi16(1))); #endif - return _mm256_cvtepi32_ps(i8); + if constexpr (is_abs) { + return _mm256_cvtepi32_ps(_mm256_sign_epi32(i8, i8)); + } else { + return _mm256_cvtepi32_ps(i8); + } } inline __m256 gen8(uint32_t val) const { auto v8 = _mm256_and_si256(next8(val), _mm256_set1_epi32(0x3f3f3f3f)); @@ -137,11 +141,14 @@ struct Trellis3 { auto dot = _mm256_maddubs_epi16(v8, _mm256_set1_epi32(0x01010101)); auto i8 = _mm256_add_epi32(_mm256_set1_epi32(-126), _mm256_madd_epi16(dot, _mm256_set1_epi16(1))); #endif - return _mm256_cvtepi32_ps(i8); + if constexpr (is_abs) { + return _mm256_cvtepi32_ps(_mm256_sign_epi32(i8, i8)); + } else { + return _mm256_cvtepi32_ps(i8); + } } - template inline __m256i next32(const uint32_t * val) const { - const __m256i offset = is_unsigned ? _mm256_setzero_si256() : _mm256_set1_epi32(-126); + const __m256i offset = _mm256_set1_epi32(-126); __m256i aux[4]; for (int i = 0; i < 4; ++i) { auto i8 = _mm256_and_si256(next8(val[2*i+0], val[2*i+1]), _mm256_set1_epi32(0x3f3f3f3f)); @@ -156,11 +163,15 @@ struct Trellis3 { aux[2] = _mm256_packs_epi32(aux[2], aux[3]); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 aux[0] = _mm256_packs_epi16(aux[0], aux[2]); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27 // 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 - return _mm256_permutevar8x32_epi32(aux[0], shuffle); + if constexpr (is_abs) { + auto result = _mm256_permutevar8x32_epi32(aux[0], shuffle); + return _mm256_sign_epi8(result, result); + } else { + return _mm256_permutevar8x32_epi32(aux[0], shuffle); + } } - template inline __m256i next32(const uint16_t * val, uint32_t v0) const { - const __m256i offset = is_unsigned ? _mm256_setzero_si256() : _mm256_set1_epi32(-126); + const __m256i offset = _mm256_set1_epi32(-126); __m256i aux[4]; for (int i = 0; i < 4; ++i) { auto i8 = _mm256_and_si256(next8(v0 + val[i]), _mm256_set1_epi32(0x3f3f3f3f)); @@ -175,11 +186,15 @@ struct Trellis3 { aux[2] = _mm256_packs_epi32(aux[2], aux[3]); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 aux[0] = _mm256_packs_epi16(aux[0], aux[2]); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27 // 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 - return _mm256_permutevar8x32_epi32(aux[0], shuffle); + if constexpr (is_abs) { + auto result = _mm256_permutevar8x32_epi32(aux[0], shuffle); + return _mm256_sign_epi8(result, result); + } else { + return _mm256_permutevar8x32_epi32(aux[0], shuffle); + } } - template inline void next64(const uint32_t * val, __m256i * result) const { - const __m256i offset = is_unsigned ? _mm256_setzero_si256() : _mm256_set1_epi32(-126); + const __m256i offset = _mm256_set1_epi32(-126); auto vka3 = _mm256_set1_epi32(ka3); __m256i aux[8]; for (int i = 0; i < 4; ++i) { @@ -203,6 +218,9 @@ struct Trellis3 { aux[4*k+0] = _mm256_packs_epi16(aux[4*k+0], aux[4*k+2]); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27 // 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 result[k] = _mm256_permutevar8x32_epi32(aux[4*k+0], shuffle); + if constexpr (is_abs) { + result[k] = _mm256_sign_epi8(result[k], result[k]); + } } } }; @@ -449,6 +467,70 @@ void mul_mat_iq2_kt_q8_2_x4_T(int n, const void * vx, size_t bx, const DataInfo& } } +void iqk_dequantize_iq3_kt_q80_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + const int nb = n/QK_K; + + Trellis3 trellis; + + auto shifts = _mm_set_epi32(0, 0, 4, 0); + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + const block_iq3_kt * x8[8]; + float dkt[8]; + float ls[8]; + float ls_all[64]; + uint32_t idx[8]; + uint32_t sign_bits[16]; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) { + const float * dptr = (const float *)((const char*)vx + (ix+k)*bx); + dkt[k] = dptr[0]; + x8[k] = (const block_iq3_kt *)(dptr + 1); + } + auto vd = _mm256_mul_ps(_mm256_set1_ps(1.01f), _mm256_loadu_ps(dkt)); + + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + auto s8 = _mm_set1_epi32(*(const uint32_t *)x8[k][i].scales); + s8 = _mm_and_si128(_mm_srlv_epi32(s8, shifts), _mm_set1_epi8(0xf)); + auto s32 = _mm256_cvtepi8_epi32(s8); + _mm256_storeu_ps(ls_all + 8*k, _mm256_cvtepi32_ps(s32)); + } + auto mask = _mm256_set1_epi8(1); + for (int ib = 0; ib < QK_K/32; ++ib) { + for (int k = 0; k < 8; ++k) ls[k] = ls_all[8*k+ib]; + auto scales = _mm256_mul_ps(vd, _mm256_loadu_ps(ls)); + _mm_storeu_si128((__m128i *)y[ib].d, _mm256_cvtps_ph(scales, _MM_FROUND_TO_NEAREST_INT)); + for (int j = 0; j < 4; ++j) { + for (int k = 0; k < 8; ++k) { + const uint16_t * ql = (const uint16_t *)x8[k][i].ql; + idx[k] = ql[4*ib+j] + 4096; + auto qh = (const uint32_t *)x8[k][i].qh; + sign_bits[k+0] = qh[2*j+0]; + sign_bits[k+8] = qh[2*j+1]; + } + __m256i packed[2]; + trellis.next64(idx, packed); + auto signs1 = _mm256_loadu_si256((const __m256i *)sign_bits+0); + auto signs2 = _mm256_loadu_si256((const __m256i *)sign_bits+1); + signs1 = _mm256_or_si256(_mm256_cmpeq_epi8(_mm256_and_si256(signs1, mask), mask), _mm256_set1_epi8(1)); + signs2 = _mm256_or_si256(_mm256_cmpeq_epi8(_mm256_and_si256(signs2, mask), mask), _mm256_set1_epi8(1)); + packed[0] = _mm256_sign_epi8(packed[0], signs1); + packed[1] = _mm256_sign_epi8(packed[1], signs2); + _mm256_storeu_si256((__m256i *)y[ib].qs+2*j+0, packed[0]); + _mm256_storeu_si256((__m256i *)y[ib].qs+2*j+1, packed[1]); + } + mask = _mm256_slli_epi16(mask, 1); + } + y += 8; // = QK_K/32; + } + } +} + inline __m256 abs_ps(__m256 vals) { // Clear sign-bit of all the 32-bit floats in vals __m256 sign_bit = _mm256_set1_ps(-0.0f); @@ -887,10 +969,10 @@ bool iqk_set_kernels_ktquants(int ne00, int typeA, int typeB, std::array