Skip to content

ggml-cpu : AVX-512 VNNI and AVX2 vec_dot for PTQ1_0 on x86 (4x token generation) #250

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lenny76 wants to merge 1 commit into
base: prism
Choose a base branch
from
Open

ggml-cpu : AVX-512 VNNI and AVX2 vec_dot for PTQ1_0 on x86 (4x token generation) #250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions ggml/src/ggml-cpu/arch-fallback.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,6 @@
#define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
#define ggml_gemm_pq2_0_4x8_q8_0_generic ggml_gemm_pq2_0_4x8_q8_0
#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
// PTQ1_0 currently has only the generic vec_dot; alias it here until a SIMD version lands
#define ggml_vec_dot_ptq1_0_q8_0_generic ggml_vec_dot_ptq1_0_q8_0
// quants.c
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
// repack.cpp
Expand Down
146 changes: 146 additions & 0 deletions ggml/src/ggml-cpu/arch/x86/quants.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,152 @@ static inline __m128i get_scale_shuffle(int i) {
# define GGML_DPBUSD_256 _mm256_dpbusd_avx_epi32
#endif

#if defined(__AVX512BW__) && defined(__AVX512VNNI__)
// per-byte x*c mod 256, c is the same for both bytes of each 16-bit lane
static inline __m512i ptq1_0_mul_u8(__m512i x, __m512i c) {
const __m512i lo = _mm512_mullo_epi16(x, c);
const __m512i hi = _mm512_mullo_epi16(_mm512_and_si512(x, _mm512_set1_epi16((short) 0xFF00)), c);
return _mm512_mask_blend_epi8(0xAAAAAAAAAAAAAAAAull, lo, hi);
}

// trit (3*v)>>8 in {0,1,2}: 1 from v >= 86, 2 from v >= 171
static inline __m512i ptq1_0_trit(__m512i v) {
const __mmask64 m1 = _mm512_cmpge_epu8_mask(v, _mm512_set1_epi8(86));
const __mmask64 m2 = _mm512_cmpge_epu8_mask(v, _mm512_set1_epi8((char) 171));
return _mm512_mask_mov_epi8(_mm512_maskz_mov_epi8(m1, _mm512_set1_epi8(1)), m2, _mm512_set1_epi8(2));
}
#elif defined(__AVX2__)
// per-byte x*c mod 256, c is the same for both bytes of each 16-bit lane
static inline __m256i ptq1_0_mul_u8(__m256i x, __m256i c) {
const __m256i lo = _mm256_and_si256(_mm256_mullo_epi16(x, c), _mm256_set1_epi16(0x00FF));
const __m256i hi = _mm256_mullo_epi16(_mm256_and_si256(x, _mm256_set1_epi16((short) 0xFF00)), c);
return _mm256_or_si256(lo, hi);
}

// trit (3*v)>>8 in {0,1,2}: unsigned v >= 86 and v >= 171 as signed compares of v ^ 0x80 (masks are -1)
static inline __m256i ptq1_0_trit(__m256i v) {
const __m256i vs = _mm256_xor_si256(v, _mm256_set1_epi8((char) 0x80));
const __m256i m1 = _mm256_cmpgt_epi8(vs, _mm256_set1_epi8((char) (85 ^ 0x80)));
const __m256i m2 = _mm256_cmpgt_epi8(vs, _mm256_set1_epi8((char) (170 ^ 0x80)));
return _mm256_sub_epi8(_mm256_setzero_si256(), _mm256_add_epi8(m1, m2));
}

// sum((t-1)*y) in 8 int32 lanes
static inline __m256i ptq1_0_dot(__m256i t, __m256i y) {
#if defined(GGML_DPBUSD_256)
return _mm256_sub_epi32(GGML_DPBUSD_256(_mm256_setzero_si256(), t, y),
GGML_DPBUSD_256(_mm256_setzero_si256(), _mm256_set1_epi8(1), y));
#else
const __m256i s16 = _mm256_sub_epi16(_mm256_maddubs_epi16(t, y), _mm256_maddubs_epi16(_mm256_set1_epi8(1), y));
return _mm256_madd_epi16(s16, _mm256_set1_epi16(1));
#endif
}
#endif

// PTQ1_0 x Q8_0. Element order of a block (see dequantize_row_ptq1_0): trit nn of qs[0..15] -> 16*nn + m,
// trit nn of qs[16..23] -> 80 + 8*nn + m, trit nn of qh[h] -> 120 + 2*nn + h. Trit nn of byte b is
// ((b*3^nn mod 256)*3)>>8. Each 128-bit lane holds one source run times one power of 3, so the unpack
// needs one in-lane shuffle and no cross-lane permute. AVX-512: 64 weights (two q8_0 sub-blocks) per
// vector; AVX2: one q8_0 sub-block per vector.
void ggml_vec_dot_ptq1_0_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) {
const int qk = QK_PTQ1_0;
const int nb = n / qk;

assert(n % qk == 0);
assert(nrc == 1);
UNUSED(nrc);
UNUSED(bx);
UNUSED(by);
UNUSED(bs);

#if defined(__AVX512BW__) && defined(__AVX512VNNI__)
const block_ptq1_0 * GGML_RESTRICT x = vx;
const block_q8_0 * GGML_RESTRICT y = vy;

// powers of 3 per 16-bit lane: weights 0..63 (qs[0..15] x 1, 3, 9, 27)
const __m512i c0 = _mm512_setr_epi32(
0x00010001, 0x00010001, 0x00010001, 0x00010001, 0x00030003, 0x00030003, 0x00030003, 0x00030003,
0x00090009, 0x00090009, 0x00090009, 0x00090009, 0x001B001B, 0x001B001B, 0x001B001B, 0x001B001B);
// weights 64..127 (qs[0..15] x 81 | qs[16..23] x 1, 3 | x 9, 27 | x 81, qh x 1, 3, 9, 27)
const __m512i c1 = _mm512_setr_epi32(
0x00510051, 0x00510051, 0x00510051, 0x00510051, 0x00010001, 0x00010001, 0x00030003, 0x00030003,
0x00090009, 0x00090009, 0x001B001B, 0x001B001B, 0x00510051, 0x00510051, 0x00030001, 0x001B0009);
// source of weights 64..127: lane 0 = qs[0..15], lanes 1..3 = block bytes 12..27 (qs[16..23] at 4..11, qh at 12..13)
const __m512i sh1 = _mm512_set_epi8(
13, 12, 13, 12, 13, 12, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4,
11, 10, 9, 8, 7, 6, 5, 4, 11, 10, 9, 8, 7, 6, 5, 4,
11, 10, 9, 8, 7, 6, 5, 4, 11, 10, 9, 8, 7, 6, 5, 4,
15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0);
const __m512i ones = _mm512_set1_epi8(1);

__m512 acc = _mm512_setzero_ps();
for (int i = 0; i < nb; i++) {
const uint8_t * GGML_RESTRICT xb = (const uint8_t *) &x[i];
const __m128i a = _mm_loadu_si128((const __m128i *) xb);
const __m128i b = _mm_loadu_si128((const __m128i *) (xb + 12));
const __m512i s0 = _mm512_broadcast_i32x4(a);
const __m512i s1 = _mm512_shuffle_epi8(_mm512_mask_broadcast_i32x4(_mm512_broadcast_i32x4(b), 0x000F, a), sh1);
const __m512i t0 = ptq1_0_trit(ptq1_0_mul_u8(s0, c0));
const __m512i t1 = ptq1_0_trit(ptq1_0_mul_u8(s1, c1));

const block_q8_0 * GGML_RESTRICT yb = &y[4*i];
const __m512i y0 = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *) yb[0].qs)),
_mm256_loadu_si256((const __m256i *) yb[1].qs), 1);
const __m512i y1 = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *) yb[2].qs)),
_mm256_loadu_si256((const __m256i *) yb[3].qs), 1);
// sum((t-1)*y) = sum(t*y) - sum(y), lanes 0..7 and 8..15 hold the two sub-blocks
const __m512i p0 = _mm512_sub_epi32(_mm512_dpbusd_epi32(_mm512_setzero_si512(), t0, y0),
_mm512_dpbusd_epi32(_mm512_setzero_si512(), ones, y0));
const __m512i p1 = _mm512_sub_epi32(_mm512_dpbusd_epi32(_mm512_setzero_si512(), t1, y1),
_mm512_dpbusd_epi32(_mm512_setzero_si512(), ones, y1));
const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
const __m512 sc0 = _mm512_insertf32x8(_mm512_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[0].d)), _mm256_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[1].d)), 1);
const __m512 sc1 = _mm512_insertf32x8(_mm512_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[2].d)), _mm256_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[3].d)), 1);
acc = _mm512_fmadd_ps(_mm512_cvtepi32_ps(p0), sc0, acc);
acc = _mm512_fmadd_ps(_mm512_cvtepi32_ps(p1), sc1, acc);
}
*s = _mm512_reduce_add_ps(acc);
#elif defined(__AVX2__)
const block_ptq1_0 * GGML_RESTRICT x = vx;
const block_q8_0 * GGML_RESTRICT y = vy;

