Commit 0efe8e4c6 for llama.cpp
commit 0efe8e4c6bb946fb71327b068b26c17d5d15121e
Author: Max Krasnyansky <maxk@qti.qualcomm.com>
Date: Wed Oct 21 00:48:39 2026 -0700
hexagon: neon-vectorized and threaded repack (tiling) (#30901)
* add vectorized q4_0 repack
* apply to more types
* hex-repack: neon-vectorize the rest of the types q1/q2/f16/...
* hex-repack: add simple chunking / threading
---------
Co-authored-by: hubezhou <hubezhou@qti.qualcomm.com>
diff --git a/ggml/src/ggml-hexagon/htp-repack.cpp b/ggml/src/ggml-hexagon/htp-repack.cpp
index 7ef2f6d0c..d2e3df2cf 100644
--- a/ggml/src/ggml-hexagon/htp-repack.cpp
+++ b/ggml/src/ggml-hexagon/htp-repack.cpp
@@ -9,6 +9,11 @@
#include <cmath>
#include <cstring>
#include <initializer_list>
+#include <thread>
+
+#if defined(__ARM_NEON) || defined(__aarch64__)
+# include <arm_neon.h>
+#endif
// --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) ---
@@ -43,8 +48,309 @@ static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const
}
}
+#if defined(__ARM_NEON) || defined(__aarch64__)
+static void repack_neon_transpose8x8(const uint8x8_t rows[8], uint8x8_t columns[8]) {
+ const uint8x8x2_t t0 = vtrn_u8(rows[0], rows[1]);
+ const uint8x8x2_t t1 = vtrn_u8(rows[2], rows[3]);
+ const uint8x8x2_t t2 = vtrn_u8(rows[4], rows[5]);
+ const uint8x8x2_t t3 = vtrn_u8(rows[6], rows[7]);
+ const uint16x4x2_t t4 = vtrn_u16(vreinterpret_u16_u8(t0.val[0]), vreinterpret_u16_u8(t1.val[0]));
+ const uint16x4x2_t t5 = vtrn_u16(vreinterpret_u16_u8(t0.val[1]), vreinterpret_u16_u8(t1.val[1]));
+ const uint16x4x2_t t6 = vtrn_u16(vreinterpret_u16_u8(t2.val[0]), vreinterpret_u16_u8(t3.val[0]));
+ const uint16x4x2_t t7 = vtrn_u16(vreinterpret_u16_u8(t2.val[1]), vreinterpret_u16_u8(t3.val[1]));
+ const uint32x2x2_t t8 = vtrn_u32(vreinterpret_u32_u16(t4.val[0]), vreinterpret_u32_u16(t6.val[0]));
+ const uint32x2x2_t t9 = vtrn_u32(vreinterpret_u32_u16(t4.val[1]), vreinterpret_u32_u16(t6.val[1]));
+ const uint32x2x2_t ta = vtrn_u32(vreinterpret_u32_u16(t5.val[0]), vreinterpret_u32_u16(t7.val[0]));
+ const uint32x2x2_t tb = vtrn_u32(vreinterpret_u32_u16(t5.val[1]), vreinterpret_u32_u16(t7.val[1]));
+ columns[0] = vreinterpret_u8_u32(t8.val[0]);
+ columns[1] = vreinterpret_u8_u32(ta.val[0]);
+ columns[2] = vreinterpret_u8_u32(t9.val[0]);
+ columns[3] = vreinterpret_u8_u32(tb.val[0]);
+ columns[4] = vreinterpret_u8_u32(t8.val[1]);
+ columns[5] = vreinterpret_u8_u32(ta.val[1]);
+ columns[6] = vreinterpret_u8_u32(t9.val[1]);
+ columns[7] = vreinterpret_u8_u32(tb.val[1]);
+}
+
+static uint8x8_t repack_neon_pack_4bit_pairs(uint8x8_t first, uint8x8_t second, bool high) {
+ const uint8x8x2_t separated = vuzp_u8(first, second);
+ const uint8x8_t mask = vdup_n_u8(0x0F);
+ uint8x8_t even = separated.val[0];
+ uint8x8_t odd = separated.val[1];
+
+ if (high) {
+ even = vshr_n_u8(even, 4);
+ odd = vshr_n_u8(odd, 4);
+ }
+
+ even = vand_u8(even, mask);
+ odd = vand_u8(odd, mask);
+ return vorr_u8(vshl_n_u8(odd, 4), even);
+}
+
+static void repack_neon_4bit_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+ uint8x8_t low_rows[8];
+ uint8x8_t high_rows[8];
+ uint8x8_t columns[8];
+
+ for (int rr = 0; rr < 8; ++rr) {
+ const uint8x16_t packed = vld1q_u8(qs[rr]);
+ const uint8x8_t first = vget_low_u8(packed);
+ const uint8x8_t second = vget_high_u8(packed);
+ low_rows[rr] = repack_neon_pack_4bit_pairs(first, second, false);
+ high_rows[rr] = repack_neon_pack_4bit_pairs(first, second, true);
+ }
+
+ repack_neon_transpose8x8(low_rows, columns);
+ for (int cp = 0; cp < 8; ++cp) {
+ vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+ }
+
+ repack_neon_transpose8x8(high_rows, columns);
+ for (int cp = 0; cp < 8; ++cp) {
+ vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+ }
+}
+
+static void repack_neon_k4_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0, bool high) {
+ uint8x8_t first_rows[8];
+ uint8x8_t second_rows[8];
+ uint8x8_t columns[8];
+
+ for (int rr = 0; rr < 8; ++rr) {
+ first_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr]), vld1_u8(qs[rr] + 8), high);
+ second_rows[rr] = repack_neon_pack_4bit_pairs(vld1_u8(qs[rr] + 16), vld1_u8(qs[rr] + 24), high);
+ }
+
+ repack_neon_transpose8x8(first_rows, columns);
+ for (int cp = 0; cp < 8; ++cp) {
+ vst1_u8(tile_dst + cp * 32 + row0, columns[cp]);
+ }
+
+ repack_neon_transpose8x8(second_rows, columns);
+ for (int cp = 0; cp < 8; ++cp) {
+ vst1_u8(tile_dst + (cp + 8) * 32 + row0, columns[cp]);
+ }
+}
+
+static void repack_neon_q8_rows(const uint8_t * qs[8], uint8_t * tile_dst, int row0) {
+ uint8x8_t rows[8];
+ uint8x8_t columns[8];
+
+ for (int cp = 0; cp < 16; cp += 4) {
+ for (int rr = 0; rr < 8; ++rr) {
+ rows[rr] = vld1_u8(qs[rr] + 2 * cp);
+ }
+ repack_neon_transpose8x8(rows, columns);
+
+ for (int pair = 0; pair < 4; ++pair) {
+ const uint8x8x2_t interleaved = vzip_u8(columns[2 * pair], columns[2 * pair + 1]);
+ uint8_t * dst = tile_dst + (cp + pair) * 64 + 2 * row0;
+ vst1_u8(dst, interleaved.val[0]);
+ vst1_u8(dst + 8, interleaved.val[1]);
+ }
+ }
+}
+
+static inline void repack_store_u32(uint8_t * dst, uint32_t v) {
+ memcpy(dst, &v, sizeof(v));
+}
+
+static inline uint32x4_t repack_neon_transpose_2bit_4x4(uint32x4_t w) {
+ uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 6)), vdupq_n_u32(0x00CC00CCu));
+ w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 6)));
+ t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 12)), vdupq_n_u32(0x0000F0F0u));
+ return veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 12)));
+}
+
+static inline void repack_neon_q2_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+ const uint32x4_t w = repack_neon_transpose_2bit_4x4(vreinterpretq_u32_u8(vld1q_u8(qs)));
+ repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+ repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+ repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+ repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_q1_0_block(const uint8_t * qs, uint8_t * tile, int row) {
+ uint32x4_t w = vreinterpretq_u32_u8(vld1q_u8(qs));
+ uint32x4_t t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 1)), vdupq_n_u32(0x22222222u));
+ w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 1)));
+ t = vandq_u32(veorq_u32(w, vshrq_n_u32(w, 3)), vdupq_n_u32(0x0A0A0A0Au));
+ w = veorq_u32(w, veorq_u32(t, vshlq_n_u32(t, 3)));
+ w = repack_neon_transpose_2bit_4x4(w);
+ repack_store_u32(tile + 0 * 128 + 4 * row, vgetq_lane_u32(w, 0));
+ repack_store_u32(tile + 1 * 128 + 4 * row, vgetq_lane_u32(w, 1));
+ repack_store_u32(tile + 2 * 128 + 4 * row, vgetq_lane_u32(w, 2));
+ repack_store_u32(tile + 3 * 128 + 4 * row, vgetq_lane_u32(w, 3));
+}
+
+static inline void repack_neon_pack_2bit_plane(uint8x16_t v0, uint8x16_t v1, uint8_t * dst, int row) {
+ const uint32x4x2_t z = vuzpq_u32(vreinterpretq_u32_u8(v0), vreinterpretq_u32_u8(v1));
+ const uint32x4_t p01 = vorrq_u32(z.val[0], vshlq_n_u32(z.val[1], 2));
+ const uint32x2x2_t u = vuzp_u32(vget_low_u32(p01), vget_high_u32(p01));
+ const uint32x2_t out = vorr_u32(u.val[0], vshl_n_u32(u.val[1], 4));
+ repack_store_u32(dst + 0 * 128 + row * 4, vget_lane_u32(out, 0));
+ repack_store_u32(dst + 1 * 128 + row * 4, vget_lane_u32(out, 1));
+}
+
+static inline void repack_neon_q2_3_k_low2(const uint8_t * qs, int kt_local, uint8_t * lo_pl, int row) {
+ const uint8_t * qs_sub = qs + (kt_local / 4) * 32;
+ const int8x16_t shift = vdupq_n_s8(-2 * (kt_local % 4));
+ const uint8x16_t mask = vdupq_n_u8(0x03);
+ const uint8x16_t v0 = vandq_u8(vshlq_u8(vld1q_u8(qs_sub + 0), shift), mask);
+ const uint8x16_t v1 = vandq_u8(vshlq_u8(vld1q_u8(qs_sub + 16), shift), mask);
+ repack_neon_pack_2bit_plane(v0, v1, lo_pl, row);
+}
+
+static inline void repack_neon_q3_k_hmask(const uint8_t * hmask, int kt_local, uint8_t * neg_pl, int row) {
+ static const uint32_t mask_words[4] = { 0x01010101u, 0x02020202u, 0x04040404u, 0x08080808u };
+ const uint8x16_t m0 = vreinterpretq_u8_u32(vld1q_u32(mask_words));
+ const uint8x16_t m1 = vshlq_n_u8(m0, 4);
+ const uint8x16_t hbit = vdupq_n_u8((uint8_t) (1u << kt_local));
+ const uint8x16_t b0 = vbicq_u8(m0, vtstq_u8(vld1q_u8(hmask + 0), hbit));
+ const uint8x16_t b1 = vbicq_u8(m1, vtstq_u8(vld1q_u8(hmask + 16), hbit));
+ const uint32x4_t b01 = vreinterpretq_u32_u8(vorrq_u8(b0, b1));
+ const uint32x2_t f0 = vorr_u32(vget_low_u32(b01), vget_high_u32(b01));
+ const uint32x2_t f1 = vpadd_u32(f0, f0);
+ repack_store_u32(neg_pl + row * 4, vget_lane_u32(f1, 0));
+}
+
+static inline void repack_neon_q6_k(const block_q6_K * b, int kt_local, uint8_t * lo_pl, uint8_t * hi_pl, int row) {
+ const int c = kt_local / 4;
+ const int g_sb = kt_local % 4;
+ const uint8_t * ql = b->ql + c * 64 + (g_sb & 1) * 32;
+ const uint8_t * qh = b->qh + c * 32;
+ const int8x16_t lshift = vdupq_n_s8((g_sb >= 2) ? -4 : 0);
+ const int8x16_t hshift = vdupq_n_s8(-2 * g_sb);
+ const uint8x16_t mask4 = vdupq_n_u8(0x0F);
+ const uint8x16_t mask2 = vdupq_n_u8(0x03);
+
+ const uint8x16_t lo0 = vandq_u8(vshlq_u8(vld1q_u8(ql + 0), lshift), mask4);
+ const uint8x16_t lo1 = vandq_u8(vshlq_u8(vld1q_u8(ql + 16), lshift), mask4);
+ const uint32x4x2_t lz = vuzpq_u32(vreinterpretq_u32_u8(lo0), vreinterpretq_u32_u8(lo1));
+ const uint32x4_t lout = vorrq_u32(lz.val[0], vshlq_n_u32(lz.val[1], 4));
+ repack_store_u32(lo_pl + 0 * 128 + row * 4, vgetq_lane_u32(lout, 0));
+ repack_store_u32(lo_pl + 1 * 128 + row * 4, vgetq_lane_u32(lout, 1));
+ repack_store_u32(lo_pl + 2 * 128 + row * 4, vgetq_lane_u32(lout, 2));
+ repack_store_u32(lo_pl + 3 * 128 + row * 4, vgetq_lane_u32(lout, 3));
+
+ const uint8x16_t hi0 = vandq_u8(vshlq_u8(vld1q_u8(qh + 0), hshift), mask2);
+ const uint8x16_t hi1 = vandq_u8(vshlq_u8(vld1q_u8(qh + 16), hshift), mask2);
+ repack_neon_pack_2bit_plane(hi0, hi1, hi_pl, row);
+}
+
+static inline void repack_neon_q5_k_qh(const uint8_t * qh, uint8_t hbit, uint8_t * plane, int row) {
+ static const uint8_t mask_bytes[16] = { 1, 2, 1, 2, 1, 2, 1, 2, 4, 8, 4, 8, 4, 8, 4, 8 };
+ const uint8x16_t m0 = vld1q_u8(mask_bytes);
+ const uint8x16_t m1 = vshlq_n_u8(m0, 4);
+ const uint8x16_t vbit = vdupq_n_u8(hbit);
+ const uint8x16_t b0 = vandq_u8(vtstq_u8(vld1q_u8(qh + 0), vbit), m0);
+ const uint8x16_t b1 = vandq_u8(vtstq_u8(vld1q_u8(qh + 16), vbit), m1);
+ const uint8x16_t b01 = vorrq_u8(b0, b1);
+ const uint8x8_t fold = vorr_u8(vget_low_u8(b01), vget_high_u8(b01));
+ const uint8x8_t out = vpadd_u8(fold, fold);
+ plane[0 * 32 + row] = vget_lane_u8(out, 0);
+ plane[1 * 32 + row] = vget_lane_u8(out, 1);
+ plane[2 * 32 + row] = vget_lane_u8(out, 2);
+ plane[3 * 32 + row] = vget_lane_u8(out, 3);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const uint16_t * src) {
+ return vld1q_u16(src);
+}
+
+#if defined(__aarch64__) || defined(__ARM_FP16_FORMAT_IEEE)
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+ const float16x4_t lo = vcvt_f16_f32(vld1q_f32(src));
+ const float16x4_t hi = vcvt_f16_f32(vld1q_f32(src + 4));
+ return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+ const uint16x8_t bf = vld1q_u16((const uint16_t *) src);
+ const float32x4_t f_lo = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(bf), 16));
+ const float32x4_t f_hi = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(bf), 16));
+ const float16x4_t lo = vcvt_f16_f32(f_lo);
+ const float16x4_t hi = vcvt_f16_f32(f_hi);
+ return vreinterpretq_u16_f16(vcombine_f16(lo, hi));
+}
+#else
+static inline uint16x8_t repack_neon_load_f16x8(const float * src) {
+ uint16_t tmp[8];
+ for (int i = 0; i < 8; ++i) {
+ tmp[i] = GGML_FP32_TO_FP16(src[i]);
+ }
+ return vld1q_u16(tmp);
+}
+
+static inline uint16x8_t repack_neon_load_f16x8(const ggml_bf16_t * src) {
+ uint16_t tmp[8];
+ for (int i = 0; i < 8; ++i) {
+ tmp[i] = GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(src[i]));
+ }
+ return vld1q_u16(tmp);
+}
+#endif
+
+template <typename SrcT>
+static void repack_neon_f16_tile(uint16_t * tile_dst, const SrcT * src_slice, int64_t row_stride, int ct, int kt) {
+ for (int row0 = 0; row0 < 32; row0 += 4) {
+ const SrcT * r0 = src_slice + (int64_t) (ct * 32 + row0 + 0) * row_stride + kt * 32;
+ const SrcT * r1 = src_slice + (int64_t) (ct * 32 + row0 + 1) * row_stride + kt * 32;
+ const SrcT * r2 = src_slice + (int64_t) (ct * 32 + row0 + 2) * row_stride + kt * 32;
+ const SrcT * r3 = src_slice + (int64_t) (ct * 32 + row0 + 3) * row_stride + kt * 32;
+ for (int cp = 0; cp < 16; cp += 4) {
+ const uint32x4_t v0 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r0 + 2 * cp));
+ const uint32x4_t v1 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r1 + 2 * cp));
+ const uint32x4_t v2 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r2 + 2 * cp));
+ const uint32x4_t v3 = vreinterpretq_u32_u16(repack_neon_load_f16x8(r3 + 2 * cp));
+
+ const uint32x2x2_t t0 = vtrn_u32(vget_low_u32(v0), vget_low_u32(v1));
+ const uint32x2x2_t t1 = vtrn_u32(vget_high_u32(v0), vget_high_u32(v1));
+ const uint32x2x2_t t2 = vtrn_u32(vget_low_u32(v2), vget_low_u32(v3));
+ const uint32x2x2_t t3 = vtrn_u32(vget_high_u32(v2), vget_high_u32(v3));
+
+ vst1q_u16(tile_dst + (cp + 0) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[0], t2.val[0])));
+ vst1q_u16(tile_dst + (cp + 1) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t0.val[1], t2.val[1])));
+ vst1q_u16(tile_dst + (cp + 2) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[0], t3.val[0])));
+ vst1q_u16(tile_dst + (cp + 3) * 64 + row0 * 2, vreinterpretq_u16_u32(vcombine_u32(t1.val[1], t3.val[1])));
+ }
+ }
+}
+#endif
+
// --- Generic 32x32 tile grid walkers for block-32 and float types ---
+template <typename Fn>
+static void repack_parallel_for_ct(size_t slice_bytes, int n_col_tiles, Fn && fn) {
+ int n_threads = 1;
+ if (slice_bytes > 128 * 1024) {
+ for (int nt : { 4, 3, 2 }) {
+ if (n_col_tiles % nt == 0) {
+ n_threads = nt;
+ break;
+ }
+ }
+ }
+
+ if (n_threads <= 1) {
+ fn(0, n_col_tiles);
+ return;
+ }
+
+ const int chunk = n_col_tiles / n_threads;
+ std::thread workers[3];
+ for (int i = 0; i < n_threads - 1; ++i) {
+ workers[i] = std::thread([&fn, i, chunk]() {
+ fn(i * chunk, (i + 1) * chunk);
+ });
+ }
+ fn((n_threads - 1) * chunk, n_col_tiles);
+ for (int i = 0; i < n_threads - 1; ++i) {
+ workers[i].join();
+ }
+}
+
template <typename SrcElemT, int ELEMS_PER_KTILE, typename PackTileFn>
static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t offset, size_t size,
size_t tile_size, PackTileFn && pack_tile) {
@@ -70,12 +376,14 @@ static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t o
const SrcElemT * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * row_stride);
uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
- for (int ct = 0; ct < n_col_tiles; ct++) {
- for (int kt = 0; kt < n_k_tiles; kt++) {
- uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
- pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+ repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+ for (int ct = ct0; ct < ct1; ct++) {
+ for (int kt = 0; kt < n_k_tiles; kt++) {
+ uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+ pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt);
+ }
}
- }
+ });
}
}
@@ -130,8 +438,24 @@ static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t
static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0,
[](uint8_t * tile_dst, const block_q4_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
- uint8_t tile_quants[32][32];
ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+ for (int row0 = 0; row0 < 32; row0 += 8) {
+ const uint8_t * qs[8];
+ for (int rr = 0; rr < 8; ++rr) {
+ qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+ }
+ repack_neon_4bit_rows(qs, tile_dst, row0);
+ }
+
+ for (int row = 0; row < 32; ++row) {
+ scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+ }
+ return;
+ }
+#endif
+ uint8_t tile_quants[32][32];
for (int row = 0; row < 32; row++) {
const int64_t r = ct * 32 + row;
if (r < ne1 && kt < ne0 / 32) {
@@ -167,8 +491,26 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset,
static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
repack_to_tiled_grid_32<block_q4_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_1,
[](uint8_t * tile_dst, const block_q4_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
- uint8_t tile_quants[32][32];
ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+ for (int row0 = 0; row0 < 32; row0 += 8) {
+ const uint8_t * qs[8];
+ for (int rr = 0; rr < 8; ++rr) {
+ qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+ }
+ repack_neon_4bit_rows(qs, tile_dst, row0);
+ }
+
+ for (int row = 0; row < 32; ++row) {
+ const block_q4_1 & src_block = src_slice[(ct * 32 + row) * row_stride + kt];
+ scale_dst[2 * row + 0] = src_block.d;
+ scale_dst[2 * row + 1] = src_block.m;
+ }
+ return;
+ }
+#endif
+ uint8_t tile_quants[32][32];
for (int row = 0; row < 32; row++) {
const int64_t r = ct * 32 + row;
if (r < ne1 && kt < ne0 / 32) {
@@ -207,8 +549,24 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset,
static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4,
[](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
- uint8_t tile_quants[32][32];
uint8_t * scale_dst = tile_dst + 512;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+ for (int row0 = 0; row0 < 32; row0 += 8) {
+ const uint8_t * qs[8];
+ for (int rr = 0; rr < 8; ++rr) {
+ qs[rr] = src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+ }
+ repack_neon_4bit_rows(qs, tile_dst, row0);
+ }
+
+ for (int row = 0; row < 32; ++row) {
+ scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].e;
+ }
+ return;
+ }
+#endif
+ uint8_t tile_quants[32][32];
for (int row = 0; row < 32; row++) {
const int64_t r = ct * 32 + row;
if (r < ne1 && kt < ne0 / 32) {
@@ -244,6 +602,23 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset
static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
repack_to_tiled_grid_32<block_q8_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q8_0,
[](uint8_t * tile_dst, const block_q8_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
+ ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ if (ct * 32 + 32 <= ne1 && kt < ne0 / 32) {
+ for (int row0 = 0; row0 < 32; row0 += 8) {
+ const uint8_t * qs[8];
+ for (int rr = 0; rr < 8; ++rr) {
+ qs[rr] = (const uint8_t *) src_slice[(ct * 32 + row0 + rr) * row_stride + kt].qs;
+ }
+ repack_neon_q8_rows(qs, tile_dst, row0);
+ }
+
+ for (int row = 0; row < 32; ++row) {
+ scale_dst[row] = src_slice[(ct * 32 + row) * row_stride + kt].d;
+ }
+ return;
+ }
+#endif
for (int cp = 0; cp < 16; cp++) {
const int col0 = cp * 2;
const int col1 = col0 + 1;
@@ -254,7 +629,6 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset,
tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0;
}
}
- ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024);
for (int row = 0; row < 32; row++) {
const int64_t r = ct * 32 + row;
scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * row_stride + kt].d : 0;
@@ -303,24 +677,33 @@ static void repack_q2_0_tiled(ggml_tensor * t, const void * data, size_t offset,
for (int64_t slice = 0; slice < t->ne[2] * t->ne[3]; ++slice) {
const block_q2_0 * src_slice = src + slice * t->ne[1] * blocks_per_row;
uint8_t * dst_slice = (uint8_t *) t->data + slice * matrix_size;
- memset(dst_slice, 0, matrix_size);
-
- for (int64_t row_index = 0; row_index < t->ne[1]; ++row_index) {
- const int64_t ct = row_index / 32;
- const int row = (int) (row_index % 32);
- for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
- const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
- uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
- for (int group = 0; group < 16; ++group) {
- for (int byte = 0; byte < 4; ++byte) {
- const int k = 4 * group + byte;
- const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
- tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+ const size_t ct_bytes = (size_t) n_weight_tiles * HTP_TENSOR_TILE_SIZE_Q2_0;
+
+ repack_parallel_for_ct(matrix_size, (int) n_col_tiles, [&](int ct0, int ct1) {
+ memset(dst_slice + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
+ const int64_t r0 = (int64_t) ct0 * 32;
+ const int64_t r1 = (std::min)((int64_t) ct1 * 32, t->ne[1]);
+ for (int64_t row_index = r0; row_index < r1; ++row_index) {
+ const int64_t ct = row_index / 32;
+ const int row = (int) (row_index % 32);
+ for (int64_t kt = 0; kt < n_weight_tiles; ++kt) {
+ const block_q2_0 * block = src_slice + row_index * blocks_per_row + kt;
+ uint8_t * tile = dst_slice + (ct * n_weight_tiles + kt) * HTP_TENSOR_TILE_SIZE_Q2_0;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ repack_neon_q2_0_block(block->qs, tile, row);
+#else
+ for (int group = 0; group < 16; ++group) {
+ for (int byte = 0; byte < 4; ++byte) {
+ const int k = 4 * group + byte;
+ const uint8_t code = (block->qs[k / 4] >> (2 * (k % 4))) & 3;
+ tile[(group >> 2) * 128 + 4 * row + byte] |= code << (2 * (group & 3));
+ }
}
+#endif
+ memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
}
- memcpy(tile + 512 + 2 * row, &block->d, sizeof(block->d));
}
- }
+ });
}
}
@@ -368,6 +751,12 @@ static void repack_float_to_f16_tiled(ggml_tensor * t, const void * data, size_t
repack_to_tiled_grid_32<SrcT, 32>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_F16,
[&](uint8_t * tile_bytes, const SrcT * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) {
uint16_t * tile_dst = (uint16_t *) tile_bytes;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ if (ct * 32 + 32 <= ne1 && kt * 32 + 32 <= ne0) {
+ repack_neon_f16_tile(tile_dst, src_slice, row_stride, ct, kt);
+ return;
+ }
+#endif
for (int cp = 0; cp < 16; cp++) {
const int64_t c0 = kt * 32 + cp * 2;
const int64_t c1 = c0 + 1;
@@ -424,20 +813,25 @@ static void repack_kquant_to_tiled_grid(ggml_tensor * t, const void * data, size
for (int64_t slice_idx = 0; slice_idx < ne2 * ne3; slice_idx++) {
const BlockT * src_slice = src_matrix + slice_idx * (ne1 * sb_per_row);
uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
-
- // Padding rows and OR-ed bit planes require zeroed tiles
- memset(matrix_dst, 0, matrix_size);
-
- for (int64_t r = 0; r < ne1; r++) {
- const int ct = (int) (r / 32);
- const int row = (int) (r % 32);
- const BlockT * src_row = src_slice + r * sb_per_row;
-
- for (int kt = 0; kt < n_k_tiles; kt++) {
- uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
- pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+ const size_t ct_bytes = (size_t) n_k_tiles * tile_size;
+
+ repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+ // Padding rows and OR-ed bit planes require zeroed tiles
+ memset(matrix_dst + (size_t) ct0 * ct_bytes, 0, (size_t) (ct1 - ct0) * ct_bytes);
+
+ const int64_t r0 = (int64_t) ct0 * 32;
+ const int64_t r1 = (std::min)((int64_t) ct1 * 32, ne1);
+ for (int64_t r = r0; r < r1; r++) {
+ const int ct = (int) (r / 32);
+ const int row = (int) (r % 32);
+ const BlockT * src_row = src_slice + r * sb_per_row;
+
+ for (int kt = 0; kt < n_k_tiles; kt++) {
+ uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+ pack_kt(tile, &src_row[kt / 8], row, kt % 8);
+ }
}
- }
+ });
}
}
@@ -493,6 +887,7 @@ static void unpack_kquant_from_tiled_grid(void * data, const ggml_tensor * t, si
// --- K-Quant Helpers & Per-Type Implementations (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) ---
+#if !defined(__ARM_NEON) && !defined(__aarch64__)
// unsigned 6-bit value (0..63) of element e of a Q6_K block, same bit layout as dequantize_row_q6_K
static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
const int c = e / 128;
@@ -510,6 +905,7 @@ static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) {
}
return (uint8_t) (lo | (hi << 4));
}
+#endif
static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
repack_kquant_to_tiled_grid<block_q6_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q6_K,
@@ -519,6 +915,9 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
uint8_t * hi_pl = tile + 512;
ggml_half * sc_pl = (ggml_half *) (tile + 768);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ repack_neon_q6_k(b, kt_local, lo_pl, hi_pl, row);
+#else
for (int lk = 0; lk < 32; lk++) {
const uint8_t q6 = q6_K_get_quant(b, kt_local * 32 + lk);
const int g = lk >> 2;
@@ -526,6 +925,7 @@ static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset,
lo_pl[(g >> 1) * 128 + pos] |= (uint8_t) ((q6 & 0xF) << ((g & 1) * 4));
hi_pl[(g >> 2) * 128 + pos] |= (uint8_t) ((q6 >> 4) << ((g & 3) * 2));
}
+#endif
for (int sub = 0; sub < 2; sub++) {
sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) b->scales[kt_local * 2 + sub]);
}
@@ -585,6 +985,7 @@ static void repack_tiled_q6_K(void * data, const ggml_tensor * t, size_t offset,
});
}
+#if !defined(__ARM_NEON) && !defined(__aarch64__)
// low 2 bits (0..3) of element e of a Q2_K or Q3_K block, same bit layout as dequantize_row_q2_K / q3_K
static inline uint8_t q2_3_K_get_low2(const uint8_t * qs, int e) {
const int c = e / 128;
@@ -600,6 +1001,7 @@ static inline bool q3_K_get_hbit(const block_q3_K * b, int e) {
const int l = e % 32;
return (b->hmask[l] >> (c * 4 + j)) & 1;
}
+#endif
// signed 6-bit scale j (-32..31) of a Q3_K block, same packing as quantize_row_q3_K_ref
static inline int q3_K_get_scale(const uint8_t * scales, int j) {
@@ -670,6 +1072,10 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
uint8_t * neg_pl = tile + 256;
ggml_half * sc_pl = (ggml_half *) (tile + 384);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+ repack_neon_q3_k_hmask(b->hmask, kt_local, neg_pl, row);
+#else
for (int lk = 0; lk < 32; lk++) {
const int e = kt_local * 32 + lk;
const int g = lk >> 2;
@@ -679,6 +1085,7 @@ static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset,
neg_pl[pos] |= (uint8_t) (1 << g);
}
}
+#endif
for (int sub = 0; sub < 2; sub++) {
sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) q3_K_get_scale(b->scales, kt_local * 2 + sub));
}
@@ -750,11 +1157,15 @@ static void repack_q2_K_tiled(ggml_tensor * t, const void * data, size_t offset,
uint8_t * sc_pl = tile + 256;
ggml_half * dd_pl = (ggml_half *) (tile + 320);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ repack_neon_q2_3_k_low2(b->qs, kt_local, lo_pl, row);
+#else
for (int lk = 0; lk < 32; lk++) {
const int g = lk >> 2;
const int pos = row * 4 + (lk & 3);
lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, kt_local * 32 + lk) << ((g & 3) * 2));
}
+#endif
for (int sub = 0; sub < 2; sub++) {
sc_pl[sub * 32 + row] = b->scales[kt_local * 2 + sub];
}
@@ -846,23 +1257,46 @@ static inline void pack_k4_scales_mins(uint8_t * scales, ggml_half * d, ggml_hal
}
}
+template <typename BlockT>
+static inline void repack_k4_quants(uint8_t * tile_dst, const BlockT * b, int row, int kt_local,
+ int64_t sb_per_row, bool use_neon) {
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ if (use_neon) {
+ if ((row & 7) == 0) {
+ const uint8_t * qs[8];
+ for (int rr = 0; rr < 8; ++rr) {
+ qs[rr] = b[rr * sb_per_row].qs + (kt_local / 2) * 32;
+ }
+ repack_neon_k4_rows(qs, tile_dst, row, (kt_local & 1) != 0);
+ }
+ return;
+ }
+#else
+ GGML_UNUSED(sb_per_row);
+ GGML_UNUSED(use_neon);
+#endif
+ const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
+ const int shift = (kt_local & 1) ? 4 : 0;
+
+ for (int cp = 0; cp < 16; cp++) {
+ const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
+ const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
+ tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
+ }
+}
+
static void repack_q4_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+ const int64_t sb_per_row = t->ne[0] / QK_K;
+ const bool use_neon = (t->ne[1] % 8 == 0);
repack_kquant_to_tiled_grid<block_q4_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_K,
- [](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
+ [sb_per_row, use_neon](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) {
const float d = GGML_FP16_TO_FP32(b->d);
const float dmin = GGML_FP16_TO_FP32(b->dmin);
uint8_t sc, m;
get_scale_min_k4(kt_local, b->scales, &sc, &m);
- const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
- const int shift = (kt_local & 1) ? 4 : 0;
-
- for (int cp = 0; cp < 16; cp++) {
- const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
- const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
- tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
- }
+ repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -900,8 +1334,10 @@ static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset,
}
static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) {
+ const int64_t sb_per_row = t->ne[0] / QK_K;
+ const bool use_neon = (t->ne[1] % 8 == 0);
repack_kquant_to_tiled_grid<block_q5_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_K,
- [](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
+ [sb_per_row, use_neon](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) {
const float d = GGML_FP16_TO_FP32(b->d);
const float dmin = GGML_FP16_TO_FP32(b->dmin);
uint8_t * plane = tile_dst + 640;
@@ -909,15 +1345,13 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
uint8_t sc, m;
get_scale_min_k4(kt_local, b->scales, &sc, &m);
- const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32;
- const int shift = (kt_local & 1) ? 4 : 0;
- const uint8_t hbit = (uint8_t) (1 << kt_local);
+ repack_k4_quants(tile_dst, b, row, kt_local, sb_per_row, use_neon);
+ const uint8_t hbit = (uint8_t) (1 << kt_local);
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ repack_neon_q5_k_qh(b->qh, hbit, plane, row);
+#else
for (int cp = 0; cp < 16; cp++) {
- const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F;
- const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F;
- tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0);
-
const int i = cp / 4;
const int lane = (cp % 4) * 32 + row;
if (b->qh[2 * cp + 0] & hbit) {
@@ -927,6 +1361,7 @@ static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset,
plane[lane] |= (uint8_t) (1 << (2 * i + 1));
}
}
+#endif
ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc);
@@ -1002,31 +1437,37 @@ static void repack_q1_0_tiled(ggml_tensor * t, const void * data, size_t offset,
const block_q1_0 * src_slice = src_matrix + (slice_idx - start_slice) * ne1 * n_k_tiles;
uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size;
- for (int ct = 0; ct < n_col_tiles; ct++) {
- for (int kt = 0; kt < n_k_tiles; kt++) {
- uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
- ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
- memset(tile_dst, 0, tile_size);
-
- for (int row = 0; row < 32; row++) {
- const int64_t r = ct * 32 + row;
- if (r >= ne1) {
- continue;
- }
+ repack_parallel_for_ct(matrix_size, n_col_tiles, [&](int ct0, int ct1) {
+ for (int ct = ct0; ct < ct1; ct++) {
+ for (int kt = 0; kt < n_k_tiles; kt++) {
+ uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size;
+ ggml_half * scale_dst = (ggml_half *) (tile_dst + 512);
+ memset(tile_dst, 0, tile_size);
+
+ for (int row = 0; row < 32; row++) {
+ const int64_t r = ct * 32 + row;
+ if (r >= ne1) {
+ continue;
+ }
- const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
- scale_dst[row] = block.d;
- for (int k = 0; k < QK1_0; k++) {
- int byte_offset;
- int bit;
- q1_0_tiled_bit_location(k, row, byte_offset, bit);
- if (block.qs[k / 8] & (1u << (k % 8))) {
- tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+ const block_q1_0 & block = src_slice[r * n_k_tiles + kt];
+ scale_dst[row] = block.d;
+#if defined(__ARM_NEON) || defined(__aarch64__)
+ repack_neon_q1_0_block(block.qs, tile_dst, row);
+#else
+ for (int k = 0; k < QK1_0; k++) {
+ int byte_offset;
+ int bit;
+ q1_0_tiled_bit_location(k, row, byte_offset, bit);
+ if (block.qs[k / 8] & (1u << (k % 8))) {
+ tile_dst[byte_offset] |= (uint8_t) (1u << bit);
+ }
}
+#endif
}
}
}
- }
+ });
}
}
@@ -1127,11 +1568,11 @@ void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t siz
break;
case GGML_TYPE_F32:
repack_float_to_f16_tiled<float>(tensor, data, 0, size,
- [](float v) { return (uint16_t) ggml_fp32_to_fp16(v); });
+ [](float v) { return (uint16_t) GGML_FP32_TO_FP16(v); });
break;
case GGML_TYPE_BF16:
repack_float_to_f16_tiled<ggml_bf16_t>(tensor, data, 0, size,
- [](ggml_bf16_t v) { return (uint16_t) ggml_fp32_to_fp16(ggml_bf16_to_fp32(v)); });
+ [](ggml_bf16_t v) { return (uint16_t) GGML_FP32_TO_FP16(GGML_BF16_TO_FP32(v)); });
break;
default:
break;