hexagon: dynamic quantizer improvements (llama/29395)

* hexagon: fix accuracy issue in Q8_0 N=1 MUL_MAT

* hex-quant: fix register spills

* hex-mm: use dma for all dyn.quant paths

Co-authored-by: Aparna M P <aparmp@qti.qualcomm.com>

* hex-mm: remove obsolete run_quant_task

* hex-mm: update tracing to properly wrap the events

* hex-mm: use act for activation data in all paths

* hex-mm: use act_ instead of src1_ to avoid confusion in fused kernels

* hex-mm: remove/reroute the rest of the non-DMA act (aka src1) logic

* hex-dma64: yet another pass at cleaning up the dma_addr_t casts

* Update ggml/src/ggml-hexagon/htp/matmul-ops.h

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

* Update ggml/src/ggml-hexagon/htp/matmul-ops.c

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

* Update ggml/src/ggml-hexagon/htp/matmul-ops.c

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

* Update ggml/src/ggml-hexagon/htp/matmul-ops.c

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

---------

Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>
Co-authored-by: Aparna M P <aparmp@qti.qualcomm.com>
Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
This commit is contained in:
kurquhar
2026-10-06 10:34:32 +03:00
committed by Georgi Gerganov
co-authored by Aparna M P Sigbjørn Skjæret Max Krasnyansky
parent e5a88776af
commit 845d1c02a8
5 changed files with 423 additions and 485 deletions
+1 -2
View File
@@ -4494,8 +4494,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params(
if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) {
// Try grouped path first
const bool use_dma_activation = (src1->nb[1]/sizeof(float) > (size_t)ne00_padded);
if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, use_dma_activation, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, ne11, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) {
use_grouped = true;
}
}
@@ -914,98 +914,6 @@ typedef struct {
// activations : fp32 -> fp16
static void transfer_activation_chunk_fp32_to_fp16(__fp16 *restrict vtcm_dst, const float *restrict src, uint32_t n_rows, uint32_t k_block, uint32_t k_stride, uint32_t k_valid) {
const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
const uint32_t n_rows_tiled = (n_rows / HTP_MM_HMX_TILE_N_ROWS) * HTP_MM_HMX_TILE_N_ROWS;
uint32_t r = 0;
#pragma unroll(2)
for (r = 0; r < n_rows_tiled; r += 2) {
uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
const float *ptr_in0 = src + (r + 0) * k_stride;
const float *ptr_in1 = src + (r + 1) * k_stride;
uint32_t c = 0;
for (; c + 32 <= k_valid; c += 32) {
HVX_Vector v0 = *(const HVX_Vector *)(ptr_in0 + c);
HVX_Vector v1 = *(const HVX_Vector *)(ptr_in1 + c);
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index
uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
tile[r1 / 2] = v_out;
}
if (c < k_block) {
HVX_Vector v0 = *(const HVX_Vector *)(ptr_in0 + c);
HVX_Vector v1 = *(const HVX_Vector *)(ptr_in1 + c);
uint32_t rem = k_valid - c;
HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0);
v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero());
v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero());
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index
uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
tile[r1 / 2] = v_out;
}
}
for (; r < n_rows_padded; r += 2) {
uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
const bool row0_valid = r < n_rows;
const bool row1_valid = (r + 1) < n_rows;
const float *ptr_in0 = row0_valid ? (src + (r + 0) * k_stride) : NULL;
const float *ptr_in1 = row1_valid ? (src + (r + 1) * k_stride) : NULL;
uint32_t c = 0;
for (; c + 32 <= k_valid; c += 32) {
HVX_Vector v0 = Q6_V_vzero();
HVX_Vector v1 = Q6_V_vzero();
if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c);
if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c);
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index
uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
tile[r1 / 2] = v_out;
}
if (c < k_block) {
HVX_Vector v0 = Q6_V_vzero();
HVX_Vector v1 = Q6_V_vzero();
if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c);
if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c);
uint32_t rem = k_valid - c;
HVX_VectorPred mask = Q6_Q_vsetq2_R(rem > 0 ? rem * sizeof(float) : 0);
v0 = Q6_V_vmux_QVV(mask, v0, Q6_V_vzero());
v1 = Q6_V_vmux_QVV(mask, v1, Q6_V_vzero());
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index
uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
tile[r1 / 2] = v_out;
}
}
}
static void transfer_activation_row_pair_fp32_to_fp16(
__fp16 *restrict vtcm_dst,
const float *restrict row0,
@@ -121,29 +121,43 @@ static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * r
HVX_Vector * vx = (HVX_Vector *) x;
HVX_Vector zero = Q6_V_vzero();
HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0]));
HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1]));
HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2]));
HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3]));
HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero);
HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero);
HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero);
HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero);
HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero);
HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero);
HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero);
HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero);
HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf)));
HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf)));
HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf)));
HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf)));
HVX_Vector vmax_hf = hvx_vec_reduce_max_f16(hvx_vec_abs_f16(vx01_hf));
vmax_hf = hvx_vec_reduce_max2_f16(hvx_vec_abs_f16(vx23_hf), vmax_hf);
HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008));
HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008));
HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16);
HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16);
HVX_Vector vd_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax_hf, Q6_Vh_vsplat_R(0x2008));
HVX_Vector vd_hf = Q6_Vhf_equals_Vqf16(vd_qf16);
HVX_Vector vd_inv_hf = hvx_vec_inverse_f16(vd_hf);
vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd_inv_hf));
vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd_inv_hf));
HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf);
HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf);
vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf));
vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf));
HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf);
HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf);
HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16);
HVX_Vector r_scale = hvx_vec_repl_f16(vd_hf);
HVX_VectorPair vp01 = Q6_W_vshuff_VVR(vd01_hf, vd01_hf, -64);
HVX_VectorPair vp23 = Q6_W_vshuff_VVR(vd23_hf, vd23_hf, -64);
static const uint8_t __attribute__((aligned(128))) repl[128] = {
0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
@@ -157,8 +171,19 @@ static inline void quantize_block_f32_q8_0_tiled(float * restrict x, uint8_t * r
};
HVX_Vector v_repl_ctrl = * (const HVX_Vector *) repl;
#pragma unroll
for (int b = 0; b < 4; b++) {
HVX_Vector v_act = Q6_V_vror_VR(vx_i8, b * 32);
HVX_Vector r_scale;
if (b == 0) {
r_scale = Q6_V_lo_W(vp01);
} else if (b == 1) {
r_scale = Q6_V_hi_W(vp01);
} else if (b == 2) {
r_scale = Q6_V_lo_W(vp23);
} else {
r_scale = Q6_V_hi_W(vp23);
}
HVX_Vector r0 = Q6_V_vdelta_VV(v_act, v_repl_ctrl);
HVX_Vector r1 = Q6_V_vdelta_VV(Q6_V_vror_VR(v_act, 4), v_repl_ctrl);
@@ -941,14 +966,9 @@ static inline void quantize_f32_q8_0_tiled_kernel(
size_t src_row_size,
size_t dst_row_size
) {
const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float));
hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float));
(void) tmp_data;
for (uint32_t i = 0; i < nrows; ++i) {
hex_l2fetch(src_data, src_row_size, src_row_size, 2);
hvx_copy_f32_aa(tmp_data, src_data, ne0);
quantize_row_f32_q8_0_tiled((float *) tmp_data, dst_data, ne0);
quantize_row_f32_q8_0_tiled((float *) src_data, dst_data, ne0);
dst_data += dst_row_size;
src_data += src_row_size;
}
@@ -963,14 +983,9 @@ static inline void quantize_f32_q8_1_tiled_kernel(
size_t src_row_size,
size_t dst_row_size
) {
const size_t src_row_size_padded = hex_round_up(src_row_size, QK_Q8_0_TILED * sizeof(float));
hvx_splat_f32_a(tmp_data, 0.0f, src_row_size_padded / sizeof(float));
(void) tmp_data;
for (uint32_t i = 0; i < nrows; ++i) {
hex_l2fetch(src_data, src_row_size, src_row_size, 2);
hvx_copy_f32_aa(tmp_data, src_data, ne0);
quantize_row_f32_q8_1_tiled((float *) tmp_data, dst_data, ne0);
quantize_row_f32_q8_1_tiled((float *) src_data, dst_data, ne0);
dst_data += dst_row_size;
src_data += src_row_size;
}
@@ -988,24 +1003,15 @@ static inline void quantize_f32_q8_0_tiled_block_kernel(
uint32_t r,
uint32_t c
) {
(void) tmp_data;
const uint32_t qk = QK_Q8_0_TILED;
const uint32_t nb = (ne0 + qk - 1) / qk;
for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
const uint8_t * restrict src_ptr = (const uint8_t *) src + r * src_row_size + c * qk * sizeof(float);
const float * restrict src_ptr = (const float *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(float));
uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 1152;
hex_l2fetch(src_ptr, qk * sizeof(float), qk * sizeof(float), 1);
if (c == nb - 1) {
uint32_t active_elements = ne0 - c * qk;
hvx_splat_f32_a(tmp_data, 0.0f, qk);
hvx_copy_f32_aa(tmp_data, src_ptr, active_elements);
} else {
hvx_copy_f32_aa(tmp_data, src_ptr, qk);
}
quantize_block_f32_q8_0_tiled((float *) tmp_data, dst_ptr);
quantize_block_f32_q8_0_tiled((float *) src_ptr, dst_ptr);
c++;
if (c == nb) {
@@ -1027,24 +1033,15 @@ static inline void quantize_f32_q8_1_tiled_block_kernel(
uint32_t r,
uint32_t c
) {
(void) tmp_data;
const uint32_t qk = QK_Q8_0_TILED;
const uint32_t nb = (ne0 + qk - 1) / qk;
for (uint32_t ib = ib_first; ib < ib_last; ++ib) {
const uint8_t * restrict src_ptr = (const uint8_t *) src + r * src_row_size + c * qk * sizeof(float);
const float * restrict src_ptr = (const float *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(float));
uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 1280;
hex_l2fetch(src_ptr, qk * sizeof(float), qk * sizeof(float), 1);
if (c == nb - 1) {
uint32_t active_elements = ne0 - c * qk;
hvx_splat_f32_a(tmp_data, 0.0f, qk);
hvx_copy_f32_aa(tmp_data, src_ptr, active_elements);
} else {
hvx_copy_f32_aa(tmp_data, src_ptr, qk);
}
quantize_block_f32_q8_1_tiled((float *) tmp_data, dst_ptr);
quantize_block_f32_q8_1_tiled((float *) src_ptr, dst_ptr);
c++;
if (c == nb) {
File diff suppressed because it is too large Load Diff
+69 -48
View File
@@ -332,6 +332,7 @@ struct htp_mm_hvx_vtcm_layout {
size_t off_src2; // vtcm_src2 (Wq / fused only)
size_t off_src3; // vtcm_src3 (Wv / fused only)
size_t off_dst; // vtcm_dst (output scratch)
size_t off_act_raw; // vtcm_act_raw (raw activation DMA staging)
// Cached sizes
size_t src0_bytes;
@@ -339,6 +340,7 @@ struct htp_mm_hvx_vtcm_layout {
size_t src2_bytes;
size_t src3_bytes;
size_t dst_bytes;
size_t act_raw_bytes;
size_t total_bytes;
};
@@ -351,7 +353,6 @@ static inline void htp_mm_hmx_vtcm_layout_build(
size_t mc,
size_t nc,
uint32_t group_size,
bool use_dma_activation,
bool pipeline,
uint32_t act_threads,
uint32_t aligned_tile_size,
@@ -366,8 +367,7 @@ static inline void htp_mm_hmx_vtcm_layout_build(
const size_t activation_area_size = hex_align_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
const size_t output_area_size = hex_align_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
const size_t scratch_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
const size_t min_f32_size = use_dma_activation
? hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128) : 0;
const size_t min_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128);
// Group A: Permanent activation tiles and scales
size_t off_group_a = 0;
@@ -388,10 +388,9 @@ static inline void htp_mm_hmx_vtcm_layout_build(
// Group C: Activation prep temporary buffer (overlaps Group B, starting at off_group_a)
const size_t max_f32_size = act_threads * 64 * k * sizeof(float);
const size_t act_f32_size = use_dma_activation
? hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128) : 0;
const size_t act_f32_size = hex_align_up(hex_smin(max_f32_size, hex_smax(min_f32_size, group_b_size)), 128);
size_t off_group_c = off_group_a;
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_f32, act_f32_size, use_dma_activation);
VTCM_LAYOUT_ALLOC(off_group_c, off_act_f32, act_f32_size);
const size_t group_c_size = off_group_c - off_group_a;
@@ -478,11 +477,12 @@ static inline void htp_mm_hvx_vtcm_layout_build(
bool is_fused_nx
) {
(void)src1_row_size;
size_t src0_sz = 0;
size_t src1_sz = 0;
size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
size_t src3_sz = 0;
size_t dst_sz = 0;
size_t src0_sz = 0;
size_t src1_sz = 0;
size_t src2_sz = src2_row_size > 0 ? htp_mm_round_up(src2_row_size, 128) : 0;
size_t src3_sz = 0;
size_t dst_sz = 0;
size_t act_raw_sz = 0;
const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
@@ -491,7 +491,6 @@ static inline void htp_mm_hvx_vtcm_layout_build(
if (is_fused_nx) {
const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
size_t weight_sz_per_thread = 0;
@@ -507,12 +506,14 @@ static inline void htp_mm_hvx_vtcm_layout_build(
size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
size_t act_sz = hex_round_up(tiled_act_row_size * src1_nrows, 128);
size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
src1_sz = act_sz; // quantized activation buffer
src2_sz = 0;
src3_sz = 0;
dst_sz = quant_scratch_size;
src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer
src1_sz = act_sz; // quantized activation buffer
src2_sz = 0;
src3_sz = 0;
dst_sz = 0;
act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
} else if (is_matmul_id) {
const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
const size_t src1_row_size_tiled = (wtype == HTP_TYPE_Q4_1 || wtype == HTP_TYPE_Q4_K) ? htp_mm_q8_1_tiled_row_size(ne10)
@@ -529,10 +530,13 @@ static inline void htp_mm_hvx_vtcm_layout_build(
src0_sz_per_thread = repacked_vtcm_size;
}
src0_sz = src0_sz_per_thread * n_threads;
dst_sz = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
src2_sz = 0;
src3_sz = 0;
size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
src0_sz = src0_sz_per_thread * n_threads;
dst_sz = 0;
src2_sz = 0;
src3_sz = 0;
act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
} else {
const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
const size_t dst_nrows = (src1_nrows > 1) ? 0 : 1;
@@ -540,16 +544,18 @@ static inline void htp_mm_hvx_vtcm_layout_build(
switch (kernel_type) {
case HTP_MM_KERNEL_HVX_F16_F16_VTCM: {
size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128);
src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
src1_sz = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * src1_nrows, 128);
break;
}
case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
src1_sz = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
act_raw_sz = 0;
break;
}
case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
@@ -569,10 +575,10 @@ static inline void htp_mm_hvx_vtcm_layout_build(
src0_sz = repacked_vtcm_size * n_threads;
}
size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
size_t dst_slice_per_thread = (dst_nrows > 0 && src1_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0;
size_t dst_size_per_thread = (dst_slice_per_thread > quant_scratch_size_per_thread) ? dst_slice_per_thread : quant_scratch_size_per_thread;
dst_sz = dst_size_per_thread * n_threads;
dst_sz = dst_slice_per_thread * n_threads;
size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
act_raw_sz = hex_round_up(raw_row_size * src1_nrows, 128);
break;
}
default:
@@ -580,19 +586,30 @@ static inline void htp_mm_hvx_vtcm_layout_build(
}
}
size_t off = 0;
VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz);
VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz);
VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz);
VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz);
VTCM_LAYOUT_ALLOC(off, off_dst, dst_sz);
// Group A: Persistent buffers across chunk compute
size_t off_group_a = 0;
VTCM_LAYOUT_ALLOC(off_group_a, off_src1, src1_sz);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src2, src2_sz, src2_sz > 0);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_a, off_src3, src3_sz, src3_sz > 0);
L->src0_bytes = src0_sz;
L->src1_bytes = src1_sz;
L->src2_bytes = src2_sz;
L->src3_bytes = src3_sz;
L->dst_bytes = dst_sz;
L->total_bytes = off;
// Group B: Compute-only buffers (starts at off_group_a)
size_t off_group_b = off_group_a;
VTCM_LAYOUT_ALLOC(off_group_b, off_src0, src0_sz);
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst, dst_sz, dst_sz > 0);
const size_t group_b_size = off_group_b - off_group_a;
// Group C: Raw activation staging buffer (overlaps Group B, starts at off_group_a)
size_t off_group_c = off_group_a;
VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_raw, act_raw_sz, act_raw_sz > 0);
const size_t group_c_size = off_group_c - off_group_a;
L->src0_bytes = src0_sz;
L->src1_bytes = src1_sz;
L->src2_bytes = src2_sz;
L->src3_bytes = src3_sz;
L->dst_bytes = dst_sz;
L->act_raw_bytes = act_raw_sz;
L->total_bytes = off_group_a + hex_smax(group_b_size, group_c_size);
}
static inline bool htp_mm_hvx_solve_vtcm_params(
@@ -642,7 +659,12 @@ static inline bool htp_mm_hvx_solve_vtcm_params(
return false;
}
uint32_t m_chunk = (uint32_t) (avail_act / row_size);
size_t eff_row_size = row_size;
if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
eff_row_size += hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
}
uint32_t m_chunk = (uint32_t) (avail_act / eff_row_size);
if (m_chunk > 1) {
m_chunk &= ~1U;
}
@@ -679,15 +701,15 @@ static inline size_t htp_mm_hmx_get_2d_vtcm_size(
int wtype, uint32_t k, size_t mc, size_t nc, bool pipeline, uint32_t act_threads, uint32_t aligned_tile_size, size_t src2_size
) {
struct htp_mm_hmx_vtcm_layout L;
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, false, pipeline, act_threads, aligned_tile_size, src2_size);
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, wtype, k, mc, nc, 1, pipeline, act_threads, aligned_tile_size, src2_size);
return L.total_bytes;
}
static inline size_t htp_mm_hmx_get_batched_vtcm_size(
int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool use_dma_activation, bool pipeline, uint32_t act_threads, size_t src2_size) {
int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t src2_size) {
(void)pipeline;
struct htp_mm_hmx_vtcm_layout L;
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, use_dma_activation, false, act_threads, 0, src2_size);
htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, src2_size);
return L.total_bytes;
}
@@ -697,7 +719,6 @@ static inline bool htp_mm_hmx_solve_batched_params(
uint32_t ne01_padded,
uint32_t ne11,
uint32_t group_size,
bool use_dma_activation,
int n_threads,
bool pipeline,
size_t src2_size,
@@ -726,7 +747,7 @@ static inline bool htp_mm_hmx_solve_batched_params(
if (htp_mm_hmx_compute_chunks(vtcm_budget, group_overhead, group_size_per_n, group_size_per_m, group_size_per_mn, hex_align_up(ne11, 32), ne01_padded,
(size_t) ne01_padded * HTP_MM_HMX_COST_W_DEQUANT, (size_t) ne11 * HTP_MM_HMX_COST_A_CONVERT,
&m_chunk_candidate, &n_chunk_candidate, &vtcm_size_candidate) == 0) {
size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, use_dma_activation, pipeline, act_threads, src2_size);
size_t exact_size = htp_mm_hmx_get_batched_vtcm_size(wtype, k, m_chunk_candidate, n_chunk_candidate, group_size, pipeline, act_threads, src2_size);
if (exact_size <= vtcm_budget) {
size_t mblocks = ((size_t) ne11 + m_chunk_candidate - 1) / m_chunk_candidate;
if (mblocks < best_mblocks || (mblocks == best_mblocks && act_threads > best_act_threads)) {