// powers of 3 per 16-bit lane, one vector per q8_0 sub-block (see the AVX-512 path for the layout)
const __m256i c0 = _mm256_setr_epi16( 1, 1, 1, 1, 1, 1, 1, 1, 3, 3, 3, 3, 3, 3, 3, 3);
const __m256i c1 = _mm256_setr_epi16( 9, 9, 9, 9, 9, 9, 9, 9, 27, 27, 27, 27, 27, 27, 27, 27);
const __m256i c2 = _mm256_setr_epi16(81, 81, 81, 81, 81, 81, 81, 81, 1, 1, 1, 1, 3, 3, 3, 3);
const __m256i c3 = _mm256_setr_epi16( 9, 9, 9, 9, 27, 27, 27, 27, 81, 81, 81, 81, 1, 3, 9, 27);
// sources of sub-blocks 2 and 3 (block bytes 12..27: qs[16..23] at 4..11, qh at 12..13)
const __m256i sh2 = _mm256_setr_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
4, 5, 6, 7, 8, 9, 10, 11, 4, 5, 6, 7, 8, 9, 10, 11);
const __m256i sh3 = _mm256_setr_epi8(4, 5, 6, 7, 8, 9, 10, 11, 4, 5, 6, 7, 8, 9, 10, 11,
4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 12, 13, 12, 13, 12, 13);

