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.
This commit is contained in:
Iwan Kawrakow
2024-12-11 13:21:16 +01:00
parent bd3bb0df85
commit 04e35a3d51

View File

@@ -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));