hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Adrien Gallouëtggml : optimize AMX (#24806)37a77fb052026-06-20

1#if defined(__GNUC__)
2#pragma GCC diagnostic ignored "-Wpedantic"
3#pragma GCC diagnostic ignored "-Wunused-local-typedefs"
4#endif
5
6#include "amx.h"
7#include "mmq.h"
8#include "ggml-impl.h"
9#include "ggml-cpu-impl.h"
10#include "simd-mappings.h"
11#include "quants.h"
12#include "ggml-quants.h"
13#include <algorithm>
14#include <type_traits>
15
16#if defined(__gnu_linux__)
17#include <sys/syscall.h>
18#include <unistd.h>
19#endif
20
21#if (defined(_WIN32) || defined(_WIN64))
22#define RESTRICT __restrict
23#else
24#define RESTRICT __restrict__
25#endif
26
27#if (defined(_WIN32) || defined(_WIN64))
28#define ALWAYS_INLINE __forceinline
29#elif __has_attribute(always_inline) || defined(__GNUC__)
30#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
31#else
32#define ALWAYS_INLINE inline
33#endif
34
35#if defined(__AMX_INT8__) && defined(__AVX512VNNI__)
36
37namespace {
38
39// Forced unrolling
40template <int n>
41struct Unroll {
42    template <typename Func, typename... Args>
43    ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
44        Unroll<n - 1>{}(f, args...);
45        f(std::integral_constant<int, n - 1>{}, args...);
46    }
47};
48
49template <>
50struct Unroll<1> {
51    template <typename Func, typename... Args>
52    ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
53        f(std::integral_constant<int, 0>{}, args...);
54    }
55};
56
57// type traits
58template <typename T> struct PackedTypes {};
59template <> struct PackedTypes<block_q4_0> { using type = int8_t; };
60template <> struct PackedTypes<block_q4_1> { using type = uint8_t; };
61template <> struct PackedTypes<block_q8_0> { using type = int8_t; };
62template <typename T> using packed_B_type = typename PackedTypes<T>::type;
63
64template <typename T>
65struct do_compensate : std::integral_constant<bool,
66    std::is_same<T, block_q8_0>::value> {};
67
68template <typename T>
69struct do_unpack : std::integral_constant<bool,
70    std::is_same<T, block_q4_0>::value ||
71    std::is_same<T, block_q4_1>::value> {};
72
73template <typename T>
74struct is_type_qkk : std::integral_constant<bool,
75    std::is_same<T, block_q4_K>::value ||
76    std::is_same<T, block_q5_K>::value ||
77    std::is_same<T, block_q6_K>::value ||
78    std::is_same<T, block_iq4_xs>::value> {};
79
80#define GGML_DISPATCH_FLOATING_TYPES(TYPE, ...)                                        \
81    [&] {                                                                              \
82        switch (TYPE) {                                                                \
83            case GGML_TYPE_F16: {                                                      \
84                using type = ggml_fp16_t;                                              \
85                constexpr int blck_size = 16;                                          \
86                return __VA_ARGS__();                                                  \
87            }                                                                          \
88            case GGML_TYPE_BF16: {                                                     \
89                using type = ggml_bf16_t;                                              \
90                constexpr int blck_size = 32;                                          \
91                return __VA_ARGS__();                                                  \
92            }                                                                          \
93            default:                                                                   \
94                fprintf(stderr, "Unsupported floating data type\n");                   \
95        }                                                                              \
96    }()
97
98#define GGML_DISPATCH_QTYPES(QT, ...)                                                  \
99    [&] {                                                                              \
100        switch (QT) {                                                                  \
101            case GGML_TYPE_Q4_0: {                                                     \
102                using type = block_q4_0;                                               \
103                using vec_dot_type = block_q8_0;                                       \
104                constexpr int blck_size = QK4_0;                                       \
105                return __VA_ARGS__();                                                  \
106            }                                                                          \
107            case GGML_TYPE_Q4_1: {                                                     \
108                using type = block_q4_1;                                               \
109                using vec_dot_type = block_q8_1;                                       \
110                constexpr int blck_size = QK4_1;                                       \
111                return __VA_ARGS__();                                                  \
112            }                                                                          \
113            case GGML_TYPE_Q8_0: {                                                     \
114                using type = block_q8_0;                                               \
115                using vec_dot_type = block_q8_0;                                       \
116                constexpr int blck_size = QK8_0;                                       \
117                return __VA_ARGS__();                                                  \
118            }                                                                          \
119            case GGML_TYPE_Q4_K: {                                                     \
120                using type = block_q4_K;                                               \
121                using vec_dot_type = block_q8_K;                                       \
122                constexpr int blck_size = QK_K;                                        \
123                return __VA_ARGS__();                                                  \
124            }                                                                          \
125            case GGML_TYPE_Q5_K: {                                                     \
126                using type = block_q5_K;                                               \
127                using vec_dot_type = block_q8_K;                                       \
128                constexpr int blck_size = QK_K;                                        \
129                return __VA_ARGS__();                                                  \
130            }                                                                          \
131            case GGML_TYPE_Q6_K: {                                                     \
132                using type = block_q6_K;                                               \
133                using vec_dot_type = block_q8_K;                                       \
134                constexpr int blck_size = QK_K;                                        \
135                return __VA_ARGS__();                                                  \
136            }                                                                          \
137            case GGML_TYPE_IQ4_XS: {                                                   \
138                using type = block_iq4_xs;                                             \
139                using vec_dot_type = block_q8_K;                                       \
140                constexpr int blck_size = QK_K;                                        \
141                return __VA_ARGS__();                                                  \
142            }                                                                          \
143            default:                                                                   \
144                fprintf(stderr, "Unsupported quantized data type: %d\n", int(TYPE));   \
145        }                                                                              \
146    }()
147
148#define GGML_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...)                                     \
149    [&] {                                                                              \
150        if (BOOL_V) {                                                                  \
151            constexpr bool BOOL_NAME = true;                                           \
152            return __VA_ARGS__();                                                      \
153        } else {                                                                       \
154            constexpr bool BOOL_NAME = false;                                          \
155            return __VA_ARGS__();                                                      \
156        }                                                                              \
157    }()
158
159// define amx tile config data structure
160struct tile_config_t{
161    uint8_t palette_id = 0;
162    uint8_t start_row = 0;
163    uint8_t reserved_0[14] = {0};
164    uint16_t colsb[16] = {0};
165    uint8_t rows[16] = {0};
166};
167
168// Notes: amx tile config
169//
170// Typically, TMUL calculates A and B of size 16 x 64 containing INT8 values,
171// and accumulate the result to a 16 x 16 matrix C containing INT32 values,
172//
173// As many GGUF quantized types as `block_size` of 32, so a 16-16-32 config is used
174// instead of the normally used 16-16-64 config.
175//
176//    Block A: {16, 32}, dtype = int8_t
177//    Block B: {16, 32}, dtype = uint8_t/int8_t
178//    Block C: {16, 16}, dtype = int32_t
179//
180// Block B needs to be prepacked to vnni format before feeding into  TMUL:
181//    packed_B: from {n, k} to {k/vnni_blk, n, vnni_blck}, viewed in 2d, we get {8, 64}
182//
183// Therefore, we get tileconfig:
184//             A    B    C
185//    rows    16    8   16
186//    colsb   32   64   16
187//
188// For tile distribution, follow a 2-2-4 pattern, e.g. A used TMM2-TMM3, B used TMM0-TMM1,
189// C used TMM4-TMM7:
190//            B TMM0  B TMM1
191//    A TMM2  C TMM4  C TMM6
192//    A TMM3  C TMM5  C TMM7
193//
194// Each `amx` kernel handles 4 blocks at a time: 2MB * 2NB, when m < 2 * BLOCK_M, unpack A
195// will be needed.
196//
197// Here another commonly used pattern 1-3-3 is skipped, as it is mostly used when m <=16;
198// and the single batch gemm (m=1) has a special fast path with `avx512-vnni`.
199//
200// ref: https://www.intel.com/content/www/us/en/developer/articles/code-sample/
201//    advanced-matrix-extensions-intrinsics-functions.html
202//
203
204inline void ggml_tile_config_init(void) {
205    static thread_local bool done = false;
206
207    if (done) {
208        return;
209    }
210
211    alignas(64) tile_config_t tc = {};
212    tc.palette_id = 1;
213    tc.start_row = 0;
214    tc.rows[0] = 8;   tc.colsb[0] = 64;
215    tc.rows[1] = 8;   tc.colsb[1] = 64;
216    tc.rows[2] = 16;  tc.colsb[2] = 32;
217    tc.rows[3] = 16;  tc.colsb[3] = 32;
218    tc.rows[4] = 16;  tc.colsb[4] = 64;
219    tc.rows[5] = 16;  tc.colsb[5] = 64;
220    tc.rows[6] = 16;  tc.colsb[6] = 64;
221    tc.rows[7] = 16;  tc.colsb[7] = 64;
222
223    _tile_loadconfig(&tc);
224    done = true;
225}
226
227// we need an extra 16 * 4B (TILE_N * int32_t) for each NB/KB block for compensation.
228// See the notes `s8s8 igemm compensation in avx512-vnni` for detail.
229template <typename TB>
230int get_tile_size() {
231    int tile_size = TILE_N * sizeof(TB);
232    if (do_compensate<TB>::value) {
233        tile_size += TILE_N * sizeof(int32_t);
234    }
235    if (std::is_same<TB, block_q4_K>::value ||
236        std::is_same<TB, block_q5_K>::value) {
237        tile_size += TILE_N * 4;
238    }
239    if (std::is_same<TB, block_iq4_xs>::value) {
240        tile_size += TILE_N * 2;
241    }
242    return tile_size;
243}
244
245template <typename TB, int BLOCK_K>
246int get_row_size(int K) {
247    int KB = K / BLOCK_K;
248    int row_size = KB * sizeof(TB);
249    if (do_compensate<TB>::value) {
250        row_size += KB * sizeof(int32_t);
251    }
252    if (std::is_same<TB, block_q4_K>::value ||
253        std::is_same<TB, block_q5_K>::value) {
254        row_size += KB * 4;
255    }
256    if (std::is_same<TB, block_iq4_xs>::value) {
257        row_size += KB * 2;
258    }
259    return row_size;
260}
261
262// transpose utils
263#define SHUFFLE_EPI32(a, b, mask) \
264    _mm256_castps_si256(_mm256_shuffle_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b), mask))
265inline void transpose_8x8_32bit(__m256i * v, __m256i * v1) {
266    // unpacking and 32-bit elements
267    v1[0] = _mm256_unpacklo_epi32(v[0], v[1]);
268    v1[1] = _mm256_unpackhi_epi32(v[0], v[1]);
269    v1[2] = _mm256_unpacklo_epi32(v[2], v[3]);
270    v1[3] = _mm256_unpackhi_epi32(v[2], v[3]);
271    v1[4] = _mm256_unpacklo_epi32(v[4], v[5]);
272    v1[5] = _mm256_unpackhi_epi32(v[4], v[5]);
273    v1[6] = _mm256_unpacklo_epi32(v[6], v[7]);
274    v1[7] = _mm256_unpackhi_epi32(v[6], v[7]);
275
276    // shuffling the 32-bit elements
277    v[0] = SHUFFLE_EPI32(v1[0], v1[2], 0x44);
278    v[1] = SHUFFLE_EPI32(v1[0], v1[2], 0xee);
279    v[2] = SHUFFLE_EPI32(v1[4], v1[6], 0x44);
280    v[3] = SHUFFLE_EPI32(v1[4], v1[6], 0xee);
281    v[4] = SHUFFLE_EPI32(v1[1], v1[3], 0x44);
282    v[5] = SHUFFLE_EPI32(v1[1], v1[3], 0xee);
283    v[6] = SHUFFLE_EPI32(v1[5], v1[7], 0x44);
284    v[7] = SHUFFLE_EPI32(v1[5], v1[7], 0xee);
285
286    // shuffling 128-bit elements
287    v1[0] = _mm256_permute2f128_si256(v[2], v[0], 0x02);
288    v1[1] = _mm256_permute2f128_si256(v[3], v[1], 0x02);
289    v1[2] = _mm256_permute2f128_si256(v[6], v[4], 0x02);
290    v1[3] = _mm256_permute2f128_si256(v[7], v[5], 0x02);
291    v1[4] = _mm256_permute2f128_si256(v[2], v[0], 0x13);
292    v1[5] = _mm256_permute2f128_si256(v[3], v[1], 0x13);
293    v1[6] = _mm256_permute2f128_si256(v[6], v[4], 0x13);
294    v1[7] = _mm256_permute2f128_si256(v[7], v[5], 0x13);
295}
296
297inline void transpose_16x4_32bit(__m512i * r, __m512i * d) {
298
299    static const __m512i index1 = _mm512_set_epi32(
300        0x0f, 0x0b, 0x07, 0x03,
301        0x0e, 0x0a, 0x06, 0x02,
302        0x0d, 0x09, 0x05, 0x01,
303        0x0c, 0x08, 0x04, 0x00);
304
305    d[0] = _mm512_permutexvar_epi32(index1, r[0]);
306    d[1] = _mm512_permutexvar_epi32(index1, r[1]);
307    d[2] = _mm512_permutexvar_epi32(index1, r[2]);
308    d[3] = _mm512_permutexvar_epi32(index1, r[3]);
309
310    r[0] = _mm512_shuffle_i32x4(d[0], d[1], 0x44);
311    r[1] = _mm512_shuffle_i32x4(d[0], d[1], 0xee);
312    r[2] = _mm512_shuffle_i32x4(d[2], d[3], 0x44);
313    r[3] = _mm512_shuffle_i32x4(d[2], d[3], 0xee);
314
315    d[0] = _mm512_shuffle_i32x4(r[0], r[2], 0x88);
316    d[1] = _mm512_shuffle_i32x4(r[0], r[2], 0xdd);
317    d[2] = _mm512_shuffle_i32x4(r[1], r[3], 0x88);
318    d[3] = _mm512_shuffle_i32x4(r[1], r[3], 0xdd);
319}
320
321inline void transpose_16x16_32bit(__m512i * v) {
322    __m512i v1[16];
323    v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
324    v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
325    v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
326    v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
327    v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
328    v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
329    v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
330    v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
331    v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
332    v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
333    v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
334    v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
335    v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
336    v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
337    v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
338    v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
339
340    v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
341    v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
342    v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
343    v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
344    v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
345    v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
346    v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
347    v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
348    v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
349    v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
350    v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
351    v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
352    v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
353    v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
354    v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
355    v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
356
357    v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
358    v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
359    v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
360    v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
361    v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
362    v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
363    v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
364    v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
365    v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
366    v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
367    v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
368    v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
369    v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
370    v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
371    v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
372    v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
373
374    v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
375    v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
376    v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
377    v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
378    v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
379    v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
380    v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
381    v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
382    v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
383    v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
384    v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
385    v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
386    v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
387    v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
388    v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
389    v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
390}
391
392void quantize_row_q8_K_vnni(const float * RESTRICT x, void * RESTRICT vy, int64_t k) {
393    assert(k % QK_K == 0);
394    const int KB = k / QK_K;
395    constexpr int kVecs = QK_K / 16;
396
397    block_q8_K * y = reinterpret_cast<block_q8_K *>(vy);
398
399    // hold 16 float vecs from x
400    __m512  v[kVecs];
401
402    // hold the quants vecs
403    __m512i vq[kVecs / 4];
404
405    // hold the packed quants vecs
406    __m512i vq_packed[kVecs / 4];
407
408    const __m512 signBit = _mm512_set1_ps(-0.f);
409
410    for (int i = 0; i < KB; ++i) {
411        // Compute max(abs(e)) for the block
412        __m512 vamax = _mm512_set1_ps(0.f);
413        for (int j = 0; j < kVecs; ++j) {
414            v[j] = _mm512_loadu_ps(x); x += 16;
415            vamax = _mm512_max_ps(vamax, _mm512_andnot_ps(signBit, v[j]));
416        }
417        const float amax = _mm512_reduce_max_ps(vamax);
418
419        // Quantize these floats
420        const float iscale = 127.f / amax;
421        y[i].d = GGML_CPU_FP32_TO_FP16(1 / iscale);
422        const float id = ( amax != 0.0f ) ? iscale : 0.f;
423        const __m512 vscale = _mm512_set1_ps(id);
424
425        // Apply multiplier and round to nearest integer
426        for (int j = 0; j < kVecs; ++j) {
427            v[j] = _mm512_mul_ps(v[j], vscale);
428            v[j] = _mm512_roundscale_ps(v[j], (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
429        }
430
431        // Pack to epi8 vecs
432        for (int j = 0; j < kVecs / 4; ++j) {
433            __m128i q8_0 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 0]));
434            __m128i q8_1 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 1]));
435            __m128i q8_2 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 2]));
436            __m128i q8_3 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 3]));
437
438            __m256i q8_01 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_0), (q8_1), 1);
439            __m256i q8_23 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_2), (q8_3), 1);
440
441            vq[j] = _mm512_inserti32x8(_mm512_castsi256_si512(q8_01), q8_23, 1);
442            _mm512_storeu_si512((__m512i *)(y[i].qs + j * 64), vq[j]);
443        }
444
445        // Compute the bsums with vnni
446        transpose_16x4_32bit(vq, vq_packed);
447
448        const __m512i one = _mm512_set1_epi8(1);
449        __m512i sum = _mm512_setzero_si512();
450        for (int k = 0; k < 4; ++k) {
451            sum = _mm512_dpbusd_epi32(sum, one, vq_packed[k]);
452        }
453        _mm256_storeu_si256((__m256i *)(y[i].bsums), _mm512_cvtepi32_epi16(sum));
454    }
455}
456
457// quantize A from float to `vec_dot_type`
458template <typename T>
459inline void from_float(const float * x, char * vy, int64_t k);
460
461template <>
462inline void from_float<block_q8_0>(const float * x, char * vy, int64_t k) {
463    quantize_row_q8_0(x, (block_q8_0 *)vy, k);
464}
465
466template <>
467inline void from_float<block_q8_1>(const float * x, char * vy, int64_t k) {
468    quantize_row_q8_1(x, (block_q8_1 *)vy, k);
469}
470
471template <>
472inline void from_float<block_q8_K>(const float * x, char * vy, int64_t k) {
473#if 1
474    // TODO: this is reference impl!
475    quantize_row_q8_K_ref(x, (block_q8_K *)vy, k);
476#else
477    quantize_row_q8_K_vnni(x, vy, k);
478#endif
479}
480
481// load A from memory to array when nrows can not fill in whole tile
482void unpack_A(int8_t * RESTRICT tile, const block_q8_0 * RESTRICT A, int lda, int nr) {
483    assert(nr != TILE_M);
484    for (int m = 0; m < nr; ++m) {
485        const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs));
486        _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v);
487    }
488}
489
490void unpack_A(int8_t * RESTRICT tile, const block_q8_1 * RESTRICT A, int lda, int nr) {
491    assert(nr != TILE_M);
492    for (int m = 0; m < nr; ++m) {
493        const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs));
494        _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v);
495    }
496}
497
498template <typename TB>
499void unpack_A(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) {
500    assert(nr <= TILE_M);
501    for (int m = 0; m < nr; ++m) {
502        const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs + k * 32));
503        _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v);
504    }
505}
506
507template <>
508void unpack_A<block_q6_K>(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) {
509    assert(nr <= TILE_M);
510    // zero padding k from 16 to 32, so that we don't have to re-config amx
511    const __m128i zero = _mm_setzero_si128();
512    for (int m = 0; m < nr; ++m) {
513        const __m128i v = _mm_loadu_si128((const __m128i *)(A[m * lda].qs + k * 16));
514        const __m256i r = _mm256_insertf128_si256(_mm256_castsi128_si256(v), zero, 1);
515        _mm256_storeu_si256((__m256i *)(tile + m * TILE_K), r);
516    }
517}
518
519#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1)
520inline __m256i bytes_from_nibbles_32(const uint8_t * rsi) {
521    const __m128i tmp = _mm_loadu_si128((const __m128i *)rsi);
522    const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp);
523    const __m256i lowMask = _mm256_set1_epi8(0xF);
524    return _mm256_and_si256(lowMask, bytes);
525}
526
527// used for block_q4_K
528inline __m512i bytes_from_nibbles_64(const uint8_t * rsi) {
529    const __m256i tmp = _mm256_loadu_si256((const __m256i *)rsi);
530    const __m256i lowMask = _mm256_set1_epi8(0xF);
531    const __m256i q4l = _mm256_and_si256(tmp, lowMask);
532    const __m256i q4h = _mm256_and_si256(_mm256_srli_epi16(tmp, 4), lowMask);
533    return _mm512_inserti32x8(_mm512_castsi256_si512(q4l), q4h, 1);
534}
535
536// used for block_q5_K
537inline __m512i bytes_from_nibbles_64(const uint8_t * qs, const uint8_t * qh, int k) {
538    const __m256i lowMask = _mm256_set1_epi8(0xF);
539    __m256i hmask = _mm256_set1_epi8(1);
540    hmask = _mm256_slli_epi16(hmask, k);
541
542    const __m256i q5bits = _mm256_loadu_si256((const __m256i *)qs);
543    const __m256i hbits = _mm256_loadu_si256((const __m256i *)qh);
544
545    const __m256i q5l_0 = _mm256_and_si256(q5bits, lowMask);
546    const __m256i q5h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 0), 4);
547    const __m256i q5_0  = _mm256_add_epi8(q5l_0, q5h_0);
548    hmask = _mm256_slli_epi16(hmask, 1);
549
550    const __m256i q5l_1 = _mm256_and_si256(_mm256_srli_epi16(q5bits, 4), lowMask);
551    const __m256i q5h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 1), 4);
552    const __m256i q5_1  = _mm256_add_epi8(q5l_1, q5h_1);
553
554    return _mm512_inserti32x8(_mm512_castsi256_si512(q5_0), q5_1, 1);
555}
556
557// used for block_q6_K
558inline void bytes_from_nibbles_128(__m512i& r0, __m512i& r1, const uint8_t * qs, const uint8_t * qh) {
559    const __m256i m4 = _mm256_set1_epi8(0xF);
560    const __m256i m2 = _mm256_set1_epi8(0x3);
561
562    const __m256i q6bits1 = _mm256_loadu_si256((const __m256i *)qs);
563    const __m256i q6bits2 = _mm256_loadu_si256((const __m256i *)(qs + 32));
564    const __m256i q6bitsH = _mm256_loadu_si256((const __m256i *)qh);
565
566    const __m256i q6h_0 = _mm256_slli_epi16(_mm256_and_si256(                  q6bitsH,     m2), 4);
567    const __m256i q6h_1 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 2), m2), 4);
568    const __m256i q6h_2 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 4), m2), 4);
569    const __m256i q6h_3 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 6), m2), 4);
570
571    const __m256i q6_0 = _mm256_or_si256(_mm256_and_si256(q6bits1, m4), q6h_0);
572    const __m256i q6_1 = _mm256_or_si256(_mm256_and_si256(q6bits2, m4), q6h_1);
573    const __m256i q6_2 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits1, 4), m4), q6h_2);
574    const __m256i q6_3 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits2, 4), m4), q6h_3);
575
576    r0 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_0), q6_1, 1);
577    r1 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_2), q6_3, 1);
578}
579
580inline __m512i packNibbles(__m512i r0, __m512i r1) {
581    return _mm512_or_si512(r0, _mm512_slli_epi16(r1, 4));
582}
583
584template <typename TB>
585inline void pack_qs(void * RESTRICT packed_B, const TB * RESTRICT B, int KB) {
586    int8_t tmp[8 * 64];
587    __m256i v[8], v2[8];
588    for (int n = 0; n < 8; ++n) {
589        v[n] = bytes_from_nibbles_32(B[n * KB].qs);
590    }
591    transpose_8x8_32bit(v, v2);
592    for (int n = 0; n < 8; ++n) {
593        _mm256_storeu_si256((__m256i *)(tmp + n * 64), v2[n]);
594    }
595    for (int n = 0; n < 8; ++n) {
596        v[n] = bytes_from_nibbles_32(B[(n + 8) * KB].qs);
597    }
598    transpose_8x8_32bit(v, v2);
599    for (int n = 0; n < 8; ++n) {
600        _mm256_storeu_si256((__m256i *)(tmp + n * 64 + 32), v2[n]);
601    }
602
603    // pack again with 128 to fully utilize vector length
604    for (int n = 0; n < 8; n += 2) {
605        __m512i r0 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64));
606        __m512i r1 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64 + 64));
607        __m512i r1r0 = packNibbles(r0, r1);
608        _mm512_storeu_si512((__m512i *)((char *)packed_B + n * 32), r1r0);
609    }
610}
611
612template <>
613inline void pack_qs<block_q8_0>(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) {
614    __m256i v[8], v2[8];
615    for (int n = 0; n < 8; ++n) {
616        v[n] = _mm256_loadu_si256((const __m256i *)(B[n * KB].qs));
617    }
618    transpose_8x8_32bit(v, v2);
619    for (int n = 0; n < 8; ++n) {
620        _mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64), v2[n]);
621    }
622    for (int n = 0; n < 8; ++n) {
623        v[n] = _mm256_loadu_si256((const __m256i *)(B[(n + 8) * KB].qs));
624    }
625    transpose_8x8_32bit(v, v2);
626    for (int n = 0; n < 8; ++n) {
627        _mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64 + 32), v2[n]);
628    }
629}
630
631template <>
632inline void pack_qs<block_q4_K>(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) {
633    __m512i v[16];
634    // QK_K 256 with 8 groups, handle 2 groups at a time
635    char * pb = (char *)packed_B;
636    for (int k = 0; k < QK_K / 64; ++k) {
637        // pack 2 groups { n, g,  k} to {g, k/4, 4n}
638        //          e.g. {16, 2, 32} to {2,   8, 64}
639        for (int n = 0; n < TILE_N; ++n) {
640            v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32);
641        }
642
643        transpose_16x16_32bit(v);
644
645        // pack again with 128 to fully utilize vector length
646        for (int n = 0; n < TILE_N; n += 2) {
647            _mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1]));
648            pb += 64;
649        }
650    }
651}
652
653template <>
654inline void pack_qs<block_q5_K>(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) {
655    __m512i v[16];
656    const __m512i lowMask = _mm512_set1_epi8(0xF);
657    // QK_K 256 with 8 groups, handle 2 groups at a time
658    char * pb = (char *)packed_B;
659    char * ph = (char *)packed_B + (QK_K / 2) * TILE_N;
660    for (int k = 0; k < QK_K / 64; ++k) {
661        // pack 2 groups { n, g,  k} to {g, k/4, 4n}
662        //          e.g. {16, 2, 32} to {2,   8, 64}
663        for (int n = 0; n < TILE_N; ++n) {
664            v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32, B[n * KB].qh, /* group */2 * k);
665        }
666
667        transpose_16x16_32bit(v);
668
669        // 1. pack lower 4bits with 2 groups
670        for (int n = 0; n < TILE_N; n += 2) {
671            // get lower 4 bits
672            const __m512i r0 = _mm512_and_si512(v[n], lowMask);
673            const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask);
674            _mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64;
675        }
676
677        // 2. pack higher 1bit with 2 groups
678        const __m512i hmask = _mm512_set1_epi8(0x10);
679        for (int g = 0; g < 2; ++g) {
680            __m512i hbits = _mm512_setzero_si512();
681            hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 0], hmask), 4));
682            hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 1], hmask), 3));
683            hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 2], hmask), 2));
684            hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 3], hmask), 1));
685            hbits = _mm512_add_epi8(hbits,                   _mm512_and_si512(v[g * 8 + 4], hmask)    );
686            hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 5], hmask), 1));
687            hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 6], hmask), 2));
688            hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 7], hmask), 3));
689            _mm512_storeu_si512((__m512i *)ph, hbits); ph += 64;
690        }
691    }
692}
693
694template <>
695inline void pack_qs<block_q6_K>(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) {
696    __m512i v[32];
697    const __m512i lowMask = _mm512_set1_epi8(0xF);
698    // QK_K 256 with 8 groups, handle 4 groups at a time
699    char * pb = (char *)packed_B;
700    char * ph = (char *)packed_B + (QK_K / 2) * TILE_N;
701    for (int k = 0; k < QK_K / 128; ++k) {
702        for (int n = 0; n < TILE_N; ++n) {
703            bytes_from_nibbles_128(v[n], v[n + 16], B[n * KB].ql + k * 64, B[n * KB].qh + k * 32);
704        }
705
706        // top half: group 0,1 or 4,5; bottom half: group 2,3 or 6,7
707        transpose_16x16_32bit(v);
708        transpose_16x16_32bit(v + 16);
709
710        // 1. pack lower 4bits with 4 groups
711        for (int n = 0; n < 32; n += 2) {
712            const __m512i r0 = _mm512_and_si512(v[n], lowMask);
713            const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask);
714            _mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64;
715        }
716
717        // 2. pack higher 2bit with 4 groups
718        const __m512i hmask = _mm512_set1_epi8(0x30);
719        for (int g = 0; g < 8; ++g) {
720            __m512i hbits = _mm512_setzero_si512();
721            hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 0], hmask), 4));
722            hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 1], hmask), 2));
723            hbits = _mm512_add_epi8(hbits,                   _mm512_and_si512(v[g * 4 + 2], hmask)    );
724            hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 4 + 3], hmask), 2));
725            _mm512_storeu_si512((__m512i *)ph, hbits); ph += 64;
726        }
727    }
728}
729
730template <>
731inline void pack_qs<block_iq4_xs>(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) {
732    __m512i v[16];
733    char * pb = (char *)packed_B;
734    for (int k = 0; k < QK_K / 64; ++k) {
735        for (int n = 0; n < TILE_N; ++n) {
736            __m256i r0 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 +  0);
737            __m256i r1 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 16);
738            v[n] = _mm512_inserti32x8(_mm512_castsi256_si512(r0), r1, 1);
739        }
740
741        transpose_16x16_32bit(v);
742
743        // pack again with 128 to fully utilize vector length
744        for (int n = 0; n < TILE_N; n += 2) {
745            _mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1]));
746            pb += 64;
747        }
748    }
749}
750
751// pack B to vnni formats in 4bits or 8 bits
752void pack_B(void * RESTRICT packed_B, const block_q4_0 * RESTRICT B, int KB) {
753    pack_qs(packed_B, B, KB);
754    ggml_half * d0 = reinterpret_cast<ggml_half *>((char *)packed_B + TILE_N * TILE_K / 2);
755    for (int n = 0; n < TILE_N; ++n) {
756        d0[n] = B[n * KB].d;
757    }
758}
759
760void pack_B(void * RESTRICT packed_B, const block_q4_1 * RESTRICT B, int KB) {
761    pack_qs(packed_B, B, KB);
762    ggml_half * d0 = reinterpret_cast<ggml_half *>((char *)packed_B + TILE_N * TILE_K / 2);
763    ggml_half * m0 = d0 + TILE_N;
764    for (int n = 0; n < TILE_N; ++n) {
765        d0[n] = B[n * KB].d;
766        m0[n] = B[n * KB].m;
767    }
768}
769
770inline void s8s8_compensation(void * RESTRICT packed_B) {
771    // packed_B layout:
772    //   quants {TILE_N, TILEK}  int8_t
773    //   d0     {TILE_N}      ggml_half
774    //   comp   {TILE_N}        int32_t
775    const int offset = TILE_N * TILE_K + TILE_N * sizeof(ggml_half);
776    __m512i vcomp = _mm512_setzero_si512();
777    const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
778    for (int k = 0; k < 8; ++k) {
779        __m512i vb = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + k * 64));
780        vcomp = _mm512_dpbusd_epi32(vcomp, off, vb);
781    }
782    _mm512_storeu_si512((__m512i *)((char *)(packed_B) + offset), vcomp);
783}
784
785void pack_B(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) {
786    pack_qs(packed_B, B, KB);
787    ggml_half * d0 = reinterpret_cast<ggml_half *>((char *)packed_B + TILE_N * TILE_K);
788    for (int n = 0; n < TILE_N; ++n) {
789        d0[n] = B[n * KB].d;
790    }
791    s8s8_compensation(packed_B);
792}
793
794// convert 8 * {min, scale} from int6 to int8
795inline void unpack_mins_and_scales(const uint8_t * scales, uint32_t * utmp) {
796    const uint32_t kmask1 = 0x3f3f3f3f;
797    const uint32_t kmask2 = 0x0f0f0f0f;
798    const uint32_t kmask3 = 0x03030303;
799
800    memcpy(utmp, scales, 12);
801    utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
802    const uint32_t uaux = utmp[1] & kmask1;
803    utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
804    utmp[2] = uaux;
805    utmp[0] &= kmask1;
806}
807
808// packed_B layout:
809//   quants {8, TILE_N, 16}  uint8
810//   scales {8, TILE_N}      uint8
811//   mins   {8, TILE_N}      uint8
812//   d      {TILE_N}     ggml_half
813//   dmin   {TILE_N}     ggml_half
814void pack_B(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) {
815    pack_qs(packed_B, B, KB);
816
817    uint8_t * scales = reinterpret_cast<uint8_t *>((char *)packed_B + (QK_K / 2) * TILE_N);
818    uint8_t * mins = scales + 8 * TILE_N;
819    ggml_half * d = reinterpret_cast<ggml_half *>(mins + 8 * TILE_N);
820    ggml_half * dmin = d + TILE_N;
821
822    union {
823        uint32_t u32[4];
824        uint8_t  u8[16];
825    } s;
826
827    for (int n = 0; n < TILE_N; ++n) {
828        unpack_mins_and_scales(B[n * KB].scales, s.u32);
829        for (int k = 0; k < 8; ++k) {
830            scales[k * TILE_N + n] = s.u8[k];
831            mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8];
832        }
833        d[n] = B[n * KB].d;
834        dmin[n] = B[n * KB].dmin;
835    }
836}
837
838// packed_B layout:
839//   quants {8, TILE_N, 16}  uint8
840//   qh     {8, TILE_N,  4}  uint8
841//   scales {8, TILE_N}      uint8
842//   mins   {8, TILE_N}      uint8
843//   d      {TILE_N}     ggml_half
844//   dmin   {TILE_N}     ggml_half
845void pack_B(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) {
846    pack_qs(packed_B, B, KB);
847
848    uint8_t * scales = reinterpret_cast<uint8_t *>((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N);
849    uint8_t * mins = scales + 8 * TILE_N;
850    ggml_half * d = reinterpret_cast<ggml_half *>(mins + 8 * TILE_N);
851    ggml_half * dmin = d + TILE_N;
852
853    union {
854        uint32_t u32[4];
855        uint8_t  u8[16];
856    } s;
857
858    for (int n = 0; n < TILE_N; ++n) {
859        unpack_mins_and_scales(B[n * KB].scales, s.u32);
860        for (int k = 0; k < 8; ++k) {
861            scales[k * TILE_N + n] = s.u8[k];
862            mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8];
863        }
864        d[n] = B[n * KB].d;
865        dmin[n] = B[n * KB].dmin;
866    }
867}
868
869// packed_B layout:
870//   quants {16, TILE_N, 8}  uint8
871//   qh     {16, TILE_N, 4}  uint8
872//   scales {16, TILE_N}      uint8
873//   d      {TILE_N}     ggml_half
874void pack_B(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) {
875    pack_qs(packed_B, B, KB);
876
877    uint8_t * scales = reinterpret_cast<uint8_t *>((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N);
878    ggml_half * d = reinterpret_cast<ggml_half *>(scales + 16 * TILE_N);
879    for (int n = 0; n < TILE_N; ++n) {
880        const int8_t * ps = B[n * KB].scales;
881        for (int k = 0; k < 16; ++k) {
882            scales[k * TILE_N + n] = ps[k];
883        }
884        d[n] = B[n * KB].d;
885    }
886}
887
888// packed_B layout:
889//   quants {8, TILE_N, 16}  uint8
890//   scales {8, TILE_N}       int8
891//   d      {TILE_N}     ggml_half
892void pack_B(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) {
893    pack_qs(packed_B, B, KB);
894
895    int8_t * scales = reinterpret_cast<int8_t *>((char *)packed_B + (QK_K / 2) * TILE_N);
896    ggml_half * d = reinterpret_cast<ggml_half *>(scales + 8 * TILE_N);
897
898    // pack the scales
899    for (int n = 0; n < TILE_N; ++n) {
900        uint16_t sh = B[n * KB].scales_h;
901        for (int k = 0; k < 8; k += 2) {
902            const int16_t ls1 = ((B[n * KB].scales_l[k / 2] & 0xf) | ((sh << 4) & 0x30)) - 32;
903            const int16_t ls2 = ((B[n * KB].scales_l[k / 2] >>  4) | ((sh << 2) & 0x30)) - 32;
904            scales[(k + 0) * TILE_N + n] = ls1;
905            scales[(k + 1) * TILE_N + n] = ls2;
906            sh >>= 4;
907        }
908        d[n] = B[n * KB].d;
909    }
910}
911
912template<typename TB, typename packed_B_t = packed_B_type<TB>>
913void unpack_B(packed_B_t * RESTRICT tile, const void * RESTRICT packed_B) {
914    GGML_UNUSED(tile);
915    GGML_UNUSED(packed_B);
916}
917
918template <>
919void unpack_B<block_q4_0>(int8_t * RESTRICT tile, const void * RESTRICT packed_B) {
920  const __m512i off = _mm512_set1_epi8(8);
921  const __m512i lowMask = _mm512_set1_epi8(0xF);
922  for (int n = 0; n < 8; n += 2) {
923    __m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32));
924    const __m512i r0 = _mm512_sub_epi8(_mm512_and_si512(bytes, lowMask), off);
925    const __m512i r1 = _mm512_sub_epi8(_mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask), off);
926    _mm512_storeu_si512((__m512i *)(tile + n * 64 +  0), r0);
927    _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
928  }
929}
930
931template <>
932void unpack_B<block_q4_1>(uint8_t * RESTRICT tile, const void * RESTRICT packed_B) {
933    const __m512i lowMask = _mm512_set1_epi8(0xF);
934    for (int n = 0; n < 8; n += 2) {
935        __m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32));
936        const __m512i r0 = _mm512_and_si512(bytes, lowMask);
937        const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
938        _mm512_storeu_si512((__m512i *)(tile + n * 64 +  0), r0);
939        _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
940    }
941}
942
943// packed_B_t for QKK is int8_t
944template <typename TB>
945void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
946    const int packed_B_group_size = QK_K / 2 * TILE_N / 8;
947    const char * packed_B_group = (const char *)packed_B + k * packed_B_group_size;
948    const __m512i lowMask = _mm512_set1_epi8(0xF);
949    for (int n = 0; n < 8; n += 2) {
950        __m512i bytes = _mm512_loadu_si512(packed_B_group + n * 32);
951        const __m512i r0 = _mm512_and_si512(bytes, lowMask);
952        const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
953        _mm512_storeu_si512((__m512i *)(tile + n * 64 +  0), r0);
954        _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
955    }
956}
957
958template <>
959void unpack_B<block_q5_K>(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
960    // lower 4bits, stride 256 bytes
961    const int packed_l4_group_size = QK_K / 2 * TILE_N / 8;
962    const char * pb = (const char *)packed_B + k * packed_l4_group_size;
963
964    // higher 1bit, stride 64 bytes
965    const int packed_h1_group_size = QK_K / 8 * TILE_N / 8;
966    const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h1_group_size;
967    const __m512i hbits = _mm512_loadu_si512(ph);
968
969    const __m512i lowMask = _mm512_set1_epi8(0xF);
970    __m512i hmask0 = _mm512_set1_epi8(0x1);
971    __m512i hmask1 = _mm512_set1_epi8(0x2);
972
973    for (int n = 0; n < 8; n += 2) {
974        __m512i bytes = _mm512_loadu_si512(pb + n * 32);
975        __m512i r0 = _mm512_and_si512(bytes, lowMask);
976        __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
977        __m512i h0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), n), 4);
978        __m512i h1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), n + 1), 4);
979
980        hmask0 = _mm512_slli_epi16(hmask0, 2);
981        hmask1 = _mm512_slli_epi16(hmask1, 2);
982        r0 = _mm512_add_epi8(r0, h0);
983        r1 = _mm512_add_epi8(r1, h1);
984        _mm512_storeu_si512((__m512i *)(tile + n * 64 +  0), r0);
985        _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
986    }
987}
988
989template <>
990void unpack_B<block_q6_K>(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
991    // lower 4bits, stride 128 bytes
992    const int packed_l4_group_size = QK_K / 2 * TILE_N / 16;
993    const char * pb = (const char *)packed_B + k * packed_l4_group_size;
994
995    // higher 2bits, stride 64 bytes
996    const int packed_h2_group_size = QK_K / 4 * TILE_N / 16;
997    const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h2_group_size;
998    const __m512i hbits = _mm512_loadu_si512(ph);
999
1000    const __m512i off = _mm512_set1_epi8(32);
1001    const __m512i lowMask = _mm512_set1_epi8(0xF);
1002    __m512i hmask0 = _mm512_set1_epi8(0x3); // 0011
1003    __m512i hmask1 = _mm512_set1_epi8(0xC); // 1100
1004
1005    // notes: skip zero padding from row4 to row7 as we have done so in `unpack_A`
1006    __m512i bytes = _mm512_loadu_si512(pb);
1007    __m512i r0 = _mm512_and_si512(bytes, lowMask);
1008    __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1009    __m512i h0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask0), 4);
1010    __m512i h1 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask1), 2);
1011    _mm512_storeu_si512((__m512i *)(tile +  0), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off));
1012    _mm512_storeu_si512((__m512i *)(tile + 64), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off));
1013
1014    hmask0 = _mm512_slli_epi16(hmask0, 4);
1015    hmask1 = _mm512_slli_epi16(hmask1, 4);
1016
1017    bytes = _mm512_loadu_si512(pb + 64);
1018    r0 = _mm512_and_si512(bytes, lowMask);
1019    r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1020    h0 =                   _mm512_and_si512(hbits, hmask0);
1021    h1 = _mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), 2);
1022    _mm512_storeu_si512((__m512i *)(tile + 128), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off));
1023    _mm512_storeu_si512((__m512i *)(tile + 192), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off));
1024}
1025
1026template <>
1027void unpack_B<block_iq4_xs>(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
1028    static const __m512i values128 = _mm512_set_epi8(
1029        113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
1030        113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
1031        113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
1032        113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127
1033    );
1034
1035    const int packed_B_group_size = QK_K / 2 * TILE_N / 8;
1036    const char * pb = (const char *)packed_B + k * packed_B_group_size;
1037    const __m512i lowMask = _mm512_set1_epi8(0xF);
1038
1039    for (int n = 0; n < 8; n += 2) {
1040        __m512i bytes = _mm512_loadu_si512(pb + n * 32);
1041        const __m512i r0 = _mm512_shuffle_epi8(values128, _mm512_and_si512(bytes, lowMask));
1042        const __m512i r1 = _mm512_shuffle_epi8(values128, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask));
1043        _mm512_storeu_si512((__m512i *)(tile + n * 64 +  0), r0);
1044        _mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
1045    }
1046}
1047
1048template <typename TA, typename TB, bool is_acc>
1049struct acc_C {};
1050
1051template <bool is_acc>
1052struct acc_C<block_q8_0, block_q4_0, is_acc> {
1053    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) {
1054        const int offset = TILE_N * TILE_K / 2;
1055        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset)));
1056
1057        for (int m = 0; m < nr; ++m) {
1058            const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d));
1059            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1060
1061            __m512 vsum;
1062            if (is_acc) {
1063                vsum = _mm512_loadu_ps(C + m * ldc);
1064            } else {
1065                vsum = _mm512_set1_ps(0.f);
1066            }
1067            vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum);
1068            _mm512_storeu_ps(C + m * ldc, vsum);
1069        }
1070    }
1071};
1072
1073template <bool is_acc>
1074struct acc_C<block_q8_1, block_q4_1, is_acc> {
1075    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_1 * A, int lda, const void * packed_B, int nr) {
1076        const int offset = TILE_N * TILE_K / 2;
1077        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset)));
1078        const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset + TILE_N * sizeof(ggml_half))));
1079
1080        for (int m = 0; m < nr; ++m) {
1081            const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d));
1082            const __m512 vs1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].s));
1083            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1084
1085            __m512 vsum;
1086            if (is_acc) {
1087                vsum = _mm512_loadu_ps(C + m * ldc);
1088            } else {
1089                vsum = _mm512_set1_ps(0.f);
1090            }
1091            vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum);
1092            vsum = _mm512_fmadd_ps(vm0, vs1, vsum);
1093            _mm512_storeu_ps(C + m * ldc, vsum);
1094        }
1095    }
1096};
1097
1098template <bool is_acc>
1099struct acc_C<block_q8_0, block_q8_0, is_acc> {
1100    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) {
1101        const int offset = TILE_N * TILE_K;
1102        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset)));
1103
1104        for (int m = 0; m < nr; ++m) {
1105            const __m512 vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[m * lda].d));
1106            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1107
1108            __m512 vsum;
1109            if (is_acc) {
1110                vsum = _mm512_loadu_ps(C + m * ldc);
1111            } else {
1112                vsum = _mm512_set1_ps(0.f);
1113            }
1114            vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum);
1115            _mm512_storeu_ps(C + m * ldc, vsum);
1116        }
1117    }
1118};
1119
1120template <bool is_acc>
1121struct acc_C<block_q8_K, block_q4_K, is_acc> {
1122    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
1123        const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N);
1124        const uint8_t * mins = scales + 8 * TILE_N;
1125        const ggml_half * d0 = reinterpret_cast<const ggml_half *>(mins + 8 * TILE_N);
1126        const ggml_half * dmin = d0 + TILE_N;
1127
1128        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
1129        const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin));
1130
1131        for (int m = 0; m < nr; ++m) {
1132            const float d1 = A[m * lda].d;
1133            const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
1134            const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin);
1135            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1136
1137            __m512 vsum;
1138            if (is_acc) {
1139                vsum = _mm512_loadu_ps(C + m * ldc);
1140            } else {
1141                vsum = _mm512_set1_ps(0.f);
1142            }
1143
1144            const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums);
1145            const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
1146
1147            __m512i acc_m = _mm512_setzero_si512();
1148            for (int k = 0; k < 4; ++k) {
1149                __m512i vmask = _mm512_set1_epi32(k);
1150                __m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s));
1151                __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32)));
1152                acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
1153            }
1154
1155            vsum = _mm512_fmadd_ps(vtile, vd, vsum);
1156            vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum);
1157            _mm512_storeu_ps(C + m * ldc, vsum);
1158        }
1159    }
1160};
1161
1162template <bool is_acc>
1163struct acc_C<block_q8_K, block_q5_K, is_acc> {
1164    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
1165        const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N);
1166        const uint8_t * mins = scales + 8 * TILE_N;
1167        const ggml_half * d0 = reinterpret_cast<const ggml_half *>(mins + 8 * TILE_N);
1168        const ggml_half * dmin = d0 + TILE_N;
1169
1170        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
1171        const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin));
1172
1173        for (int m = 0; m < nr; ++m) {
1174            const float d1 = A[m * lda].d;
1175            const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
1176            const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin);
1177            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1178
1179            __m512 vsum;
1180            if (is_acc) {
1181                vsum = _mm512_loadu_ps(C + m * ldc);
1182            } else {
1183                vsum = _mm512_set1_ps(0.f);
1184            }
1185
1186            const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums);
1187            const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
1188
1189            __m512i acc_m = _mm512_setzero_si512();
1190            for (int k = 0; k < 4; ++k) {
1191                __m512i vmask = _mm512_set1_epi32(k);
1192                __m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s));
1193                __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32)));
1194                acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
1195            }
1196
1197            vsum = _mm512_fmadd_ps(vtile, vd, vsum);
1198            vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum);
1199            _mm512_storeu_ps(C + m * ldc, vsum);
1200        }
1201    }
1202};
1203
1204template <bool is_acc>
1205struct acc_C<block_q8_K, block_q6_K, is_acc> {
1206    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
1207        const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N);
1208        const ggml_half * d0 = reinterpret_cast<const ggml_half *>(scales + 16 * TILE_N);
1209
1210        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
1211
1212        for (int m = 0; m < nr; ++m) {
1213            const float d1 = A[m * lda].d;
1214            const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
1215            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1216
1217            __m512 vsum;
1218            if (is_acc) {
1219                vsum = _mm512_loadu_ps(C + m * ldc);
1220            } else {
1221                vsum = _mm512_set1_ps(0.f);
1222            }
1223
1224            vsum = _mm512_fmadd_ps(vtile, vd, vsum);
1225            _mm512_storeu_ps(C + m * ldc, vsum);
1226        }
1227    }
1228};
1229
1230template <bool is_acc>
1231struct acc_C<block_q8_K, block_iq4_xs, is_acc> {
1232    static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
1233        const int8_t * scales = reinterpret_cast<const int8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N);
1234        const ggml_half * d0 = reinterpret_cast<const ggml_half *>(scales + 8 * TILE_N);
1235
1236        const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
1237
1238        for (int m = 0; m < nr; ++m) {
1239            const float d1 = A[m * lda].d;
1240            const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
1241            const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
1242
1243            __m512 vsum;
1244            if (is_acc) {
1245                vsum = _mm512_loadu_ps(C + m * ldc);
1246            } else {
1247                vsum = _mm512_set1_ps(0.f);
1248            }
1249
1250            vsum = _mm512_fmadd_ps(vtile, vd, vsum);
1251            _mm512_storeu_ps(C + m * ldc, vsum);
1252        }
1253    }
1254};
1255
1256template <typename TB> constexpr int get_quants_size();
1257template <> constexpr int get_quants_size<block_q4_K>() { return (QK_K / 2) * TILE_N; }
1258template <> constexpr int get_quants_size<block_q5_K>() { return (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N; }
1259template <> constexpr int get_quants_size<block_q6_K>() { return (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N; }
1260template <> constexpr int get_quants_size<block_iq4_xs>() { return (QK_K / 2) * TILE_N; }
1261
1262// used for QKK format
1263template <typename TB, bool is_acc,
1264          typename std::enable_if<is_type_qkk<TB>::value, int>::type = 0>
1265inline void scale_C(const int32_t * RESTRICT tile, int32_t * RESTRICT sumi, const void * packed_B, int k, int nr) {
1266    const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + get_quants_size<TB>());
1267    const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(scales + k * TILE_N)));
1268
1269    for (int m = 0; m < nr; ++m) {
1270        __m512i vsumi;
1271        if (is_acc) {
1272            vsumi = _mm512_loadu_si512(sumi + m * TILE_N);
1273        } else {
1274            vsumi = _mm512_setzero_si512();
1275        }
1276        __m512i vtile = _mm512_loadu_si512(tile + m * TILE_N);
1277        vsumi = _mm512_add_epi32(vsumi, _mm512_mullo_epi32(vtile, vscale));
1278        _mm512_storeu_si512((__m512i *)(sumi + m * TILE_N), vsumi);
1279    }
1280}
1281
1282template <typename TA, typename TB, typename TC, int BLOCK_M, int BLOCK_N, int BLOCK_K>
1283struct tinygemm_kernel_avx {
1284    static void apply(int K, const TA * RESTRICT A, const TB * RESTRICT B, TC * RESTRICT C, int ldc) {
1285        GGML_UNUSED(K);
1286        GGML_UNUSED(A);
1287        GGML_UNUSED(B);
1288        GGML_UNUSED(C);
1289        GGML_UNUSED(ldc);
1290    }
1291};
1292
1293template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1294struct tinygemm_kernel_avx<float, ggml_fp16_t, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1295    static void apply(int K, const float * RESTRICT A, const ggml_fp16_t * RESTRICT B, float * RESTRICT C, int ldc) {
1296        constexpr int ROWS = BLOCK_M;
1297        constexpr int COLS = BLOCK_N;
1298        assert(BLOCK_K == 16);
1299
1300        __m512 va;
1301        __m512 vb[COLS];
1302        __m512 vc[ROWS * COLS];
1303
1304        auto loadc = [&](auto idx) {
1305            vc[idx] = _mm512_setzero_ps();
1306        };
1307        Unroll<ROWS * COLS>{}(loadc);
1308
1309        auto compute = [&](auto idx, auto k) {
1310            constexpr int row = idx / COLS;
1311            constexpr int col = idx % COLS;
1312
1313            if constexpr (col == 0) {
1314                va = _mm512_loadu_ps(A + row * K + k);
1315            }
1316            if constexpr (row == 0) {
1317                vb[col] =  _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(B + col * K + k)));
1318            }
1319            vc[idx] = _mm512_fmadd_ps(va, vb[col], vc[idx]);
1320        };
1321
1322        for (int k = 0; k < K; k += 16) {
1323            Unroll<ROWS * COLS>{}(compute, k);
1324        }
1325
1326        auto storec = [&](auto idx) {
1327            constexpr int row = idx / COLS;
1328            constexpr int col = idx % COLS;
1329            C[row * ldc + col] = _mm512_reduce_add_ps(vc[idx]);
1330        };
1331        Unroll<ROWS * COLS>{}(storec);
1332    }
1333};
1334
1335#define LAUNCH_TINYGEMM_KERNEL_AVX(MB_SIZE, NB_SIZE)                                \
1336    tinygemm_kernel_avx<float, type, float, MB_SIZE, NB_SIZE, blck_size>::apply(    \
1337        K, (const float *)src1->data + src1_offset + mb_start * K,                  \
1338        (const type *)src0->data + src0_offset + nb_start * K,                      \
1339        (float *)dst->data + dst_offset + mb_start * ldc + nb_start, ldc)
1340
1341
1342// re-organize in the format {NB, KB, TILE_SIZE}:
1343#define PACKED_INDEX(n, k, KB, tile_size) (n * KB + k) * tile_size
1344
1345template<typename TB, int BLOCK_K>
1346void convert_B_packed_format(void * RESTRICT packed_B, const TB * RESTRICT B, int N, int K) {
1347    const int NB = N / TILE_N;
1348    const int KB = K / BLOCK_K;
1349    const int TILE_SIZE = get_tile_size<TB>();
1350
1351    // parallel on NB should be enough
1352    parallel_for(NB, [&](int begin, int end) {
1353        for (int n = begin; n < end; ++n) {
1354            for (int k = 0; k < KB; ++k) {
1355                int n0 = n * TILE_N;
1356                pack_B((char *)packed_B + PACKED_INDEX(n, k, KB, TILE_SIZE), &B[n0 * KB + k], KB);
1357            }
1358        }
1359    });
1360}
1361
1362template <typename TA, typename TB, typename TC, int BLOCK_M, int BLOCK_N, int BLOCK_K>
1363struct tinygemm_kernel_vnni {};
1364
1365template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1366struct tinygemm_kernel_vnni<block_q8_0, block_q4_0, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1367    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1368
1369        constexpr int COLS = BLOCK_N / 16;
1370        const int TILE_SIZE = TILE_N * sizeof(block_q4_0);
1371
1372        const block_q8_0 * RESTRICT A = static_cast<const block_q8_0 *>(_A);
1373        const char * RESTRICT B = static_cast<const char *>(_B);
1374
1375        __m512i va[8];
1376        __m512 vc[COLS];
1377        __m512 vd1;
1378
1379        // sum of offsets, shared across COLS
1380        //
1381        // avx512-vnni does not have `_mm512_dpbssd_epi32`,
1382        // need to transform ss to us:
1383        //   a * (b - 8) is equivalent to b * a - 8 * a
1384        //   s    u   u                   u   s   u   s
1385        //
1386        __m512i vcomp;
1387
1388        const __m512i off = _mm512_set1_epi8(8);
1389        const __m512i lowMask = _mm512_set1_epi8(0xF);
1390
1391        auto loadc = [&](auto col) {
1392            vc[col] = _mm512_setzero_ps();
1393        };
1394        Unroll<COLS>{}(loadc);
1395
1396        auto compute = [&](auto col, auto i) {
1397            // load a and compute compensation
1398            if constexpr (col == 0) {
1399                const int32_t * a_ptr = reinterpret_cast<const int32_t *>(A[0 * KB + i].qs);
1400                vcomp = _mm512_setzero_si512();
1401                for (int k = 0; k < 8; ++k) {
1402                    va[k] = _mm512_set1_epi32(a_ptr[k]);
1403                    vcomp = _mm512_dpbusd_epi32(vcomp, off, va[k]);
1404                }
1405                vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d));
1406            }
1407
1408            // load b
1409            __m512i vsum = _mm512_setzero_si512();
1410            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1411            for (int k = 0; k < 8; k += 2) {
1412                __m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32));
1413                __m512i vb0 = _mm512_and_si512(bytes, lowMask);
1414                vsum = _mm512_dpbusd_epi32(vsum, vb0, va[k + 0]);
1415                __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1416                vsum = _mm512_dpbusd_epi32(vsum, vb1, va[k + 1]);
1417            }
1418            const int offset = TILE_N * TILE_K / 2;
1419            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset)));
1420            vsum = _mm512_sub_epi32(vsum, vcomp);
1421
1422            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]);
1423        };
1424
1425        for (int i = 0; i < KB; ++i) {
1426            Unroll<COLS>{}(compute, i);
1427        }
1428
1429        //store to C
1430        auto storec = [&](auto col) {
1431            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1432        };
1433        Unroll<COLS>{}(storec);
1434    }
1435};
1436
1437template <int BLOCK_N, int BLOCK_K>
1438struct tinygemm_kernel_vnni<block_q8_1, block_q4_1, float, 1, BLOCK_N, BLOCK_K> {
1439    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1440
1441        constexpr int COLS = BLOCK_N / 16;
1442        const int TILE_SIZE = TILE_N * sizeof(block_q4_1);
1443
1444        const block_q8_1 * RESTRICT A = static_cast<const block_q8_1 *>(_A);
1445        const char * RESTRICT B = static_cast<const char *>(_B);
1446
1447        __m512i va[8];
1448        __m512i vb[8];
1449        __m512 vc[COLS];
1450        __m512 vd1, vs1;
1451
1452        const __m512i lowMask = _mm512_set1_epi8(0xF);
1453
1454        auto loadc = [&](auto col) {
1455            vc[col] = _mm512_setzero_ps();
1456        };
1457        Unroll<COLS>{}(loadc);
1458
1459        auto compute = [&](auto col, auto i) {
1460            // load a
1461            if constexpr (col == 0) {
1462                const int32_t * a_ptr = reinterpret_cast<const int32_t *>(A[0 * KB + i].qs);
1463                for (int k = 0; k < 8; ++k) {
1464                    va[k] = _mm512_set1_epi32(a_ptr[k]);
1465                }
1466                vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d));
1467                vs1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].s));
1468            }
1469
1470            // load b
1471            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1472            for (int k = 0; k < 8; k += 2) {
1473                __m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32));
1474                vb[k + 0] = _mm512_and_si512(bytes, lowMask);
1475                vb[k + 1] = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1476            }
1477            const int offset = TILE_N * TILE_K / 2;
1478            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset)));
1479            const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset + TILE_N * sizeof(ggml_half))));
1480
1481            __m512i vsum = _mm512_setzero_si512();
1482            for (int k = 0; k < 8; ++k) {
1483                vsum = _mm512_dpbusd_epi32(vsum, vb[k], va[k]);
1484            }
1485
1486            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]);
1487            vc[col] = _mm512_fmadd_ps(vm0, vs1, vc[col]);
1488        };
1489
1490        for (int i = 0; i < KB; ++i) {
1491            Unroll<COLS>{}(compute, i);
1492        }
1493
1494        //store to C
1495        auto storec = [&](auto col) {
1496            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1497        };
1498        Unroll<COLS>{}(storec);
1499    }
1500};
1501
1502template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1503struct tinygemm_kernel_vnni<block_q8_0, block_q8_0, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1504    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1505
1506        constexpr int COLS = BLOCK_N / 16;
1507        const int TILE_SIZE = TILE_N * sizeof(block_q8_0) + TILE_N * sizeof(int32_t);
1508
1509        const block_q8_0 * RESTRICT A = static_cast<const block_q8_0 *>(_A);
1510        const char * RESTRICT B = static_cast<const char *>(_B);
1511
1512        __m512i va[8];
1513        __m512i vb[8];
1514        __m512 vc[COLS];
1515        __m512 vd1;
1516
1517        // Notes: s8s8 igemm compensation in avx512-vnni
1518        // change s8s8 to u8s8 with compensate
1519        //   a * b = (a + 128) * b - 128 * b
1520        //   s   s       u       s    u    s
1521        //
1522        // (128 * b is pre-computed when packing B to vnni formats)
1523        //
1524        const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
1525
1526        auto loadc = [&](auto col) {
1527            vc[col] = _mm512_setzero_ps();
1528        };
1529        Unroll<COLS>{}(loadc);
1530
1531        auto compute = [&](auto col, auto i) {
1532            // load a and add offset 128
1533            if constexpr (col == 0) {
1534                const int32_t * a_ptr = reinterpret_cast<const int32_t *>(A[0 * KB + i].qs);
1535                for (int k = 0; k < 8; ++k) {
1536                    va[k] = _mm512_set1_epi32(a_ptr[k]);
1537                    va[k] = _mm512_add_epi8(va[k], off);
1538                }
1539                vd1 = _mm512_set1_ps(GGML_CPU_FP16_TO_FP32(A[0 * KB + i].d));
1540            }
1541
1542            // load b
1543            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1544            for (int k = 0; k < 8; ++k) {
1545                vb[k] = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 64));
1546            }
1547            const int offset = TILE_N * TILE_K;
1548            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset)));
1549            const int offset2 = TILE_N * TILE_K + TILE_N * sizeof(ggml_half);
1550            const __m512i vcomp = _mm512_loadu_si512((const __m512i *)(b_ptr + offset2));
1551
1552            __m512i vsum = _mm512_setzero_si512();
1553            for (int k = 0; k < 8; ++k) {
1554                vsum = _mm512_dpbusd_epi32(vsum, va[k], vb[k]);
1555            }
1556            vsum = _mm512_sub_epi32(vsum, vcomp);
1557
1558            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]);
1559        };
1560
1561        for (int i = 0; i < KB; ++i) {
1562            Unroll<COLS>{}(compute, i);
1563        }
1564
1565        //store to C
1566        auto storec = [&](auto col) {
1567            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1568        };
1569        Unroll<COLS>{}(storec);
1570    }
1571};
1572
1573template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1574struct tinygemm_kernel_vnni<block_q8_K, block_q4_K, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1575    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1576
1577        constexpr int COLS = BLOCK_N / 16;
1578        const int TILE_SIZE = TILE_N * sizeof(block_q4_K) + TILE_N * 4;
1579
1580        const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
1581        const char * RESTRICT B = static_cast<const char *>(_B);
1582
1583        // a.qs:   8 groups, 32 bytes each group (m256i)
1584        __m512i va[8];
1585        // a.bsum: 8 groups,  2 bytes each group (m128i)
1586        __m512i va_bsum;
1587        __m512 vc[COLS];
1588        __m512 vd1;
1589
1590        // packed_B:
1591        const int offset_scales = (QK_K / 2) * TILE_N;
1592        const int offset_mins   = (QK_K / 2) * TILE_N +  8 * TILE_N;
1593        const int offset_d0     = (QK_K / 2) * TILE_N + 16 * TILE_N;
1594        const int offset_dmin   = (QK_K / 2) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half);
1595
1596        const __m512i lowMask = _mm512_set1_epi8(0xF);
1597
1598        auto loadc = [&](auto col) {
1599            vc[col] = _mm512_setzero_ps();
1600        };
1601        Unroll<COLS>{}(loadc);
1602
1603        // Notes: vnni formats in QK_K
1604        //   a) quants vnni format
1605        //     int8  {k/4, n, 4}, viewed as 2d {k/4, 4n}, k = 32
1606        //     from {16, 32} to {8, 64}
1607        //
1608        //   b) min vnni format
1609        //     int16 {k/2, n, 2}, viewed as 2d {k/2, 2n}, k = 8
1610        //     from {16,  8} to {4, 32}
1611        //
1612        auto compute = [&](auto col, auto i) {
1613            // load a
1614            if constexpr (col == 0) {
1615                for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
1616                    va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32)));
1617                }
1618                const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
1619                const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
1620                va_bsum = _mm512_castsi128_si512(q8s);
1621                vd1 = _mm512_set1_ps(A[0 * KB + i].d);
1622            }
1623
1624            // step 1: accumultate the quants
1625            __m512i acc = _mm512_setzero_si512();
1626            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1627            const char * b_qs  = b_ptr;
1628            for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
1629                __m512i vsum = _mm512_setzero_si512();
1630                for (int k = 0; k < 8; k += 2) {
1631                    __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]);
1632                    __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]);
1633
1634                    __m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs);
1635                    __m512i vb0 = _mm512_and_si512(bytes, lowMask);
1636                    vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
1637                    __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1638                    vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
1639
1640                    b_qs += 64;
1641                }
1642                // vacc += scale * (q8 @ q4)
1643                const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
1644                acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
1645            }
1646            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
1647            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
1648
1649            // step 2: accumulate the mins
1650            __m512i acc_m = _mm512_setzero_si512();
1651            for (int k = 0; k < 4; ++k) {
1652                __m512i vmask = _mm512_set1_epi32(k);
1653                __m512i va = _mm512_permutexvar_epi32(vmask, va_bsum);
1654                __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32)));
1655                acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
1656            }
1657            const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin)));
1658            vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]);
1659        };
1660
1661        for (int i = 0; i < KB; ++i) {
1662            Unroll<COLS>{}(compute, i);
1663        }
1664
1665        //store to C
1666        auto storec = [&](auto col) {
1667            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1668        };
1669        Unroll<COLS>{}(storec);
1670    }
1671};
1672
1673template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1674struct tinygemm_kernel_vnni<block_q8_K, block_q5_K, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1675    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1676
1677        constexpr int COLS = BLOCK_N / 16;
1678        const int TILE_SIZE = TILE_N * sizeof(block_q5_K) + TILE_N * 4;
1679
1680        const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
1681        const char * RESTRICT B = static_cast<const char *>(_B);
1682
1683        // a.qs:   8 groups, 32 bytes each group (m256i)
1684        __m512i va[8];
1685        // a.bsum: 8 groups,  2 bytes each group (m128i)
1686        __m512i va_bsum;
1687        __m512 vc[COLS];
1688        __m512 vd1;
1689
1690        // packed_B:
1691        const int offset_qh     = (QK_K / 2) * TILE_N;
1692        const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N;
1693        const int offset_mins   = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N +  8 * TILE_N;
1694        const int offset_d0     = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N;
1695        const int offset_dmin   = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half);
1696
1697        const __m512i lowMask = _mm512_set1_epi8(0xF);
1698
1699        auto loadc = [&](auto col) {
1700            vc[col] = _mm512_setzero_ps();
1701        };
1702        Unroll<COLS>{}(loadc);
1703
1704        // Q5_K and Q4_K shares the same vnni formats, refer to notes above.
1705        auto compute = [&](auto col, auto i) {
1706            // load a
1707            if constexpr (col == 0) {
1708                for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
1709                    va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32)));
1710                }
1711                const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
1712                const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
1713                va_bsum = _mm512_castsi128_si512(q8s);
1714                vd1 = _mm512_set1_ps(A[0 * KB + i].d);
1715            }
1716
1717            // step 1: accumultate the quants
1718            __m512i acc = _mm512_setzero_si512();
1719            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1720            const char * b_qs  = b_ptr;
1721            const char * b_qh  = b_ptr + offset_qh;
1722            for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
1723                __m512i vsum = _mm512_setzero_si512();
1724                __m512i hmask0 = _mm512_set1_epi8(0x1);
1725                __m512i hmask1 = _mm512_set1_epi8(0x2);
1726                __m512i hbits = _mm512_loadu_si512((const __m512i *)(b_qh + k_group * 64));
1727                for (int k = 0; k < 8; k += 2) {
1728                    __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]);
1729                    __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]);
1730
1731                    __m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs);
1732                    __m512i vb0 = _mm512_and_si512(bytes, lowMask);
1733                    __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1734
1735                    __m512i vh0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), k), 4);
1736                    __m512i vh1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), k + 1), 4);
1737
1738                    hmask0 = _mm512_slli_epi16(hmask0, 2);
1739                    hmask1 = _mm512_slli_epi16(hmask1, 2);
1740                    vb0 = _mm512_add_epi8(vb0, vh0);
1741                    vb1 = _mm512_add_epi8(vb1, vh1);
1742
1743                    vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
1744                    vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
1745
1746                    b_qs += 64;
1747                }
1748                // vacc += scale * (q8 @ q5)
1749                const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
1750                acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
1751            }
1752            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
1753            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
1754
1755            // step 2: accumulate the mins
1756            __m512i acc_m = _mm512_setzero_si512();
1757            for (int k = 0; k < 4; ++k) {
1758                __m512i vmask = _mm512_set1_epi32(k);
1759                __m512i va = _mm512_permutexvar_epi32(vmask, va_bsum);
1760                __m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32)));
1761                acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
1762            }
1763            const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin)));
1764            vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]);
1765        };
1766
1767        for (int i = 0; i < KB; ++i) {
1768            Unroll<COLS>{}(compute, i);
1769        }
1770
1771        //store to C
1772        auto storec = [&](auto col) {
1773            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1774        };
1775        Unroll<COLS>{}(storec);
1776    }
1777};
1778
1779template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1780struct tinygemm_kernel_vnni<block_q8_K, block_q6_K, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1781    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1782
1783        constexpr int COLS = BLOCK_N / 16;
1784        const int TILE_SIZE = TILE_N * sizeof(block_q6_K);
1785
1786        const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
1787        const char * RESTRICT B = static_cast<const char *>(_B);
1788
1789        // load the 256 bytes from A to 4 avx512 vectors
1790        __m512i va[4];
1791        __m512 vc[COLS];
1792        __m512 vd1;
1793
1794        // packed_B:
1795        const int offset_qh     = (QK_K / 2) * TILE_N;
1796        const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N;
1797        const int offset_d0     = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N + 16 * TILE_N;
1798
1799        // compensation
1800        __m512i vcomp;
1801
1802        const __m512i m32s = _mm512_set1_epi32(32);
1803        const __m512i lowMask = _mm512_set1_epi8(0xF);
1804
1805        auto loadc = [&](auto col) {
1806            vc[col] = _mm512_setzero_ps();
1807        };
1808        Unroll<COLS>{}(loadc);
1809
1810        auto compute = [&](auto col, auto i) {
1811            if constexpr (col == 0) {
1812                // load a
1813                va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs +   0));
1814                va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs +  64));
1815                va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128));
1816                va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192));
1817
1818                const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
1819                vcomp = _mm512_mullo_epi32(_mm512_cvtepi16_epi32(q8sums), m32s);
1820                vd1 = _mm512_set1_ps(A[0 * KB + i].d);
1821            }
1822
1823            // accmulate the quants
1824            __m512i acc = _mm512_setzero_si512();
1825            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1826            const char * b_qs = b_ptr;
1827            const char * b_qh = b_ptr + offset_qh;
1828            int mask = 0;
1829            for (int k_group = 0; k_group < QK_K / 16; ++k_group) {
1830                int r = k_group >> 2;
1831                __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
1832                __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
1833
1834                __m512i vsum = _mm512_setzero_si512();
1835                __m512i hmask = _mm512_set1_epi8(0x3);
1836
1837                __m512i bytes = _mm512_loadu_si512(b_qs);
1838                __m512i hbits = _mm512_loadu_si512(b_qh);
1839                __m512i vb0 = _mm512_and_si512(bytes, lowMask);
1840                __m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1841                __m512i vh0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask), 4);
1842                __m512i vh1 = _mm512_slli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 2)), 2);
1843
1844                vb0 = _mm512_add_epi8(vb0, vh0);
1845                vb1 = _mm512_add_epi8(vb1, vh1);
1846                vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
1847                vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
1848                b_qs += 64;
1849
1850                va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
1851                va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
1852
1853                bytes = _mm512_loadu_si512(b_qs);
1854                vb0 = _mm512_and_si512(bytes, lowMask);
1855                vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
1856                vh0 =                   _mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 4));
1857                vh1 = _mm512_srli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 6)), 2);
1858                vb0 = _mm512_add_epi8(vb0, vh0);
1859                vb1 = _mm512_add_epi8(vb1, vh1);
1860                vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
1861                vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
1862                b_qs += 64;
1863                b_qh += 64;
1864
1865                // B * A - 32 * A
1866                __m512i vmask = _mm512_set1_epi32(k_group);
1867                vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp));
1868
1869                // vacc += scale * (q8 @ q6)
1870                const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
1871                acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
1872            }
1873            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
1874            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
1875        };
1876
1877        for (int i = 0; i < KB; ++i) {
1878            Unroll<COLS>{}(compute, i);
1879        }
1880
1881        //store to C
1882        auto storec = [&](int col) {
1883            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1884        };
1885        Unroll<COLS>{}(storec);
1886    }
1887};
1888
1889template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
1890struct tinygemm_kernel_vnni<block_q8_K, block_iq4_xs, float, BLOCK_M, BLOCK_N, BLOCK_K> {
1891    static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
1892
1893        constexpr int COLS = BLOCK_N / 16;
1894        const int TILE_SIZE = TILE_N * sizeof(block_iq4_xs) + TILE_N * 2;
1895
1896        const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
1897        const char * RESTRICT B = static_cast<const char *>(_B);
1898
1899        // load the 256 bytes from A to 4 avx512 vectors
1900        __m512i va[4];
1901        __m512 vc[COLS];
1902        __m512 vd1;
1903
1904        // packed_B:
1905        const int offset_scales = (QK_K / 2) * TILE_N ;
1906        const int offset_d0     = (QK_K / 2) * TILE_N + 8 * TILE_N;
1907
1908        // compensation
1909        __m512i vcomp;
1910
1911        const __m256i m128s = _mm256_set1_epi16(128);
1912        const __m512i lowMask = _mm512_set1_epi8(0xF);
1913
1914        const __m512i values128 = _mm512_set_epi8(
1915            113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
1916            113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
1917            113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
1918            113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127
1919        );
1920        const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
1921        const __m512i values256 = _mm512_add_epi8(values128, off);
1922
1923        auto loadc = [&](auto col) {
1924            vc[col] = _mm512_setzero_ps();
1925        };
1926        Unroll<COLS>{}(loadc);
1927
1928        auto compute = [&](auto col, auto i) {
1929            if constexpr (col == 0) {
1930                // load a
1931                va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs +   0));
1932                va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs +  64));
1933                va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128));
1934                va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192));
1935
1936                // compensation: 128 * A
1937                const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
1938                vcomp = _mm512_castsi256_si512(_mm256_madd_epi16(q8sums, m128s));
1939                vd1 = _mm512_set1_ps(A[0 * KB + i].d);
1940            }
1941
1942            // accmulate the quants
1943            __m512i acc = _mm512_setzero_si512();
1944            const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
1945            const char * b_qs = b_ptr;
1946            int mask = 0;
1947            for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
1948                int r = k_group >> 1;
1949                __m512i vmask = _mm512_set1_epi32(k_group);
1950                __m512i vsum = _mm512_setzero_si512();
1951                for (int k = 0; k < 8; k += 2) {
1952                    __m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
1953                    __m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
1954
1955                    __m512i bytes = _mm512_loadu_si512(b_qs);
1956                    __m512i vb0 = _mm512_shuffle_epi8(values256, _mm512_and_si512(bytes, lowMask));
1957                    __m512i vb1 = _mm512_shuffle_epi8(values256, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask));
1958
1959                    vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
1960                    vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
1961                    b_qs += 64;
1962                }
1963                // (B + 128) * A - 128 * A
1964                vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp));
1965
1966                // vacc += scale * (q8 @ q4)
1967                const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
1968                acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
1969            }
1970            const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
1971            vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
1972        };
1973
1974        for (int i = 0; i < KB; ++i) {
1975            Unroll<COLS>{}(compute, i);
1976        }
1977
1978        //store to C
1979        auto storec = [&](auto col) {
1980            _mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
1981        };
1982        Unroll<COLS>{}(storec);
1983    }
1984};
1985
1986#define LAUNCH_TINYGEMM_KERNEL_VNNI(NB_SIZE)                                                   \
1987    tinygemm_kernel_vnni<vec_dot_type, type, float, 1, NB_SIZE, blck_size>::apply(             \
1988        KB, wdata_batch,                                                                       \
1989        (const char *)src0->data + src0_offset + PACKED_INDEX(nb * kTilesN, 0, KB, TILE_SIZE), \
1990        (float *) dst->data + dst_offset + nb_start, ldc)
1991
1992template <typename TA, typename TB, typename TC, int BLOCK_K,
1993          typename std::enable_if<!is_type_qkk<TB>::value, int>::type = 0>
1994void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, TC * RESTRICT C, int ldc) {
1995    using packed_B_t = packed_B_type<TB>;
1996    const int TILE_SIZE = get_tile_size<TB>();
1997    const bool need_unpack = do_unpack<TB>::value;
1998
1999    GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N);
2000    const TA * RESTRICT A = static_cast<const TA *>(_A);
2001    const char * RESTRICT B = static_cast<const char *>(_B);
2002
2003    const int m0 = std::min(M, TILE_M);
2004    const int m1 = std::max(M - TILE_M, 0);
2005    const int lda = KB * sizeof(TA);
2006    //const int ldb = KB * sizeof(TB);
2007
2008    alignas(64) static thread_local packed_B_t Tile0[TILE_N * TILE_K];
2009    alignas(64) static thread_local packed_B_t Tile1[TILE_N * TILE_K];
2010    alignas(64) static thread_local int8_t Tile23[TILE_M * TILE_K];
2011
2012    alignas(64) static thread_local int32_t TileC0[TILE_M * TILE_N * 4];
2013    alignas(64) static thread_local int32_t TileC1[TILE_M * TILE_N * 4];
2014
2015    // double buffering C to interleave avx512 and amx
2016    int32_t * C_cur = TileC0;
2017    int32_t * C_pre = TileC1;
2018
2019    auto Tile4 = [&](int32_t * base) { return base; };
2020    auto Tile5 = [&](int32_t * base) { return base + TILE_M * TILE_N; };
2021    auto Tile6 = [&](int32_t * base) { return base + 2 * TILE_M * TILE_N; };
2022    auto Tile7 = [&](int32_t * base) { return base + 3 * TILE_M * TILE_N; };
2023
2024    if (M == 2 * TILE_M) {
2025        // i = 0
2026        const char * B_blk0 = B + PACKED_INDEX(0, 0, KB, TILE_SIZE);
2027        const char * B_blk1 = B + PACKED_INDEX(1, 0, KB, TILE_SIZE);
2028        if (need_unpack) {
2029            unpack_B<TB>(Tile0, B_blk0);
2030            _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
2031        } else {
2032            _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK);
2033        }
2034
2035        _tile_zero(TMM4);
2036        _tile_loadd(TMM2, A[0].qs, lda);
2037        _tile_dpbssd(TMM4, TMM2, TMM0);
2038        _tile_stored(TMM4, Tile4(C_pre), TILE_N * sizeof(int32_t));
2039
2040        _tile_zero(TMM5);
2041        _tile_loadd(TMM3, A[TILE_M * KB + 0].qs, lda);
2042        _tile_dpbssd(TMM5, TMM3, TMM0);
2043        _tile_stored(TMM5, Tile5(C_pre), TILE_N * sizeof(int32_t));
2044
2045        if (need_unpack) {
2046            unpack_B<TB>(Tile1, B_blk1);
2047            _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
2048        } else {
2049            _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK);
2050        }
2051
2052        _tile_zero(TMM6);
2053        _tile_dpbssd(TMM6, TMM2, TMM1);
2054        _tile_stored(TMM6, Tile6(C_pre), TILE_N * sizeof(int32_t));
2055
2056        _tile_zero(TMM7);
2057        _tile_dpbssd(TMM7, TMM3, TMM1);
2058        _tile_stored(TMM7, Tile7(C_pre), TILE_N * sizeof(int32_t));
2059
2060        for (int i = 1; i < KB; ++i) {
2061            // index of previous iter
2062            const int ii = i - 1;
2063            const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE);
2064            const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE);
2065            GGML_DISPATCH_BOOL(ii > 0, is_acc, [&] {
2066                if (need_unpack) {
2067                    unpack_B<TB>(Tile0, B_blk0);
2068                    _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
2069                } else {
2070                    _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK);
2071                }
2072                _tile_zero(TMM4);
2073                _tile_loadd(TMM2, A[i].qs, lda);
2074                acc_C<TA, TB, is_acc>::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
2075
2076                _tile_dpbssd(TMM4, TMM2, TMM0);
2077                _tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t));
2078
2079                _tile_zero(TMM5);
2080                _tile_loadd(TMM3, A[TILE_M * KB + i].qs, lda);
2081                acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
2082
2083                _tile_dpbssd(TMM5, TMM3, TMM0);
2084                _tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t));
2085
2086                if (need_unpack) {
2087                    unpack_B<TB>(Tile1, B_blk1);
2088                    _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
2089                } else {
2090                    _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK);
2091                }
2092                _tile_zero(TMM6);
2093                acc_C<TA, TB, is_acc>::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
2094
2095                _tile_dpbssd(TMM6, TMM2, TMM1);
2096                _tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t));
2097
2098                _tile_zero(TMM7);
2099                acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
2100
2101                _tile_dpbssd(TMM7, TMM3, TMM1);
2102                _tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t));
2103
2104                std::swap(C_cur, C_pre);
2105            });
2106        }
2107        // final accumulation
2108        {
2109            int ii = KB - 1;
2110            acc_C<TA, TB, true>::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
2111            acc_C<TA, TB, true>::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
2112            acc_C<TA, TB, true>::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
2113            acc_C<TA, TB, true>::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
2114        }
2115    } else {
2116        for (int i = 0; i < KB; ++i) {
2117            _tile_zero(TMM4);
2118            _tile_zero(TMM6);
2119            if (m1 != 0) {
2120                _tile_zero(TMM5);
2121                _tile_zero(TMM7);
2122            }
2123
2124            const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE);
2125            const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE);
2126            if (need_unpack) {
2127                unpack_B<TB>(Tile0, B_blk0);
2128                _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
2129            } else {
2130                _tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK);
2131            }
2132
2133            if (need_unpack) {
2134                unpack_B<TB>(Tile1, B_blk1);
2135                _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
2136            } else {
2137                _tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK);
2138            }
2139
2140            if (m0 == TILE_M) {
2141                _tile_loadd(TMM2, A[i].qs, lda);
2142            } else {
2143                unpack_A(Tile23, &A[i], KB, m0);
2144                _tile_loadd(TMM2, Tile23, TILE_K);
2145            }
2146
2147            _tile_dpbssd(TMM4, TMM2, TMM0);
2148            _tile_dpbssd(TMM6, TMM2, TMM1);
2149
2150            _tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t));
2151            _tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t));
2152
2153            GGML_DISPATCH_BOOL(i > 0, is_acc, [&] {
2154                acc_C<TA, TB, is_acc>::apply(C,          ldc, Tile4(C_cur), &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0);
2155                acc_C<TA, TB, is_acc>::apply(C + TILE_N, ldc, Tile6(C_cur), &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0);
2156            });
2157
2158            if (m1 != 0) {
2159                unpack_A(Tile23, &A[TILE_M * KB + i], KB, m1);
2160                _tile_loadd(TMM3, Tile23, TILE_K);
2161
2162                _tile_dpbssd(TMM5, TMM3, TMM0);
2163                _tile_dpbssd(TMM7, TMM3, TMM1);
2164                _tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t));
2165                _tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t));
2166                GGML_DISPATCH_BOOL(i > 0, is_acc, [&] {
2167                    acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc,          ldc, Tile5(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1);
2168                    acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1);
2169                });
2170            }
2171        }
2172    }
2173    return;
2174}
2175
2176template <typename TA, typename TB, typename TC, int BLOCK_K,
2177          typename std::enable_if<is_type_qkk<TB>::value, int>::type = 0>
2178void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
2179    static_assert(std::is_same<TA, block_q8_K>::value);
2180    const int TILE_SIZE = get_tile_size<TB>();
2181
2182    GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N);
2183    const TA * RESTRICT A = static_cast<const TA *>(_A);
2184    const char * RESTRICT B = static_cast<const char *>(_B);
2185
2186    const int m0 = std::min(M, TILE_M);
2187    const int m1 = std::max(M - TILE_M, 0);
2188    //const int lda = KB * sizeof(TA);
2189
2190    alignas(64) static thread_local int8_t Tile0[TILE_N * TILE_K];
2191    alignas(64) static thread_local int8_t Tile1[TILE_N * TILE_K];
2192    alignas(64) static thread_local int8_t Tile23[TILE_M * TILE_K];
2193
2194    // mat mul result for each group
2195    alignas(64) static thread_local int32_t Tile4[TILE_M * TILE_N];
2196    alignas(64) static thread_local int32_t Tile5[TILE_M * TILE_N];
2197    alignas(64) static thread_local int32_t Tile6[TILE_M * TILE_N];
2198    alignas(64) static thread_local int32_t Tile7[TILE_M * TILE_N];
2199
2200    // sum of each QK_K block, contains 8 groups, int32
2201    alignas(64) static thread_local int32_t Sumi4[TILE_M * TILE_N];
2202    alignas(64) static thread_local int32_t Sumi5[TILE_M * TILE_N];
2203    alignas(64) static thread_local int32_t Sumi6[TILE_M * TILE_N];
2204    alignas(64) static thread_local int32_t Sumi7[TILE_M * TILE_N];
2205
2206    const int k_group_size = std::is_same<TB, block_q6_K>::value ? 16 : 32;
2207    for (int i = 0; i < KB; ++i) {
2208        // step 1: accumulate the quants across 8 groups, each group with 32
2209        for (int k = 0; k < QK_K / k_group_size; ++k) {
2210            GGML_DISPATCH_BOOL(k > 0, is_acc, [&] {
2211                _tile_zero(TMM4);
2212                _tile_zero(TMM6);
2213
2214                unpack_B<TB>(Tile0, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k);
2215                _tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
2216
2217                unpack_B<TB>(Tile1, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k);
2218                _tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
2219
2220                unpack_A<TB>(Tile23, &A[i], KB, k, m0);
2221                _tile_loadd(TMM2, Tile23, TILE_K);
2222
2223                _tile_dpbssd(TMM4, TMM2, TMM0);
2224                _tile_dpbssd(TMM6, TMM2, TMM1);
2225
2226                _tile_stored(TMM4, Tile4, TILE_N * sizeof(int32_t));
2227                _tile_stored(TMM6, Tile6, TILE_N * sizeof(int32_t));
2228
2229                scale_C<TB, is_acc>(Tile4, Sumi4, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m0);
2230                scale_C<TB, is_acc>(Tile6, Sumi6, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m0);
2231
2232                if (m1 != 0) {
2233                    _tile_zero(TMM5);
2234                    _tile_zero(TMM7);
2235
2236                    unpack_A<TB>(Tile23, &A[TILE_M * KB + i], KB, k, m1);
2237                    _tile_loadd(TMM3, Tile23, TILE_K);
2238
2239                    _tile_dpbssd(TMM5, TMM3, TMM0);
2240                    _tile_dpbssd(TMM7, TMM3, TMM1);
2241
2242                    _tile_stored(TMM5, Tile5, TILE_N * sizeof(int32_t));
2243                    _tile_stored(TMM7, Tile7, TILE_N * sizeof(int32_t));
2244
2245                    scale_C<TB, is_acc>(Tile5, Sumi5, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m1);
2246                    scale_C<TB, is_acc>(Tile7, Sumi7, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m1);
2247                }
2248            });
2249        }
2250
2251        // step 2: accmulate the mins
2252        GGML_DISPATCH_BOOL(i > 0, is_acc, [&] {
2253            acc_C<TA, TB, is_acc>::apply(C,          ldc, Sumi4, &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0);
2254            acc_C<TA, TB, is_acc>::apply(C + TILE_N, ldc, Sumi6, &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0);
2255            if (m1 != 0) {
2256                acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc,          ldc, Sumi5, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1);
2257                acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc + TILE_N, ldc, Sumi7, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1);
2258            }
2259        });
2260    }
2261    return;
2262}
2263
2264} // anonymous namespace
2265
2266// get the packed tensor size for quantized weights
2267size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor) {
2268    const enum ggml_type TYPE = tensor->type;
2269
2270    const int K = tensor->ne[0]; // ne0: in_features
2271    const int N = tensor->ne[1]; // ne1: out_features
2272
2273    auto get_tensor_size = [&] {
2274        size_t row_size_B{0};
2275        GGML_DISPATCH_QTYPES(TYPE, [&] {
2276            row_size_B = get_row_size<type, blck_size>(K);
2277        });
2278        return N * row_size_B;
2279    };
2280
2281    if (qtype_has_amx_kernels(TYPE)) {
2282        return get_tensor_size();
2283    } else {
2284        // for f16, bf16 we don't do packing
2285        return ggml_nbytes(tensor);
2286    }
2287}
2288
2289// pack weight to vnni format
2290void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
2291    GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor)); // only full tensor conversion is supported for now
2292
2293    const enum ggml_type TYPE = tensor->type;
2294
2295    const int K = tensor->ne[0]; // ne0: in_features
2296    const int N = tensor->ne[1]; // ne1: out_features
2297
2298    GGML_DISPATCH_QTYPES(TYPE, [&] {
2299        convert_B_packed_format<type, blck_size>((void *)((char *)tensor->data + offset), (const type *)data, N, K);
2300    });
2301}
2302
2303// ne2 is passed explicitly to help compiler optimize repeated calls
2304inline int64_t ggml_batch_offset(const ggml_tensor * t, int64_t batch_idx, int64_t ne2) {
2305    const int64_t i2 = batch_idx % ne2;
2306    const int64_t i3 = batch_idx / ne2;
2307    return i3 * t->nb[3] + i2 * t->nb[2];
2308}
2309
2310size_t ggml_backend_amx_desired_wsize(const struct ggml_tensor * dst) {
2311    struct ggml_tensor * src0 = dst->src[0];
2312
2313    const enum ggml_type TYPE = src0->type;
2314
2315    const bool is_floating_type = TYPE == GGML_TYPE_F16;
2316    if (is_floating_type) {
2317        return 0;
2318    }
2319
2320    const int M = dst->ne[1];
2321    const int K = src0->ne[0];
2322    const int64_t n_batch = dst->ne[2] * dst->ne[3];
2323
2324    size_t desired_wsize = 0;
2325
2326    GGML_DISPATCH_QTYPES(TYPE, [&] {
2327        const size_t row_size_A = K / blck_size * sizeof(vec_dot_type);
2328        desired_wsize = n_batch * M * row_size_A;
2329    });
2330
2331    return desired_wsize;
2332}
2333
2334// NB: mixed dtype gemm with Advanced Matrix Extensions (Intel AMX)
2335//
2336// src0: weight in shape of {N, K}, quantized
2337// src1: input  in shape of {M, K}, float32
2338// dst:  output in shape of {M, N}, float32
2339//
2340// the function performs: dst = src1 @ src0.T for each batch
2341//
2342void ggml_backend_amx_mul_mat(const ggml_compute_params * params, struct ggml_tensor * dst) {
2343    struct ggml_tensor * src0 = dst->src[0];
2344    struct ggml_tensor * src1 = dst->src[1];
2345
2346    const enum ggml_type TYPE = src0->type;
2347
2348    // f16 only has avx512 kernels for now,
2349    // amx kernels will be added once 6th gen xeon is released.
2350    const bool is_floating_type = TYPE == GGML_TYPE_F16;
2351
2352    const int M = dst->ne[1];
2353    const int N = dst->ne[0];
2354    const int K = src0->ne[0];
2355    const int ldc = dst->nb[1] / dst->nb[0];
2356
2357    const int64_t ne2 = dst->ne[2];
2358    const int64_t n_batch = ne2 * dst->ne[3];
2359
2360    if (is_floating_type) {
2361        constexpr int BLOCK_M = 4;
2362        constexpr int BLOCK_N = 6;
2363        const int MB = div_up(M, BLOCK_M);
2364        const int NB = div_up(N, BLOCK_N);
2365
2366        parallel_for_ggml(params, n_batch * MB * NB, [&](int begin, int end) {
2367            GGML_DISPATCH_FLOATING_TYPES(TYPE, [&] {
2368                for (int i = begin; i < end; ++i) {
2369                    int batch_idx = i / (MB * NB);
2370                    int remaining = i % (MB * NB);
2371                    int mb = remaining / NB;
2372                    int nb = remaining % NB;
2373
2374                    int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2);
2375                    int64_t src1_offset = ggml_batch_offset(src1, batch_idx, ne2);
2376                    int64_t dst_offset  = ggml_batch_offset(dst,  batch_idx, ne2);
2377
2378                    int mb_start = mb * BLOCK_M;
2379                    int mb_size = std::min(BLOCK_M, M - mb_start);
2380                    int nb_start = nb * BLOCK_N;
2381                    int nb_size = std::min(BLOCK_N, N - nb_start);
2382
2383                    switch (mb_size << 4 | nb_size) {
2384                        case 0x12: LAUNCH_TINYGEMM_KERNEL_AVX(1, 2); break;
2385                        case 0x14: LAUNCH_TINYGEMM_KERNEL_AVX(1, 4); break;
2386                        case 0x16: LAUNCH_TINYGEMM_KERNEL_AVX(1, 6); break;
2387                        case 0x22: LAUNCH_TINYGEMM_KERNEL_AVX(2, 2); break;
2388                        case 0x24: LAUNCH_TINYGEMM_KERNEL_AVX(2, 4); break;
2389                        case 0x26: LAUNCH_TINYGEMM_KERNEL_AVX(2, 6); break;
2390                        case 0x32: LAUNCH_TINYGEMM_KERNEL_AVX(3, 2); break;
2391                        case 0x34: LAUNCH_TINYGEMM_KERNEL_AVX(3, 4); break;
2392                        case 0x36: LAUNCH_TINYGEMM_KERNEL_AVX(3, 6); break;
2393                        case 0x42: LAUNCH_TINYGEMM_KERNEL_AVX(4, 2); break;
2394                        case 0x44: LAUNCH_TINYGEMM_KERNEL_AVX(4, 4); break;
2395                        case 0x46: LAUNCH_TINYGEMM_KERNEL_AVX(4, 6); break;
2396                        default: fprintf(stderr, "Unexpected block size!\n");
2397                    }
2398                }
2399            });
2400        });
2401        return;
2402    }
2403
2404    // pointer to work space, used convert A from float to quantized type
2405    void * wdata = params->wdata;
2406
2407    //TODO: performance improvement: merge quant A
2408 // if (params->ith == 0) {
2409        GGML_DISPATCH_QTYPES(TYPE, [&] {
2410            const size_t row_size_A = K / blck_size * sizeof(vec_dot_type);
2411            const size_t desired_wsize = n_batch * M * row_size_A;
2412            if (params->wsize < desired_wsize) {
2413                GGML_ABORT("insufficient work space size");
2414            }
2415
2416            // Q4_0, Q4_1, Q8_0 handles 1 TILE_K per blck_size
2417            // Q4_K, Q5_K, Q6_K, IQ4_XS handles 8 TILE_K per blck_size
2418            GGML_ASSERT(TILE_K == blck_size || TILE_K * 8 == blck_size);
2419
2420            parallel_for_ggml(params, n_batch * M, [&](int begin, int end) {
2421                for (int idx = begin; idx < end; ++idx) {
2422                    int batch_idx = idx / M;
2423                    int m         = idx % M;
2424                    int64_t src1_offset = ggml_batch_offset(src1, batch_idx, ne2);
2425                    const float * A_data = (const float *)((const char *)src1->data + src1_offset);
2426                    char * wdata_batch = (char *)wdata + batch_idx * M * row_size_A;
2427                    from_float<vec_dot_type>(A_data + m * K, wdata_batch + m * row_size_A, K);
2428                }
2429            });
2430        });
2431 // }
2432
2433    ggml_barrier(params->threadpool);
2434
2435    if (M == 1) {
2436        // MB = 1 and handle 8 tiles in each block
2437        constexpr int kTilesN = 4;
2438        constexpr int BLOCK_N = TILE_N * kTilesN;
2439        const int NB = div_up(N, BLOCK_N);
2440
2441        parallel_for_ggml(params, n_batch * NB, [&](int begin, int end) {
2442            GGML_DISPATCH_QTYPES(TYPE, [&] {
2443                const int KB = K / blck_size;
2444                const int TILE_SIZE = get_tile_size<type>();
2445                const int row_size_A = KB * sizeof(vec_dot_type);
2446                for (int i = begin; i < end; ++i) {
2447                    int batch_idx = i / NB;
2448                    int nb = i % NB;
2449
2450                    int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2);
2451                    int64_t dst_offset  = ggml_batch_offset(dst,  batch_idx, ne2);
2452                    const char * wdata_batch = (const char *)wdata + batch_idx * row_size_A;
2453
2454                    int nb_start = nb * BLOCK_N;
2455                    int nb_size = std::min(BLOCK_N, N - nb_start); // 32, 64, 96
2456
2457                    switch (nb_size) {
2458                        //case 160: LAUNCH_TINYGEMM_KERNEL_VNNI(160); break;
2459                        case 128: LAUNCH_TINYGEMM_KERNEL_VNNI(128); break;
2460                        case 96: LAUNCH_TINYGEMM_KERNEL_VNNI(96); break;
2461                        case 64: LAUNCH_TINYGEMM_KERNEL_VNNI(64); break;
2462                        case 32: LAUNCH_TINYGEMM_KERNEL_VNNI(32); break;
2463                        default: fprintf(stderr, "Unexpected n block size!\n");
2464                    }
2465                }
2466            });
2467        });
2468        return;
2469    }
2470
2471    // handle 4 tiles at a tile
2472    constexpr int BLOCK_M = TILE_M * 2;
2473    constexpr int BLOCK_N = TILE_N * 2;
2474    const int MB = div_up(M, BLOCK_M);
2475    const int NB = div_up(N, BLOCK_N);
2476
2477    parallel_for_ggml(params, n_batch * MB * NB, [&](int begin, int end) {
2478        // init tile config for each thread
2479        ggml_tile_config_init();
2480
2481        GGML_DISPATCH_QTYPES(TYPE, [&] {
2482            const int KB = K / blck_size;
2483            const int TILE_SIZE = get_tile_size<type>();
2484            const int row_size_A = KB * sizeof(vec_dot_type);
2485
2486            for (int i = begin; i < end; ++i) {
2487                int batch_idx = i / (MB * NB);
2488                int remaining = i % (MB * NB);
2489                int mb = remaining / NB;
2490                int nb = remaining % NB;
2491
2492                int64_t src0_offset = ggml_batch_offset(src0, batch_idx, ne2);
2493                int64_t dst_offset  = ggml_batch_offset(dst,  batch_idx, ne2);
2494                const char * wdata_batch = (const char *)wdata + batch_idx * M * row_size_A;
2495
2496                int mb_start = mb * BLOCK_M;
2497                int mb_size = std::min(BLOCK_M, M - mb_start);
2498                int nb_start = nb * BLOCK_N;
2499                int nb_size = BLOCK_N;
2500
2501                tinygemm_kernel_amx<vec_dot_type, type, float, blck_size>(
2502                    mb_size, nb_size, KB,
2503                    wdata_batch + mb_start * row_size_A,
2504                    (const char *)src0->data + src0_offset + PACKED_INDEX(nb * 2, 0, KB, TILE_SIZE),
2505                    (float *) dst->data + dst_offset + mb_start * N + nb_start, ldc);
2506            }
2507        });
2508    });
2509}
2510
2511#endif // if defined(__AMX_INT8__) && defined(__AVX512VNNI__)