Skip to content

Commit

Permalink
ggml-quants : ARM NEON vec_dot for q2_2 and q1_3
Browse files Browse the repository at this point in the history
  • Loading branch information
@compilade
compilade committed Jun 25, 2024
1 parent 5a64b71 commit 5fefd47
Show file tree
Hide file tree
Showing 2 changed files with 182 additions and 18 deletions.
188 changes: 170 additions & 18 deletions ggml-quants.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -686,6 +686,10 @@ void quantize_row_q2_2_reference(const float * restrict x, block_q2_2 * restrict
}
}

void quantize_row_q2_2(const float * restrict x, void * restrict y, int64_t k) {
quantize_row_q2_2_reference(x, y, k);
}

// reference implementation for deterministic creation of model files
void quantize_row_q4_0_reference(const float * restrict x, block_q4_0 * restrict y, int64_t k) {
static const int qk = QK4_0;
Expand Down Expand Up @@ -3900,17 +3904,81 @@ void ggml_vec_dot_q2_2_q8_0(int n, float * restrict s, size_t bs, const void * r
}

*s = hsum_float_8(acc);
#elif defined(__ARM_NEON)
float sumf0 = 0.0f;
float sumf1 = 0.0f;

const uint8x8_t mask = vdup_n_u8(3);
const int8x8_t offset = vdup_n_s8(2);

const int leftovers = nb % 2;

for (int i = 0; i < nb - leftovers; i += 2) {
const uint8x8_t xq8_0 = vld1_u8(x[0].qs);
const uint8x8_t xq8_1 = vld1_u8(x[1].qs);

const int8x8_t xq8_0_0 = vsub_s8(vreinterpret_s8_u8(vand_u8(xq8_0, mask)), offset);
const int8x8_t xq8_0_1 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8_0, 2), mask)), offset);
const int8x8_t xq8_0_2 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8_0, 4), mask)), offset);
const int8x8_t xq8_0_3 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8_0, 6), mask)), offset);
const int8x8_t xq8_1_0 = vsub_s8(vreinterpret_s8_u8(vand_u8(xq8_1, mask)), offset);
const int8x8_t xq8_1_1 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8_1, 2), mask)), offset);
const int8x8_t xq8_1_2 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8_1, 4), mask)), offset);
const int8x8_t xq8_1_3 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8_1, 6), mask)), offset);

const int8x16_t xq8_0_l = vcombine_s8(xq8_0_0, xq8_0_1);
const int8x16_t xq8_0_h = vcombine_s8(xq8_0_2, xq8_0_3);
const int8x16_t xq8_1_l = vcombine_s8(xq8_1_0, xq8_1_1);
const int8x16_t xq8_1_h = vcombine_s8(xq8_1_2, xq8_1_3);

const int8x16_t yq8_0_l = vld1q_s8(y[0].qs + 0);
const int8x16_t yq8_0_h = vld1q_s8(y[0].qs + 16);
const int8x16_t yq8_1_l = vld1q_s8(y[1].qs + 0);
const int8x16_t yq8_1_h = vld1q_s8(y[1].qs + 16);

const int16x8_t dot0 = vaddq_s16(vpaddlq_s8(vmulq_s8(xq8_0_l, yq8_0_l)), vpaddlq_s8(vmulq_s8(xq8_0_h, yq8_0_h)));
const int16x8_t dot1 = vaddq_s16(vpaddlq_s8(vmulq_s8(xq8_1_l, yq8_1_l)), vpaddlq_s8(vmulq_s8(xq8_1_h, yq8_1_h)));

sumf0 += GGML_FP16_TO_FP32(y[0].d) * (float) vaddlvq_s16(dot0);
sumf1 += GGML_FP16_TO_FP32(y[1].d) * (float) vaddlvq_s16(dot1);
x += 2;
y += 2;
}

// one block at a time
for (int i = nb - leftovers; i < nb; ++i) {
const uint8x8_t xq8 = vld1_u8(x->qs);
const int8x8_t xq8_0 = vsub_s8(vreinterpret_s8_u8(vand_u8(xq8, mask)), offset);
const int8x8_t xq8_1 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8, 2), mask)), offset);
const int8x8_t xq8_2 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8, 4), mask)), offset);
const int8x8_t xq8_3 = vsub_s8(vreinterpret_s8_u8(vand_u8(vshr_n_u8(xq8, 6), mask)), offset);

