mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-09 05:55:54 +02:00
Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Max Krasnyansky
|
81ff7abe50 |
@@ -2028,10 +2028,10 @@ static bool ggml_hexagon_precompute_flash_attn_params(
|
||||
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]);
|
||||
@@ -2385,31 +2385,30 @@ static void ggml_hexagon_precompute_hvx_mm_params(
|
||||
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) {
|
||||
@@ -2420,37 +2419,36 @@ static void ggml_hexagon_precompute_hvx_mm_params(
|
||||
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
|
||||
@@ -2458,17 +2456,17 @@ static void ggml_hexagon_precompute_hvx_mm_params(
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2478,19 +2476,19 @@ static void ggml_hexagon_precompute_hvx_mm_params(
|
||||
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) {
|
||||
@@ -2499,14 +2497,14 @@ static void ggml_hexagon_precompute_hvx_mm_params(
|
||||
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 {
|
||||
@@ -2514,31 +2512,31 @@ static void ggml_hexagon_precompute_hvx_mm_params(
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -2608,80 +2606,57 @@ static void ggml_hexagon_precompute_fused_qkv_params(
|
||||
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);
|
||||
|
||||
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;
|
||||
|
||||
size_t tiled_vtcm_size = src0_sz + src1_sz + src2_sz + src3_sz + quant_scratch_size;
|
||||
struct htp_mm_hvx_vtcm_layout L;
|
||||
bool try_tiled = (opt_mm_select >= 2);
|
||||
if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
|
||||
|
||||
// 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
|
||||
);
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -2701,72 +2676,55 @@ static void ggml_hexagon_precompute_fused_ffn_params(
|
||||
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);
|
||||
|
||||
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 tiled_vtcm_size = src0_sz + src1_sz + src2_sz + quant_scratch_size;
|
||||
struct htp_mm_hvx_vtcm_layout L;
|
||||
bool try_tiled = (opt_mm_select >= 2);
|
||||
if (try_tiled && tiled_vtcm_size <= sess->vtcm_size) {
|
||||
|
||||
// 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
|
||||
);
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,6 +20,7 @@ add_library(${HTP_LIB} SHARED
|
||||
worker-pool.c
|
||||
hex-dma.c
|
||||
hmx-queue.c
|
||||
gated-delta-net-ops.c
|
||||
binary-ops.c
|
||||
unary-ops.c
|
||||
sum-rows-ops.c
|
||||
@@ -37,10 +38,9 @@ add_library(${HTP_LIB} SHARED
|
||||
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
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
#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>
|
||||
|
||||
@@ -8,6 +8,7 @@
|
||||
#include <HAP_perf.h>
|
||||
#include <math.h>
|
||||
#include <stdbool.h>
|
||||
#include <stdatomic.h>
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
@@ -22,7 +23,7 @@
|
||||
#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
|
||||
@@ -142,6 +143,10 @@ struct hmx_fa_context {
|
||||
__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;
|
||||
};
|
||||
@@ -463,7 +468,7 @@ typedef struct {
|
||||
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;
|
||||
@@ -483,19 +488,19 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data)
|
||||
|
||||
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 {
|
||||
@@ -507,7 +512,8 @@ typedef struct {
|
||||
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;
|
||||
@@ -526,11 +532,11 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
|
||||
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));
|
||||
}
|
||||
@@ -538,7 +544,8 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data)
|
||||
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;
|
||||
@@ -547,7 +554,7 @@ static void fa_phase_v_interleave(struct hmx_fa_context * factx,
|
||||
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 {
|
||||
@@ -563,6 +570,9 @@ typedef struct {
|
||||
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) {
|
||||
@@ -587,9 +597,8 @@ 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);
|
||||
@@ -643,72 +652,63 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) {
|
||||
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];
|
||||
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;
|
||||
|
||||
assert(factx->DK == factx->DV);
|
||||
|
||||
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));
|
||||
|
||||
__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;
|
||||
|
||||
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 {
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
}
|
||||
|
||||
// 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 *) 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;
|
||||
|
||||
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);
|
||||
|
||||
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;
|
||||
|
||||
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_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);
|
||||
|
||||
__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;
|
||||
|
||||
*pv_out0 = Q6_V_lo_W(vp);
|
||||
*pv_out1 = Q6_V_hi_W(vp);
|
||||
}
|
||||
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));
|
||||
@@ -726,7 +726,18 @@ static void fa_phase_q_load(struct hmx_fa_context * factx,
|
||||
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 {
|
||||
@@ -798,11 +809,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
|
||||
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);
|
||||
|
||||
@@ -831,10 +841,10 @@ static void fa_o_store_thread_f16(unsigned int n, unsigned int i, void * data) {
|
||||
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 {
|
||||
@@ -957,14 +967,14 @@ static inline void fa_softmax_impl(
|
||||
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);
|
||||
@@ -974,14 +984,14 @@ static inline void fa_softmax_impl(
|
||||
}
|
||||
} 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;
|
||||
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);
|
||||
my_row_buf0[ci] = Q6_V_lo_W(vp_s_dual_row);
|
||||
my_row_buf1[ci] = Q6_V_hi_W(vp_s_dual_row);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1118,9 +1128,9 @@ static inline void fa_softmax_impl(
|
||||
|
||||
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);
|
||||
@@ -1150,7 +1160,7 @@ static inline void fa_softmax_impl(
|
||||
}
|
||||
|
||||
// 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);
|
||||
|
||||
@@ -1160,7 +1170,7 @@ static inline void fa_softmax_impl(
|
||||
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);
|
||||
@@ -1331,14 +1341,17 @@ static void hmx_fa_qk_dot_worker(void * data) {
|
||||
__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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1373,17 +1386,21 @@ static void hmx_fa_o_update_worker(void * data) {
|
||||
__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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1409,14 +1426,17 @@ static void hmx_fa_o_norm_worker(void * data) {
|
||||
__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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1475,7 +1495,7 @@ static void fa_push_mask_dma_gqa(
|
||||
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) {
|
||||
@@ -1484,7 +1504,7 @@ static void fa_push_mask_dma_gqa(
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1582,62 +1602,57 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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);
|
||||
// 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);
|
||||
|
||||
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) {
|
||||
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));
|
||||
@@ -1655,11 +1670,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
|
||||
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;
|
||||
@@ -1669,28 +1679,44 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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);
|
||||
|
||||
// Prefetch first KV block
|
||||
// 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);
|
||||
}
|
||||
|
||||
// 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 ----
|
||||
@@ -1709,12 +1735,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
const size_t k_src_stride = size_k_row_padded / sizeof(__fp16);
|
||||
const size_t v_src_stride = size_v_row_padded / sizeof(__fp16);
|
||||
|
||||
if (factx.pipeline) {
|
||||
// ==================================================================
|
||||
// Pipeline path
|
||||
// ==================================================================
|
||||
struct hmx_queue * hmx_q = ctx->hmx_queue;
|
||||
struct hmx_queue * hmx_q = ctx->hmx_queue;
|
||||
|
||||
if (factx.pipeline) {
|
||||
// Pipeline 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);
|
||||
@@ -1724,15 +1748,22 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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) {
|
||||
@@ -1749,7 +1780,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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;
|
||||
@@ -1762,16 +1796,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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);
|
||||
@@ -1838,24 +1865,21 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
}
|
||||
|
||||
} 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);
|
||||
@@ -1865,31 +1889,29 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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);
|
||||
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;
|
||||
|
||||
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);
|
||||
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) {
|
||||
@@ -1922,33 +1944,23 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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);
|
||||
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;
|
||||
|
||||
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_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
|
||||
hmx_queue_pop(ctx->hmx_queue);
|
||||
|
||||
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);
|
||||
hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
|
||||
}
|
||||
|
||||
@@ -1962,37 +1974,15 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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 ----
|
||||
@@ -2001,12 +1991,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
}
|
||||
}
|
||||
|
||||
if (factx.pipeline) {
|
||||
hmx_queue_suspend(ctx->hmx_queue);
|
||||
} else {
|
||||
HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
|
||||
}
|
||||
|
||||
return HTP_STATUS_OK;
|
||||
}
|
||||
|
||||
@@ -2040,10 +2024,10 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
|
||||
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;
|
||||
|
||||
@@ -7,19 +7,23 @@
|
||||
|
||||
#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 {
|
||||
@@ -54,6 +58,11 @@ struct htp_fa_kernel_params {
|
||||
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;
|
||||
@@ -69,10 +78,6 @@ struct htp_fa_kernel_params {
|
||||
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;
|
||||
};
|
||||
@@ -81,39 +86,124 @@ struct htp_fa_kernel_params {
|
||||
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
|
||||
@@ -153,23 +243,8 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
|
||||
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;
|
||||
|
||||
@@ -185,51 +260,26 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
|
||||
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);
|
||||
// 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;
|
||||
|
||||
// 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) {
|
||||
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;
|
||||
}
|
||||
|
||||
// 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) {
|
||||
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) {
|
||||
@@ -237,7 +287,7 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
|
||||
}
|
||||
}
|
||||
|
||||
if (best_Br == 0) {
|
||||
if (best_Br == 0 || best_Bc == 0) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
#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
|
||||
|
||||
@@ -47,22 +48,75 @@ static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) =
|
||||
};
|
||||
// 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,
|
||||
@@ -70,17 +124,71 @@ static inline void hmx_fa_o_update_tile(
|
||||
__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(
|
||||
@@ -88,9 +196,360 @@ 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 */
|
||||
|
||||
@@ -506,7 +506,8 @@ static void dequantize_tiled_weight_to_fp16_task_q8_0(
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
|
||||
@@ -543,17 +544,13 @@ static void convert_f16_weight_to_fp16_tiles_task(
|
||||
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) {
|
||||
|
||||
@@ -594,120 +591,178 @@ static void quantize_f32_weight_to_fp16_tiles_task(
|
||||
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;
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
__fp16 *out_tile = output + (r * n_col_tiles + c) * HTP_MM_HMX_TILE_N_ELMS;
|
||||
Q6_mxmem_AR_after_hf(out_tile, 0);
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// C += AB
|
||||
static void core_mma_chunk_fp16(__fp16 *restrict c, const __fp16 *restrict a, const __fp16 *restrict b,
|
||||
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;
|
||||
|
||||
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) {
|
||||
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;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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;
|
||||
asm volatile(HMX_SET_BIAS("%0") :: "r"((unsigned int)col_scales));
|
||||
|
||||
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);
|
||||
}
|
||||
const size_t dot_tile_stride = n_dot_tiles * HTP_MM_HMX_TILE_N_ELMS;
|
||||
|
||||
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;
|
||||
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
|
||||
@@ -901,148 +956,55 @@ static void transfer_activation_chunk_fp32_to_fp16(__fp16 *restrict vtcm_dst, co
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
bool row0_valid,
|
||||
bool row1_valid) {
|
||||
|
||||
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);
|
||||
uint32_t r0 = r / HTP_MM_HMX_TILE_N_ROWS; // tile row index
|
||||
uint32_t r1 = r % HTP_MM_HMX_TILE_N_ROWS; // intra-tile row idx
|
||||
|
||||
const uint32_t n_steps = n_rows_padded / R;
|
||||
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);
|
||||
|
||||
// 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);
|
||||
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
|
||||
|
||||
uint32_t c0 = c / HTP_MM_HMX_TILE_N_COLS; // tile column index
|
||||
uint32_t tile_idx = r0 * (k_block / HTP_MM_HMX_TILE_N_COLS) + c0;
|
||||
|
||||
HVX_Vector *tile = (HVX_Vector *) (vtcm_dst + tile_idx * HTP_MM_HMX_TILE_N_ELMS);
|
||||
tile[r1 / 2] = v_out;
|
||||
}
|
||||
if (c < k_block) {
|
||||
HVX_Vector v0 = Q6_V_vzero();
|
||||
HVX_Vector v1 = Q6_V_vzero();
|
||||
if (row0_valid) v0 = *(const HVX_Vector *)(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;
|
||||
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());
|
||||
|
||||
if (r < n_rows) {
|
||||
dma_queue_pop(dma_q);
|
||||
}
|
||||
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
|
||||
|
||||
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;
|
||||
|
||||
uint32_t c = 0;
|
||||
for (; c + 32 <= k_valid; c += 32) {
|
||||
HVX_Vector v0 = Q6_V_vzero();
|
||||
HVX_Vector v1 = Q6_V_vzero();
|
||||
if (row0_valid) v0 = *(const HVX_Vector *)(ptr_in0 + c);
|
||||
if (row1_valid) v1 = *(const HVX_Vector *)(ptr_in1 + c);
|
||||
|
||||
HVX_Vector v_out = hvx_vec_f32_to_f16_shuff(v0, v1);
|
||||
|
||||
uint32_t 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;
|
||||
}
|
||||
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,
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
|
||||
#include <qurt_thread.h>
|
||||
#include <qurt_futex.h>
|
||||
#include <qurt_hvx.h>
|
||||
|
||||
#include <HAP_compute_res.h>
|
||||
|
||||
@@ -42,6 +43,7 @@ static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) {
|
||||
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);
|
||||
@@ -70,9 +72,14 @@ static void hmx_queue_thread(void * arg) {
|
||||
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;
|
||||
|
||||
@@ -18,13 +18,19 @@ extern "C" {
|
||||
#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
|
||||
};
|
||||
@@ -97,7 +103,7 @@ static inline uint32_t hmx_queue_capacity(struct hmx_queue * q) {
|
||||
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;
|
||||
|
||||
@@ -120,13 +126,28 @@ static inline struct hmx_queue_desc hmx_queue_pop(struct hmx_queue * q) {
|
||||
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
|
||||
|
||||
@@ -197,4 +197,26 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
|
||||
}
|
||||
}
|
||||
|
||||
// --- 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
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
#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
|
||||
@@ -948,6 +948,8 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
|
||||
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;
|
||||
|
||||
@@ -976,6 +978,8 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
|
||||
}
|
||||
}
|
||||
|
||||
hmx_queue_suspend(ctx->hmx_queue);
|
||||
|
||||
struct htp_opbatch_rsp rsp;
|
||||
rsp.id = req.id;
|
||||
rsp.status = op_status;
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
#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,
|
||||
@@ -1514,37 +1514,26 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
|
||||
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],
|
||||
@@ -1556,10 +1545,10 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) {
|
||||
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;
|
||||
@@ -1948,14 +1937,95 @@ static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void
|
||||
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);
|
||||
@@ -1964,18 +2034,55 @@ static void transfer_activation_chunk_worker_fn(unsigned int n, unsigned int i,
|
||||
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;
|
||||
@@ -2112,32 +2219,89 @@ static void transfer_activation_chunk_threaded(
|
||||
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,
|
||||
@@ -2198,48 +2362,33 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
|
||||
|
||||
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);
|
||||
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);
|
||||
|
||||
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);
|
||||
|
||||
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);
|
||||
@@ -2254,107 +2403,118 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
|
||||
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]));
|
||||
|
||||
// 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);
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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);
|
||||
@@ -2365,7 +2525,6 @@ static int hmx_mm_2d_f32(struct htp_context *ctx,
|
||||
}
|
||||
}
|
||||
}
|
||||
HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
|
||||
}
|
||||
|
||||
return 0;
|
||||
@@ -2458,37 +2617,34 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
|
||||
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);
|
||||
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);
|
||||
|
||||
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) {
|
||||
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) {
|
||||
@@ -2505,58 +2661,59 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
|
||||
// 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;
|
||||
|
||||
const size_t nc_next = nc + n_chunk_n_cols;
|
||||
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 * 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2565,8 +2722,6 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_
|
||||
}
|
||||
}
|
||||
|
||||
HAP_compute_res_hmx_unlock(ctx->vtcm_rctx);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -2758,7 +2913,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
|
||||
|
||||
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);
|
||||
@@ -2768,37 +2923,52 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx,
|
||||
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) {
|
||||
@@ -2960,22 +3130,14 @@ static int hvx_mm_matmul_id(
|
||||
}
|
||||
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],
|
||||
@@ -2989,11 +3151,11 @@ static int hvx_mm_matmul_id(
|
||||
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;
|
||||
@@ -3003,10 +3165,10 @@ static int hvx_mm_matmul_id(
|
||||
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;
|
||||
@@ -3181,19 +3343,11 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
|
||||
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",
|
||||
@@ -3201,12 +3355,12 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
|
||||
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;
|
||||
@@ -3219,11 +3373,11 @@ int op_matmul_qkv(struct htp_ops_context * octx) {
|
||||
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;
|
||||
@@ -3331,28 +3485,22 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
|
||||
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;
|
||||
@@ -3363,10 +3511,10 @@ int op_matmul_ffn(struct htp_ops_context * octx) {
|
||||
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;
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
#include "htp-ops.h"
|
||||
#include "hex-fastdiv.h"
|
||||
#include "hex-common.h"
|
||||
#include "htp-vtcm.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
@@ -44,7 +45,7 @@ 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,
|
||||
@@ -295,197 +296,351 @@ static inline void htp_mm_hmx_get_batched_chunk_costs(
|
||||
*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
|
||||
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;
|
||||
};
|
||||
|
||||
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;
|
||||
};
|
||||
|
||||
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
|
||||
) {
|
||||
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);
|
||||
size_t off = 0;
|
||||
|
||||
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;
|
||||
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);
|
||||
}
|
||||
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);
|
||||
|
||||
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;
|
||||
|
||||
return weight_area_size + act_area_size + act_f32_size + output_area_size +
|
||||
scratch0_size + scratch1_size + scratch2_size + 256;
|
||||
L->total_bytes = off;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
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);
|
||||
size_t src0_sz_per_thread = 0;
|
||||
size_t src2_sz_per_thread = 0;
|
||||
size_t src3_sz_per_thread = 0;
|
||||
|
||||
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);
|
||||
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;
|
||||
|
||||
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 = 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 (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;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
vtcm_dst_size = dst_size_per_thread * n_threads;
|
||||
break;
|
||||
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);
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
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;
|
||||
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);
|
||||
|
||||
return vtcm_src0_size + vtcm_src1_size + vtcm_dst_size;
|
||||
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);
|
||||
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;
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
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
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
#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
|
||||
@@ -1,6 +1,9 @@
|
||||
#include "worker-pool.h"
|
||||
#include "hex-utils.h"
|
||||
|
||||
#include <qurt.h>
|
||||
#include <qurt_hvx.h>
|
||||
|
||||
#include <stdatomic.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
@@ -9,7 +12,6 @@
|
||||
|
||||
#include "HAP_farf.h"
|
||||
|
||||
#define WORKER_THREAD_STACK_SZ (2 * 16384)
|
||||
#define LOWEST_USABLE_QURT_PRIO (254)
|
||||
|
||||
struct worker_pool_s;
|
||||
@@ -42,17 +44,27 @@ static void worker_pool_main(void * context) {
|
||||
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) {
|
||||
|
||||
@@ -24,9 +24,17 @@ typedef struct {
|
||||
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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user