__m256 acc = _mm256_setzero_ps();
for (int i = 0; i < nb; i++) {
const uint8_t * GGML_RESTRICT xb = (const uint8_t *) &x[i];
const __m128i a = _mm_loadu_si128((const __m128i *) xb);
const __m128i b = _mm_loadu_si128((const __m128i *) (xb + 12));
const __m256i aa = _mm256_broadcastsi128_si256(a);
const __m256i ab = _mm256_shuffle_epi8(_mm256_inserti128_si256(_mm256_castsi128_si256(a), b, 1), sh2);
const __m256i bb = _mm256_shuffle_epi8(_mm256_broadcastsi128_si256(b), sh3);

const block_q8_0 * GGML_RESTRICT yb = &y[4*i];
const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
const __m256i p0 = ptq1_0_dot(ptq1_0_trit(ptq1_0_mul_u8(aa, c0)), _mm256_loadu_si256((const __m256i *) yb[0].qs));
const __m256i p1 = ptq1_0_dot(ptq1_0_trit(ptq1_0_mul_u8(aa, c1)), _mm256_loadu_si256((const __m256i *) yb[1].qs));
const __m256i p2 = ptq1_0_dot(ptq1_0_trit(ptq1_0_mul_u8(ab, c2)), _mm256_loadu_si256((const __m256i *) yb[2].qs));
const __m256i p3 = ptq1_0_dot(ptq1_0_trit(ptq1_0_mul_u8(bb, c3)), _mm256_loadu_si256((const __m256i *) yb[3].qs));
acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(p0), _mm256_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[0].d)), acc);
acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(p1), _mm256_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[1].d)), acc);
acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(p2), _mm256_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[2].d)), acc);
acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(p3), _mm256_set1_ps(d0*GGML_CPU_FP16_TO_FP32(yb[3].d)), acc);
}
*s = hsum_float_8(acc);
#else
ggml_vec_dot_ptq1_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
#endif
}

void ggml_vec_dot_pq2_0_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) {
const int qk = QK_PQ2_0;
const int nb = n / qk;
Expand Down