From 04e35a3d51ae0aacc98d2427600fb58e83e09f08 Mon Sep 17 00:00:00 2001 From: Iwan Kawrakow Date: Wed, 11 Dec 2024 13:21:16 +0100 Subject: [PATCH] q4_k_r4: Better ARM implementation PP-512(LLaMA-3.1-8B) is now 122.1 t/s up from 110 t/s. I.e., q4_k_r4 is now (nearly) on par with q4_0_r4. --- ggml/src/iqk/iqk_mul_mat.cpp | 105 +++++++++++++++++------------------ 1 file changed, 52 insertions(+), 53 deletions(-) diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 1b469afe..64d5ccc5 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -8439,8 +8439,8 @@ void mul_mat_q4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& inf auto m3 = vdupq_n_u8(0x30); int nbl = n / QK_K; int8x16_t qx[8]; - int8x16x4_t iscales; - float32x4x4_t scales; + int8x16x2_t iscales; + int32x4x4_t scales; float32x4_t acc[nrc_y] = {}; for (int ix = 0; ix < nrc_x; ix += 4) { const block_q4_k_r4 * iq4 = (const block_q4_k_r4 *)((const char *)vx + ix*bx); @@ -8448,50 +8448,49 @@ void mul_mat_q4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& inf auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq4[ibl].d)); auto m4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq4[ibl].d+4)); m4 = vmulq_f32(m4, vdupq_n_f32(-1.f)); - if constexpr (nrc_y == 1) { - d4 = vmulq_f32(d4, vdupq_n_f32(q8.scale(0, ibl))); - } auto sl = vld1q_u8_x2(iq4[ibl].scales_l); auto sh = vld1q_u8(iq4[ibl].scales_h); - iscales.val[0] = vorrq_u8(vandq_u8(sl.val[0], mf), vandq_u8(vshlq_n_u8(sh, 4), m3)); - iscales.val[1] = vorrq_u8(vandq_u8(sl.val[1], mf), vandq_u8(sh, m3)); - iscales.val[2] = vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(vshlq_n_u8(sh, 2), m3)); - iscales.val[3] = vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m3)); + iscales.val[0] = vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(vshlq_n_u8(sh, 2), m3)); + iscales.val[1] = vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m3)); for (int is = 0; is < 2; ++is) { - auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is+2])); - auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is+2])); - scales.val[0] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); - scales.val[1] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); - scales.val[2] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); - scales.val[3] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + float32x4x4_t fscales; + fscales.val[0] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); + fscales.val[1] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); + fscales.val[2] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); + fscales.val[3] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); for (int iy = 0; iy < nrc_y; ++iy) { auto m8 = vld1q_f32((const float *)q8.y[iy][ibl].bsums + 4*is); - acc[iy] = vmlaq_laneq_f32(acc[iy], scales.val[0], m8, 0); - acc[iy] = vmlaq_laneq_f32(acc[iy], scales.val[1], m8, 1); - acc[iy] = vmlaq_laneq_f32(acc[iy], scales.val[2], m8, 2); - acc[iy] = vmlaq_laneq_f32(acc[iy], scales.val[3], m8, 3); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[0], m8, 0); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[1], m8, 1); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[2], m8, 2); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[3], m8, 3); } - iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); - iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); - scales.val[0] = vmulq_f32(d4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); - scales.val[1] = vmulq_f32(d4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); - scales.val[2] = vmulq_f32(d4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); - scales.val[3] = vmulq_f32(d4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); + } + iscales.val[0] = vorrq_u8(vandq_u8(sl.val[0], mf), vandq_u8(vshlq_n_u8(sh, 4), m3)); + iscales.val[1] = vorrq_u8(vandq_u8(sl.val[1], mf), vandq_u8(sh, m3)); + int32x4_t isum[nrc_y] = {}; + for (int is = 0; is < 2; ++is) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + scales.val[0] = vmovl_s16(vget_low_s16(iscales16_1)); + scales.val[1] = vmovl_s16(vget_high_s16(iscales16_1)); + scales.val[2] = vmovl_s16(vget_low_s16(iscales16_2)); + scales.val[3] = vmovl_s16(vget_high_s16(iscales16_2)); for (int ib = 0; ib < 4; ++ib) { auto bits = vld1q_u8_x4(iq4[ibl].qs + 256*is + 64*ib); prepare_q4_k_quants(mf, bits, qx); for (int iy = 0; iy < nrc_y; ++iy) { auto y = vld1q_s8_x2(q8.y[iy][ibl].qs+128*is+32*ib); auto sumi = interleaved_dotq(qx, y); - if constexpr (nrc_y == 1) { - acc[iy] = vfmaq_f32(acc[iy], scales.val[ib], vcvtq_f32_s32(sumi)); - } else { - auto d4d8 = vmulq_f32(scales.val[ib], vdupq_n_f32(q8.scale(iy, ibl))); - acc[iy] = vfmaq_f32(acc[iy], d4d8, vcvtq_f32_s32(sumi)); - } + isum[iy] = vmlaq_s32(isum[iy], scales.val[ib], sumi); } } } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(d4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(isum[iy])); + } } for (int iy = 0; iy < nrc_y; ++iy) { info.store(ix, iy, acc[iy]); @@ -8509,7 +8508,7 @@ void mul_mat_q5_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& inf auto m10 = vdupq_n_u8(0x10); int nbl = n / QK_K; int8x16_t qx[8]; - int8x16x4_t iscales; + int8x16x2_t iscales; int32x4x4_t scales; float32x4_t acc[nrc_y] = {}; for (int ix = 0; ix < nrc_x; ix += 4) { @@ -8518,32 +8517,32 @@ void mul_mat_q5_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& inf auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq5[ibl].d)); auto m4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq5[ibl].d+4)); m4 = vmulq_f32(m4, vdupq_n_f32(-1.f)); - int32x4_t isum[nrc_y] = {}; auto sl = vld1q_u8_x2(iq5[ibl].scales_l); auto sh = vld1q_u8(iq5[ibl].scales_h); + iscales.val[0] = vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(vshlq_n_u8(sh, 2), m30)); + iscales.val[1] = vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m30)); + for (int is = 0; is < 2; ++is) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + float32x4x4_t fscales; + fscales.val[0] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); + fscales.val[1] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); + fscales.val[2] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); + fscales.val[3] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto m8 = vld1q_f32((const float *)q8.y[iy][ibl].bsums + 4*is); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[0], m8, 0); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[1], m8, 1); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[2], m8, 2); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[3], m8, 3); + } + } iscales.val[0] = vorrq_u8(vandq_u8(sl.val[0], mf), vandq_u8(vshlq_n_u8(sh, 4), m30)); iscales.val[1] = vorrq_u8(vandq_u8(sl.val[1], mf), vandq_u8(sh, m30)); - iscales.val[2] = vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(vshlq_n_u8(sh, 2), m30)); - iscales.val[3] = vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m30)); + int32x4_t isum[nrc_y] = {}; for (int is = 0; is < 2; ++is) { - auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is+2])); - auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is+2])); - { - float32x4x4_t fscales; - fscales.val[0] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); - fscales.val[1] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); - fscales.val[2] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); - fscales.val[3] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); - for (int iy = 0; iy < nrc_y; ++iy) { - auto m8 = vld1q_f32((const float *)q8.y[iy][ibl].bsums + 4*is); - acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[0], m8, 0); - acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[1], m8, 1); - acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[2], m8, 2); - acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[3], m8, 3); - } - } - iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); - iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); scales.val[0] = vmovl_s16(vget_low_s16(iscales16_1)); scales.val[1] = vmovl_s16(vget_high_s16(iscales16_1)); scales.val[2] = vmovl_s16(vget_low_s16(iscales16_2));