kparams->u.hvx.size_v_row_padded = size_v_row_padded;
kparams->u.hvx.src0_div21 = init_fastdiv_values(q->ne[2] * q->ne[1]);
kparams->u.hvx.src0_div1 = init_fastdiv_values(q->ne[1]);
- kparams->u.hvx.broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]);
- kparams->u.hvx.broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]);
- kparams->u.hvx.broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]);
- kparams->u.hvx.broadcast_rv3 = init_fastdiv_values(q->ne[3]/v->ne[3]);
+ kparams->broadcast_rk2 = init_fastdiv_values(q->ne[2]/k->ne[2]);
+ kparams->broadcast_rk3 = init_fastdiv_values(q->ne[3]/k->ne[3]);
+ kparams->broadcast_rv2 = init_fastdiv_values(q->ne[2]/v->ne[2]);
+ kparams->broadcast_rv3 = init_fastdiv_values(q->ne[3]/v->ne[3]);
if (mask) {
kparams->src3_div2 = init_fastdiv_values(mask->ne[2]);
kparams->src3_div3 = init_fastdiv_values(mask->ne[3]);
kparams->kernel_type = (src1_nrows < (int) sess->n_threads) ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW;
kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
- size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
+ struct htp_mm_hvx_vtcm_layout L;
uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
uint32_t best_n_prefetch = 2;
- size_t total_size = 0;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
- total_size = htp_mm_hvx_id_get_vtcm_sizes(
- wtype, ne10, src1_nrows, sess->n_threads, src0->nb[1], d,
- &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0->nb[1], 0, d, true, false, false
);
- if (total_size <= vtcm_budget) {
+ if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
break;
}
}
- if (best_n_prefetch == 2 && total_size > vtcm_budget) {
- total_size = htp_mm_hvx_id_get_vtcm_sizes(
- wtype, ne10, src1_nrows, sess->n_threads, src0->nb[1], 2,
- &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0->nb[1], 0, 2, true, false, false
);
}
kparams->n_prefetch = best_n_prefetch;
- kparams->vtcm_size = total_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
} else {
bool try_tiled = (k_align && opt_mm_select >= 2);
if (try_tiled) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
}
+ struct htp_mm_hvx_vtcm_layout L;
uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
uint32_t best_n_prefetch = 2;
- size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
- size_t total_size = 0;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
- total_size = htp_mm_hvx_get_vtcm_sizes(
- kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], d, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], d, false, false, false
);
- if (total_size <= vtcm_budget) {
+ if (L.total_bytes <= vtcm_budget) {
best_n_prefetch = d;
break;
}
}
- if (best_n_prefetch == 2 && total_size > vtcm_budget) {
- total_size = htp_mm_hvx_get_vtcm_sizes(
- kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 2, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) {
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], 2, false, false, false
);
}
kparams->n_prefetch = best_n_prefetch;
- if (total_size <= vtcm_budget) {
- kparams->vtcm_size = total_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ if (L.total_bytes <= vtcm_budget) {
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
goto done_quant;
}
- HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), total_size, vtcm_budget);
+ HEX_VERBOSE("ggml-hex: %s HVX tiled path VTCM size needed (%zu) > budget (%zu), falling back to HVX flat\n", sess->name.c_str(), L.total_bytes, vtcm_budget);
}
// Flat HVX fallback
kparams->src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
- size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
- size_t total_size = htp_mm_hvx_get_vtcm_sizes(
- kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
);
kparams->n_prefetch = 16;
- kparams->vtcm_size = total_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
}
}
const bool is_batched = (ne02 > 1) || (ne03 > 1);
const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1);
- size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
- size_t vtcm_size = htp_mm_hvx_get_vtcm_sizes(
- HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
);
- if (!is_batched && !is_permuted && vtcm_size <= vtcm_budget) {
+ if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_VTCM;
kparams->src1_row_size = hex_round_up(ne10 * 2, 128);
- kparams->vtcm_size = vtcm_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
kparams->n_prefetch = 16;
} else {
if (src1->type == GGML_TYPE_F32) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_F16_F16_DDR;
}
kparams->src1_row_size = src1->nb[1];
- size_t ddr_size = htp_mm_hvx_get_vtcm_sizes(
- kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
);
- kparams->vtcm_size = ddr_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
kparams->n_prefetch = 16;
}
} else {
const bool is_batched = (ne02 > 1) || (ne03 > 1);
const bool is_permuted = ggml_is_permuted(src0) || ggml_is_permuted(src1);
- size_t vtcm_src0_size = 0, vtcm_src1_size = 0, vtcm_dst_size = 0;
- size_t vtcm_size = htp_mm_hvx_get_vtcm_sizes(
- HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
);
- if (!is_batched && !is_permuted && vtcm_size <= vtcm_budget) {
+ if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_VTCM;
kparams->src1_row_size = hex_round_up(ne10 * 4, 128);
- kparams->vtcm_size = vtcm_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
kparams->n_prefetch = 16;
} else {
kparams->kernel_type = HTP_MM_KERNEL_HVX_F32_F32_DDR;
kparams->src1_row_size = src1->nb[1];
- size_t ddr_size = htp_mm_hvx_get_vtcm_sizes(
- kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
- dst->nb[1], src0->nb[1], src1->nb[1], 16, &vtcm_src0_size, &vtcm_src1_size, &vtcm_dst_size
+ htp_mm_hvx_vtcm_layout_build(
+ &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads,
+ dst->nb[1], src0->nb[1], src1->nb[1], 16, false, false, false
);
- kparams->vtcm_size = ddr_size;
- kparams->vtcm_src0_size = vtcm_src0_size;
- kparams->vtcm_src1_size = vtcm_src1_size;
- kparams->vtcm_dst_size = vtcm_dst_size;
+ kparams->vtcm_size = L.total_bytes;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
kparams->n_prefetch = 16;
}
}
const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
const size_t src0_row_size = src0->nb[1];
- const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
- size_t src0_sz_per_thread = 0;
- size_t src2_sz_per_thread = 0;
- size_t src3_sz_per_thread = 0;
uint32_t best_n_prefetch = 16;
- size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * sess->n_threads;
-
if (is_repack) {
- uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
- uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
- uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
- size_t src1_sz = src1_sz_per_thread;
-
const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
best_n_prefetch = 2;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
- size_t repacked_vtcm_size = hex_round_up(d * tile_row_size, 128);
- size_t src0_sz = repacked_vtcm_size * sess->n_threads;
- size_t src2_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
- size_t src3_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
-
- if (tiled_vtcm_size <= sess->vtcm_size) {
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0_row_size, src1_row_size, d, false, true, false
+ );
+ if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
- src0_sz_per_thread = repacked_vtcm_size;
- src2_sz_per_thread = hex_round_up(d * tile_row_size, 128);
- src3_sz_per_thread = hex_round_up(d * tile_row_size, 128);
break;
}
}
- if (best_n_prefetch == 2 && src0_sz_per_thread == 0) {
- size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128);
- src0_sz_per_thread = repacked_vtcm_size;
- src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
- src3_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
- }
- } else {
- best_n_prefetch = 16;
- src0_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
- src2_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
- src3_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
}
- size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
+ struct htp_mm_hvx_vtcm_layout L;
+ bool try_tiled = (opt_mm_select >= 2);
- size_t src0_sz = src0_sz_per_thread * sess->n_threads;
- size_t src1_sz = src1_sz_per_thread;
- size_t src2_sz = src2_sz_per_thread * sess->n_threads;
- size_t src3_sz = src3_sz_per_thread * sess->n_threads;
+ // Test tiled first
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0_row_size, src1_row_size, best_n_prefetch, false, true, false
+ );
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
- bool try_tiled = (opt_mm_select >= 2);
- if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
+ if (try_tiled && L.total_bytes <= sess->vtcm_size) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
- kparams->vtcm_src0_size = src0_sz;
- kparams->vtcm_src1_size = src1_sz;
- kparams->vtcm_src2_size = src2_sz;
- kparams->vtcm_src3_size = src3_sz;
- kparams->vtcm_dst_size = quant_scratch_size;
- kparams->vtcm_size = tiled_vtcm_size;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_src2_size = L.src2_bytes;
+ kparams->vtcm_src3_size = L.src3_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
+ kparams->vtcm_size = L.total_bytes;
kparams->n_prefetch = best_n_prefetch;
} else {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
- size_t flat_src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
- kparams->vtcm_src0_size = src0_sz;
- kparams->vtcm_src1_size = flat_src1_sz;
- kparams->vtcm_src2_size = src2_sz;
- kparams->vtcm_src3_size = src3_sz;
- kparams->vtcm_dst_size = quant_scratch_size;
- kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + src3_sz + quant_scratch_size;
+
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, true, false
+ );
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_src2_size = L.src2_bytes;
+ kparams->vtcm_src3_size = L.src3_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
+ kparams->vtcm_size = L.total_bytes;
kparams->n_prefetch = best_n_prefetch;
}
}
const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3];
const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
const size_t src0_row_size = src0->nb[1];
- const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128);
- size_t src0_sz_per_thread = 0;
- size_t src2_sz_per_thread = 0;
uint32_t best_n_prefetch = 16;
- size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * sess->n_threads;
-
if (is_repack) {
- uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
- uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
- uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
- size_t src1_sz = src1_sz_per_thread;
-
const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16;
best_n_prefetch = 2;
for (uint32_t d = max_prefetch; d >= 2; d /= 2) {
- size_t repacked_vtcm_size = hex_round_up(d * tile_row_size, 128);
- size_t src0_sz = repacked_vtcm_size * sess->n_threads;
- size_t src2_sz = hex_round_up(d * tile_row_size, 128) * sess->n_threads;
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
-
- if (tiled_vtcm_size <= sess->vtcm_size) {
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0_row_size, src1_row_size, d, false, false, true
+ );
+ if (L.total_bytes <= sess->vtcm_size) {
best_n_prefetch = d;
- src0_sz_per_thread = repacked_vtcm_size;
- src2_sz_per_thread = hex_round_up(d * tile_row_size, 128);
break;
}
}
- if (best_n_prefetch == 2 && src0_sz_per_thread == 0) {
- size_t repacked_vtcm_size = hex_round_up(2 * tile_row_size, 128);
- src0_sz_per_thread = repacked_vtcm_size;
- src2_sz_per_thread = hex_round_up(2 * tile_row_size, 128);
- }
- } else {
- best_n_prefetch = 16;
- src0_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
- src2_sz_per_thread = hex_round_up(best_n_prefetch * src0_row_size_padded, 128);
}
- size_t src1_sz_per_thread = hex_round_up(src1_row_size * src1_nrows, 128);
+ struct htp_mm_hvx_vtcm_layout L;
+ bool try_tiled = (opt_mm_select >= 2);
- size_t src0_sz = src0_sz_per_thread * sess->n_threads;
- size_t src1_sz = src1_sz_per_thread;
- size_t src2_sz = src2_sz_per_thread * sess->n_threads;
+ // Test tiled first
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0_row_size, src1_row_size, best_n_prefetch, false, false, true
+ );
- size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
- bool try_tiled = (opt_mm_select >= 2);
- if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
+ if (try_tiled && L.total_bytes <= sess->vtcm_size) {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW;
- kparams->vtcm_src0_size = src0_sz;
- kparams->vtcm_src1_size = src1_sz;
- kparams->vtcm_src2_size = src2_sz;
- kparams->vtcm_dst_size = quant_scratch_size;
- kparams->vtcm_size = tiled_vtcm_size;
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_src2_size = L.src2_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
+ kparams->vtcm_size = L.total_bytes;
kparams->n_prefetch = best_n_prefetch;
} else {
kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT;
size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
- size_t flat_src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
- kparams->vtcm_src0_size = src0_sz;
- kparams->vtcm_src1_size = flat_src1_sz;
- kparams->vtcm_src2_size = src2_sz;
- kparams->vtcm_dst_size = quant_scratch_size;
- kparams->vtcm_size = src0_sz + flat_src1_sz + src2_sz + quant_scratch_size;
+
+ htp_mm_hvx_vtcm_layout_build(
+ &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads,
+ 0, src0_row_size, flat_src1_row_size, best_n_prefetch, false, false, true
+ );
+ kparams->vtcm_src0_size = L.src0_bytes;
+ kparams->vtcm_src1_size = L.src1_bytes;
+ kparams->vtcm_src2_size = L.src2_bytes;
+ kparams->vtcm_dst_size = L.dst_bytes;
+ kparams->vtcm_size = L.total_bytes;
kparams->n_prefetch = best_n_prefetch;
}
}
worker-pool.c
hex-dma.c
hmx-queue.c
+ gated-delta-net-ops.c
binary-ops.c
unary-ops.c
sum-rows-ops.c
concat-ops.c
diag-ops.c
solve-tri-ops.c
- gated-delta-net-ops.c
pad-ops.c
- matmul-ops.c
flash-attn-ops.c
+ matmul-ops.c
)
target_compile_definitions(${HTP_LIB} PRIVATE
#include "hexagon_protos.h"
#include "hvx_hexagon_protos.h"
#include "hex-dma.h"
-#include "vtcm-utils.h"
+#include "htp-vtcm.h"
#include "hvx-utils.h"
#include "hex-fastdiv.h"
#include <string.h>
#include <HAP_perf.h>
#include <math.h>
#include <stdbool.h>
+#include <stdatomic.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include "hvx-copy.h"
#include "hvx-reduce.h"
#include "hvx-flash-attn.h"
-#include "vtcm-utils.h"
+#include "htp-vtcm.h"
#include "worker-pool.h"
#define GGML_COMMON_DECL_C
__fp16 * vtcm_slopes; // ALiBi slopes [g_br]
size_t row_buf_stride; // HVX vectors per row buffer (Bc/64)
size_t mask_buf_row_stride; // elements (__fp16) per row in mask buffer
+ size_t q_tile_bytes;
+ size_t o_tile_bytes;
+ size_t col_vec_bytes;
+ size_t d_tile_bytes;
bool mask_broadcast; // true when mask->ne[2] == 1 (head-independent, single 2D DMA)
dma_cache m_cache;
};
struct hmx_fa_context * factx;
uint32_t kv_rows;
size_t src_stride;
- size_t buf_idx;
+ void * curr_k;
uint32_t kv_start;
uint32_t rows_per_t;
} fa_k_int_args_t;
struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
- hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, factx->vtcm_k_fp16[args->buf_idx], total_rows, factx->DK,
+ hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles, (const __fp16 *) args->curr_k, total_rows, factx->DK,
args->src_stride, start, end);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start));
}
-static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, size_t buf_idx, uint32_t kv_start) {
+static void fa_phase_k_interleave(struct hmx_fa_context * factx, uint32_t kv_rows, size_t src_stride, void * curr_k, uint32_t kv_start) {
worker_pool_context_t wp = factx->octx->ctx->worker_pool;
uint32_t n = 1;
if (factx->n_threads > 1 && kv_rows >= factx->n_threads * 2) {
n = factx->n_threads;
}
uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
- fa_k_int_args_t args = { factx, kv_rows, src_stride, buf_idx, kv_start, rows_per_t };
+ fa_k_int_args_t args = { factx, kv_rows, src_stride, curr_k, kv_start, rows_per_t };
if (n > 1) {
worker_pool_run_func(wp, fa_k_interleave_thread, &args, n);
} else {
struct hmx_fa_context * factx;
uint32_t kv_rows;
size_t src_stride;
- size_t buf_idx;
+ void * v_src;
+ void * v_tiles_dst;
size_t n_col_tiles;
uint32_t kv_start;
uint32_t rows_per_t;
return;
}
- __fp16 * v_tiles_dest = factx->pipeline ? factx->vtcm_v_tiles[args->buf_idx] : factx->vtcm_v_tiles[0];
+ __fp16 * v_tiles_dst = (__fp16 *) args->v_tiles_dst;
struct htp_thread_trace * tr = factx->octx->ctx ? &factx->octx->ctx->trace[i] : NULL;
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
- hmx_interleave_cols_to_tiles(v_tiles_dest, factx->vtcm_v_fp16[args->buf_idx], total_rows, factx->DV,
+ hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV,
args->src_stride, (uint32_t) args->n_col_tiles, start, end);
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start));
}
static void fa_phase_v_interleave(struct hmx_fa_context * factx,
uint32_t kv_rows,
size_t src_stride,
- size_t buf_idx,
+ void * v_src,
+ void * v_tiles_dst,
size_t n_col_tiles,
uint32_t kv_start) {
worker_pool_context_t wp = factx->octx->ctx->worker_pool;
n = factx->n_threads;
}
uint32_t rows_per_t = hex_align_up(hmx_ceil_div(kv_rows, n), 2);
- fa_v_int_args_t args = { factx, kv_rows, src_stride, buf_idx, n_col_tiles, kv_start, rows_per_t };
+ fa_v_int_args_t args = { factx, kv_rows, src_stride, v_src, v_tiles_dst, n_col_tiles, kv_start, rows_per_t };
if (n > 1) {
worker_pool_run_func(wp, fa_v_interleave_thread, &args, n);
} else {
uint32_t ib3;
size_t n_rows_g;
size_t rows_per_t;
+ size_t n_rows_q;
+ bool q_transposed;
+ atomic_uint barrier;
} fa_q_load_args_t;
static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
const uint32_t g_br = factx->g_br;
const uint32_t DV = factx->DV;
- const size_t col_vec_bytes = hex_align_up(g_br * sizeof(float), 256);
- const size_t d_tile_bytes = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
- const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
+ const size_t col_vec_bytes = factx->col_vec_bytes;
+ const size_t d_tile_bytes = factx->d_tile_bytes;
// Initialize vtcm_l_vec & vtcm_m_vec
const size_t l_bytes_per_t = hex_align_up(col_vec_bytes / n, 128);
if (d_start < d_tile_bytes) {
hvx_splat_u8_a((char *) factx->vtcm_d_tiles + d_start, 0, d_end - d_start);
}
-
- // Initialize vtcm_o_tiles[0] to 0
- __fp16 * o_tile_prev = factx->vtcm_o_tiles[0];
- const size_t o_bytes_per_t = hex_align_up(o_tile_bytes / n, 128);
- const size_t o_start = i * o_bytes_per_t;
- const size_t o_end = hex_smin(o_start + o_bytes_per_t, o_tile_bytes);
- if (o_start < o_tile_bytes) {
- hvx_splat_u8_a((char *) o_tile_prev + o_start, 0, o_end - o_start);
- }
}
- if (start >= factx->g_br) {
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
- return;
- }
-
- const struct htp_tensor * q = args->q;
- const uint32_t q_start = args->q_start;
- const uint32_t kv_head = args->kv_head;
- const uint32_t ib3 = args->ib3;
+ if (start < factx->g_br) {
+ const struct htp_tensor * q = args->q;
+ const uint32_t q_start = args->q_start;
+ const uint32_t kv_head = args->kv_head;
+ const uint32_t ib3 = args->ib3;
- for (size_t r = start; r < end; r += 2) {
- const size_t q_idx0 = fastdiv(r + 0, &factx->div_G);
- const size_t h_idx0 = fastmodulo(r + 0, G, &factx->div_G);
- const size_t q_idx1 = fastdiv(r + 1, &factx->div_G);
- const size_t h_idx1 = fastmodulo(r + 1, G, &factx->div_G);
+ assert(factx->DK == factx->DV);
- const uint8_t * q_ptr0 = (r + 0 < n_rows_g) ? ((const uint8_t *) q->data + (q_start + q_idx0) * q->nb[1] +
- (kv_head * G + h_idx0) * q->nb[2] + ib3 * q->nb[3]) :
- NULL;
- const uint8_t * q_ptr1 = (r + 1 < n_rows_g) ? ((const uint8_t *) q->data + (q_start + q_idx1) * q->nb[1] +
- (kv_head * G + h_idx1) * q->nb[2] + ib3 * q->nb[3]) :
- NULL;
-
- size_t r0 = r / HMX_FP16_TILE_N_ROWS;
- size_t r1 = r % HMX_FP16_TILE_N_ROWS;
- __fp16 * out_base = factx->vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
-
- if (factx->is_q_fp32) {
- const HVX_Vector * pv_in0 = q_ptr0 ? (const HVX_Vector *) q_ptr0 : NULL;
- const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
+ const size_t o_tile_bytes = factx->o_tile_bytes;
+ const bool use_q_dma = (2 * o_tile_bytes >= factx->g_br * DK * (factx->is_q_fp32 ? 4 : 2));
- for (uint32_t d = 0; d < DK / 32; ++d) {
- HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
- HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
- HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+ __fp16 * q_tiles = factx->vtcm_q_tiles;
+ if (use_q_dma) {
+ const size_t g_rows_end = hex_smin(end, n_rows_g);
+ const uint32_t d_limit = factx->is_q_fp32 ? DK / 32 : DK / 64;
- HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
- out_tile[r1 / 2] = v_hf;
+ uint8_t * q_flat = (uint8_t *) factx->vtcm_o_tiles[0];
+ if (factx->is_q_fp32) {
+ switch (d_limit) {
+ case 2: hmx_fa_q_prep_fp32_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+ case 4: hmx_fa_q_prep_fp32_d4(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+ default: hmx_fa_q_prep_fp32( q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); break;
+ }
+ } else {
+ switch (d_limit) {
+ case 1: hmx_fa_q_prep_fp16_d1(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+ case 2: hmx_fa_q_prep_fp16_d2(q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, args->q_transposed); break;
+ default: hmx_fa_q_prep_fp16( q_tiles, q_flat, start, end, g_rows_end, DK, G, args->n_rows_q, &factx->div_G, d_limit, args->q_transposed); break;
+ }
}
} else {
- const HVX_Vector * pv_in0 = q_ptr0 ? (const HVX_Vector *) q_ptr0 : NULL;
- const HVX_Vector * pv_in1 = q_ptr1 ? (const HVX_Vector *) q_ptr1 : NULL;
-
- for (uint32_t d = 0; d < DK / 64; ++d) {
- HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
- HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
- HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+ // Fallback: direct-from-DDR/L2 path
+ hmx_fa_q_prep_fallback(q_tiles, q->data, q->nb[1], q->nb[2], q->nb[3],
+ q_start, kv_head, ib3, start, end, n_rows_g, G, DK, factx->is_q_fp32, &factx->div_G);
+ }
+ }
- __fp16 * out_dual_tile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
- HVX_Vector * pv_out0 = ((HVX_Vector *) out_dual_tile) + r1 / 2;
- HVX_Vector * pv_out1 = pv_out0 + 16;
+ // Synchronize threads before zeroing out vtcm_o_tiles[0] to prevent race condition
+ if (n > 1) {
+ atomic_fetch_sub(&args->barrier, 1);
+ while (atomic_load(&args->barrier) > 0) {
+ // spin wait
+ }
+ }
- *pv_out0 = Q6_V_lo_W(vp);
- *pv_out1 = Q6_V_hi_W(vp);
- }
+ // Zero out vtcm_o_tiles[0] as it was used as temp_q_vtcm
+ {
+ const uint32_t g_br = factx->g_br;
+ const uint32_t DV = factx->DV;
+ const size_t o_tile_bytes = factx->o_tile_bytes;
+ const size_t o_bytes_per_t = hex_align_up(o_tile_bytes / n, 128);
+ const size_t o_start = i * o_bytes_per_t;
+ const size_t o_end = hex_smin(o_start + o_bytes_per_t, o_tile_bytes);
+ if (o_start < o_tile_bytes) {
+ hvx_splat_u8_a((char *) factx->vtcm_o_tiles[0] + o_start, 0, o_end - o_start);
}
}
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_Q_PREP, (uint16_t) (args->q_start * G + start));
n = factx->n_threads;
}
size_t rows_per_t = hex_align_up(hmx_ceil_div(factx->g_br, n), 2);
- fa_q_load_args_t args = { factx, q, q_start, kv_head, ib3, n_rows_g, rows_per_t };
+ const uint32_t n_rows_q = hex_smin(factx->Br, factx->neq1 - q_start);
+ fa_q_load_args_t args;
+ args.factx = factx;
+ args.q = q;
+ args.q_start = q_start;
+ args.kv_head = kv_head;
+ args.ib3 = ib3;
+ args.n_rows_g = n_rows_g;
+ args.rows_per_t = rows_per_t;
+ args.n_rows_q = n_rows_q;
+ args.q_transposed = q->nb[1] < q->nb[2];
+ atomic_init(&args.barrier, n);
if (n > 1) {
worker_pool_run_func(wp, fa_q_load_thread, &args, n);
} else {
fa_o_store_args_t * args = (fa_o_store_args_t *) data;
struct hmx_fa_context * factx = args->factx;
- const size_t n_rows_g = args->n_rows_g;
- const size_t G = factx->G;
- const size_t DV = factx->DV;
-
+ const size_t n_rows_g = args->n_rows_g;
const size_t rows_per_t = args->rows_per_t;
+ const size_t G = factx->G;
+ const size_t DV = factx->DV;
const size_t start = (size_t) i * rows_per_t;
const size_t end = hex_smin(start + rows_per_t, n_rows_g);
const __fp16 * tile_row_base = o_tile_src + r0 * HMX_FP16_TILE_N_ROWS * DV;
for (uint32_t d = 0; d < DV / 64; ++d) {
- const __fp16 * in_dual_tile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
- const HVX_Vector * pv_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
- const HVX_Vector * pv_in1 = pv_in0 + 16;
- HVX_VectorPair vp = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
+ const __fp16 * in_dtile = tile_row_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
+ const HVX_Vector * pv_in1 = pv_in0 + 16;
+ HVX_VectorPair vp = Q6_W_vdeal_VVR(*pv_in1, *pv_in0, -2);
if (r1 % 2 == 0) {
*(HVX_UVector *) (out + d * 64) = Q6_V_lo_W(vp);
} else {
if (has_softcap) {
const HVX_Vector v_cap = hvx_vec_splat_f16(factx->logit_softcap);
for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 64;
- const __fp16 * in_dual_tile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
- const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
- const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
+ size_t ci = c / 64;
+ const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
+ const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
- HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
- HVX_Vector v_s_row0 = Q6_V_lo_W(vp_s_dual_row);
- HVX_Vector v_s_row1 = Q6_V_hi_W(vp_s_dual_row);
+ HVX_VectorPair vp_s_drow = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
+ HVX_Vector v_s_row0 = Q6_V_lo_W(vp_s_drow);
+ HVX_Vector v_s_row1 = Q6_V_hi_W(vp_s_drow);
HVX_Vector t0 = hvx_vec_tanh_f16(v_s_row0);
my_row_buf0[ci] = hvx_vec_mul_f16_f16(t0, v_cap);
}
} else {
for (size_t c = 0; c < kv_rows; c += 64) {
- size_t ci = c / 64;
- const __fp16 * in_dual_tile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
- const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dual_tile) + r1 / 2;
- const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
-
- HVX_VectorPair vp_s_dual_row = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
- my_row_buf0[ci] = Q6_V_lo_W(vp_s_dual_row);
- my_row_buf1[ci] = Q6_V_hi_W(vp_s_dual_row);
+ size_t ci = c / 64;
+ const __fp16 * in_dtile = s_ld_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+ const HVX_Vector * pv_s_in0 = ((const HVX_Vector *) in_dtile) + r1 / 2;
+ const HVX_Vector * pv_s_in1 = pv_s_in0 + 16;
+
+ HVX_VectorPair vp_s_drow = Q6_W_vdeal_VVR(*pv_s_in1, *pv_s_in0, -2);
+ my_row_buf0[ci] = Q6_V_lo_W(vp_s_drow);
+ my_row_buf1[ci] = Q6_V_hi_W(vp_s_drow);
}
}
HVX_Vector v_p_row0_hf = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m0));
HVX_Vector v_p_row1_hf = hvx_vec_exp2_f16(Q6_Vhf_equals_Vqf16(v_s_minus_m1));
- __fp16 * out_dual_tile = p_st_base + (c / 64) * HMX_FP16_TILE_N_ELMS * 2;
- HVX_Vector * pv_p_out0 = ((HVX_Vector *) out_dual_tile) + r1 / 2;
- HVX_Vector * pv_p_out1 = pv_p_out0 + 16;
+ __fp16 * out_dtile = p_st_base + ci * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_p_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_p_out1 = pv_p_out0 + 16;
HVX_VectorPair vp_p_dual = Q6_W_vshuff_VVR(v_p_row1_hf, v_p_row0_hf, -2);
*pv_p_out0 = Q6_V_lo_W(vp_p_dual);
}
// Inline fa_ml_update_and_build_d for this vector (lock-free and in parallel)
- HVX_VectorPair rowmax_acc_pair = hvx_vec_f16_to_f32(rowmax_acc_v);
+ HVX_VectorPair rowmax_acc_pair = hvx_vec_f16_to_f32(rowmax_acc_v);
HVX_Vector v_rowmax_acc_f32_0 = Q6_V_lo_W(rowmax_acc_pair);
HVX_Vector v_rowmax_acc_f32_1 = Q6_V_hi_W(rowmax_acc_pair);
HVX_Vector v_m_diff0 = HVX_OP_SUB_F32(m_prev_v0, v_m_curr0);
HVX_Vector v_m_diff1 = HVX_OP_SUB_F32(m_prev_v1, v_m_curr1);
- HVX_Vector v_m_diff_f16 = hvx_vec_f32_to_f16(v_m_diff0, v_m_diff1);
+ HVX_Vector v_m_diff_f16 = hvx_vec_f32_to_f16(v_m_diff0, v_m_diff1);
HVX_Vector exp_m_diff_f16 = hvx_vec_exp2_f16(v_m_diff_f16);
HVX_VectorPair exp_m_diff_pair = hvx_vec_f16_to_f32(exp_m_diff_f16);
__builtin_assume(n_col_tiles > 0);
__builtin_assume(n_dot_tiles > 0);
- Q6_bias_mxmem2_A((void *) job->hmx_scales);
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
+ const size_t dot_stride = n_dot_tiles * HMX_FP16_TILE_N_ELMS;
for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < n_col_tiles; ++c) {
- const __fp16 * row_tiles = q_tiles + r * HMX_FP16_TILE_N_ROWS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
- const __fp16 * col_tiles = k_tiles + c * HMX_FP16_TILE_N_COLS * n_dot_tiles * HMX_FP16_TILE_N_COLS;
- __fp16 * out_tile = s_tiles + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * row_tiles = q_tiles + r * dot_stride;
+ const __fp16 * col_tiles = k_tiles;
+ __fp16 * out_tile = s_tiles + r * n_tiles_per_bc * HMX_FP16_TILE_N_ELMS;
+ for (size_t c = 0; c < n_col_tiles; ++c) {
hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles);
+ col_tiles += dot_stride;
+ out_tile += HMX_FP16_TILE_N_ELMS;
}
}
}
__builtin_assume(n_col_tiles > 0);
__builtin_assume(DV_tiles > 0);
- Q6_bias_mxmem2_A((void *) job->hmx_scales);
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
+ const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS;
+ const size_t v_stride = n_tiles_per_bc * HMX_FP16_TILE_N_ELMS;
for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- // D[r,r] @ O_prev[r,c] — only the diagonal tile
- const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * p_tile_in = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * v_tile_in = v_tiles + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- __fp16 * o_tile_out = o_curr + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * p_tile_in = p_tiles + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * o_rc = o_prev + r * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * v_tile_in = v_tiles;
+ __fp16 * o_tile_out = o_curr + r * HMX_FP16_TILE_N_ELMS;
+ for (size_t c = 0; c < DV_tiles; ++c) {
hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles);
+ o_rc += o_stride;
+ v_tile_in += v_stride;
+ o_tile_out += o_stride;
}
}
}
__builtin_assume(n_row_tiles > 0);
__builtin_assume(DV_tiles > 0);
- Q6_bias_mxmem2_A((void *) job->hmx_scales);
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)job->hmx_scales));
+ const size_t o_stride = n_row_tiles_g_br * HMX_FP16_TILE_N_ELMS;
for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = o_prev + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- __fp16 * o_out = o_curr + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * d_diag = d_tiles + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
+ const __fp16 * o_rc = o_prev + r * HMX_FP16_TILE_N_ELMS;
+ __fp16 * o_out = o_curr + r * DV_tiles * HMX_FP16_TILE_N_ELMS;
+ for (size_t c = 0; c < DV_tiles; ++c) {
hmx_fa_o_norm_tile(d_diag, o_rc, o_out);
+ o_rc += o_stride;
+ o_out += HMX_FP16_TILE_N_ELMS;
}
}
}
uint32_t G,
uint32_t m_line_bytes,
uint32_t kv_rows,
- uint32_t n_q_rows,
+ uint32_t n_rows_q,
struct hmx_fa_context * factx
) {
for (uint32_t g = 0; g < G; ++g) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] +
im2 * mask->nb[2] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
uint8_t * ms_dst = (uint8_t *) factx->vtcm_mask_buf + g * m_line_bytes;
- dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+ dma_queue_push(dma, dma_make_ptr(ms_dst, ms_src), G * m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
}
}
const uint32_t G = factx.G;
// ======== VTCM allocation (GQA-aware) ========
+ // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used
+ // throughout the KV loop below.
const size_t size_k_row = DK * sizeof(__fp16);
const size_t size_v_row = DV * sizeof(__fp16);
const size_t size_k_row_padded = hex_round_up(size_k_row, 128);
const size_t size_v_row_padded = hex_round_up(size_v_row, 128);
- const size_t q_tile_bytes = hex_align_up(g_br * DK * sizeof(__fp16), 4096);
- const size_t o_tile_bytes = hex_align_up(g_br * DV * sizeof(__fp16), 4096);
- const size_t k_dma_bytes = hex_align_up(Bc * size_k_row_padded, 4096);
- const size_t v_dma_bytes = hex_align_up(Bc * size_v_row_padded, 4096);
- const size_t k_tile_bytes = hex_align_up(Bc * DK * sizeof(__fp16), 4096);
- const size_t v_tile_bytes = hex_align_up(Bc * DV * sizeof(__fp16), 4096);
- const size_t s_tile_bytes = hex_align_up(g_br * Bc * sizeof(__fp16), 4096);
- const size_t d_tile_bytes = hex_align_up(g_br * g_br * sizeof(__fp16), 4096);
- const size_t col_vec_bytes = hex_align_up(g_br * sizeof(float), 256);
- const size_t row_vec_bytes = hex_align_up(Bc * sizeof(__fp16), 256);
- const size_t m_line_bytes = hex_align_up(Bc * sizeof(__fp16), 128);
- const size_t m_buf_bytes = hex_align_up(Br * m_line_bytes, 4096) * HMX_FA_DMA_CACHE_SIZE;
- const size_t slopes_bytes = hex_align_up(g_br * sizeof(__fp16), 128);
-
- uint8_t * vtcm_cur = ctx->vtcm_base;
-
- factx.vtcm_q_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, q_tile_bytes);
- factx.vtcm_o_tiles[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
- factx.vtcm_o_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, o_tile_bytes);
- factx.vtcm_k_fp16[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
- factx.vtcm_k_fp16[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_dma_bytes);
- factx.vtcm_v_fp16[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
- factx.vtcm_v_fp16[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_dma_bytes);
- factx.vtcm_k_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, k_tile_bytes);
- factx.vtcm_v_tiles[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
- if (pipeline) {
- factx.vtcm_v_tiles[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, v_tile_bytes);
- } else {
- factx.vtcm_v_tiles[1] = NULL;
- }
- factx.vtcm_s_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
- factx.vtcm_p_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, s_tile_bytes);
- factx.vtcm_d_tiles = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, d_tile_bytes);
- factx.vtcm_m_vec = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_l_vec = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_s_rowmax = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_p_rowsum = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, col_vec_bytes);
- factx.vtcm_row_bufs = (HVX_Vector *) vtcm_seq_alloc(&vtcm_cur, row_vec_bytes * 2 * n_threads);
- factx.row_buf_stride = row_vec_bytes / sizeof(HVX_Vector);
- factx.vtcm_hmx_scales_id = vtcm_seq_alloc(&vtcm_cur, 256);
- factx.vtcm_hmx_scales_qk = vtcm_seq_alloc(&vtcm_cur, 256);
- factx.vtcm_mask_buf = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, m_buf_bytes);
- factx.mask_buf_row_stride = m_line_bytes / sizeof(__fp16);
- factx.vtcm_slopes = (__fp16 *) vtcm_seq_alloc(&vtcm_cur, slopes_bytes);
-
- dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, hex_align_up(Br * m_line_bytes, 4096), HMX_FA_DMA_CACHE_SIZE);
-
- if ((size_t) (vtcm_cur - ctx->vtcm_base) > ctx->vtcm_size) {
+ // Build the VTCM layout once (shared with the host estimator) and place every
+ // scratch buffer at its computed offset.
+ struct hmx_fa_vtcm_layout L;
+ hmx_fa_vtcm_layout_build(&L, G, DK, DV, Br, Bc, n_threads, pipeline);
+
+ if (L.total_bytes > ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL;
}
+ uint8_t * const base = ctx->vtcm_base;
+
+ factx.vtcm_q_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_q_tiles);
+ factx.vtcm_o_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[0]);
+ factx.vtcm_o_tiles[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_o_tiles[1]);
+ factx.vtcm_k_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[0]);
+ factx.vtcm_k_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_fp16[1]);
+ factx.vtcm_v_fp16[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[0]);
+ factx.vtcm_v_fp16[1] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_fp16[1]);
+ factx.vtcm_k_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_k_tiles);
+ factx.vtcm_v_tiles[0] = VTCM_LAYOUT_PTR(__fp16, base, L.off_v_tiles[0]);
+ factx.vtcm_v_tiles[1] = VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_v_tiles[1], pipeline);
+ factx.vtcm_s_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_s_tiles);
+ factx.vtcm_p_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_p_tiles);
+ factx.vtcm_d_tiles = VTCM_LAYOUT_PTR(__fp16, base, L.off_d_tiles);
+ factx.vtcm_m_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_m_vec);
+ factx.vtcm_l_vec = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_l_vec);
+ factx.vtcm_s_rowmax = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_s_rowmax);
+ factx.vtcm_p_rowsum = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_p_rowsum);
+ factx.vtcm_row_bufs = VTCM_LAYOUT_PTR(HVX_Vector, base, L.off_row_bufs);
+ factx.row_buf_stride = L.row_buf_stride;
+ factx.vtcm_hmx_scales_id = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_id);
+ factx.vtcm_hmx_scales_qk = VTCM_LAYOUT_PTR(uint8_t, base, L.off_hmx_scales_qk);
+ factx.vtcm_mask_buf = VTCM_LAYOUT_PTR(__fp16, base, L.off_mask_buf);
+ factx.mask_buf_row_stride = L.mask_buf_row_stride;
+ factx.q_tile_bytes = L.q_tile_bytes;
+ factx.o_tile_bytes = L.o_tile_bytes;
+ factx.col_vec_bytes = L.col_vec_bytes;
+ factx.d_tile_bytes = L.d_tile_bytes;
+ factx.vtcm_slopes = VTCM_LAYOUT_PTR(__fp16, base, L.off_slopes);
+
+ const size_t m_line_bytes = L.m_line_bytes; // used by the mask DMAs in the KV loop
+
+ dma_cache_init(&factx.m_cache, (uint8_t *) factx.vtcm_mask_buf, L.m_buf_slot_bytes, HMX_FA_DMA_CACHE_SIZE);
+
// ======== Initialize HMX output scales ========
hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0
hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale));
const size_t qo_element_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
- // ======== HMX lock strategy ========
- if (!factx.pipeline) {
- HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
- }
-
// ======== Reusable job descriptors for pipeline ========
hmx_fa_qk_job_t qk_job;
hmx_fa_o_update_job_t ou_job;
for (uint32_t ib3 = 0; ib3 < neq3; ++ib3) {
const uint32_t im3 = mask ? fastmodulo(ib3, mask->ne[3], &factx.src3_div3) : 0;
for (uint32_t q_start = 0; q_start < neq1; q_start += Br) {
- const uint32_t n_q_rows = hex_smin(Br, neq1 - q_start);
- const size_t n_rows_g = n_q_rows * G;
+ const uint32_t n_rows_q = hex_smin(Br, neq1 - q_start);
+ const size_t n_rows_g = n_rows_q * G;
const size_t g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
const size_t n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
for (uint32_t kv_head = 0; kv_head < n_kv_heads; ++kv_head) {
const uint32_t ik2 = kv_head;
- const uint32_t ik3 = ib3 / (neq3 / k->ne[3]);
+ const uint32_t ik3 = fastdiv(ib3, &kparams->broadcast_rk3);
const uint32_t iv2 = kv_head;
- const uint32_t iv3 = ib3 / (neq3 / v->ne[3]);
+ const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);
+
+ // 1. Push Q DMA (if Q DMA is used)
+ const size_t o_tile_bytes = factx.o_tile_bytes;
+ const bool use_q_dma = (2 * o_tile_bytes >= factx.g_br * factx.DK * (factx.is_q_fp32 ? 4 : 2));
+ if (use_q_dma) {
+ const bool q_transposed = q->nb[1] < q->nb[2];
+ const uint8_t * q_ptr = (const uint8_t *) q->data + q_start * q->nb[1] + (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
+ const size_t el_size = factx.is_q_fp32 ? sizeof(float) : sizeof(__fp16);
+ const size_t q_row_bytes = q_transposed ? n_rows_q * factx.DK * el_size : factx.G * factx.DK * el_size;
+ const size_t src_stride = q_transposed ? q->nb[2] : q->nb[1];
+ const size_t n_rows = q_transposed ? factx.G : n_rows_q;
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_o_tiles[0], q_ptr), q_row_bytes, hex_smax(src_stride, q_row_bytes), q_row_bytes, n_rows);
+ }
- // Prefetch first KV block
+ // 2. Prefetch first KV block
if (factx.n_kv_blocks > 0) {
const uint32_t kv_rows0 = hex_smin(Bc, nek1);
const uint8_t * k_src = (const uint8_t *) k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1],
- size_k_row, kv_rows0);
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
const uint8_t * v_src = (const uint8_t *) v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1],
- size_v_row, kv_rows0);
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+ }
+
+ // 3. Pop Q DMA (blocks until Q is loaded)
+ if (use_q_dma) {
+ dma_queue_pop(dma);
}
// ---- Load Q block & Initialize per-block state ----
const size_t k_src_stride = size_k_row_padded / sizeof(__fp16);
const size_t v_src_stride = size_v_row_padded / sizeof(__fp16);
+ struct hmx_queue * hmx_q = ctx->hmx_queue;
+
if (factx.pipeline) {
- // ==================================================================
// Pipeline path
- // ==================================================================
- struct hmx_queue * hmx_q = ctx->hmx_queue;
-
for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
const uint32_t kv_start = kv_blk * Bc;
const uint32_t kv_rows = hex_smin(Bc, nek1 - kv_start);
if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
- dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+ dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
} else {
- fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx);
+ fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
}
}
- // Wait for current KV DMA
- dma_queue_pop(dma); // K
- dma_queue_pop(dma); // V
+ // Prefetch next KV block early
+ if (kv_blk + 1 < factx.n_kv_blocks) {
+ const uint32_t prefetch_start = (kv_blk + 1) * Bc;
+ const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
+ const size_t prefetch_buf = 1 - buf_idx;
+ const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
+ const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+ dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
+ }
// ---- Phase 1: K_int ----
if (kv_blk > 0) {
ou_job.DV = DV;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
}
- fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start);
+
+ // Wait for current K DMA and interleave
+ void * curr_k = dma_queue_pop(dma).dst;
+ fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
// ---- Phase 2: qk_dot ----
qk_job.q_tiles = factx.vtcm_q_tiles;
qk_job.hmx_scales = factx.vtcm_hmx_scales_qk;
hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
- if (kv_blk + 1 < factx.n_kv_blocks) {
- const uint32_t prefetch_start = (kv_blk + 1) * Bc;
- const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
- const size_t prefetch_buf = 1 - buf_idx;
- const uint8_t * k_prefetch_src = (const uint8_t *) k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_k_fp16[prefetch_buf], k_prefetch_src), size_k_row_padded, k->nb[1], size_k_row, prefetch_rows);
- const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
- dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
- }
- fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc, kv_start);
+ // Wait for current V DMA and interleave
+ void * curr_v = dma_queue_pop(dma).dst;
+ fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[buf_idx], n_tiles_per_bc, kv_start);
if (kv_blk > 0) {
hmx_queue_pop(hmx_q);
}
} else {
- // ==================================================================
// Fallback path
- // ==================================================================
for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
const uint32_t kv_start = kv_blk * Bc;
const uint32_t kv_rows = hex_smin(Bc, nek1 - kv_start);
const size_t n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
- dma_queue_pop(dma); // K
- dma_queue_pop(dma); // V
if (mask) {
if (__builtin_expect(factx.mask_broadcast, true)) {
const uint8_t * ms_src = (const uint8_t *) mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + kv_start * sizeof(__fp16);
- dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_q_rows);
+ dma_cache_push(dma, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows * sizeof(__fp16), n_rows_q);
} else {
- fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_q_rows, &factx);
+ fa_push_mask_dma_gqa(dma, mask, q_start, im3, kv_start, kv_head, G, m_line_bytes, kv_rows, n_rows_q, &factx);
}
}
+
if (kv_blk + 1 < factx.n_kv_blocks) {
const uint32_t prefetch_start = (kv_blk + 1) * Bc;
const uint32_t prefetch_rows = hex_smin(Bc, nek1 - prefetch_start);
const uint8_t * v_prefetch_src = (const uint8_t *) v->data + prefetch_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
dma_queue_push(dma, dma_make_ptr(factx.vtcm_v_fp16[prefetch_buf], v_prefetch_src), size_v_row_padded, v->nb[1], size_v_row, prefetch_rows);
}
- fa_phase_k_interleave(&factx, kv_rows, k_src_stride, buf_idx, kv_start);
+
+ // Wait for current K DMA and interleave
+ void * curr_k = dma_queue_pop(dma).dst;
+ fa_phase_k_interleave(&factx, kv_rows, k_src_stride, curr_k, kv_start);
{
- const size_t n_dot_tiles = (size_t) (DK / 32);
- const __fp16 * restrict q_base = factx.vtcm_q_tiles;
- const __fp16 * restrict k_base = factx.vtcm_k_tiles;
- __fp16 * restrict s_base = factx.vtcm_s_tiles;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(n_col_tiles > 0);
- __builtin_assume(n_dot_tiles > 0);
-
- htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
- Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_qk);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < n_col_tiles; ++c) {
- const __fp16 * row_tiles = q_base + r * HMX_FP16_TILE_N_ROWS * DK;
- const __fp16 * col_tiles = k_base + c * HMX_FP16_TILE_N_COLS * DK;
- __fp16 * out_tile = s_base + (r * n_tiles_per_bc + c) * HMX_FP16_TILE_N_ELMS;
-
- hmx_fa_qk_dot_tile(row_tiles, col_tiles, out_tile, n_dot_tiles);
- }
- }
- htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ qk_job.q_tiles = factx.vtcm_q_tiles;
+ qk_job.k_tiles = factx.vtcm_k_tiles;
+ qk_job.s_tiles = factx.vtcm_s_tiles;
+ qk_job.n_row_tiles = n_row_tiles;
+ qk_job.n_col_tiles = n_col_tiles;
+ qk_job.n_dot_tiles = (size_t) (DK / 32);
+ qk_job.n_tiles_per_bc = n_tiles_per_bc;
+ qk_job.hmx_scales = factx.vtcm_hmx_scales_qk;
+
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job));
+ hmx_queue_pop(ctx->hmx_queue);
}
+ // Wait for current V DMA and interleave
+ void * curr_v = dma_queue_pop(dma).dst;
+ fa_phase_v_interleave(&factx, kv_rows, v_src_stride, curr_v, factx.vtcm_v_tiles[0], n_tiles_per_bc, kv_start);
+
// ---- Phase 3: softmax + build_D ----
__fp16 * current_mask_vtcm = NULL;
if (mask) {
sargs.mask_vtcm_row_stride = factx.mask_buf_row_stride;
sargs.slopes = factx.vtcm_slopes;
fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
- fa_phase_v_interleave(&factx, kv_rows, v_src_stride, buf_idx, n_tiles_per_bc, kv_start);
{
- const size_t DV_tiles = (size_t) (DV / 32);
- const __fp16 * restrict d_base = factx.vtcm_d_tiles;
- const __fp16 * restrict p_base = factx.vtcm_p_tiles;
- const __fp16 * restrict v_base = factx.vtcm_v_tiles[0];
- const __fp16 * restrict op_base = o_tile_prev;
- __fp16 * restrict oc_base = o_tile_curr;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(n_col_tiles > 0);
- __builtin_assume(DV_tiles > 0);
-
- htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
- Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * p_tile_in = p_base + (r * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * v_tile_in = v_base + (c * n_tiles_per_bc) * HMX_FP16_TILE_N_ELMS;
- __fp16 * o_tile_out = oc_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
-
- hmx_fa_o_update_tile(d_diag, o_rc, p_tile_in, v_tile_in, o_tile_out, n_col_tiles);
- }
- }
- htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
+ ou_job.o_curr = o_tile_curr;
+ ou_job.o_prev = o_tile_prev;
+ ou_job.p_tiles = factx.vtcm_p_tiles;
+ ou_job.v_tiles = factx.vtcm_v_tiles[0];
+ ou_job.d_tiles = factx.vtcm_d_tiles;
+ ou_job.hmx_scales = factx.vtcm_hmx_scales_id;
+ ou_job.n_row_tiles = n_row_tiles;
+ ou_job.n_col_tiles = n_col_tiles;
+ ou_job.n_row_tiles_g_br = n_row_tiles_g_br;
+ ou_job.n_tiles_per_bc = n_tiles_per_bc;
+ ou_job.DV = DV;
+
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
+ hmx_queue_pop(ctx->hmx_queue);
+
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
}
fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br);
htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
- if (factx.pipeline) {
- on_job.o_curr = o_tile_curr;
- on_job.o_prev = o_tile_prev;
- on_job.d_tiles = factx.vtcm_d_tiles;
- on_job.hmx_scales = factx.vtcm_hmx_scales_id;
- on_job.n_row_tiles = n_row_tiles;
- on_job.n_row_tiles_g_br = n_row_tiles_g_br;
- on_job.DV = DV;
- hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
- hmx_queue_pop(ctx->hmx_queue);
- } else {
- const size_t DV_tiles = (size_t) (DV / 32);
- const __fp16 * restrict d_base = factx.vtcm_d_tiles;
- const __fp16 * restrict op_base = o_tile_prev;
- __fp16 * restrict oc_base = o_tile_curr;
- __builtin_assume(n_row_tiles > 0);
- __builtin_assume(DV_tiles > 0);
-
- htp_trace_event_start(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
- Q6_bias_mxmem2_A((void *) factx.vtcm_hmx_scales_id);
- for (size_t r = 0; r < n_row_tiles; ++r) {
- for (size_t c = 0; c < DV_tiles; ++c) {
- const __fp16 * d_diag = d_base + r * (n_row_tiles_g_br + 1) * HMX_FP16_TILE_N_ELMS;
- const __fp16 * o_rc = op_base + (c * n_row_tiles_g_br + r) * HMX_FP16_TILE_N_ELMS;
- __fp16 * o_out = oc_base + (r * DV_tiles + c) * HMX_FP16_TILE_N_ELMS;
-
- hmx_fa_o_norm_tile(d_diag, o_rc, o_out);
- }
- }
- htp_trace_event_stop(tr_hmx, HTP_TRACE_EVT_HMX_COMP, (uint16_t) q_start);
- }
+ on_job.o_curr = o_tile_curr;
+ on_job.o_prev = o_tile_prev;
+ on_job.d_tiles = factx.vtcm_d_tiles;
+ on_job.hmx_scales = factx.vtcm_hmx_scales_id;
+ on_job.n_row_tiles = n_row_tiles;
+ on_job.n_row_tiles_g_br = n_row_tiles_g_br;
+ on_job.DV = DV;
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_norm_worker, &on_job));
+ hmx_queue_pop(ctx->hmx_queue);
}
// ---- Store O block ----
}
}
- if (factx.pipeline) {
- hmx_queue_suspend(ctx->hmx_queue);
- } else {
- HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
- }
-
return HTP_STATUS_OK;
}
factx.src0_div21 = kparams->u.hvx.src0_div21;
factx.src0_div1 = kparams->u.hvx.src0_div1;
- factx.broadcast_rk2 = kparams->u.hvx.broadcast_rk2;
- factx.broadcast_rk3 = kparams->u.hvx.broadcast_rk3;
- factx.broadcast_rv2 = kparams->u.hvx.broadcast_rv2;
- factx.broadcast_rv3 = kparams->u.hvx.broadcast_rv3;
+ factx.broadcast_rk2 = kparams->broadcast_rk2;
+ factx.broadcast_rk3 = kparams->broadcast_rk3;
+ factx.broadcast_rv2 = kparams->broadcast_rv2;
+ factx.broadcast_rv3 = kparams->broadcast_rv3;
if (mask) {
factx.src3_div2 = kparams->src3_div2;
#include "hex-fastdiv.h"
#include "hex-common.h"
+#include "htp-vtcm.h"
#ifdef __cplusplus
extern "C" {
#endif
// Tile constants (mirrored from hmx-utils.h for use on host side if needed)
+#define HTP_FA_HMX_TILE_SIZE 2048
+#define HMX_FP16_TILE_SIZE 2048
#define HMX_FP16_TILE_N_ROWS 32
#define HMX_FP16_TILE_N_COLS 32
#define HMX_FP16_TILE_N_ELMS 1024
-#define HMX_FP16_TILE_SIZE 2048
+
#define HVX_FA_DMA_CACHE_SIZE 128
#define HMX_FA_DMA_CACHE_SIZE 4
+
#define HTP_FA_M_INITIAL_VAL -10000.0f
enum htp_fa_kernel_type {
struct fastdiv_values src3_div2;
struct fastdiv_values src3_div3;
+ struct fastdiv_values broadcast_rk2;
+ struct fastdiv_values broadcast_rk3;
+ struct fastdiv_values broadcast_rv2;
+ struct fastdiv_values broadcast_rv3;
+
union {
struct {
uint32_t g_br;
uint32_t size_v_row_padded;
struct fastdiv_values src0_div21;
struct fastdiv_values src0_div1;
- struct fastdiv_values broadcast_rk2;
- struct fastdiv_values broadcast_rk3;
- struct fastdiv_values broadcast_rv2;
- struct fastdiv_values broadcast_rv3;
} hvx;
} u;
};
static_assert(sizeof(struct htp_fa_kernel_params) <= 128, "htp_fa_kernel_params is too large for kernel_params blob");
#endif
-// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
-// g_br = hex_align_up(gqa_factor * Br, 32) replaces Br for all Q/O/S/P/D dimensions.
-// Layout: Q + O_ping + O_pong + K_dma*2 + V_dma*2 + K_tile + V_tile + S + P + D + vectors + scales
-// Mask is DMA'd into a VTCM buffer (Br rows per KV block) to avoid DDR reads in softmax.
-static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+// VTCM region layout for the HMX flash-attention kernel.
+//
+// Single source of truth for both the host (which needs the total size to pick a
+// (Br, Bc) tiling that fits the VTCM budget) and the device (which needs the actual
+// byte offsets to place each scratch buffer). Building the layout once and reading
+// offsets/total from it makes host estimate and device allocation impossible to
+// desync -- previously they were duplicated formulas in two files and drifted.
+//
+// All fields are byte offsets / byte sizes -- no HVX_Vector type is named here so the
+// header stays host-includable. The device casts (base + off_*) to the proper type.
+// An offset of 0 marks a region that is not allocated for this configuration (only
+// off_v_tiles[1], which exists only when pipelining); the device sets such pointers NULL.
+struct hmx_fa_vtcm_layout {
+ // Byte offsets from vtcm_base for each region.
+ size_t off_q_tiles;
+ size_t off_o_tiles[2];
+ size_t off_k_fp16[2];
+ size_t off_v_fp16[2];
+ size_t off_k_tiles;
+ size_t off_v_tiles[2]; // [1] allocated only when pipeline, else 0
+ size_t off_s_tiles;
+ size_t off_p_tiles;
+ size_t off_d_tiles;
+ size_t off_m_vec;
+ size_t off_l_vec;
+ size_t off_s_rowmax;
+ size_t off_p_rowsum;
+ size_t off_row_bufs;
+ size_t off_hmx_scales_id;
+ size_t off_hmx_scales_qk;
+ size_t off_mask_buf;
+ size_t off_slopes;
+
+ // Region byte sizes reused by the device at runtime (not just for allocation).
+ size_t q_tile_bytes;
+ size_t o_tile_bytes;
+ size_t s_tile_bytes; // S and P tiles (same size)
+ size_t d_tile_bytes;
+ size_t m_line_bytes; // one mask row
+ size_t m_buf_slot_bytes; // one dma_cache slot = align_up(Br * m_line_bytes, 4096)
+ size_t col_vec_bytes;
+
+ // Derived strides.
+ size_t row_buf_stride; // HVX vectors (128B) per row buffer
+ size_t mask_buf_row_stride; // __fp16 elements per row in the mask buffer
+
+ bool pipeline;
+ size_t total_bytes;
+};
+
+// Build the VTCM layout.
+
+static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
+ size_t gqa_factor, size_t DK, size_t DV,
+ size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
const size_t g_br = hex_align_up(gqa_factor * Br, HMX_FP16_TILE_N_ROWS);
- const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), 4096); // Q: [g_br, DK]
- const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), 4096); // O: [g_br, DV] x2 ping-pong
- const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 4096); // K DMA: [Bc, DK] x2 double-buf
- const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 4096); // V DMA: [Bc, DV] x2 double-buf
- const size_t k_tile_size = hex_align_up(Bc * DK * sizeof(__fp16), 4096); // K tiles: [Bc, DK] interleaved
- const size_t v_tile_size = hex_align_up(Bc * DV * sizeof(__fp16), 4096); // V tiles: [Bc, DV] interleaved
- const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), 4096); // S/P:[g_br, Bc]
- const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), 4096); // D: [g_br, g_br]
- const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256); // m, l, etc.
- const size_t row_vec_size = hex_align_up(Bc * sizeof(__fp16), 256);
- const size_t m_line_size = hex_align_up(Bc * sizeof(__fp16), 128);
- const size_t m_buf_size = hex_align_up(Br * m_line_size, 4096) * HMX_FA_DMA_CACHE_SIZE;
+ const size_t q_tile_size = hex_align_up(g_br * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+ const size_t o_tile_size = hex_align_up(g_br * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+ const size_t k_tile_size = hex_align_up(Bc * DK * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+ const size_t v_tile_size = hex_align_up(Bc * DV * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+ const size_t s_tile_size = hex_align_up(g_br * Bc * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+ const size_t d_tile_size = hex_align_up(g_br * g_br * sizeof(__fp16), HTP_FA_HMX_TILE_SIZE);
+
+ const size_t k_dma_size = hex_align_up(Bc * hex_round_up(DK * sizeof(__fp16), 128), 128);
+ const size_t v_dma_size = hex_align_up(Bc * hex_round_up(DV * sizeof(__fp16), 128), 128);
+ const size_t col_vec_size = hex_align_up(g_br * sizeof(float), 256);
+ const size_t row_vec_size = hex_align_up(Bc * sizeof(__fp16), 256);
+ const size_t m_line_size = hex_align_up(Bc * sizeof(__fp16), 128);
+ const size_t m_buf_slot = hex_align_up(Br * m_line_size, 256);
+ const size_t m_buf_size = m_buf_slot * HMX_FA_DMA_CACHE_SIZE;
const size_t slopes_size = hex_align_up(g_br * sizeof(__fp16), 128);
- return q_tile_size * 1 // Q tiles
- + o_tile_size * 2 // O ping-pong
- + k_dma_size * 2 // K DMA x2
- + v_dma_size * 2 // V DMA x2
- + k_tile_size * 1 // K tiles
- + v_tile_size * (pipeline ? 2 : 1) // V tiles (double-buffered if pipelining)
- + s_tile_size * 2 // S + P
- + d_tile_size * 1 // D (diagonal matrix)
- + col_vec_size * 4 // m_vec, l_vec, s_rowmax, p_rowsum
- + row_vec_size * 2 * n_threads // per-thread softmax row scratch
- + m_buf_size * 1 // mask VTCM buffer [Br rows]
- + slopes_size // Slopes
- + 256 * 2; // HMX scales (id + qk)
+ size_t off = 0;
+
+ // Section 1: HMX Tiled Buffers (FA_HMX_TILE_SIZE = 2KB Aligned)
+ VTCM_LAYOUT_ALLOC(off, off_q_tiles, q_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_o_tiles[0], o_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_o_tiles[1], o_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_k_tiles, k_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_v_tiles[0], v_tile_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_v_tiles[1], v_tile_size, pipeline);
+ VTCM_LAYOUT_ALLOC(off, off_s_tiles, s_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_p_tiles, s_tile_size);
+ VTCM_LAYOUT_ALLOC(off, off_d_tiles, d_tile_size);
+
+ // Section 2: HVX/DMA flat and vector buffers (128B / 256B Aligned)
+ VTCM_LAYOUT_ALLOC(off, off_k_fp16[0], k_dma_size);
+ VTCM_LAYOUT_ALLOC(off, off_k_fp16[1], k_dma_size);
+ VTCM_LAYOUT_ALLOC(off, off_v_fp16[0], v_dma_size);
+ VTCM_LAYOUT_ALLOC(off, off_v_fp16[1], v_dma_size);
+ VTCM_LAYOUT_ALLOC(off, off_m_vec, col_vec_size);
+ VTCM_LAYOUT_ALLOC(off, off_l_vec, col_vec_size);
+ VTCM_LAYOUT_ALLOC(off, off_s_rowmax, col_vec_size);
+ VTCM_LAYOUT_ALLOC(off, off_p_rowsum, col_vec_size);
+ VTCM_LAYOUT_ALLOC(off, off_row_bufs, row_vec_size * 2 * n_threads);
+ VTCM_LAYOUT_ALLOC(off, off_hmx_scales_id, 256);
+ VTCM_LAYOUT_ALLOC(off, off_hmx_scales_qk, 256);
+ VTCM_LAYOUT_ALLOC(off, off_mask_buf, m_buf_size);
+ VTCM_LAYOUT_ALLOC(off, off_slopes, slopes_size);
+
+ L->q_tile_bytes = q_tile_size;
+ L->o_tile_bytes = o_tile_size;
+ L->col_vec_bytes = col_vec_size;
+ L->s_tile_bytes = s_tile_size;
+ L->d_tile_bytes = d_tile_size;
+ L->m_line_bytes = m_line_size;
+ L->m_buf_slot_bytes = m_buf_slot;
+ L->row_buf_stride = row_vec_size / 128;
+ L->mask_buf_row_stride = m_line_size / sizeof(__fp16);
+ L->pipeline = pipeline;
+ L->total_bytes = off;
+}
+
+// Exact VTCM usage for a given (gqa_factor, DK, DV, Br, Bc) configuration.
+static inline size_t hmx_fa_compute_vtcm_usage(size_t gqa_factor, size_t DK, size_t DV, size_t Br, size_t Bc, size_t n_threads, bool pipeline) {
+ struct hmx_fa_vtcm_layout L;
+ hmx_fa_vtcm_layout_build(&L, gqa_factor, DK, DV, Br, Bc, n_threads, pipeline);
+ return L.total_bytes;
}
#define FA_HVX_BLOCK_SIZE 64
const size_t T = HMX_FP16_TILE_N_ROWS; // 32
const size_t br_unit = hmx_ceil_div(T, gqa_factor);
const size_t bc_unit = HMX_FP16_TILE_N_COLS * 2; // 64
- const size_t fp16 = sizeof(__fp16);
const bool can_pipeline = (kv_len >= FA_MIN_KV_BLOCKS * bc_unit && n_threads >= 2);
- // Approximate per-unit VTCM costs (without per-buffer alignment padding).
- const size_t per_gbr = (DK + 2 * DV) * fp16 + 4 * sizeof(float); // Q + O*2 + 4 col vectors
- const size_t per_gbr2 = fp16; // D diagonal matrix
- const size_t per_bc =
- 3 * DK * fp16 + (can_pipeline ? 4 : 3) * DV * fp16 + 2 * n_threads * fp16; // K/V DMA x2 + tiles + row bufs
- const size_t per_gbr_bc = 2 * fp16; // S + P
-
- const size_t overhead = 256 * 2 + 13 * 4096;
-
- if (vtcm_budget <= overhead) {
- return -1;
- }
- const size_t usable = vtcm_budget - overhead;
-
// Br_max: largest Br aligned to br_unit that does not exceed qo_len.
const size_t Br_max = qo_len >= br_unit ? hex_align_down(qo_len, br_unit) : br_unit;
size_t best_Br = 0, best_Bc = 0;
for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
- const size_t g_br = hex_align_up(gqa_factor * Br, T);
-
- // g_br-dependent VTCM cost: g_br * per_gbr + g_br*g_br * per_gbr2
- const size_t gbr_cost = g_br * per_gbr + g_br * g_br * per_gbr2;
- if (gbr_cost >= usable) {
- if (Br == br_unit) {
- break;
- }
- continue;
- }
-
- // Analytically solve for max Bc:
- // remain >= Bc * (per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE)
- // The Br * fp16 term accounts for the VTCM mask buffer [Br * Bc].
- const size_t remain = usable - gbr_cost;
- const size_t bc_denom = per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE;
- size_t Bc = hex_smin(hex_align_down(remain / bc_denom, bc_unit), Bc_limit);
- if (Bc < bc_unit) {
- if (Br == br_unit) {
+ // Try all Bc candidates from Bc_limit down to bc_unit
+ for (size_t Bc = Bc_limit; Bc >= bc_unit; Bc -= bc_unit) {
+ size_t vtcm_needed = hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline);
+ if (vtcm_needed <= vtcm_budget) {
+ // This Bc fits for this Br!
+ const size_t q_blocks = (qo_len + Br - 1) / Br;
+ const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
+ const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
+ const size_t mn = Br * Bc;
+
+ if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
+ best_cost = cost;
+ best_mn = mn;
+ best_Br = Br;
+ best_Bc = Bc;
+ }
+ // Since we iterate Bc from largest to smallest, this is the largest Bc that fits
+ // for this Br. We can break to the next Br.
break;
}
- continue;
- }
-
- // Exact VTCM verification (alignment padding may push over budget)
- while (Bc >= bc_unit && hmx_fa_compute_vtcm_usage(gqa_factor, DK, DV, Br, Bc, n_threads, can_pipeline) > vtcm_budget) {
- Bc -= bc_unit;
- }
- if (Bc < bc_unit) {
- if (Br == br_unit) {
- break;
- }
- continue;
- }
-
- const size_t q_blocks = (qo_len + Br - 1) / Br;
- const size_t kv_blocks = (kv_len + Bc - 1) / Bc;
- const size_t cost = q_blocks * (c_q_fixed + kv_blocks * c_iter_fixed);
- const size_t mn = Br * Bc;
-
- if (cost < best_cost || (cost == best_cost && mn > best_mn)) {
- best_cost = cost;
- best_mn = mn;
- best_Br = Br;
- best_Bc = Bc;
}
if (Br == br_unit) {
}
}
- if (best_Br == 0) {
+ if (best_Br == 0 || best_Bc == 0) {
return -1;
}
#include <stdbool.h>
#include "hvx-utils.h"
#include "hmx-utils.h"
+#include "hex-fastdiv.h"
// HMX-specific parameters, offsets and inner kernels for Flash Attention
};
// Inner HMX tile computation kernels
-static inline void hmx_fa_qk_dot_tile(
+static void hmx_fa_qk_dot_tile(
const __fp16 * row_tiles,
const __fp16 * col_tiles,
__fp16 * out_tile,
size_t n_dot_tiles
) {
- for (size_t k = 0; k < n_dot_tiles; ++k) {
- Q6_activation_hf_mxmem_RR((unsigned int) row_tiles, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) col_tiles, 2047);
- row_tiles += HMX_FP16_TILE_N_ELMS;
- col_tiles += HMX_FP16_TILE_N_ELMS;
+ if (n_dot_tiles == 2) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ HMX_LOAD_MPY_F16("%3", "%4", "%0")
+ :
+ : "r"(2047),
+ "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS)
+ );
+ } else if (n_dot_tiles == 4) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ HMX_LOAD_MPY_F16("%3", "%4", "%0")
+ HMX_LOAD_MPY_F16("%5", "%6", "%0")
+ HMX_LOAD_MPY_F16("%7", "%8", "%0")
+ :
+ : "r"(2047),
+ "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS)
+ );
+ } else if (n_dot_tiles == 8) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ HMX_LOAD_MPY_F16("%3", "%4", "%0")
+ HMX_LOAD_MPY_F16("%5", "%6", "%0")
+ HMX_LOAD_MPY_F16("%7", "%8", "%0")
+ HMX_LOAD_MPY_F16("%9", "%10", "%0")
+ HMX_LOAD_MPY_F16("%11", "%12", "%0")
+ HMX_LOAD_MPY_F16("%13", "%14", "%0")
+ HMX_LOAD_MPY_F16("%15", "%16", "%0")
+ :
+ : "r"(2047),
+ "r"(row_tiles + 0 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 0 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 1 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 1 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 2 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 2 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 3 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 3 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 4 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 4 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 5 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 5 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 6 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 6 * HMX_FP16_TILE_N_ELMS),
+ "r"(row_tiles + 7 * HMX_FP16_TILE_N_ELMS), "r"(col_tiles + 7 * HMX_FP16_TILE_N_ELMS)
+ );
+ } else {
+ for (size_t k = 0; k < n_dot_tiles; ++k) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ :
+ : "r"(2047), "r"(row_tiles), "r"(col_tiles)
+ );
+ row_tiles += HMX_FP16_TILE_N_ELMS;
+ col_tiles += HMX_FP16_TILE_N_ELMS;
+ }
}
- Q6_mxmem_AR_after_hf(out_tile, 0);
+ asm volatile(
+ HMX_STORE_AFTER_F16("%0", "%1")
+ :
+ : "r"(out_tile), "r"(0)
+ : "memory"
+ );
}
-static inline void hmx_fa_o_update_tile(
+static void hmx_fa_o_update_tile(
const __fp16 * d_diag,
const __fp16 * o_rc,
const __fp16 * p_tile_in,
__fp16 * o_tile_out,
size_t n_col_tiles
) {
- Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
-
- for (size_t k = 0; k < n_col_tiles; ++k) {
- Q6_activation_hf_mxmem_RR((unsigned int) p_tile_in, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) v_tile_in, 2047);
- p_tile_in += HMX_FP16_TILE_N_ELMS;
- v_tile_in += HMX_FP16_TILE_N_ELMS;
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ :
+ : "r"(2047), "r"(d_diag), "r"(o_rc)
+ );
+ if (n_col_tiles == 2) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ HMX_LOAD_MPY_F16("%3", "%4", "%0")
+ :
+ : "r"(2047),
+ "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS)
+ );
+ } else if (n_col_tiles == 4) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ HMX_LOAD_MPY_F16("%3", "%4", "%0")
+ HMX_LOAD_MPY_F16("%5", "%6", "%0")
+ HMX_LOAD_MPY_F16("%7", "%8", "%0")
+ :
+ : "r"(2047),
+ "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS)
+ );
+ } else if (n_col_tiles == 8) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ HMX_LOAD_MPY_F16("%3", "%4", "%0")
+ HMX_LOAD_MPY_F16("%5", "%6", "%0")
+ HMX_LOAD_MPY_F16("%7", "%8", "%0")
+ HMX_LOAD_MPY_F16("%9", "%10", "%0")
+ HMX_LOAD_MPY_F16("%11", "%12", "%0")
+ HMX_LOAD_MPY_F16("%13", "%14", "%0")
+ HMX_LOAD_MPY_F16("%15", "%16", "%0")
+ :
+ : "r"(2047),
+ "r"(p_tile_in + 0 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 0 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 1 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 1 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 2 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 2 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 3 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 3 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 4 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 4 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 5 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 5 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 6 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 6 * HMX_FP16_TILE_N_ELMS),
+ "r"(p_tile_in + 7 * HMX_FP16_TILE_N_ELMS), "r"(v_tile_in + 7 * HMX_FP16_TILE_N_ELMS)
+ );
+ } else {
+ for (size_t k = 0; k < n_col_tiles; ++k) {
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ :
+ : "r"(2047), "r"(p_tile_in), "r"(v_tile_in)
+ );
+ p_tile_in += HMX_FP16_TILE_N_ELMS;
+ v_tile_in += HMX_FP16_TILE_N_ELMS;
+ }
}
-
- Q6_mxmem_AR_after_hf(o_tile_out, 0);
+ asm volatile(
+ HMX_STORE_AFTER_F16("%0", "%1")
+ :
+ : "r"(o_tile_out), "r"(0)
+ : "memory"
+ );
}
static inline void hmx_fa_o_norm_tile(
const __fp16 * o_rc,
__fp16 * o_out
) {
- Q6_activation_hf_mxmem_RR((unsigned int) d_diag, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int) o_rc, 2047);
- Q6_mxmem_AR_after_hf(o_out, 0);
+ asm volatile(
+ HMX_LOAD_MPY_F16("%1", "%2", "%0")
+ :
+ : "r"(2047), "r"(d_diag), "r"(o_rc)
+ );
+ asm volatile(
+ HMX_STORE_AFTER_F16("%0", "%1")
+ :
+ : "r"(o_out), "r"(0)
+ : "memory"
+ );
+}
+
+static inline void hmx_fa_q_prep_fp32_d2(
+ __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+ size_t start, size_t end, size_t g_rows_end,
+ size_t DK, size_t G, size_t n_rows_q,
+ const struct fastdiv_values * div_G, bool q_transposed
+) {
+ for (size_t r = start; r < end; r += 2) {
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (r >= g_rows_end) {
+ ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+ ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+ continue;
+ }
+
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+ const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+ const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float));
+ const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+ ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float))
+ : NULL;
+
+ {
+ HVX_Vector v0 = pv_in0[0];
+ HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero();
+ HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+ ((HVX_Vector *) (out_base + 0 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf;
+ }
+ {
+ HVX_Vector v0 = pv_in0[1];
+ HVX_Vector v1 = pv_in1 ? pv_in1[1] : Q6_V_vzero();
+ HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+ ((HVX_Vector *) (out_base + 1 * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf;
+ }
+ }
+}
+
+static inline void hmx_fa_q_prep_fp32_d4(
+ __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+ size_t start, size_t end, size_t g_rows_end,
+ size_t DK, size_t G, size_t n_rows_q,
+ const struct fastdiv_values * div_G, bool q_transposed
+) {
+ for (size_t r = start; r < end; r += 2) {
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (r >= g_rows_end) {
+ for (uint32_t d = 0; d < 4; ++d) {
+ ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+ }
+ continue;
+ }
+
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+ const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+ const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float));
+ const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+ ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float))
+ : NULL;
+
+ for (uint32_t d = 0; d < 4; ++d) {
+ HVX_Vector v0 = pv_in0[d];
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+ ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = v_hf;
+ }
+ }
+}
+
+static inline void hmx_fa_q_prep_fp32(
+ __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+ size_t start, size_t end, size_t g_rows_end,
+ size_t DK, size_t G, size_t n_rows_q,
+ const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed
+) {
+ for (size_t r = start; r < end; r += 2) {
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (r >= g_rows_end) {
+ for (uint32_t d = 0; d < d_limit; ++d) {
+ ((HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS))[r1 / 2] = Q6_V_vzero();
+ }
+ continue;
+ }
+
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+ const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+ const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(float));
+ const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+ ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(float))
+ : NULL;
+
+ for (uint32_t d = 0; d < d_limit; ++d) {
+ HVX_Vector v0 = pv_in0[d];
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+ HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+ out_tile[r1 / 2] = v_hf;
+ }
+ }
+}
+
+static inline void hmx_fa_q_prep_fp16_d1(
+ __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+ size_t start, size_t end, size_t g_rows_end,
+ size_t DK, size_t G, size_t n_rows_q,
+ const struct fastdiv_values * div_G, bool q_transposed
+) {
+ for (size_t r = start; r < end; r += 2) {
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (r >= g_rows_end) {
+ __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+ *pv_out0 = Q6_V_vzero();
+ *pv_out1 = Q6_V_vzero();
+ continue;
+ }
+
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+ const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+ const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16));
+ const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+ ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16))
+ : NULL;
+
+ HVX_Vector v0 = pv_in0[0];
+ HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero();
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+ __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+
+ *pv_out0 = Q6_V_lo_W(vp);
+ *pv_out1 = Q6_V_hi_W(vp);
+ }
+}
+
+static inline void hmx_fa_q_prep_fp16_d2(
+ __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+ size_t start, size_t end, size_t g_rows_end,
+ size_t DK, size_t G, size_t n_rows_q,
+ const struct fastdiv_values * div_G, bool q_transposed
+) {
+ for (size_t r = start; r < end; r += 2) {
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (r >= g_rows_end) {
+ for (uint32_t d = 0; d < 2; ++d) {
+ __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+ *pv_out0 = Q6_V_vzero();
+ *pv_out1 = Q6_V_vzero();
+ }
+ continue;
+ }
+
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+ const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+ const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16));
+ const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+ ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16))
+ : NULL;
+
+ {
+ HVX_Vector v0 = pv_in0[0];
+ HVX_Vector v1 = pv_in1 ? pv_in1[0] : Q6_V_vzero();
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+ __fp16 * out_dtile = out_base + 0 * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+
+ *pv_out0 = Q6_V_lo_W(vp);
+ *pv_out1 = Q6_V_hi_W(vp);
+ }
+ {
+ HVX_Vector v0 = pv_in0[1];
+ HVX_Vector v1 = pv_in1 ? pv_in1[1] : Q6_V_vzero();
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+ __fp16 * out_dtile = out_base + 1 * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+
+ *pv_out0 = Q6_V_lo_W(vp);
+ *pv_out1 = Q6_V_hi_W(vp);
+ }
+ }
+}
+
+static inline void hmx_fa_q_prep_fp16(
+ __fp16 * vtcm_q_tiles, const uint8_t * temp_q_vtcm,
+ size_t start, size_t end, size_t g_rows_end,
+ size_t DK, size_t G, size_t n_rows_q,
+ const struct fastdiv_values * div_G, uint32_t d_limit, bool q_transposed
+) {
+ for (size_t r = start; r < end; r += 2) {
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (r >= g_rows_end) {
+ for (uint32_t d = 0; d < d_limit; ++d) {
+ __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+ *pv_out0 = Q6_V_vzero();
+ *pv_out1 = Q6_V_vzero();
+ }
+ continue;
+ }
+
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const size_t offset0 = q_transposed ? (h_idx0 * n_rows_q + q_idx0) : (q_idx0 * G + h_idx0);
+ const size_t offset1 = q_transposed ? (h_idx1 * n_rows_q + q_idx1) : (q_idx1 * G + h_idx1);
+
+ const HVX_Vector * pv_in0 = (const HVX_Vector *) (temp_q_vtcm + offset0 * DK * sizeof(__fp16));
+ const HVX_Vector * pv_in1 = (r + 1 < g_rows_end)
+ ? (const HVX_Vector *) (temp_q_vtcm + offset1 * DK * sizeof(__fp16))
+ : NULL;
+
+ for (uint32_t d = 0; d < d_limit; ++d) {
+ HVX_Vector v0 = pv_in0[d];
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+ __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+
+ *pv_out0 = Q6_V_lo_W(vp);
+ *pv_out1 = Q6_V_hi_W(vp);
+ }
+ }
+}
+
+
+static inline void hmx_fa_q_prep_fallback(
+ __fp16 * vtcm_q_tiles, uintptr_t q_data,
+ size_t q_nb1, size_t q_nb2, size_t q_nb3,
+ uint32_t q_start, uint32_t kv_head, uint32_t ib3,
+ size_t start, size_t end, size_t n_rows_g,
+ size_t G, size_t DK, bool is_q_fp32,
+ const struct fastdiv_values * div_G
+) {
+ for (size_t r = start; r < end; r += 2) {
+ const size_t q_idx0 = fastdiv(r + 0, div_G);
+ const size_t h_idx0 = fastmodulo(r + 0, G, div_G);
+ const size_t q_idx1 = fastdiv(r + 1, div_G);
+ const size_t h_idx1 = fastmodulo(r + 1, G, div_G);
+
+ const uint8_t * q_ptr0 = (r + 0 < n_rows_g) ? ((const uint8_t *) q_data + (q_start + q_idx0) * q_nb1 +
+ (kv_head * G + h_idx0) * q_nb2 + ib3 * q_nb3) :
+ NULL;
+ const uint8_t * q_ptr1 = (r + 1 < n_rows_g) ? ((const uint8_t *) q_data + (q_start + q_idx1) * q_nb1 +
+ (kv_head * G + h_idx1) * q_nb2 + ib3 * q_nb3) :
+ NULL;
+
+ size_t r0 = r / HMX_FP16_TILE_N_ROWS;
+ size_t r1 = r % HMX_FP16_TILE_N_ROWS;
+ __fp16 * out_base = vtcm_q_tiles + r0 * HMX_FP16_TILE_N_ROWS * DK;
+
+ if (is_q_fp32) {
+ const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL;
+ const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) q_ptr1 : NULL;
+
+ for (uint32_t d = 0; d < DK / 32; ++d) {
+ HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_Vector v_hf = hvx_vec_f32_to_f16_shuff(v0, v1);
+
+ HVX_Vector * out_tile = (HVX_Vector *) (out_base + d * HMX_FP16_TILE_N_ELMS);
+ out_tile[r1 / 2] = v_hf;
+ }
+ } else {
+ const HVX_UVector * pv_in0 = q_ptr0 ? (const HVX_UVector *) q_ptr0 : NULL;
+ const HVX_UVector * pv_in1 = q_ptr1 ? (const HVX_UVector *) q_ptr1 : NULL;
+
+ for (uint32_t d = 0; d < DK / 64; ++d) {
+ HVX_Vector v0 = pv_in0 ? pv_in0[d] : Q6_V_vzero();
+ HVX_Vector v1 = pv_in1 ? pv_in1[d] : Q6_V_vzero();
+ HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
+
+ __fp16 * out_dtile = out_base + d * HMX_FP16_TILE_N_ELMS * 2;
+ HVX_Vector * pv_out0 = ((HVX_Vector *) out_dtile) + r1 / 2;
+ HVX_Vector * pv_out1 = pv_out0 + 16;
+
+ *pv_out0 = Q6_V_lo_W(vp);
+ *pv_out1 = Q6_V_hi_W(vp);
+ }
+ }
+ }
}
#endif /* HMX_FA_KERNELS_H */
}
}
-static void convert_f16_weight_to_fp16_tiles_task(
+static __attribute__((noinline))
+void convert_f16_weight_to_fp16_tiles_task(
const tiled_dequantize_state_t *state,
uint32_t start_tile, uint32_t end_tile) {
Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v1);
v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step);
}
- (void) *(volatile HVX_Vector *)(tile_base);
}
++t; ++kt;
}
-
- if (start_tile < end_tile) {
- (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS);
- }
}
-static void quantize_f32_weight_to_fp16_tiles_task(
+static __attribute__((noinline))
+void quantize_f32_weight_to_fp16_tiles_task(
const tiled_dequantize_state_t *state,
uint32_t start_tile, uint32_t end_tile) {
Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out_hi);
v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step);
}
- (void) *(volatile HVX_Vector *)(tile_base);
}
++t; ++kt;
}
-
- if (start_tile < end_tile) {
- (void) *(volatile HVX_Vector *)(state->dst + (end_tile - 1) * HTP_MM_HMX_TILE_N_ELMS);
- }
}
// --- End tiled dequantizers ---
-// requires external HMX lock
-static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation, const __fp16 *restrict weight, const __fp16 *restrict scales,
+// dot-chunk functions require external HMX lock
+
+static void core_dot_chunk_fp16_short(__fp16 *restrict output, const __fp16 *restrict activation,
+ const __fp16 *restrict weight, const __fp16 *restrict scales,
uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) {
__builtin_assume(n_row_tiles > 0);
__builtin_assume(n_col_tiles > 0);
__builtin_assume(n_dot_tiles > 0);
+ __builtin_assume(n_dot_tiles <= 32);
+
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales));
+
+ const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
+ const uint32_t range = 2048u * n_dot_tiles - 1;
+
+ for (uint32_t r = 0; r < n_row_tiles; ++r) {
+ const __fp16 *row_base = activation + r * dot_stride;
+ const __fp16 *col_base = weight;
+ __fp16 *out_tile = output + r * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+
+ for (size_t c = 0; c < n_col_tiles; ++c) {
+ asm volatile(HMX_CLRACC_F16());
+ asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_base), "r"(col_base));
+ asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory");
+ col_base += dot_stride;
+ out_tile += HTP_MM_HMX_TILE_N_ELMS;
+ }
+ }
+}
+
+static void core_dot_chunk_fp16(__fp16 *restrict output, const __fp16 *restrict activation,
+ const __fp16 *restrict weight, const __fp16 *restrict scales,
+ uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles) {
+ if (n_dot_tiles <= 32) {
+ core_dot_chunk_fp16_short(output, activation, weight, scales, n_row_tiles, n_col_tiles, n_dot_tiles);
+ return;
+ }
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(n_col_tiles > 0);
+ __builtin_assume(n_dot_tiles > 32);
+
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)scales));
+
+ const size_t dot_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
- Q6_bias_mxmem2_A((void *)scales);
for (uint32_t r = 0; r < n_row_tiles; ++r) {
+ const __fp16 *row_base = activation + r * dot_stride;
+ const __fp16 *col_base = weight;
+ __fp16 *out_tile = output + r * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+
for (size_t c = 0; c < n_col_tiles; ++c) {
- Q6_mxclracc_hf();
-
- const __fp16 *row_tiles = activation + r * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
- const __fp16 *col_tiles = weight + c * n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
-
- for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) {
- k_block = hex_smin(n_dot_tiles - k, 32);
- const uint32_t range = 2048u * (uint32_t)k_block - 1;
- Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range);
- Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range);
- row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
- col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
+ const __fp16 *row_tiles = row_base;
+ const __fp16 *col_tiles = col_base;
+
+ asm volatile(HMX_CLRACC_F16());
+
+ const uint32_t n_loops = n_dot_tiles / 32;
+ const uint32_t rem = n_dot_tiles % 32;
+
+ for (uint32_t l = 0; l < n_loops; ++l) {
+ asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(65535), "r"(row_tiles), "r"(col_tiles));
+ row_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+ col_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
}
- __fp16 *out_tile = output + (r * n_col_tiles + c) * HTP_MM_HMX_TILE_N_ELMS;
- Q6_mxmem_AR_after_hf(out_tile, 0);
+ if (rem > 0) {
+ const uint32_t range = 2048u * rem - 1;
+ asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles));
+ }
+
+ asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(out_tile), "r"(0) : "memory");
+
+ col_base += dot_stride;
+ out_tile += HTP_MM_HMX_TILE_N_ELMS;
}
}
}
-// C += AB
-static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
+static void core_mma_chunk_fp16_short(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile,
uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) {
__builtin_assume(n_row_tiles > 0);
__builtin_assume(n_col_tiles > 0);
__builtin_assume(n_dot_tiles > 0);
+ __builtin_assume(n_dot_tiles <= 32);
- Q6_bias_mxmem2_A((void *)col_scales);
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales));
const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
+ const uint32_t range = 2048u * n_dot_tiles - 1;
+
for (size_t i = 0; i < n_row_tiles; ++i) {
const __fp16 *row_base = a + i * dot_tile_stride;
__fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+ const __fp16 *col_base = b;
+ __fp16 *accum_tile = res_base;
+
for (size_t j = 0; j < n_col_tiles; ++j) {
- Q6_mxclracc_hf();
+ asm volatile(HMX_CLRACC_F16());
- const __fp16 *col_tiles = b + j * dot_tile_stride;
- const __fp16 *row_tiles = row_base;
- __fp16 *accum_tile = res_base + j * HTP_MM_HMX_TILE_N_ELMS;
if (!zero_init) {
- Q6_activation_hf_mxmem_RR((unsigned int)accum_tile, 2047);
- Q6_weight_hf_mxmem_RR((unsigned int)eye_tile, 2047);
+ asm volatile(HMX_LOAD_MPY_F16("%1", "%2", "%0") : : "r"(2047), "r"(accum_tile), "r"(eye_tile));
}
- for (uint32_t k = 0, k_block; k < n_dot_tiles; k += k_block) {
- k_block = hex_smin(n_dot_tiles - k, 32);
- const uint32_t range = 2048u * k_block - 1;
- Q6_activation_hf_mxmem_RR_deep((unsigned int)row_tiles, range);
- Q6_weight_hf_mxmem_RR((unsigned int)col_tiles, range);
- row_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
- col_tiles += k_block * HTP_MM_HMX_TILE_N_ELMS;
- }
+ asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_base), "r"(col_base));
- Q6_mxmem_AR_after_hf(accum_tile, 0);
+ asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(accum_tile), "r"(0) : "memory");
+
+ col_base += dot_tile_stride;
+ accum_tile += HTP_MM_HMX_TILE_N_ELMS;
}
}
}
-// --- Async HMX matmul job (for pipeline overlap) ---
+static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
+ const __fp16 *restrict col_scales, const __fp16 *restrict eye_tile,
+ uint32_t n_row_tiles, uint32_t n_col_tiles, uint32_t n_dot_tiles, bool zero_init) {
+ if (n_dot_tiles <= 32) {
+ core_mma_chunk_fp16_short(c, a, b, col_scales, eye_tile, n_row_tiles, n_col_tiles, n_dot_tiles, zero_init);
+ return;
+ }
+ __builtin_assume(n_row_tiles > 0);
+ __builtin_assume(n_col_tiles > 0);
+ __builtin_assume(n_dot_tiles > 32);
-typedef struct {
- __fp16 * output;
- const __fp16 * activation;
- const __fp16 * weight;
- const __fp16 * scales;
- uint32_t n_row_tiles;
- uint32_t n_col_tiles;
- uint32_t n_dot_tiles;
-} hmx_matmul_job_t;
-
-static void hmx_matmul_worker_fn(void * data) {
- hmx_matmul_job_t * job = (hmx_matmul_job_t *) data;
- FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
- core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
-}
+ asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales));
-static inline void hmx_matmul_job_init(hmx_matmul_job_t * job,
- __fp16 * output,
- const __fp16 * activation,
- const __fp16 * weight,
- const __fp16 * scales,
- uint32_t n_row_tiles,
- uint32_t n_col_tiles,
- uint32_t n_dot_tiles) {
- job->output = output;
- job->activation = activation;
- job->weight = weight;
- job->scales = scales;
- job->n_row_tiles = n_row_tiles;
- job->n_col_tiles = n_col_tiles;
- job->n_dot_tiles = n_dot_tiles;
+ const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
+
+ for (size_t i = 0; i < n_row_tiles; ++i) {
+ const __fp16 *row_base = a + i * dot_tile_stride;
+ __fp16 *res_base = c + i * n_col_tiles * HTP_MM_HMX_TILE_N_ELMS;
+ const __fp16 *col_base = b;
+ __fp16 *accum_tile = res_base;
+
+ for (size_t j = 0; j < n_col_tiles; ++j) {
+ const __fp16 *col_tiles = col_base;
+ const __fp16 *row_tiles = row_base;
+
+ asm volatile(HMX_CLRACC_F16());
+
+ if (!zero_init) {
+ asm volatile(HMX_LOAD_MPY_F16("%1", "%2", "%0") : : "r"(2047), "r"(accum_tile), "r"(eye_tile));
+ }
+
+ const uint32_t n_loops = n_dot_tiles / 32;
+ const uint32_t rem = n_dot_tiles % 32;
+
+ for (uint32_t l = 0; l < n_loops; ++l) {
+ asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(65535), "r"(row_tiles), "r"(col_tiles));
+ row_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+ col_tiles += 32 * HTP_MM_HMX_TILE_N_ELMS;
+ }
+
+ if (rem > 0) {
+ const uint32_t range = 2048u * rem - 1;
+ asm volatile(HMX_LOAD_MPY_DEEP_F16("%1", "%2", "%0") : : "r"(range), "r"(row_tiles), "r"(col_tiles));
+ }
+
+ asm volatile(HMX_STORE_AFTER_F16("%0", "%1") : : "r"(accum_tile), "r"(0) : "memory");
+
+ col_base += dot_tile_stride;
+ accum_tile += HTP_MM_HMX_TILE_N_ELMS;
+ }
+ }
}
// output : fp16 -> f32p
}
}
-typedef struct {
- __fp16 *dst;
- const float *src;
- uint32_t n_tasks;
- uint32_t n_tot_chunks;
- uint32_t n_chunks_per_task;
- uint32_t k_block;
- uint32_t k_stride;
- uint32_t k_valid;
- struct htp_thread_trace * traces;
- struct htp_context * ctx;
- float * vtcm_f32_act;
-} activation_transfer_task_state_t;
-
-static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
- dma_queue *dma_q,
+static void transfer_activation_row_pair_fp32_to_fp16(
__fp16 *restrict vtcm_dst,
- const float *restrict src,
- uint32_t n_rows,
+ const float *restrict row0,
+ const float *restrict row1,
+ uint32_t r,
uint32_t k_block,
- uint32_t k_stride,
uint32_t k_valid,
- float *thread_f32_act) {
-
- const uint32_t R = HTP_MM_DMA_ACT_ROWS_PER_STEP;
- const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+ bool row0_valid,
+ bool row1_valid) {
- const uint32_t n_steps = n_rows_padded / R;
+ 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
- // pre-fetch step 0
- if (n_steps > 0 && n_rows > 0) {
- uint32_t nrows_to_fetch = hex_smin(n_rows, R);
- dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src),
- k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
- }
+ 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 *)(row0 + c);
+ if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c);
- for (uint32_t s = 0; s < n_steps; ++s) {
- uint32_t r = R * s;
- float *curr_buf = thread_f32_act + (s % 2) * R * k_block;
+ HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
- if (r < n_rows) {
- dma_queue_pop(dma_q);
- }
-
- uint32_t next_s = s + 1;
- uint32_t next_r = R * next_s;
- if (next_r < n_rows) {
- uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
- const float *next_src = src + next_r * k_stride;
- float *next_buf = thread_f32_act + (next_s % 2) * R * k_block;
- dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src),
- k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
- }
+ 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;
- #pragma unroll
- for (uint32_t i = 0; i < HTP_MM_DMA_ACT_ROWS_PER_STEP; i += 2) {
- uint32_t curr_r = r + i;
- const bool row0_valid = (curr_r < n_rows);
- const bool row1_valid = (curr_r + 1) < n_rows;
-
- const float *ptr_in0 = curr_buf + i * k_block;
- const float *ptr_in1 = curr_buf + (i + 1) * k_block;
+ 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 *)(row0 + c);
+ if (row1_valid) v1 = *(const HVX_Vector *)(row1 + c);
- 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);
+ 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);
+ HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
- uint32_t r0 = curr_r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
- uint32_t r1 = curr_r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
- 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;
+ 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 r0 = curr_r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
- uint32_t r1 = curr_r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
- 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;
- }
- }
+ HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
+ tile[r1 / 2] = v_out;
}
}
-typedef struct {
- const struct mmid_row_mapping *matrix_rows;
- __fp16 *dst;
- const float *src;
- uint32_t n_tasks;
- uint32_t n_tot_chunks;
- uint32_t n_chunks_per_task;
- uint32_t k_block;
- uint32_t cur_a;
- uint32_t mapping_stride;
- uint32_t ne11;
- struct fastdiv_values ne11_div;
- size_t nb11;
- size_t nb12;
- uint32_t start_row;
- uint32_t cne1;
- uint32_t k_valid;
- struct htp_thread_trace *traces;
-} activation_transfer_gathered_task_state_t;
-
-typedef struct {
- const struct mmid_row_mapping *matrix_rows;
- const __fp16 *vtcm_src;
- float *dst;
- uint32_t n_tasks;
- uint32_t n_tot_chunks;
- uint32_t n_chunks_per_task;
- uint32_t n_cols;
- uint32_t cur_a;
- uint32_t mapping_stride;
- size_t dst_nb1;
- size_t dst_nb2;
- uint32_t start_row;
- uint32_t cne1;
- struct htp_thread_trace *traces;
-} output_transfer_scattered_task_state_t;
-
static void transfer_activation_chunk_fp32_to_fp16_gathered(
__fp16 *restrict vtcm_dst,
const float *restrict src,
#include <qurt_thread.h>
#include <qurt_futex.h>
+#include <qurt_hvx.h>
#include <HAP_compute_res.h>
case HMX_QUEUE_NOOP: /* noop */; break;
case HMX_QUEUE_KILL: *killed = true; break;
case HMX_QUEUE_SUSPEND: hmx_unlock(q); break;
+ case HMX_QUEUE_WAKEUP: hmx_lock(q); break;
default:
hmx_lock(q);
htp_trace_event_start(q->trace, HTP_TRACE_EVT_HMX_COMP, ir);
while (!killed) {
unsigned int seqn = atomic_load(&q->seqn);
if (seqn == prev_seqn) {
+ // drop HVX context while spinning
+ if (poll_cnt > 1 && poll_cnt == HMX_QUEUE_POLL_COUNT) {
+ qurt_hvx_unlock();
+ }
if (--poll_cnt) { hex_pause(); continue; }
FARF(HIGH, "hmx-queue-thread: sleeping");
qurt_futex_wait(&q->seqn, prev_seqn);
+ poll_cnt = HMX_QUEUE_POLL_COUNT;
continue;
}
prev_seqn = seqn;
#endif
#define HMX_QUEUE_THREAD_STACK_SIZE (16 * 1024)
-#define HMX_QUEUE_POLL_COUNT 2000
+
+#if __HVX_ARCH__ > 79
+#define HMX_QUEUE_POLL_COUNT 2000
+#else
+#define HMX_QUEUE_POLL_COUNT 1
+#endif
typedef void (*hmx_queue_func)(void *);
// Dummy funcs used as signals
enum hmx_queue_signal {
HMX_QUEUE_NOOP = 0, // aka NULL
+ HMX_QUEUE_WAKEUP,
HMX_QUEUE_SUSPEND,
HMX_QUEUE_KILL
};
return q->capacity;
}
-static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
+static inline struct hmx_queue_desc hmx_queue_pop_one(struct hmx_queue * q) {
unsigned int ip = q->idx_pop;
unsigned int iw = q->idx_write;
return rd;
}
+static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
+ while (1) {
+ struct hmx_queue_desc d = hmx_queue_pop_one(q);
+
+ uint32_t sig = (uint32_t) d.func;
+ if (sig && sig <= HMX_QUEUE_KILL)
+ continue;
+
+ return d;
+ }
+}
+
static inline void hmx_queue_flush(struct hmx_queue * q) {
- while (hmx_queue_pop(q).func != NULL) ;
+ while (hmx_queue_pop_one(q).func != NULL) ;
+}
+
+static inline void hmx_queue_wakeup(struct hmx_queue * q) {
+ hmx_queue_signal(q, HMX_QUEUE_WAKEUP);
}
static inline void hmx_queue_suspend(struct hmx_queue *q) {
hmx_queue_signal(q, HMX_QUEUE_SUSPEND);
- hmx_queue_flush(q);
}
#ifdef __cplusplus
}
}
+// --- HMX inline asm macros for load-store packetization ---
+#define HMX_LOAD_MPY_F16(act, wt, range) \
+ "{\n" \
+ " activation.hf = mxmem(" act ", " range ")\n" \
+ " weight.hf = mxmem(" wt ", " range ")\n" \
+ "}\n"
+
+#define HMX_LOAD_MPY_DEEP_F16(act, wt, range) \
+ "{\n" \
+ " activation.hf = mxmem(" act ", " range "):deep\n" \
+ " weight.hf = mxmem(" wt ", " range ")\n" \
+ "}\n"
+
+#define HMX_STORE_AFTER_F16(out, scale_reg) \
+ "mxmem(" out ", " scale_reg "):after.hf = acc\n"
+
+#define HMX_SET_BIAS(scales) \
+ "bias = mxmem2(" scales ")\n"
+
+#define HMX_CLRACC_F16() \
+ "mxclracc.hf\n"
+
#endif // HMX_UTILS_H
--- /dev/null
+#ifndef HTP_VTCM_H
+#define HTP_VTCM_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) {
+ uint8_t *p = *vtcm_ptr;
+ *vtcm_ptr += size;
+ return p;
+}
+
+#define VTCM_LAYOUT_ALLOC(off, field, sz) do { (L)->field = (off); (off) += (sz); } while (0)
+#define VTCM_LAYOUT_ALLOC_OPTIONAL(off, field, sz, cond) do { if (cond) { VTCM_LAYOUT_ALLOC(off, field, sz); } else { (L)->field = 0; } } while (0)
+
+#define VTCM_LAYOUT_PTR(type, base, offset) ((type *)((uint8_t *)(base) + (offset)))
+#define VTCM_LAYOUT_PTR_OPTIONAL(type, base, offset, cond) ((cond) ? VTCM_LAYOUT_PTR(type, base, offset) : NULL)
+
+#endif // HTP_VTCM_H
int op_status = HTP_STATUS_OK;
uint32_t op_wakeup = n_ops / 2; // half-way throgh the batch
+ hmx_queue_wakeup(ctx->hmx_queue);
+
for (uint32_t i=0; i < n_ops; i++) {
struct profile_data prof;
}
}
+ hmx_queue_suspend(ctx->hmx_queue);
+
struct htp_opbatch_rsp rsp;
rsp.id = req.id;
rsp.status = op_status;
#include "htp-ctx.h"
#include "htp-ops.h"
#include "matmul-ops.h"
-#include "vtcm-utils.h"
+#include "htp-vtcm.h"
static void hvx_tensor_add_f32_grid(
const struct htp_tensor * restrict dst,
break;
}
- size_t src0_sz = 0, src1_sz = 0, dst_sz = 0;
- if (kparams->vtcm_src0_size > 0 || kparams->vtcm_src1_size > 0 || kparams->vtcm_dst_size > 0) {
- src0_sz = kparams->vtcm_src0_size;
- src1_sz = kparams->vtcm_src1_size;
- dst_sz = kparams->vtcm_dst_size;
- } else {
- const uint32_t n_prefetch = kparams->n_prefetch;
- assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0);
- htp_mm_hvx_get_vtcm_sizes(
- kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
- dst_row_size, src0_row_size, src1_row_size, n_prefetch,
- &src0_sz, &src1_sz, &dst_sz
- );
- }
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
+ dst_row_size, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, false);
if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM ||
kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM ||
kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW ||
kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) {
- mmctx->vtcm_src1_size_per_thread = src1_sz;
+ mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
} else {
- mmctx->vtcm_src1_size_per_thread = src1_sz / octx->n_threads;
+ mmctx->vtcm_src1_size_per_thread = L.src1_bytes / octx->n_threads;
}
- mmctx->vtcm_src0_size_per_thread = src0_sz / octx->n_threads;
- mmctx->vtcm_dst_size_per_thread = dst_sz / octx->n_threads;
+ mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+ mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads;
- size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : (src1_sz + src0_sz + dst_sz);
+ size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
FARF(HIGH, "matmul-%s : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu (%zu)\n", mmctx->type,
- src0_sz, src1_sz, dst_sz, vtcm_size);
+ L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
FARF(HIGH, "matmul-%s : %ux%ux%ux%u * %ux%ux%ux%u-> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type, src0->ne[0],
src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], dst->ne[0],
return HTP_STATUS_VTCM_TOO_SMALL;
}
- uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
- mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
- mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
- mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+ uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+ mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+ mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+ mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL;
octx->src0_spad.src = NULL;
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_O_PROC, start_chunk_idx);
}
+typedef struct {
+ const struct mmid_row_mapping *matrix_rows;
+ __fp16 *dst;
+ const float *src;
+ uint32_t n_tasks;
+ uint32_t n_tot_chunks;
+ uint32_t n_chunks_per_task;
+ uint32_t k_block;
+ uint32_t k_stride;
+ uint32_t k_valid;
+ struct htp_thread_trace * traces;
+ struct htp_context * ctx;
+ float * vtcm_f32_act;
+ size_t vtcm_f32_act_bytes_per_thread;
+ uint32_t dma_step_rows;
+ uint32_t dma_step_rows_shift;
+} activation_transfer_task_state_t;
+
+static void transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
+ dma_queue *dma_q,
+ __fp16 *restrict vtcm_dst,
+ const float *restrict src,
+ uint32_t n_rows,
+ uint32_t k_block,
+ uint32_t k_stride,
+ uint32_t k_valid,
+ float *thread_f32_act,
+ struct htp_thread_trace *tr,
+ uint32_t dma_step_rows,
+ uint32_t dma_step_rows_shift) {
+
+ const uint32_t R = dma_step_rows;
+ const uint32_t n_rows_padded = hex_align_up(n_rows, HTP_MM_HMX_TILE_N_ROWS);
+
+ const uint32_t n_steps = n_rows_padded >> dma_step_rows_shift;
+
+ // Push step 0
+ if (n_steps > 0 && n_rows > 0) {
+ uint32_t nrows_to_fetch = hex_smin(n_rows, R);
+ dma_queue_push(dma_q, dma_make_ptr(thread_f32_act, src),
+ k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
+ }
+ // Push step 1 (if valid)
+ if (n_steps > 1) {
+ uint32_t next_r = R * 1;
+ if (next_r < n_rows) {
+ uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
+ const float *next_src = src + next_r * k_stride;
+ float *next_buf = thread_f32_act + 1 * R * k_block;
+ dma_queue_push(dma_q, dma_make_ptr(next_buf, next_src),
+ k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
+ }
+ }
+ for (uint32_t s = 0; s < n_steps; ++s) {
+ uint32_t r = s << dma_step_rows_shift;
+ float *curr_buf = thread_f32_act;
+
+ if (r < n_rows) {
+ curr_buf = (float *) dma_queue_pop(dma_q).dst;
+ }
+
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, r);
+ for (uint32_t p = 0; p < (R >> 1); ++p) {
+ uint32_t row_idx = r + (p << 1);
+ float *pair_buf = curr_buf + (p << 1) * k_block;
+ bool r0_valid = ((row_idx + 0) < n_rows);
+ bool r1_valid = ((row_idx + 1) < n_rows);
+
+ transfer_activation_row_pair_fp32_to_fp16(vtcm_dst, pair_buf, pair_buf + k_block, row_idx, k_block, k_valid, r0_valid, r1_valid);
+ }
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, r);
+
+ // Push step s + 2
+ uint32_t next_s = s + 2;
+ uint32_t next_r = next_s << dma_step_rows_shift;
+ if (next_r < n_rows) {
+ uint32_t nrows_to_fetch = hex_smin(n_rows - next_r, R);
+ const float *next_src = src + next_r * k_stride;
+ dma_queue_push(dma_q, dma_make_ptr(curr_buf, next_src),
+ k_block * sizeof(float), k_stride * sizeof(float), k_valid * sizeof(float), nrows_to_fetch);
+ }
+ }
+}
+
static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i, void *data) {
activation_transfer_task_state_t *st = (activation_transfer_task_state_t *) data;
struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
- int start_chunk_idx = i * st->n_chunks_per_task;
- htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx);
-
for (unsigned int task_id = i; task_id < (unsigned int)st->n_tasks; task_id += n) {
int chunk_idx = task_id * st->n_chunks_per_task;
size_t chunk_size = hex_smin(st->n_tot_chunks - chunk_idx, st->n_chunks_per_task);
const float *src = st->src + chunk_idx * st->k_stride;
if (st->vtcm_f32_act) {
- float *thread_f32_act = st->vtcm_f32_act + i * HTP_MM_DMA_ACT_MULTIPLIER * st->k_block;
+ float *thread_f32_act = (float *)((char *)st->vtcm_f32_act + i * st->vtcm_f32_act_bytes_per_thread);
transfer_activation_chunk_fp32_to_fp16_dma_pipelined(
- st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act
+ st->ctx->dma[i], dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid, thread_f32_act, tr, st->dma_step_rows, st->dma_step_rows_shift
);
} else {
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
transfer_activation_chunk_fp32_to_fp16(dst, src, chunk_size, st->k_block, st->k_stride, st->k_valid);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, chunk_idx);
}
}
-
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, start_chunk_idx);
}
+typedef struct {
+ const struct mmid_row_mapping *matrix_rows;
+ __fp16 *dst;
+ const float *src;
+ uint32_t n_tasks;
+ uint32_t n_tot_chunks;
+ uint32_t n_chunks_per_task;
+ uint32_t k_block;
+ uint32_t cur_a;
+ uint32_t mapping_stride;
+ uint32_t ne11;
+ struct fastdiv_values ne11_div;
+ size_t nb11;
+ size_t nb12;
+ uint32_t start_row;
+ uint32_t cne1;
+ uint32_t k_valid;
+ struct htp_thread_trace *traces;
+} activation_transfer_gathered_task_state_t;
+
+typedef struct {
+ const struct mmid_row_mapping *matrix_rows;
+ const __fp16 *vtcm_src;
+ float *dst;
+ uint32_t n_tasks;
+ uint32_t n_tot_chunks;
+ uint32_t n_chunks_per_task;
+ uint32_t n_cols;
+ uint32_t cur_a;
+ uint32_t mapping_stride;
+ size_t dst_nb1;
+ size_t dst_nb2;
+ uint32_t start_row;
+ uint32_t cne1;
+ struct htp_thread_trace *traces;
+} output_transfer_scattered_task_state_t;
+
static void transfer_activation_chunk_gathered_worker_fn(unsigned int n, unsigned int i, void *data) {
activation_transfer_gathered_task_state_t *st = data;
struct htp_thread_trace * tr = st->traces ? &st->traces[i] : NULL;
int k_stride,
int n_threads,
int k_valid,
- float *vtcm_f32_act) {
+ float *vtcm_f32_act,
+ size_t vtcm_f32_act_bytes) {
+ if (n_rows <= 0) {
+ return;
+ }
+
assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0 && k_stride % HTP_MM_HMX_TILE_N_COLS == 0);
size_t n_tot_chunks = n_rows;
size_t n_chunks_per_task = (n_threads == 1) ? n_tot_chunks : 32; // must be multiple of 32 to ensure correct destination address
+ uint32_t dma_step_rows = 2;
+ uint32_t dma_step_rows_shift = 1;
+ if (vtcm_f32_act && vtcm_f32_act_bytes > 0 && k_block > 0) {
+ size_t thread_scratch_elements = vtcm_f32_act_bytes / (n_threads * sizeof(float));
+ size_t dma_step_rows_max = (thread_scratch_elements / 2) / k_block;
+ if (dma_step_rows_max >= 4) {
+ dma_step_rows = 4;
+ dma_step_rows_shift = 2;
+ } else {
+ dma_step_rows = 2;
+ dma_step_rows_shift = 1;
+ }
+ }
+
activation_transfer_task_state_t state;
- state.n_tasks = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task;
- state.n_tot_chunks = n_tot_chunks;
- state.n_chunks_per_task = n_chunks_per_task;
- state.dst = dst;
- state.src = src;
- state.k_block = k_block;
- state.k_stride = k_stride;
- state.k_valid = k_valid;
- state.traces = ctx->trace;
- state.ctx = ctx;
- state.vtcm_f32_act = vtcm_f32_act;
+ state.n_tasks = (n_tot_chunks + n_chunks_per_task - 1) / n_chunks_per_task;
+ state.n_tot_chunks = n_tot_chunks;
+ state.n_chunks_per_task = n_chunks_per_task;
+ state.dst = dst;
+ state.src = src;
+ state.k_block = k_block;
+ state.k_stride = k_stride;
+ state.k_valid = k_valid;
+ state.traces = ctx->trace;
+ state.ctx = ctx;
+ state.vtcm_f32_act = vtcm_f32_act;
+
+ int active_threads = hex_smin(n_threads, (int)state.n_tasks);
+ state.vtcm_f32_act_bytes_per_thread = (vtcm_f32_act_bytes / active_threads) & ~127u;
+ state.dma_step_rows = dma_step_rows;
+ state.dma_step_rows_shift = dma_step_rows_shift;
if (state.n_tasks == 1 || n_threads == 1) {
transfer_activation_chunk_worker_fn(1, 0, &state);
} else {
- int n_tasks = hex_smin((int) state.n_tasks, n_threads);
- worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, n_tasks);
+ worker_pool_run_func(ctx->worker_pool, transfer_activation_chunk_worker_fn, &state, active_threads);
}
}
+// --- Async HMX matmul job (for pipeline overlap) ---
+
+typedef struct {
+ __fp16 * output;
+ const __fp16 * activation;
+ const __fp16 * weight;
+ const __fp16 * scales;
+ uint32_t n_row_tiles;
+ uint32_t n_col_tiles;
+ uint32_t n_dot_tiles;
+} hmx_matmul_job_t;
+
+static void hmx_matmul_worker_fn(void * data) {
+ hmx_matmul_job_t * job = (hmx_matmul_job_t *) data;
+ FARF(HIGH, "hmx-mm-job: n_row_tiles %u n_col_tiles %u n_dot_tiles %u", job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
+ core_dot_chunk_fp16(job->output, job->activation, job->weight, job->scales, job->n_row_tiles, job->n_col_tiles, job->n_dot_tiles);
+}
+
+static inline void hmx_matmul_job_init(hmx_matmul_job_t * job,
+ __fp16 * output,
+ const __fp16 * activation,
+ const __fp16 * weight,
+ const __fp16 * scales,
+ uint32_t n_row_tiles,
+ uint32_t n_col_tiles,
+ uint32_t n_dot_tiles) {
+ job->output = output;
+ job->activation = activation;
+ job->weight = weight;
+ job->scales = scales;
+ job->n_row_tiles = n_row_tiles;
+ job->n_col_tiles = n_col_tiles;
+ job->n_dot_tiles = n_dot_tiles;
+}
static int hmx_mm_2d_f32(struct htp_context *ctx,
float *restrict dst,
const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0;
- const size_t act_f32_size = hex_align_up((size_t)act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), HTP_MM_HMX_TILE_SIZE);
-
- const size_t weight_area_size = is_quant
- ? hex_align_up((n_chunk_n_cols / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE)
- : hex_align_up(n_chunk_n_cols * row_stride, HTP_MM_HMX_TILE_SIZE);
- const size_t act_area_size = hex_align_up(m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- const size_t output_area_size = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE);
-
- size_t scratch0_size, scratch1_size, scratch2_size;
- scratch0_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE); // dequant buf 0
- scratch1_size = pipeline ? scratch0_size : 0; // dequant buf 1
- scratch2_size = pipeline ? output_area_size : 0; // output buf 1
-
- uint8_t *vtcm_ptr = (uint8_t *) ctx->vtcm_base;
- __fp16 *vtcm_weight_raw[2] = { NULL, NULL };
- if (weight_area_size) {
- if (pipeline) {
- vtcm_weight_raw[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
- vtcm_weight_raw[1] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
- } else {
- vtcm_weight_raw[0] = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
- }
- }
-
- __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, act_area_size);
- float *vtcm_f32_act = (float *) vtcm_seq_alloc(&vtcm_ptr, act_f32_size);
- __fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
- void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch0_size);
- void *vtcm_scratch1 = scratch1_size ? vtcm_seq_alloc(&vtcm_ptr, scratch1_size) : NULL;
- void *vtcm_scratch2 = scratch2_size ? vtcm_seq_alloc(&vtcm_ptr, scratch2_size) : NULL;
- __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256);
+ struct htp_mm_hmx_vtcm_layout L;
+ htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, false, pipeline, act_threads, aligned_tile_size);
- vtcm_used = vtcm_ptr - (uint8_t *) ctx->vtcm_base;
+ vtcm_used = L.total_bytes;
if (vtcm_used > vtcm_budget) {
FARF(ERROR, "hmx-mm-2d-precomputed: VTCM overflow: used %zu budget %zu, m %d k %d n %d mc %zu nc %zu",
vtcm_used, vtcm_budget, m, k, n, m_chunk_n_rows, n_chunk_n_cols);
return -1;
}
+ uint8_t * const base = (uint8_t *) ctx->vtcm_base;
+ __fp16 *vtcm_weight_raw[2] = {
+ VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]),
+ VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_weight[1], pipeline)
+ };
+
+ __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act);
+ float *vtcm_f32_act = VTCM_LAYOUT_PTR(float, base, L.off_act_f32);
+ __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]);
+ void *vtcm_scratch0 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]);
+ void *vtcm_scratch1 = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_scratch[1], pipeline);
+ void *vtcm_scratch2 = VTCM_LAYOUT_PTR_OPTIONAL(void, base, L.off_dst[1], pipeline);
+ __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales);
+
hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16
- FARF(HIGH, "hmx-mm-2d-precomputed: standard : m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu",
+ FARF(HIGH, "hmx-mm-2d: m %d k %d n %d wtype %d mc %zu nc %zu vtcm %zu/%zu",
m, k, n, weight_type, m_chunk_n_rows, n_chunk_n_cols, vtcm_used, vtcm_budget);
int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols);
void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 };
void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 };
- transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act);
+ transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes);
// Prologue: push A0 and optionally A1 (if n_chunk_cnt > 1)
- const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols);
+ const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols);
const uint32_t height_A0 = is_quant ? (n_cols_A0 / 32) * n_k_tiles : n_cols_A0;
dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight),
dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0);
if (1 < n_chunk_cnt) {
- const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols);
+ const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols);
const uint32_t height_A1 = is_quant ? (n_cols_A1 / 32) * n_k_tiles : n_cols_A1;
dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride),
dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1);
}
- // pop A0 -> dequantize A0 -> submit C0
- dma_queue_pop(ctx->dma[0]);
- dequantize_tiled_weight_chunk_to_fp16_tiles(
- ctx, vtcm_weight_bufs[0], vtcm_weight_raw[0],
- n_cols_A0, k, row_stride, weight_type,
- n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
-
- hmx_matmul_job_init(&job_slots[0], (__fp16 *) vtcm_output_bufs[0], (__fp16 *) vtcm_f16_act,
- (__fp16 *) vtcm_weight_bufs[0], vtcm_scales,
- hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS),
- hmx_ceil_div(n_cols_A0, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS);
- hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[0]));
-
- // Main loop: pop/dequantize A_{i+1} -> push A_{i+2} -> submit C_{i+1} -> wait C_i and store D_i
+ // Main loop: pop A_i -> dequantize A_i -> push A_{i+2} -> submit C_i -> wait C_{i-1} and store D_{i-1}
for (int i = 0; i < n_chunk_cnt; ++i) {
const size_t nc = i * n_chunk_n_cols;
- const size_t nc_p1 = nc + 1 * n_chunk_n_cols;
const size_t nc_p2 = nc + 2 * n_chunk_n_cols;
const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols);
- const size_t n_cols_p1 = hex_smin(n - nc_p1, n_chunk_n_cols);
const size_t n_cols_p2 = hex_smin(n - nc_p2, n_chunk_n_cols);
- // 1. pop A_{i+1} and dequantize it (if i+1 < n_chunk_cnt)
- if (i + 1 < n_chunk_cnt) {
- dma_queue_pop(ctx->dma[0]);
- dequantize_tiled_weight_chunk_to_fp16_tiles(
- ctx, vtcm_weight_bufs[(i + 1) % 2], vtcm_weight_raw[(i + 1) % 2],
- n_cols_p1, k, row_stride, weight_type,
- n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
- }
+ // 1. pop A_i
+ void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
- // 2. push A_{i+2} (if i+2 < n_chunk_cnt)
+ // 2. dequantize A_i
+ dequantize_tiled_weight_chunk_to_fp16_tiles(
+ ctx, vtcm_weight_bufs[i % 2], curr_raw,
+ n_cols, k, row_stride, weight_type,
+ n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
+
+ // 3. push A_{i+2} (if i+2 < n_chunk_cnt)
if (i + 2 < n_chunk_cnt) {
const uint32_t height_p2 = is_quant ? (n_cols_p2 / 32) * n_k_tiles : n_cols_p2;
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[(i + 2) % 2], weight + nc_p2 * weight_stride),
+ dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_p2 * weight_stride),
dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2);
}
- // 3. submit C_{i+1} (if i+1 < n_chunk_cnt)
- if (i + 1 < n_chunk_cnt) {
- hmx_matmul_job_init(&job_slots[(i + 1) % 2], (__fp16 *) vtcm_output_bufs[(i + 1) % 2],
- (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[(i + 1) % 2],
- vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS),
- hmx_ceil_div(n_cols_p1, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS);
- hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[(i + 1) % 2]));
+ // 4. submit C_i
+ hmx_matmul_job_init(&job_slots[i % 2], (__fp16 *) vtcm_output_bufs[i % 2],
+ (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2],
+ vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS),
+ hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS);
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[i % 2]));
+
+ // 5. wait C_{i-1} and store D_{i-1} (multi-thread HVX, parallel with C_i)
+ if (i > 0) {
+ hmx_queue_pop(ctx->hmx_queue);
+ const size_t nc_prev = (i - 1) * n_chunk_n_cols;
+ const size_t n_cols_prev = hex_smin(n - nc_prev, n_chunk_n_cols);
+ float *output_chunk = dst + (mr * dst_stride + nc_prev);
+ const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_prev) : NULL;
+ int chunk_dst_cols = dst_cols - (int)nc_prev;
+ if (chunk_dst_cols > 0) {
+ transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(i - 1) % 2], n_rows, n_cols_prev, dst_stride, src2_stride, chunk_dst_cols, n_threads);
+ }
}
+ }
- // 4. wait C_i and store D_i (multi-thread HVX, parallel with C_{i+1})
- hmx_queue_pop(ctx->hmx_queue);
- float *output_chunk = dst + (mr * dst_stride + nc);
- const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc) : NULL;
- int chunk_dst_cols = dst_cols - (int)nc;
- if (chunk_dst_cols > 0) {
- transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[i % 2], n_rows, n_cols, dst_stride, src2_stride, chunk_dst_cols, n_threads);
- }
+ // Epilogue: wait C_{last} and store D_{last}
+ hmx_queue_pop(ctx->hmx_queue);
+ const size_t nc_last = (n_chunk_cnt - 1) * n_chunk_n_cols;
+ const size_t n_cols_last = hex_smin(n - nc_last, n_chunk_n_cols);
+ float *output_chunk = dst + (mr * dst_stride + nc_last);
+ const float *src2_chunk = src2 ? (src2 + mr * src2_stride + nc_last) : NULL;
+ int chunk_dst_cols = dst_cols - (int)nc_last;
+ if (chunk_dst_cols > 0) {
+ transfer_output_chunk_threaded(ctx, output_chunk, src2_chunk, vtcm_output_bufs[(n_chunk_cnt - 1) % 2], n_rows, n_cols_last, dst_stride, src2_stride, chunk_dst_cols, n_threads);
}
}
- hmx_queue_suspend(ctx->hmx_queue);
} else {
- // --- Synchronous Un-pipelined loop (m <= 32 or fallback) ---
- HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+ // --- Synchronous loop (m <= 32 or fallback) ---
+ hmx_matmul_job_t job;
for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) {
const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows);
- transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act);
+ transfer_activation_chunk_threaded(ctx, vtcm_f16_act, activation + mr * act_stride, n_rows, k, act_stride, act_threads, k_valid, vtcm_f32_act, L.act_f32_bytes);
+
+ // A0: Pre-fetch the first weight chunk (nc = 0)
+ if (n > 0) {
+ const size_t n_cols = hex_smin(n, n_chunk_n_cols);
+ const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height);
+ }
for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) {
const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols);
const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS);
const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS);
- // A: Weight DMA (Synchronous)
- const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight_raw[0], weight + nc * weight_stride),
- dma_dst_stride, dma_src_stride, dma_width_bytes, height);
- dma_queue_pop(ctx->dma[0]);
+ // A: Wait for weight DMA
+ void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
// B: Weight Dequantize (Threaded)
dequantize_tiled_weight_chunk_to_fp16_tiles(
- ctx, vtcm_scratch0, vtcm_weight_raw[0],
+ ctx, vtcm_scratch0, curr_raw,
n_cols, k, row_stride, weight_type,
n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads);
- // C: HMX Compute (Synchronous)
- core_dot_chunk_fp16(vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
+ // Start weight DMA for the next chunk early
+ const size_t nc_next = nc + n_chunk_n_cols;
+ if (nc_next < n) {
+ const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols);
+ const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next);
+ }
+
+ // C: HMX Compute (Queue-based)
+ hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job));
+ hmx_queue_pop(ctx->hmx_queue);
// D: Output Store
float *output_chunk = dst + (mr * dst_stride + nc);
}
}
}
- HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
}
return 0;
size_t n_chunk_n_cols = n_chunk;
size_t vtcm_used = vtcm_size;
- const size_t act_head_stride = m_chunk_n_rows * (size_t) params->k; // fp16 elements between heads
- const size_t weight_area_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- const size_t activation_area_size = hex_align_up(group_size * m_chunk_n_rows * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- const size_t output_area_size = hex_align_up(m_chunk_n_rows * n_chunk_n_cols * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE);
- const size_t scratch_area_size = hex_align_up(n_chunk_n_cols * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-
- uint8_t *vtcm_ptr = (uint8_t *) ctx->vtcm_base;
- __fp16 *vtcm_weight = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, weight_area_size);
- __fp16 *vtcm_f16_act = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, activation_area_size);
- __fp16 *vtcm_output = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, output_area_size);
- void *vtcm_scratch0 = vtcm_seq_alloc(&vtcm_ptr, scratch_area_size);
- void *vtcm_scratch1 = vtcm_seq_alloc(&vtcm_ptr, scratch_area_size);
- __fp16 *vtcm_scales = (__fp16 *) vtcm_seq_alloc(&vtcm_ptr, 256);
- float *vtcm_f32_act = use_dma_activation ? (float *) vtcm_seq_alloc(&vtcm_ptr, f32_scratch_size) : NULL;
-
- if ((size_t) (vtcm_ptr - (uint8_t *) ctx->vtcm_base) > vtcm_budget) {
+ struct htp_mm_hmx_vtcm_layout L;
+ htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, use_dma_activation, false, act_threads, 0);
+
+ if (L.total_bytes > vtcm_budget) {
FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__);
return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size);
}
+ uint8_t * const base = (uint8_t *) ctx->vtcm_base;
+ __fp16 *vtcm_weight = VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]);
+ __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act);
+ __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]);
+ void *vtcm_scratch0 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[0]);
+ void *vtcm_scratch1 = VTCM_LAYOUT_PTR(void, base, L.off_scratch[1]);
+ __fp16 *vtcm_scales = VTCM_LAYOUT_PTR(__fp16, base, L.off_scales);
+ float *vtcm_f32_act = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_act_f32, use_dma_activation);
+
hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00)); // scale: 1.0, bias: 0.0 in FP16
FARF(HIGH, "%s: grouped path m=%d k=%d n=%d group=%d streams=%d mc=%zu nc=%zu vtcm=%zu/%zu",
__func__, params->m, params->k, params->n, group_size, params->ne13,
m_chunk_n_rows, n_chunk_n_cols,
- (size_t) (vtcm_ptr - (uint8_t *) ctx->vtcm_base), vtcm_budget);
+ L.total_bytes, vtcm_budget);
const size_t fp16_row_bytes = (size_t) params->k * sizeof(__fp16);
const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16);
- HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+ hmx_matmul_job_t job;
for (int b3 = 0; b3 < params->ne13; ++b3) {
for (int b2_base = 0; b2_base < params->ne12; b2_base += group_size) {
// thrashing from HVX loads at large strides.
for (int g = 0; g < group_size; ++g) {
const float *activation_chunk = hmx_mm_activation_batch_ptr(params, b2_base + g, b3) + mr * params->act_stride;
- __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * act_head_stride;
+ __fp16 *vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride;
transfer_activation_chunk_threaded(ctx, vtcm_act_g,
activation_chunk, (int) n_rows,
- params->k, params->act_stride, act_threads, params->k, vtcm_f32_act);
+ params->k, params->act_stride, act_threads, params->k, vtcm_f32_act, L.act_f32_bytes);
}
- void *buf_curr = vtcm_scratch0;
- void *buf_next = vtcm_scratch1;
-
+ // Prologue: Push A0 and A1 (if exists)
{
const size_t n_cols_first = hex_smin((size_t) params->n, n_chunk_n_cols);
- dma_queue_push(ctx->dma[0], dma_make_ptr(buf_curr, weight_group),
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch0, weight_group),
fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_first);
}
+ if (n_chunk_n_cols < (size_t) params->n) {
+ const size_t n_cols_second = hex_smin((size_t) params->n - n_chunk_n_cols, n_chunk_n_cols);
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_scratch1, weight_group + params->weight_stride),
+ fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_second);
+ }
for (size_t nc = 0; nc < (size_t) params->n; nc += n_chunk_n_cols) {
- const size_t n_cols = hex_smin((size_t) params->n - nc, n_chunk_n_cols);
+ const size_t n_cols = hex_smin((size_t) params->n - nc, n_chunk_n_cols);
const size_t n_col_tiles = hmx_ceil_div((int) n_cols, HTP_MM_HMX_TILE_N_COLS);
{
- dma_queue_pop(ctx->dma[0]);
+ void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
+
+ hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) curr_raw, n_cols, params->k, params->k, 0, n_cols);
- const size_t nc_next = nc + n_chunk_n_cols;
+ const size_t nc_next = nc + n_chunk_n_cols * 2;
if (nc_next < (size_t) params->n) {
const size_t n_cols_next = hex_smin((size_t) params->n - nc_next, n_chunk_n_cols);
const __fp16 *next_weight_chunk = weight_group + nc_next * params->weight_stride;
- dma_queue_push(ctx->dma[0], dma_make_ptr(buf_next, next_weight_chunk),
+ dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, next_weight_chunk),
fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_next);
}
-
- hmx_interleave_rows_to_tiles(vtcm_weight, (const __fp16 *) buf_curr, n_cols, params->k, params->k, 0, n_cols);
- hex_swap_ptr(&buf_curr, &buf_next);
}
// Reuse the interleaved weight for every q_head in this GQA group
for (int g = 0; g < group_size; ++g) {
- struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS];
- htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, g);
{
- const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * act_head_stride;
- core_dot_chunk_fp16(vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles,
- params->k / 32);
+ const __fp16 * vtcm_act_g = vtcm_f16_act + (size_t) g * L.act_head_stride;
+ hmx_matmul_job_init(&job, vtcm_output, vtcm_act_g, vtcm_weight, vtcm_scales, n_row_tiles, n_col_tiles, params->k / 32);
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job));
+ hmx_queue_pop(ctx->hmx_queue);
}
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, g);
{
float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc;
const float *src2_chunk = params->src2 ? (hmx_mm_src2_batch_ptr(params, b2_base + g, b3) + mr * params->src2_stride + nc) : NULL;
int chunk_dst_cols = params->n - (int)nc;
if (chunk_dst_cols > 0) {
- transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
+ transfer_output_chunk_threaded(ctx, output, src2_chunk, vtcm_output, (int) n_rows, (int) n_cols,
+ params->dst_stride, params->src2_stride, chunk_dst_cols, ctx->n_threads);
}
}
}
}
}
- HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
-
return 0;
}
hmx_init_column_scales(vtcm_scales, Q6_V_vsplat_R(0x3c00));
- HAP_compute_res_hmx_lock(ctx->vtcm_rctx);
+ hmx_matmul_job_t job;
for (size_t mr = 0; mr < (size_t) m_padded; mr += m_chunk_n_rows) {
const size_t n_rows = hex_smin(m_padded - mr, m_chunk_n_rows);
ctx, vtcm_f16_act, activation, (int) mr, (int) n_rows, k,
matrix_rows, cur_a, mapping_stride, ne11, act_nb1, act_nb2, cne1, n_threads, k_valid);
+ // A0: Pre-fetch the first weight chunk (nc = 0)
+ if (n > 0) {
+ const size_t n_cols = hex_smin((size_t) n, n_chunk_n_cols);
+ const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height);
+ }
+
for (size_t nc = 0; nc < (size_t) n; nc += n_chunk_n_cols) {
const size_t n_cols = hex_smin((size_t) n - nc, n_chunk_n_cols);
const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS);
- const uint32_t height = is_quant ? (n_cols / 32) * n_k_tiles : n_cols;
- dma_queue_push(ctx->dma[0], dma_make_ptr(vtcm_weight, weight + nc * weight_stride),
- dma_dst_stride, dma_src_stride, dma_width_bytes, height);
- dma_queue_pop(ctx->dma[0]);
+ // A: Wait for weight DMA
+ void * curr_raw = dma_queue_pop(ctx->dma[0]).dst;
+ // B: Weight Dequantize (Threaded)
dequantize_tiled_weight_chunk_to_fp16_tiles(
- ctx, vtcm_scratch0, vtcm_weight,
+ ctx, vtcm_scratch0, curr_raw,
n_cols, k, row_stride, weight_type,
n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads
);
- struct htp_thread_trace * tr = &ctx->trace[HTP_MAX_NTHREADS];
- htp_trace_event_start(tr, HTP_TRACE_EVT_HMX_COMP, nc);
- core_dot_chunk_fp16(vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
- htp_trace_event_stop(tr, HTP_TRACE_EVT_HMX_COMP, nc);
+ // Start weight DMA for the next chunk early
+ const size_t nc_next = nc + n_chunk_n_cols;
+ if (nc_next < (size_t) n) {
+ const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols);
+ const uint32_t height_next = is_quant ? (n_cols_next / 32) * n_k_tiles : n_cols_next;
+ dma_queue_push(ctx->dma[0], dma_make_ptr(curr_raw, weight + nc_next * weight_stride),
+ dma_dst_stride, dma_src_stride, dma_width_bytes, height_next);
+ }
+ // C: HMX Compute (Queue-based)
+ hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS);
+ hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job));
+ hmx_queue_pop(ctx->hmx_queue);
+
+ // D: Output Store
transfer_output_chunk_scattered_threaded(
ctx, dst + nc, vtcm_output, (int) mr, (int) n_rows, (int) n_cols,
matrix_rows, cur_a, mapping_stride, dst_nb1, dst_nb2, cne1, n_threads);
}
}
- HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
return 0;
}
-
// --- Dispatchers and Public Entry Points ---
static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_kernel_params * kparams) {
}
size_t src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
- // Scratchpad sizes are computed on the host (htp_mm_hvx_id_get_vtcm_sizes) and passed in.
- // The ID layout is routing-independent, so the host has exact visibility -- consume it here
- // rather than recomputing, to keep host budgeting and device allocation in lockstep.
- size_t src0_sz = kparams->vtcm_src0_size;
- size_t src1_sz = kparams->vtcm_src1_size;
- size_t src2_sz = 0; // mapping lives in DDR
- size_t dst_sz = kparams->vtcm_dst_size;
- size_t vtcm_size = kparams->vtcm_size;
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads,
+ 0, src0_row_size, src1_row_size, kparams->n_prefetch, true, false, false);
- size_t src0_sz_per_thread = src0_sz / octx->n_threads;
- size_t src1_sz_per_thread = src1_sz;
- size_t src2_sz_per_thread = 0;
- size_t dst_sz_per_thread = dst_sz / octx->n_threads;
+ size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
- FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size %zu dst-spad-size %zu (%zu)\n", mmctx->type,
- src0_sz, src1_sz, src2_sz, dst_sz, vtcm_size);
+ FARF(HIGH, "matmul-id-%s : src0-spad-size %zu src1-spad-size %zu src2-spad-size 0 dst-spad-size %zu (%zu)\n", mmctx->type,
+ L.src0_bytes, L.src1_bytes, L.dst_bytes, vtcm_size);
FARF(HIGH, "matmul-id-%s : %ux%ux%ux%u * %ux%ux%ux%u (%ux%ux%ux%u) -> %ux%ux%ux%u (0x%p, 0x%p, 0x%p)\n", mmctx->type,
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3],
return HTP_STATUS_VTCM_TOO_SMALL;
}
- uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
- mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
- mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
- mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
- mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+ uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+ mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+ mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+ mmctx->vtcm_src2 = NULL;
+ mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL;
octx->src0_spad.src = NULL;
mmctx->vtcm_src0_stride = src0_row_size_padded;
mmctx->vtcm_src1_stride = src1_row_size;
- mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
- mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
- mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
- mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread;
+ mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+ mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
+ mmctx->vtcm_src2_size_per_thread = 0;
+ mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads;
mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks;
mmctx->quant_task_func = quant_task_func;
src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]);
}
- // Set up scratchpads using precomputed sizes from the host
- size_t src0_sz = kparams->vtcm_src0_size;
- size_t src1_sz = kparams->vtcm_src1_size;
- size_t src2_sz = kparams->vtcm_src2_size;
- size_t src3_sz = kparams->vtcm_src3_size;
- size_t dst_sz = kparams->vtcm_dst_size;
- size_t vtcm_size = kparams->vtcm_size;
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
+ 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, true, false);
- size_t src0_sz_per_thread = src0_sz / octx->n_threads;
- size_t src1_sz_per_thread = src1_sz;
- size_t src2_sz_per_thread = src2_sz / octx->n_threads;
- size_t src3_sz_per_thread = src3_sz / octx->n_threads;
- size_t dst_sz_per_thread = dst_sz / octx->n_threads;
+ size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-qkv: current VTCM reservation %zu is too small, needed %zu\n",
return HTP_STATUS_VTCM_TOO_SMALL;
}
- uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
- mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
- mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
- mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
- mmctx->vtcm_src3 = vtcm_seq_alloc(&vtcm_ptr, src3_sz);
- mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+ uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+ mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+ mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+ mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
+ mmctx->vtcm_src3 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src3);
+ mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL;
octx->src0_spad.src = NULL;
mmctx->vtcm_src3_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src1_stride = src1_row_size;
- mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
- mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
- mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
- mmctx->vtcm_src3_size_per_thread = src3_sz_per_thread;
- mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread;
+ mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+ mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
+ mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads;
+ mmctx->vtcm_src3_size_per_thread = L.src3_bytes / octx->n_threads;
+ mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads;
if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
return HTP_STATUS_OK;
src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]);
}
- // Set up scratchpads using precomputed sizes from the host
- size_t src0_sz = kparams->vtcm_src0_size;
- size_t src1_sz = kparams->vtcm_src1_size;
- size_t src2_sz = kparams->vtcm_src2_size;
- size_t dst_sz = kparams->vtcm_dst_size;
- size_t vtcm_size = kparams->vtcm_size;
+ struct htp_mm_hvx_vtcm_layout L;
+ htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads,
+ 0, src0_row_size, src1_row_size, kparams->n_prefetch, false, false, true);
- size_t src0_sz_per_thread = src0_sz / octx->n_threads;
- size_t src1_sz_per_thread = src1_sz;
- size_t src2_sz_per_thread = src2_sz / octx->n_threads;
- size_t dst_sz_per_thread = dst_sz / octx->n_threads;
+ size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes;
if (octx->ctx->vtcm_size < vtcm_size) {
FARF(ERROR, "matmul-ffn: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size);
return HTP_STATUS_VTCM_TOO_SMALL;
}
- uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
- mmctx->vtcm_src1 = vtcm_seq_alloc(&vtcm_ptr, src1_sz);
- mmctx->vtcm_src0 = vtcm_seq_alloc(&vtcm_ptr, src0_sz);
- mmctx->vtcm_src2 = vtcm_seq_alloc(&vtcm_ptr, src2_sz);
- mmctx->vtcm_dst = vtcm_seq_alloc(&vtcm_ptr, dst_sz);
+ uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
+ mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
+ mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
+ mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2);
+ mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
octx->src1_spad.src = NULL;
octx->src0_spad.src = NULL;
mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded;
mmctx->vtcm_src1_stride = src1_row_size;
- mmctx->vtcm_src0_size_per_thread = src0_sz_per_thread;
- mmctx->vtcm_src1_size_per_thread = src1_sz_per_thread;
- mmctx->vtcm_src2_size_per_thread = src2_sz_per_thread;
- mmctx->vtcm_dst_size_per_thread = dst_sz_per_thread;
+ mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads;
+ mmctx->vtcm_src1_size_per_thread = L.src1_bytes;
+ mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads;
+ mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads;
if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)
return HTP_STATUS_OK;
#include "htp-ops.h"
#include "hex-fastdiv.h"
#include "hex-common.h"
+#include "htp-vtcm.h"
#ifdef __cplusplus
extern "C" {
// --- DMA Activation Transfer Configuration ---
#define HTP_MM_DMA_ACT_ROWS_PER_STEP 2
-#define HTP_MM_DMA_ACT_MULTIPLIER 4
+#define HTP_MM_DMA_ACT_MULTIPLIER (2 * HTP_MM_DMA_ACT_ROWS_PER_STEP)
enum htp_mm_kernel_type {
HTP_MM_KERNEL_UNSUPPORTED = 0,
*size_per_mn_out = sizeof(uint16_t);
}
-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
-) {
- const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
- const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32);
- const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k);
- const size_t vec_dot_size = k * sizeof(uint16_t);
-
- const size_t act_f32_size = htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE);
- size_t weight_area_size = is_quant
- ? htp_mm_round_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE)
- : htp_mm_round_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE);
- if (pipeline) {
- weight_area_size *= 2;
- }
- const size_t act_area_size = htp_mm_round_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- const size_t output_area_size = htp_mm_round_up(mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
+struct htp_mm_hmx_vtcm_layout {
+ // Byte offsets from vtcm_base for each region
+ size_t off_weight[2]; // [1] is only used when pipelined
+ size_t off_act;
+ size_t off_act_f32; // fp32 activation conversion scratch
+ size_t off_dst[2]; // [1] is only used when pipelined
+ size_t off_scratch[2]; // dequantization scratch pads
+ size_t off_scales; // HMX scales (256 bytes)
+
+ // Cached sizes of regions for HMX kernel use
+ size_t weight_area_bytes;
+ size_t act_area_bytes;
+ size_t act_f32_bytes;
+ size_t output_area_bytes;
+ size_t scratch_bytes[2];
+ size_t act_head_stride;
+
+ size_t total_bytes;
+};
- size_t scratch0_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- size_t scratch1_size = pipeline ? scratch0_size : 0;
- size_t scratch2_size = pipeline ? output_area_size : 0;
+struct htp_mm_hvx_vtcm_layout {
+ // Byte offsets from vtcm_base for each region
+ size_t off_src1; // vtcm_src1 (activation)
+ size_t off_src0; // vtcm_src0 (weight/Wk)
+ 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)
+
+ // Cached sizes
+ size_t src0_bytes;
+ size_t src1_bytes;
+ size_t src2_bytes;
+ size_t src3_bytes;
+ size_t dst_bytes;
+
+ size_t total_bytes;
+};
- return weight_area_size + act_area_size + act_f32_size + output_area_size +
- scratch0_size + scratch1_size + scratch2_size + 256;
-}
+static inline void htp_mm_hmx_vtcm_layout_build(
+ struct htp_mm_hmx_vtcm_layout * L,
+ int kernel_type,
+ 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,
+ uint32_t aligned_tile_size
+) {
+ size_t off = 0;
+
+ if (kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) {
+ const size_t vec_dot_size = k * sizeof(uint16_t);
+ const size_t act_head_stride = mc * k;
+ const size_t weight_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+ 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;
+
+ // Group A: Permanent activation tiles and scales
+ size_t off_group_a = 0;
+ VTCM_LAYOUT_ALLOC(off_group_a, off_act, activation_area_size);
+ VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
+
+ // 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_weight[0], weight_area_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_weight[1], weight_area_size, false);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_dst[0], output_area_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst[1], output_area_size, false);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[0], scratch_area_size);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[1], scratch_area_size);
+
+ const size_t group_b_size = off_group_b - off_group_a;
+
+ // 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;
+ size_t off_group_c = off_group_a;
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_c, off_act_f32, act_f32_size, use_dma_activation);
+
+ const size_t group_c_size = off_group_c - off_group_a;
+
+ L->weight_area_bytes = weight_area_size;
+ L->act_area_bytes = activation_area_size;
+ L->act_f32_bytes = act_f32_size;
+ L->output_area_bytes = output_area_size;
+ L->scratch_bytes[0] = scratch_area_size;
+ L->scratch_bytes[1] = scratch_area_size;
+ L->act_head_stride = act_head_stride;
+
+ off = off_group_a + hex_smax(group_b_size, group_c_size);
+ } else {
+ // HTP_MM_KERNEL_HMX_2D
+ const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32);
+ const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k);
+ const size_t vec_dot_size = k * sizeof(uint16_t);
+ const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS;
+
+ const size_t min_f32_size = hex_align_up(act_threads * HTP_MM_DMA_ACT_MULTIPLIER * k * sizeof(float), 128);
+ const size_t weight_area_size = is_quant
+ ? hex_align_up((nc / 32) * n_k_tiles * aligned_tile_size, HTP_MM_HMX_TILE_SIZE)
+ : hex_align_up(nc * row_stride, HTP_MM_HMX_TILE_SIZE);
+ const size_t act_area_size = hex_align_up(mc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+ const size_t output_area_size = hex_align_up(mc * nc * sizeof(__fp16), HTP_MM_HMX_TILE_SIZE);
+
+ const size_t scratch0_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
+ const size_t scratch1_size = pipeline ? scratch0_size : 0;
+
+ // Group A: Scales and activation tiles (must not overlap with Group B or C)
+ size_t off_group_a = 0;
+ VTCM_LAYOUT_ALLOC(off_group_a, off_scales, HTP_MM_HMX_TILE_SIZE); // Padded to 2K for alignment and future persistent data
+ VTCM_LAYOUT_ALLOC(off_group_a, off_act, act_area_size);
+
+ // 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_weight[0], weight_area_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_weight[1], weight_area_size, pipeline);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_dst[0], output_area_size);
+ VTCM_LAYOUT_ALLOC(off_group_b, off_scratch[0], scratch0_size);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_scratch[1], scratch0_size, pipeline);
+ VTCM_LAYOUT_ALLOC_OPTIONAL(off_group_b, off_dst[1], output_area_size, pipeline);
+
+ const size_t group_b_size = off_group_b - off_group_a;
+
+ // 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 = 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(off_group_c, off_act_f32, act_f32_size);
+
+ const size_t group_c_size = off_group_c - off_group_a;
+
+ L->weight_area_bytes = weight_area_size;
+ L->act_area_bytes = act_area_size;
+ L->act_f32_bytes = act_f32_size;
+ L->output_area_bytes = output_area_size;
+ L->scratch_bytes[0] = scratch0_size;
+ L->scratch_bytes[1] = scratch1_size;
+ L->act_head_stride = 0;
+
+ off = off_group_a + hex_smax(group_b_size, group_c_size);
+ }
-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) {
- (void)wtype;
- (void)pipeline;
- const size_t vec_dot_size = k * sizeof(uint16_t);
- const size_t f32_scratch_size = use_dma_activation
- ? htp_mm_round_up(act_threads * 4 * k * sizeof(float), HTP_MM_HMX_TILE_SIZE) : 0;
-
- const size_t act_head_stride = mc * k;
- const size_t weight_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
- const size_t act_area_size = htp_mm_round_up(group_size * act_head_stride * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
- const size_t output_area_size = htp_mm_round_up(group_size * mc * nc * sizeof(uint16_t), HTP_MM_HMX_TILE_SIZE);
- const size_t scratch_area_size = htp_mm_round_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE);
-
- return weight_area_size + act_area_size + output_area_size +
- 2 * scratch_area_size + 256 + f32_scratch_size;
+ L->total_bytes = off;
}
-static inline size_t htp_mm_hvx_get_vtcm_sizes(
+static inline void htp_mm_hvx_vtcm_layout_build(
+ struct htp_mm_hvx_vtcm_layout * L,
int kernel_type,
int wtype,
uint32_t ne10, // k
- uint32_t src1_nrows, // m_total (or act_nrows)
+ uint32_t src1_nrows, // m_total
uint32_t n_threads,
size_t dst_row_size,
size_t src0_row_size,
size_t src1_row_size,
uint32_t n_prefetch,
- size_t * vtcm_src0_size_out,
- size_t * vtcm_src1_size_out,
- size_t * vtcm_dst_size_out
+ bool is_matmul_id,
+ bool is_fused_qkv,
+ bool is_fused_ffn
) {
- size_t vtcm_src0_size = 0;
- size_t vtcm_src1_size = 0;
- size_t vtcm_dst_size = 0;
+ size_t src0_sz = 0;
+ size_t src1_sz = 0;
+ size_t src2_sz = 0;
+ size_t src3_sz = 0;
+ size_t dst_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 ||
wtype == HTP_TYPE_MXFP4);
- 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;
+ if (is_fused_qkv || is_fused_ffn) {
+ 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;
- switch (kernel_type) {
- case HTP_MM_KERNEL_HVX_F16_F16_VTCM: {
- size_t f16_src1_row_size = htp_mm_round_up(ne10 * 2, 128);
- vtcm_src1_size = htp_mm_round_up(f16_src1_row_size * src1_nrows, 256);
- vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
- vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
- break;
- }
- case HTP_MM_KERNEL_HVX_F16_F32_DDR:
- case HTP_MM_KERNEL_HVX_F16_F16_DDR:
- case HTP_MM_KERNEL_HVX_F32_F32_DDR:
- case HTP_MM_KERNEL_HVX_F32_F16_DDR: {
- vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads;
- vtcm_src1_size = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads;
- vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
- break;
- }
- case HTP_MM_KERNEL_HVX_F32_F32_VTCM: {
- size_t f32_src1_row_size = htp_mm_round_up(ne10 * 4, 128);
- vtcm_src1_size = htp_mm_round_up(f32_src1_row_size * src1_nrows, 256);
- vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads;
- vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
- break;
- }
- case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
- case HTP_MM_KERNEL_HVX_QUANT_ROW: {
- size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
-
- vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
- vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+ size_t src0_sz_per_thread = 0;
+ size_t src2_sz_per_thread = 0;
+ size_t src3_sz_per_thread = 0;
- vtcm_src0_size = vtcm_src0_size * n_threads;
+ if (is_repack) {
+ uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+ uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32;
+ uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- if (is_repack) {
- uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
- uint32_t n_k_tiles = ne10 / 32;
- uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
- vtcm_src0_size = repacked_vtcm_size * n_threads;
+ src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
+ src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
+ if (is_fused_qkv) {
+ src3_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128);
}
-
- size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
- size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
- if (dst_size_per_thread < quant_scratch_size_per_thread) {
- dst_size_per_thread = quant_scratch_size_per_thread;
+ } else {
+ src0_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
+ src2_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
+ if (is_fused_qkv) {
+ src3_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128);
}
- vtcm_dst_size = dst_size_per_thread * n_threads;
- break;
}
- case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
- size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
- vtcm_src0_size = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
- vtcm_src1_size = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+ size_t flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
+ size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
- vtcm_src0_size = vtcm_src0_size * n_threads;
+ if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) {
+ src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128);
+ } else {
+ src1_sz = hex_round_up(tiled_src1_row_size * src1_nrows, 128);
+ }
- if (is_repack) {
- uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
- uint32_t n_k_tiles = ne10 / 32;
- uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
- vtcm_src0_size = repacked_vtcm_size * n_threads;
- }
+ src0_sz = src0_sz_per_thread * n_threads;
+ src2_sz = src2_sz_per_thread * n_threads;
+ src3_sz = src3_sz_per_thread * n_threads;
+ dst_sz = quant_scratch_size;
+ } 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) ? htp_mm_q8_1_tiled_row_size(ne10)
+ : htp_mm_q8_0_tiled_row_size(ne10);
+
+ size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+ src1_sz = htp_mm_round_up(src1_row_size_tiled * src1_nrows, 256);
+
+ if (is_repack) {
+ const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+ const uint32_t n_k_tiles = ne10 / 32;
+ const uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+ size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
+ src0_sz_per_thread = repacked_vtcm_size;
+ }
- size_t quant_scratch_size_per_thread = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float));
- size_t dst_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
- if (dst_size_per_thread < quant_scratch_size_per_thread) {
- dst_size_per_thread = quant_scratch_size_per_thread;
+ 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;
+ } 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;
+
+ 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;
+ break;
+ }
+ case HTP_MM_KERNEL_HVX_F16_F32_DDR:
+ case HTP_MM_KERNEL_HVX_F16_F16_DDR:
+ case HTP_MM_KERNEL_HVX_F32_F32_DDR:
+ case HTP_MM_KERNEL_HVX_F32_F16_DDR: {
+ src0_sz = htp_mm_round_up(n_prefetch * src0_row_size, 256) * n_threads;
+ src1_sz = htp_mm_round_up(n_prefetch * src1_row_size, 256) * n_threads;
+ dst_sz = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) * n_threads : 0;
+ 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;
+ break;
+ }
+ case HTP_MM_KERNEL_HVX_QUANT_BLOCK:
+ case HTP_MM_KERNEL_HVX_QUANT_ROW: {
+ size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
+
+ src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+ src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+
+ src0_sz = src0_sz * n_threads;
+
+ if (is_repack) {
+ uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+ uint32_t n_k_tiles = ne10 / 32;
+ uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+ size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
+ 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_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+ if (dst_size_per_thread < quant_scratch_size_per_thread) {
+ dst_size_per_thread = quant_scratch_size_per_thread;
+ }
+ dst_sz = dst_size_per_thread * n_threads;
+ break;
}
- vtcm_dst_size = dst_size_per_thread * n_threads;
- break;
+ case HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT: {
+ size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10);
+
+ src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
+ src1_sz = htp_mm_round_up(q_src1_row_size * src1_nrows, 256);
+
+ src0_sz = src0_sz * n_threads;
+
+ if (is_repack) {
+ uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
+ uint32_t n_k_tiles = ne10 / 32;
+ uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
+ size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
+ 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_size_per_thread = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 0;
+ if (dst_size_per_thread < quant_scratch_size_per_thread) {
+ dst_size_per_thread = quant_scratch_size_per_thread;
+ }
+ dst_sz = dst_size_per_thread * n_threads;
+ break;
+ }
+ default:
+ break;
}
- default:
- break;
}
- *vtcm_src0_size_out = vtcm_src0_size;
- *vtcm_src1_size_out = vtcm_src1_size;
- *vtcm_dst_size_out = vtcm_dst_size;
-
- return vtcm_src0_size + vtcm_src1_size + vtcm_dst_size;
+ size_t off = 0;
+ VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz);
+ VTCM_LAYOUT_ALLOC(off, off_src0, src0_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);
+
+ 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;
}
-static inline size_t htp_mm_hvx_id_get_vtcm_sizes(
- int wtype,
- uint32_t ne10, // k
- uint32_t src1_nrows,
- uint32_t n_threads,
- size_t src0_row_size, // nb01
- uint32_t n_prefetch,
- size_t * vtcm_src0_size_out,
- size_t * vtcm_src1_size_out,
- size_t * vtcm_dst_size_out
+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
) {
- const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
- wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
- wtype == HTP_TYPE_MXFP4);
-
- const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128);
- const size_t src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10)
- : htp_mm_q8_0_tiled_row_size(ne10);
-
- size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256);
- size_t src1_sz = htp_mm_round_up(src1_row_size * src1_nrows, 256);
-
- if (is_repack) {
- const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype);
- const uint32_t n_k_tiles = ne10 / 32;
- const uint32_t tile_row_size = n_k_tiles * aligned_tile_size;
- size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256);
- src0_sz_per_thread = repacked_vtcm_size;
- }
-
- const size_t vtcm_src0_size = src0_sz_per_thread * n_threads;
- const size_t vtcm_dst_size = htp_mm_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads;
-
- *vtcm_src0_size_out = vtcm_src0_size;
- *vtcm_src1_size_out = src1_sz;
- *vtcm_dst_size_out = vtcm_dst_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);
+ return L.total_bytes;
+}
- return vtcm_src0_size + src1_sz + vtcm_dst_size;
+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) {
+ (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);
+ return L.total_bytes;
}
#ifdef __cplusplus
+++ /dev/null
-#ifndef VTCM_UTILS_H
-#define VTCM_UTILS_H
-
-#include "hex-utils.h"
-
-#include <assert.h>
-#include <stdint.h>
-#include <hexagon_types.h>
-
-static inline uint8_t *vtcm_seq_alloc(uint8_t **vtcm_ptr, size_t size) {
- uint8_t *p = *vtcm_ptr;
- *vtcm_ptr += size;
- return p;
-}
-
-#endif // VTCM_UTILS_H
#include "worker-pool.h"
+#include "hex-utils.h"
#include <qurt.h>
+#include <qurt_hvx.h>
+
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include "HAP_farf.h"
-#define WORKER_THREAD_STACK_SZ (2 * 16384)
#define LOWEST_USABLE_QURT_PRIO (254)
struct worker_pool_s;
FARF(HIGH, "worker-pool: thread %u started", me->id);
unsigned int prev_seqn = 0;
+ unsigned int poll_cnt = WORKER_POOL_POLL_COUNT;
while (!atomic_load(&pool->killed)) {
unsigned int seqn = atomic_load(&pool->seqn);
if (seqn == prev_seqn) {
- // Nothing to do
+ // drop HVX context while spinning
+ if (poll_cnt > 1 && poll_cnt == WORKER_POOL_POLL_COUNT) {
+ qurt_hvx_unlock();
+ }
+ if (--poll_cnt) {
+ hex_pause();
+ continue;
+ }
qurt_futex_wait(&pool->seqn, prev_seqn);
+ poll_cnt = WORKER_POOL_POLL_COUNT;
continue;
}
- // New job
prev_seqn = seqn;
+ poll_cnt = WORKER_POOL_POLL_COUNT;
+ // New job
unsigned int n = atomic_load(&pool->n_jobs);
unsigned int i = atomic_fetch_add(&pool->next_job, 1);
if (i >= n) {
void * data;
} worker_pool_job_t;
+#define WORKER_THREAD_STACK_SZ (2 * 16384)
+
/// Maximum supported number of worker threads.
#define MAX_NUM_WORKERS 10
+#if __HVX_ARCH__ > 79
+#define WORKER_POOL_POLL_COUNT 2000
+#else
+#define WORKER_POOL_POLL_COUNT 1
+#endif
+
// Initialize worker pool.
WORKERPOOL_API AEEResult worker_pool_init(worker_pool_context_t * context, uint32_t n_threads);
overview / pkg / ggml / sources / llama.cpp / commitdiff