Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
0a50d9909a3478e82679f505bf8595d1eee4b0a8
llama.cpp/ggml/src/ggml-hexagon/htp/hmx-fa-kernels.h
T
Max Krasnyansky 81ff7abe50 hexagon: new vtcm layouts and improved pipelines for MUL_MAT, MUL_MAT_ID and FLASH_ATTN_EXT (#25425) 
* hex-fa: refactor kernel param compute to use common layout builder

* hmx: add explicit compiler barriers to make hmx funcs more robust

* hex-vtcm: more generic vtcm layout builder for mm and flash-attn kernels

* hex-hmx: unroll inner kernels

* hex-hmx: use inline asm instead of intrinsics to avoid compiler issues

* hex-hmx: define inline asm macros and simplify code

* hex-hmx: replace leftover intrinsics

* hmx-fa: minor cleanup for hmx asm

* hmx-mm: move per-task stucts out of the kernels header

* hmx-mm: simplify core_dot_chunk

* hmx-mm: simplify inner loops that call hmx instructions

* hmx-mm: proper instrumentation for activation prep work for dma pipelined version

* hmx-mm: update a-prep loop for better prefetch

* hex-vtcm: improved vtcm layout alloc for mm to support overlapping areas

* hmx-mm: reduce the number of act fetch tows to 4 for now, going larger doesnt help here

* hex-hmx: always use hmx-queue in all modes

* hmx-mm: update comments and minor formatting

* hmx-mm: further improve synchro fallback path to prefetch the weights earlier

* hex-fa: further pipeline improvements (earlier prefetch)

* hmx-mm: cleanup dma pipelines to use dst cached in the queue

* hmx-fa: minor cleanup and opts for fa dma pipelines

* hmx-fa: optimize q-prep stage with dma and unrolling

* hmx-fa: use o_tile size from layout instead of computing it

* hmx-mm: cleanup types and size handling

* hmx-mm: replace divs with fastdiv in qprep loops

* hmx-fa: minor update/formatting to q_tile handling

* hmx-fa: cleanup the layout to avoid overpadding

* hmx-fa: simplified and improved cost mode for hmx fa solver that uses vtcm layout funcs

* hmx-queue: add support queue wakeup and make suspend async to avoid hmx-lock latency

* hex-hmx: move queue wakeup / suspend to the op-batch level

* hex-threads: add hybrid polling to workpool

* hex-mm: fix trailing spaces
2026-07-08 07:38:27 -07:00

556 lines
22 KiB
C
Raw Blame History

#ifndef HMX_FA_KERNELS_H
#define HMX_FA_KERNELS_H
#include <stdint.h>
#include <stddef.h>
#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
// Scatter offsets for diagonal tile: entry[2i] = i*136, entry[2i+1] = i*136+6
// 136 = 4 * 32 + 8 = byte offset to diagonal in a 32x32 fp16 interleaved tile
static const int16_t d_tile_scatter_offsets[64] __attribute__((aligned(128))) = {
0 * 136, 0 * 136 + 6,
1 * 136, 1 * 136 + 6,
2 * 136, 2 * 136 + 6,
3 * 136, 3 * 136 + 6,
4 * 136, 4 * 136 + 6,
5 * 136, 5 * 136 + 6,
6 * 136, 6 * 136 + 6,
7 * 136, 7 * 136 + 6,
8 * 136, 8 * 136 + 6,
9 * 136, 9 * 136 + 6,
10 * 136, 10 * 136 + 6,
11 * 136, 11 * 136 + 6,
12 * 136, 12 * 136 + 6,
13 * 136, 13 * 136 + 6,
14 * 136, 14 * 136 + 6,
15 * 136, 15 * 136 + 6,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
0, 0,
};
// Inner HMX tile computation kernels
static void hmx_fa_qk_dot_tile(
const __fp16 * row_tiles,
const __fp16 * col_tiles,
__fp16 * out_tile,
size_t n_dot_tiles
) {
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;
}
}
asm volatile(
HMX_STORE_AFTER_F16("%0", "%1")
:
: "r"(out_tile), "r"(0)
: "memory"
);
}
static void hmx_fa_o_update_tile(
const __fp16 * d_diag,
const __fp16 * o_rc,
const __fp16 * p_tile_in,
const __fp16 * v_tile_in,
__fp16 * o_tile_out,
size_t n_col_tiles
) {
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;
}
}
asm volatile(
HMX_STORE_AFTER_F16("%0", "%1")
:
: "r"(o_tile_out), "r"(0)
: "memory"
);
}
static inline void hmx_fa_o_norm_tile(
const __fp16 * d_diag,
const __fp16 * o_rc,
__fp16 * o_out
) {
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 */
Reference in New Issue View Git Blame Copy Permalink