const int8x16_t xq8_l = vcombine_s8(xq8_0, xq8_1);
const int8x16_t xq8_h = vcombine_s8(xq8_2, xq8_3);

const int8x16_t yq8_l = vld1q_s8(y->qs + 0);
const int8x16_t yq8_h = vld1q_s8(y->qs + 16);

const int16x8_t dot0 = vpaddlq_s8(vmulq_s8(xq8_l, yq8_l));
const int16x8_t dot1 = vpaddlq_s8(vmulq_s8(xq8_h, yq8_h));

sumf0 += GGML_FP16_TO_FP32(y->d) * (float) vaddlvq_s16(vaddq_s16(dot0, dot1));
x += 1;
y += 1;
}

*s = sumf0 + sumf1;
#else

float sumf = 0.0;
float sumf = 0.0f;
for (int i = 0; i < nb; i++) {
int sumi = 0;
for (int j = 0; j < qk / 4; j++) {
const uint8_t weight = x[i].qs[j];
sumi += (int)y[i].qs[j + 0*qk/4] * ((weight >> 0) & 3) - 2;
sumi += (int)y[i].qs[j + 1*qk/4] * ((weight >> 2) & 3) - 2;
sumi += (int)y[i].qs[j + 2*qk/4] * ((weight >> 4) & 3) - 2;
sumi += (int)y[i].qs[j + 3*qk/4] * ((weight >> 6) & 3) - 2;
sumi += (int)y[i].qs[j + 0*qk/4] * (((weight >> 0) & 3) - 2);
sumi += (int)y[i].qs[j + 1*qk/4] * (((weight >> 2) & 3) - 2);
sumi += (int)y[i].qs[j + 2*qk/4] * (((weight >> 4) & 3) - 2);
sumi += (int)y[i].qs[j + 3*qk/4] * (((weight >> 6) & 3) - 2);
}
sumf += (float)(sumi)*(GGML_FP16_TO_FP32(y[i].d));
}
Expand Down Expand Up @@ -11314,27 +11382,27 @@ void ggml_vec_dot_q1_3_q8_0(int n, float * restrict s, size_t bs, const void * r

for (int i = 0; i < nb; ++i) {
// const __m128i x12a = _mm_maskload_epi32((const int32_t *) x, _mm_set_epi32(0, -1, -1, -1));
// const __m128i x12b = _mm_insert_epi8(x12a, x->qs[0], 12);
// const __m128i x13b = _mm_insert_epi8(x12a, x->qs[0], 12);
// WARNING: reading 3 bytes further than necessary.
// It's measurably faster than a masked load on an Intel Core m3-8100Y
const __m128i x12b = _mm_loadu_si128((const __m128i_u *) x);
const __m256i x12 = MM256_SET_M128I(x12b, x12b);
const __m128i x13b = _mm_loadu_si128((const __m128i_u *) x);
const __m256i x13 = MM256_SET_M128I(x13b, x13b);

{
// pre-shift the values by 8 bits, and prepare the layout for later packing
__m256i x0l = _mm256_shuffle_epi8(x12, _mm256_set_epi8(5, -1, 5, -1, 5, -1, 5, -1,
__m256i x0l = _mm256_shuffle_epi8(x13, _mm256_set_epi8(5, -1, 5, -1, 5, -1, 5, -1,
4, -1, 4, -1, 4, -1, 4, -1,
1, -1, 1, -1, 1, -1, 1, -1,
0, -1, 0, -1, 0, -1, 0, -1));
__m256i x0h = _mm256_shuffle_epi8(x12, _mm256_set_epi8(7, -1, 7, -1, 7, -1, 7, -1,
__m256i x0h = _mm256_shuffle_epi8(x13, _mm256_set_epi8(7, -1, 7, -1, 7, -1, 7, -1,
6, -1, 6, -1, 6, -1, 6, -1,
3, -1, 3, -1, 3, -1, 3, -1,
2, -1, 2, -1, 2, -1, 2, -1));
__m256i x1l = _mm256_shuffle_epi8(x12, _mm256_set_epi8(7, -1, 6, -1, 5, -1, 4, -1,
__m256i x1l = _mm256_shuffle_epi8(x13, _mm256_set_epi8(7, -1, 6, -1, 5, -1, 4, -1,
3, -1, 2, -1, 1, -1, 0, -1,
9, -1, 9, -1, 9, -1, 9, -1,
8, -1, 8, -1, 8, -1, 8, -1));
__m256i x1h = _mm256_shuffle_epi8(x12, _mm256_set_epi8(12, -1, 12, -1, 12, -1, 12, -1,
__m256i x1h = _mm256_shuffle_epi8(x13, _mm256_set_epi8(12, -1, 12, -1, 12, -1, 12, -1,
11, -1, 10, -1, 9, -1, 8, -1,
11, -1, 11, -1, 11, -1, 11, -1,
10, -1, 10, -1, 10, -1, 10, -1));
Expand Down Expand Up @@ -11385,6 +11453,88 @@ void ggml_vec_dot_q1_3_q8_0(int n, float * restrict s, size_t bs, const void * r
}

*s = hsum_float_8(accumf);
#elif defined(__ARM_NEON)

static const uint8_t k_mask0[16] = {0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3};
static const uint8_t k_mask1[16] = {4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7, 7};
static const uint8_t k_mask2[16] = {8, 8, 8, 8, 9, 9, 9, 9, 10, 10, 10, 10, 11, 11, 11, 11};
static const uint8_t k_mask3[16] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 12, 12, 12};

static const uint8_t k_shift0[16] = {81, 27, 9, 3, 81, 27, 9, 3, 81, 27, 9, 3, 81, 27, 9, 3};
static const uint8_t k_shift3[16] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 81, 27, 9, 3};

// float32x4_t sumv0 = vdupq_n_f32(0.0f);
// float32x4_t sumv1 = vdupq_n_f32(0.0f);

float sumf0 = 0.0f;
float sumf1 = 0.0f;

const uint8x16_t mask0 = vld1q_u8(k_mask0);
const uint8x16_t mask1 = vld1q_u8(k_mask1);
const uint8x16_t mask2 = vld1q_u8(k_mask2);
const uint8x16_t mask3 = vld1q_u8(k_mask3);

const uint8x16_t shift0 = vld1q_u8(k_shift0);
const uint8x16_t shift3 = vld1q_u8(k_shift3);

const int8x16_t one = vdupq_n_s8(1);

for (int i = 0; i < nb; ++i) {
// WARNING: reading 3 bytes further than necessary
const uint8x16_t x13b = vld1q_u8((const uint8_t *) x);

uint8x16_t x0 = vqtbl1q_u8(x13b, mask0);
uint8x16_t x1 = vqtbl1q_u8(x13b, mask1);
uint8x16_t x2 = vqtbl1q_u8(x13b, mask2);
uint8x16_t x3 = vqtbl1q_u8(x13b, mask3);

x0 = vmulq_u8(x0, shift0);
x1 = vmulq_u8(x1, shift0);
x2 = vmulq_u8(x2, shift0);
x3 = vmulq_u8(x3, shift3);

// multiply by 3 and keep the 2 bits above 8 bits
x0 = vshrq_n_u8(vhaddq_u8(x0, vshrq_n_u8(x0, 1)), 6);
x1 = vshrq_n_u8(vhaddq_u8(x1, vshrq_n_u8(x1, 1)), 6);
x2 = vshrq_n_u8(vhaddq_u8(x2, vshrq_n_u8(x2, 1)), 6);
x3 = vshrq_n_u8(vhaddq_u8(x3, vshrq_n_u8(x3, 1)), 6);

// 0, 1, 2 => -1, 0, 1
int8x16_t x0i = vsubq_s8(vreinterpretq_s8_u8(x0), one);
int8x16_t x1i = vsubq_s8(vreinterpretq_s8_u8(x1), one);
int8x16_t x2i = vsubq_s8(vreinterpretq_s8_u8(x2), one);
int8x16_t x3i = vsubq_s8(vreinterpretq_s8_u8(x3), one);

const int8x16_t y0 = vld1q_s8(y[0].qs + 0);
const int8x16_t y1 = vld1q_s8(y[0].qs + 16);
const int8x16_t y2 = vld1q_s8(y[1].qs + 0);
const int8x16_t y3 = vld1q_s8(y[1].qs + 16);

// const int32x4_t p0 = vpaddlq_s16(vaddq_s16(vpaddlq_s8(x0i), vpaddlq_s8(x1i)));
// const int32x4_t p1 = vpaddlq_s16(vaddq_s16(vpaddlq_s8(x2i), vpaddlq_s8(x3i)));

// there's no direct equivalent to _mm_sign_epi8, unfortunately
x0i = vmulq_s8(x0i, y0);
x1i = vmulq_s8(x1i, y1);
x2i = vmulq_s8(x2i, y2);
x3i = vmulq_s8(x3i, y3);

// overall 18.5% faster than with vector sums on a cortex-A72
sumf0 += GGML_FP16_TO_FP32(y[0].d) * (float) vaddlvq_s16(vaddq_s16(vpaddlq_s8(x0i), vpaddlq_s8(x1i)));
sumf1 += GGML_FP16_TO_FP32(y[1].d) * (float) vaddlvq_s16(vaddq_s16(vpaddlq_s8(x2i), vpaddlq_s8(x3i)));

// const int32x4_t p0 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), x0i, y0), x1i, y1);
// const int32x4_t p1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), x2i, y2), x3i, y3);

// sumv0 = vmlaq_n_f32(sumv0, vcvtq_f32_s32(p0), GGML_FP16_TO_FP32(y[0].d));
// sumv1 = vmlaq_n_f32(sumv1, vcvtq_f32_s32(p1), GGML_FP16_TO_FP32(y[1].d));

y += 2;
x += 1;
}

// *s = vaddvq_f32(sumv0) + vaddvq_f32(sumv1);
*s = sumf0 + sumf1;
#else
float sumf = 0.0f;

Expand All @@ -11393,34 +11543,36 @@ void ggml_vec_dot_q1_3_q8_0(int n, float * restrict s, size_t bs, const void * r
for (int j = 0; j < 8; ++j) {
const int8_t * xj = (const int8_t *) (q1_3_grid + x[i].q[j]);
for (int k = 0; k < 4; ++k) {
sum += xj[k] * (int16_t) y[2*i].qs[4*j + k];
sum += xj[k] * (int16_t) y->qs[4*j + k];
}
}

sumf += GGML_FP16_TO_FP32(y[2*i].d) * sum;
sumf += GGML_FP16_TO_FP32(y->d) * sum;
y += 1;
sum = 0;

for (int j = 0; j < 4; ++j) {
const int8_t * xj = (const int8_t *) (q1_3_grid + x[i].q[8 + j]);
for (int k = 0; k < 4; ++k) {
sum += xj[k] * (int16_t) y[2*i + 1].qs[4*j + k];
sum += xj[k] * (int16_t) y->qs[4*j + k];
}
}

for (size_t j = 0; j < 12; ++j) {
uint16_t xj = x[i].q[j];
xj = (xj * 3) >> 8;
sum += ((int16_t) xj - 1) * (int16_t) y[2*i + 1].qs[16 + j];
sum += ((int16_t) xj - 1) * (int16_t) y->qs[16 + j];
}

{
const int8_t * xj = (const int8_t *) (q1_3_grid + x[i].qs[0]);
for (int k = 0; k < 4; ++k) {
sum += (int16_t) xj[k] * (int16_t) y[2*i + 1].qs[28 + k];
sum += (int16_t) xj[k] * (int16_t) y->qs[28 + k];
}
}

sumf += GGML_FP16_TO_FP32(y[2*i + 1].d) * sum;
sumf += GGML_FP16_TO_FP32(y->d) * sum;
y += 1;
}

*s = sumf;
Expand Down
12 changes: 12 additions & 0 deletions ggml.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -836,6 +836,18 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = {
.vec_dot_type = GGML_TYPE_Q8_K,
.nrows = 1,
},
[GGML_TYPE_Q2_2] = {
.type_name = "q2_2",
.blck_size = QK2_2,
.type_size = sizeof(block_q2_2),
.is_quantized = true,
.to_float = (ggml_to_float_t) dequantize_row_q2_2,
.from_float = quantize_row_q2_2,
.from_float_reference = (ggml_from_float_t) quantize_row_q2_2_reference,
.vec_dot = ggml_vec_dot_q2_2_q8_0,
.vec_dot_type = GGML_TYPE_Q8_0,
.nrows = 1,
},
[GGML_TYPE_Q1_3] = {
.type_name = "q1_3",
.blck_size = QK1_3,
Expand Down

0 comments on commit 5fefd47

Please sign in to comment.