hexagon: flash attention rework (optimizations, accuracy improvements, etc) (#25085)
* hex-mm: fold mm quant tasks into the main matmul threads * hex-mm: minor formatting fixes * hex-mm: cleanup is_quant checks in dma dispatch * hex-mm: fix dst-spad alignment * hex-mm: move fp kernels in the hvx-mm-kernels header * hex-mm: fuse with ADD * hex-fa: factor out ukernels into separate headers and unify the rest * hex-fa: move kernel-params compute into the host * hex-fa: refactor vtcm alloc for consistency * hex-fa: add support for FA_SELECT * hex-fa: update tracing insrumentation to cover all functions * hex-fa: update hvx fallback thresholds to recover t/g regressions * hex-fa: update tracing instrumentation * hex-fa: improved tracing with additional events * hex-fa: optimize mask processing (fastdiv, etc) * hex-fa: improve mask dma caching * hmx-fa: change loop order to maximize mask cache hits * hex-fa: remove over instrumentation * hex-fa: breakdown QKV prep trace events * hmx-fa: further mask proc optimizations * hex-fa: mask broadcast is the common case, optimize for that * hex-fa: use aligned loads where possible * hex-fa: update loops to use uint32_t indices * hmx-fa: fold vtcm init into q prep task * hex-fa: update rest of the hmx funcs to use uint32_t * hmx-fa: fold build_d into the main softmax loop * hmx-fa: start kv dmas earlier * hmx-fa: start mask dma a bit earlier * hex-fa: precompute rows per task to avoid divs * hmx-fa: specialize fa_o_store for f16 and f32 * hmx-fa: prelim support for Sinks * hmx-fa: keep softmax accumulators in fp32 * hex-fa: add tanh_f16 and exp2_f16 and use that in FA * hex-fa: use fp16 math in the hvx kernel * hex-fa: avoid expensive float -> __fp16 cast for slopes and softcap * hex-fa: replace most vec_exp_f32 with vec_exp2_f16 * hmx-fa: vectorize sinks update * hex-fa: minor formatting * hmx-fa: fold softcap loop into the tile load * hmx-fa: use vectoralias to populate sinks * hex-fa: remove redudant check * hex-fa: fix vtcm size compute to use fp32 for accumulators * hex-mm: fix trailing spaces * hmx-fa: dont use -inf to init mask to avoid conversion overflows * hex-fa: no need to explicitly guard -inf in the f16->f32 converter now * hmx-fa: cleanup fa sinks handling * hex-mm: fixed src2 stride handling when mm is fused with add * hex-fa: make lto happy
This commit is contained in:
Max Krasnyansky
@@ -20,9 +20,6 @@ add_library(${HTP_LIB} SHARED
|
||||
worker-pool.c
|
||||
hex-dma.c
|
||||
hmx-queue.c
|
||||
flash-attn-ops.c
|
||||
hmx-flash-attn-ops.c
|
||||
matmul-ops.c
|
||||
binary-ops.c
|
||||
unary-ops.c
|
||||
sum-rows-ops.c
|
||||
@@ -42,16 +39,14 @@ add_library(${HTP_LIB} SHARED
|
||||
solve-tri-ops.c
|
||||
gated-delta-net-ops.c
|
||||
pad-ops.c
|
||||
matmul-ops.c
|
||||
flash-attn-ops.c
|
||||
)
|
||||
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE
|
||||
$<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,HTP_DEBUG=1,NDEBUG=1>
|
||||
$<IF:$<BOOL:${HEXAGON_HTP_DEBUG}>,FARF_HIGH=1,>)
|
||||
|
||||
if (GGML_HEXAGON_FA_EXP2_HF)
|
||||
message(STATUS "ggml-htp: HMX_FA_USE_EXP2_HF=1 (use FP16 exp2 polynomial in FA softmax)")
|
||||
target_compile_definitions(${HTP_LIB} PRIVATE HMX_FA_USE_EXP2_HF=1)
|
||||
endif()
|
||||
|
||||
build_idl(htp_iface.idl ${HTP_LIB})
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,253 @@
|
||||
#ifndef HTP_FLASH_ATTN_OPS_H
|
||||
#define HTP_FLASH_ATTN_OPS_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stddef.h>
|
||||
#include <stdbool.h>
|
||||
|
||||
#include "hex-fastdiv.h"
|
||||
#include "hex-common.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
// Tile constants (mirrored from hmx-utils.h for use on host side if needed)
|
||||
#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 {
|
||||
HTP_FA_KERNEL_UNSUPPORTED = 0,
|
||||
HTP_FA_KERNEL_HVX,
|
||||
HTP_FA_KERNEL_HMX
|
||||
};
|
||||
|
||||
struct htp_fa_kernel_params {
|
||||
uint8_t kernel_type; // enum htp_fa_kernel_type
|
||||
uint8_t is_q_fp32; // 1 = Q type is F32, 0 = F16
|
||||
uint8_t is_dst_fp32; // 1 = dst type is F32, 0 = F16
|
||||
uint8_t n_threads; // Number of threads to run
|
||||
|
||||
// Common parameters
|
||||
uint16_t Br;
|
||||
uint16_t Bc;
|
||||
uint16_t n_kv_blocks; // also HVX's n_blocks
|
||||
uint16_t G; // GQA factor (n_heads / n_kv_heads)
|
||||
|
||||
float scale;
|
||||
float max_bias;
|
||||
float logit_softcap;
|
||||
uint32_t vtcm_size;
|
||||
|
||||
uint32_t qrows;
|
||||
uint32_t qrows_per_thread;
|
||||
float m0;
|
||||
float m1;
|
||||
uint32_t n_head_log2;
|
||||
|
||||
struct fastdiv_values src3_div2;
|
||||
struct fastdiv_values src3_div3;
|
||||
|
||||
union {
|
||||
struct {
|
||||
uint32_t g_br;
|
||||
uint32_t row_buf_stride;
|
||||
uint32_t mask_buf_row_stride;
|
||||
int32_t mask_broadcast;
|
||||
int32_t pipeline;
|
||||
struct fastdiv_values div_G;
|
||||
} hmx;
|
||||
struct {
|
||||
uint32_t size_q_row_padded;
|
||||
uint32_t size_k_row_padded;
|
||||
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;
|
||||
};
|
||||
|
||||
#if defined(__cplusplus)
|
||||
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) {
|
||||
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 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)
|
||||
}
|
||||
|
||||
#define FA_HVX_BLOCK_SIZE 64
|
||||
|
||||
static inline size_t hvx_fa_compute_vtcm_usage(size_t DK, size_t DV, bool is_q_fp32, bool has_mask, size_t n_threads) {
|
||||
const size_t size_q_row_padded = hex_round_up(DK * (is_q_fp32 ? 4 : 2), 128);
|
||||
const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128);
|
||||
const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128);
|
||||
|
||||
const size_t size_q_block = size_q_row_padded * 1;
|
||||
const size_t size_k_block = size_k_row_padded * FA_HVX_BLOCK_SIZE;
|
||||
const size_t size_v_block = size_v_row_padded * FA_HVX_BLOCK_SIZE;
|
||||
const size_t size_m_block = hex_round_up(FA_HVX_BLOCK_SIZE * sizeof(__fp16), 128);
|
||||
const size_t size_vkq_acc = hex_round_up(DV * sizeof(float), 128);
|
||||
|
||||
const size_t size_per_thread = size_q_block * 1
|
||||
+ size_k_block * 2
|
||||
+ size_v_block * 2
|
||||
+ (has_mask ? size_m_block * HVX_FA_DMA_CACHE_SIZE : 0)
|
||||
+ size_vkq_acc;
|
||||
|
||||
return size_per_thread * n_threads;
|
||||
}
|
||||
|
||||
#define FA_MIN_KV_BLOCKS 3
|
||||
|
||||
// Cost-based (Br, Bc) search for flash attention with pipeline constraint.
|
||||
static inline int hmx_fa_find_chunk_size(size_t * Br_out,
|
||||
size_t * Bc_out,
|
||||
size_t gqa_factor,
|
||||
size_t DK,
|
||||
size_t DV,
|
||||
size_t qo_len,
|
||||
size_t kv_len,
|
||||
size_t vtcm_budget,
|
||||
size_t n_threads) {
|
||||
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;
|
||||
|
||||
// Pipeline constraint: cap Bc so n_kv_blocks >= FA_MIN_KV_BLOCKS.
|
||||
// Only relax when kv_len is too short to form enough blocks.
|
||||
const size_t Bc_limit = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
|
||||
(kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
|
||||
// Cost coefficients calibrated from profiling
|
||||
const size_t c_q_fixed = 1400; // per-Q-block: q_load + epilogue o_update + o_norm + o_store
|
||||
const size_t c_iter_fixed = 200; // per-KV-iter: HMX queue push/pop + DMA pop + barriers
|
||||
|
||||
size_t best_cost = SIZE_MAX, best_mn = 0;
|
||||
size_t best_Br = 0, best_Bc = 0;
|
||||
|
||||
for (size_t Br = Br_max; Br >= br_unit; Br -= br_unit) {
|
||||
const size_t g_br = hex_align_up(gqa_factor * Br, T);
|
||||
|
||||
// g_br-dependent VTCM cost: g_br * per_gbr + g_br*g_br * per_gbr2
|
||||
const size_t gbr_cost = g_br * per_gbr + g_br * g_br * per_gbr2;
|
||||
if (gbr_cost >= usable) {
|
||||
if (Br == br_unit) {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Analytically solve for max Bc:
|
||||
// remain >= Bc * (per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE)
|
||||
// The Br * fp16 term accounts for the VTCM mask buffer [Br * Bc].
|
||||
const size_t remain = usable - gbr_cost;
|
||||
const size_t bc_denom = per_bc + g_br * per_gbr_bc + Br * fp16 * HMX_FA_DMA_CACHE_SIZE;
|
||||
size_t Bc = hex_smin(hex_align_down(remain / bc_denom, bc_unit), Bc_limit);
|
||||
if (Bc < bc_unit) {
|
||||
if (Br == br_unit) {
|
||||
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) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (best_Br == 0) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
*Br_out = best_Br;
|
||||
*Bc_out = best_Bc;
|
||||
return 0;
|
||||
}
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif /* HTP_FLASH_ATTN_OPS_H */
|
||||
@@ -138,27 +138,28 @@ static inline bool dma_queue_push_single_1d(dma_queue * q, dma_ptr dptr, size_t
|
||||
}
|
||||
|
||||
dma_descriptor_1d * desc = (dma_descriptor_1d *) &q->desc[q->push_idx];
|
||||
desc->next = NULL;
|
||||
desc->desc_size = 0; // 1D mode
|
||||
desc->src_bypass = dma_src_l2_bypass_on;
|
||||
desc->dst_bypass = dma_dst_l2_bypass_on;
|
||||
desc->order = 0;
|
||||
desc->done = 0;
|
||||
desc->src = (void *) dptr.src;
|
||||
desc->dst = (void *) dptr.dst;
|
||||
desc->size = size;
|
||||
desc->src = (void *) dptr.src;
|
||||
desc->dst = (void *) dptr.dst;
|
||||
desc->size = size;
|
||||
|
||||
q->dptr[q->push_idx] = dptr;
|
||||
|
||||
if (size) {
|
||||
desc->next = NULL;
|
||||
desc->desc_size = 0; // 1D mode
|
||||
desc->src_bypass = dma_src_l2_bypass_on;
|
||||
desc->dst_bypass = dma_dst_l2_bypass_on;
|
||||
desc->order = 0;
|
||||
desc->done = 0;
|
||||
|
||||
htp_trace_event_start(q->trace, HTP_TRACE_EVT_DMA, q->push_idx);
|
||||
dmlink(q->tail, desc);
|
||||
q->tail = (dma_descriptor_2d *) desc;
|
||||
} else {
|
||||
desc->done = 1;
|
||||
desc->desc_size = 0;
|
||||
desc->done = 1;
|
||||
}
|
||||
|
||||
// FARF(ERROR, "dma-push: i %u row-size %u nrows %d dst %p src %p\n", q->push_idx, row_size, nrows, dptr.dst, dptr.src);
|
||||
q->push_idx = (q->push_idx + 1) & q->idx_mask;
|
||||
return true;
|
||||
}
|
||||
@@ -320,7 +321,7 @@ static inline bool dma_queue_push_vtcm_to_ddr(dma_queue * q, dma_ptr dptr, size_
|
||||
return dma_queue_push(q, dptr, dst_row_size, src_row_size, dst_row_size, nrows);
|
||||
}
|
||||
|
||||
#define DMA_CACHE_MAX_SIZE 64U
|
||||
#define DMA_CACHE_MAX_SIZE 256U
|
||||
|
||||
typedef struct {
|
||||
uint8_t *base;
|
||||
@@ -352,20 +353,19 @@ static inline bool dma_cache_push(dma_queue *q, dma_cache *c, const uint8_t * sr
|
||||
if (c->src[i] == (uint32_t) src) {
|
||||
c->age[i] = 0;
|
||||
dst = c->base + (i * c->line_size); nrows = 0; // dummy dma
|
||||
// FARF(ERROR, "dma-cache: found %p", src);
|
||||
} else {
|
||||
c->age[i]++;
|
||||
if (c->age[i] > o_age) { o_age = c->age[i]; o_idx = i; }
|
||||
}
|
||||
}
|
||||
if (!dst) {
|
||||
// FARF(ERROR, "dma-cache: replacing #%u : age %u %p -> %p", o_idx, c->age[o_idx], (void *) c->src[o_idx], src);
|
||||
c->age[o_idx] = 0;
|
||||
c->src[o_idx] = (uint32_t) src;
|
||||
dst = c->base + o_idx * c->line_size; // normal nrows dma
|
||||
return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
|
||||
}
|
||||
|
||||
return dma_queue_push(q, dma_make_ptr(dst, src), dst_stride, src_stride, row_size, nrows);
|
||||
return dma_queue_push_single_1d(q, dma_make_ptr(dst, src), 0);
|
||||
}
|
||||
|
||||
#ifdef __cplusplus
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
#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"
|
||||
|
||||
// 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 inline 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;
|
||||
}
|
||||
Q6_mxmem_AR_after_hf(out_tile, 0);
|
||||
}
|
||||
|
||||
static inline 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
|
||||
) {
|
||||
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;
|
||||
}
|
||||
|
||||
Q6_mxmem_AR_after_hf(o_tile_out, 0);
|
||||
}
|
||||
|
||||
static inline void hmx_fa_o_norm_tile(
|
||||
const __fp16 * d_diag,
|
||||
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);
|
||||
}
|
||||
|
||||
#endif /* HMX_FA_KERNELS_H */
|
||||
File diff suppressed because it is too large
Load Diff
@@ -712,7 +712,17 @@ static inline void hmx_matmul_job_init(hmx_matmul_job_t * job,
|
||||
|
||||
// output : fp16 -> f32p
|
||||
|
||||
static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16 *restrict vtcm_src, uint32_t start_row, uint32_t n_rows, uint32_t n_cols, uint32_t dst_stride, uint32_t dst_cols) {
|
||||
static void transfer_output_chunk_fp16_to_fp32(
|
||||
float *restrict dst,
|
||||
const float *restrict src2,
|
||||
const __fp16 *restrict vtcm_src,
|
||||
uint32_t start_row,
|
||||
uint32_t n_rows,
|
||||
uint32_t n_cols,
|
||||
uint32_t dst_stride,
|
||||
uint32_t src2_stride,
|
||||
uint32_t dst_cols
|
||||
) {
|
||||
assert(n_cols % HTP_MM_HMX_TILE_N_COLS == 0);
|
||||
const size_t tile_row_stride = (n_cols / HTP_MM_HMX_TILE_N_COLS) * HTP_MM_HMX_TILE_N_ELMS;
|
||||
|
||||
@@ -727,6 +737,7 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
|
||||
const size_t r1 = (r_idx0 % HTP_MM_HMX_TILE_N_ROWS) / 2; // index of the row pair within the tile
|
||||
const __fp16 *row_base = vtcm_src + r0 * tile_row_stride;
|
||||
float *output_row_base = dst + r * dst_stride; // global memory row base for row r (and r+1)
|
||||
const float *src2_row_base = src2 ? (src2 + r * src2_stride) : NULL;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (size_t c = 0; c < limit_c_aligned; c += HTP_MM_HMX_TILE_N_COLS) {
|
||||
@@ -738,9 +749,20 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
|
||||
HVX_Vector *pv_out0 = (HVX_Vector *) (output_row_base + c + 0);
|
||||
HVX_Vector *pv_out1 = (HVX_Vector *) (output_row_base + c + dst_stride);
|
||||
|
||||
*pv_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
|
||||
HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
|
||||
if (src2_row_base) {
|
||||
HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0);
|
||||
v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0);
|
||||
}
|
||||
*pv_out0 = v_out0;
|
||||
|
||||
if (r + 1 < n_rows) {
|
||||
*pv_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
|
||||
HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
|
||||
if (src2_row_base) {
|
||||
HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride);
|
||||
v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1);
|
||||
}
|
||||
*pv_out1 = v_out1;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -752,9 +774,20 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
|
||||
HVX_Vector v = ((const HVX_Vector *) tile)[r1];
|
||||
HVX_VectorPair vp = Q6_Wqf32_vmpy_VhfVhf(v, one);
|
||||
|
||||
hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp)));
|
||||
HVX_Vector v_out0 = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(vp));
|
||||
if (src2_row_base) {
|
||||
HVX_Vector v_src2_0 = hvx_vmemu(src2_row_base + c + 0);
|
||||
v_out0 = hvx_vec_add_f32_f32(v_out0, v_src2_0);
|
||||
}
|
||||
hvx_vec_store_u(output_row_base + c, valid_c * sizeof(float), v_out0);
|
||||
|
||||
if (r + 1 < n_rows) {
|
||||
hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp)));
|
||||
HVX_Vector v_out1 = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(vp));
|
||||
if (src2_row_base) {
|
||||
HVX_Vector v_src2_1 = hvx_vmemu(src2_row_base + c + src2_stride);
|
||||
v_out1 = hvx_vec_add_f32_f32(v_out1, v_src2_1);
|
||||
}
|
||||
hvx_vec_store_u(output_row_base + c + dst_stride, valid_c * sizeof(float), v_out1);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -763,11 +796,13 @@ static void transfer_output_chunk_fp16_to_fp32(float *restrict dst, const __fp16
|
||||
typedef struct {
|
||||
const __fp16 *vtcm_src;
|
||||
float *dst;
|
||||
const float *src2;
|
||||
uint32_t n_tasks;
|
||||
uint32_t n_tot_chunks;
|
||||
uint32_t n_chunks_per_task;
|
||||
uint32_t n_cols;
|
||||
uint32_t dst_stride; // DDR row stride
|
||||
uint32_t src2_stride; // DDR row stride for residual
|
||||
uint32_t dst_cols; // Actual output columns
|
||||
struct htp_thread_trace * traces;
|
||||
} output_transfer_task_state_t;
|
||||
|
||||
@@ -42,14 +42,14 @@ static const int32_t hmx_transpose_scatter_offsets[32] __attribute__((aligned(VL
|
||||
// Full range: start_row=0, end_row=n_cols.
|
||||
static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
const __fp16 * restrict vtcm_src,
|
||||
int n_cols,
|
||||
int k,
|
||||
int src_stride,
|
||||
int start_row,
|
||||
int end_row) {
|
||||
uint32_t n_cols,
|
||||
uint32_t k,
|
||||
size_t src_stride,
|
||||
uint32_t start_row,
|
||||
uint32_t end_row) {
|
||||
assert(k % HMX_FP16_TILE_N_COLS == 0);
|
||||
|
||||
const int n_k_tiles = k / HMX_FP16_TILE_N_COLS;
|
||||
const uint32_t n_k_tiles = k / HMX_FP16_TILE_N_COLS;
|
||||
const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets);
|
||||
const HVX_Vector v_scat_step = Q6_V_vsplat_R(4);
|
||||
const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64);
|
||||
@@ -65,14 +65,14 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
|
||||
if (pair_scatter) {
|
||||
// Step c by 64 fp16 (two K-tiles per scatter), advance dst by 2 tiles per iter.
|
||||
const int c_step = 2 * HMX_FP16_TILE_N_COLS;
|
||||
const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
|
||||
const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS;
|
||||
const int n_c_iters = k / c_step;
|
||||
const uint32_t c_step = 2 * HMX_FP16_TILE_N_COLS;
|
||||
const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
|
||||
const size_t dst_step = 2 * (size_t) HMX_FP16_TILE_N_ELMS;
|
||||
const uint32_t n_c_iters = k / c_step;
|
||||
|
||||
for (int r = start_row; r < end_row; r += 2) {
|
||||
const int ct = r / HMX_FP16_TILE_N_ROWS;
|
||||
const int local_r = r % HMX_FP16_TILE_N_ROWS;
|
||||
for (uint32_t r = start_row; r < end_row; r += 2) {
|
||||
const uint32_t ct = r / HMX_FP16_TILE_N_ROWS;
|
||||
const uint32_t local_r = r % HMX_FP16_TILE_N_ROWS;
|
||||
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
|
||||
const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
|
||||
const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
|
||||
@@ -86,7 +86,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
assert(c_byte_step % 128 == 0);
|
||||
|
||||
if (p1) {
|
||||
for (int i = 0; i < n_c_iters; ++i) {
|
||||
for (uint32_t i = 0; i < n_c_iters; ++i) {
|
||||
HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
|
||||
HVX_Vector v1 = hvx_vmem(p1); p1 += c_byte_step;
|
||||
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
|
||||
@@ -95,7 +95,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
}
|
||||
} else {
|
||||
const HVX_Vector vzero = Q6_V_vzero();
|
||||
for (int i = 0; i < n_c_iters; ++i) {
|
||||
for (uint32_t i = 0; i < n_c_iters; ++i) {
|
||||
HVX_Vector v0 = hvx_vmem(p0); p0 += c_byte_step;
|
||||
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off0, v0);
|
||||
Q6_vscatter_RMVwV((size_t) tile_base, pair_region, v_off1, vzero);
|
||||
@@ -105,14 +105,14 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
}
|
||||
} else {
|
||||
// Fallback: scatter one K-tile per call (region 2047, masked).
|
||||
const int c_step = HMX_FP16_TILE_N_COLS;
|
||||
const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
|
||||
const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS;
|
||||
const int n_c_iters = k / c_step;
|
||||
const uint32_t c_step = HMX_FP16_TILE_N_COLS;
|
||||
const size_t c_byte_step = (size_t) c_step * sizeof(__fp16);
|
||||
const size_t dst_step = (size_t) HMX_FP16_TILE_N_ELMS;
|
||||
const uint32_t n_c_iters = k / c_step;
|
||||
|
||||
for (int r = start_row; r < end_row; r += 2) {
|
||||
const int ct = r / HMX_FP16_TILE_N_ROWS;
|
||||
const int local_r = r % HMX_FP16_TILE_N_ROWS;
|
||||
for (uint32_t r = start_row; r < end_row; r += 2) {
|
||||
const uint32_t ct = r / HMX_FP16_TILE_N_ROWS;
|
||||
const uint32_t local_r = r % HMX_FP16_TILE_N_ROWS;
|
||||
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_cols;
|
||||
const HVX_Vector v_off0 = Q6_Vw_vadd_VwVw(v_scat_base, Q6_V_vsplat_R(local_r * 4));
|
||||
const HVX_Vector v_off1 = Q6_Vw_vadd_VwVw(v_off0, v_scat_step);
|
||||
@@ -122,7 +122,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
const uint8_t * p1 = next_row_valid ? (const uint8_t *) (vtcm_src + (r + 1) * src_stride) : NULL;
|
||||
|
||||
if (p1) {
|
||||
for (int i = 0; i < n_c_iters; ++i) {
|
||||
for (uint32_t i = 0; i < n_c_iters; ++i) {
|
||||
HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
|
||||
HVX_Vector v1 = hvx_vmemu(p1); p1 += c_byte_step;
|
||||
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
|
||||
@@ -131,7 +131,7 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
}
|
||||
} else {
|
||||
const HVX_Vector vzero = Q6_V_vzero();
|
||||
for (int i = 0; i < n_c_iters; ++i) {
|
||||
for (uint32_t i = 0; i < n_c_iters; ++i) {
|
||||
HVX_Vector v0 = hvx_vmemu(p0); p0 += c_byte_step;
|
||||
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off0, v0);
|
||||
Q6_vscatter_QRMVwV(q_mask64, (size_t) tile_base, single_region, v_off1, vzero);
|
||||
@@ -148,24 +148,24 @@ static inline void hmx_interleave_rows_to_tiles(__fp16 * restrict vtcm_dst,
|
||||
// Full range: start_row=0, end_row=n_rows.
|
||||
static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
|
||||
const __fp16 * restrict src,
|
||||
int n_rows,
|
||||
int head_dim,
|
||||
int src_stride,
|
||||
int n_row_tiles,
|
||||
int start_row,
|
||||
int end_row) {
|
||||
uint32_t n_rows,
|
||||
uint32_t head_dim,
|
||||
size_t src_stride,
|
||||
uint32_t n_row_tiles,
|
||||
uint32_t start_row,
|
||||
uint32_t end_row) {
|
||||
__builtin_assume(head_dim > 0);
|
||||
const size_t tile_stride_elms = (size_t) n_row_tiles * HMX_FP16_TILE_N_ELMS;
|
||||
|
||||
for (int r = start_row; r < end_row; r += 2) {
|
||||
for (uint32_t r = start_row; r < end_row; r += 2) {
|
||||
const bool next_row_valid = (r + 1) < end_row && (r + 1) < n_rows;
|
||||
|
||||
const HVX_Vector * pv_in0 = (const HVX_Vector *) (src + r * src_stride);
|
||||
const HVX_Vector * pv_in1 = next_row_valid ? (const HVX_Vector *) (src + (r + 1) * src_stride) : NULL;
|
||||
|
||||
// Row-pair invariants hoisted out of the c loop.
|
||||
const int r0 = r / HMX_FP16_TILE_N_ROWS;
|
||||
const int r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2;
|
||||
const uint32_t r0 = r / HMX_FP16_TILE_N_ROWS;
|
||||
const uint32_t r1_half = (r % HMX_FP16_TILE_N_ROWS) / 2;
|
||||
|
||||
// tb0 starts at tile (c0=0, r0); tb1 at the adjacent dim-tile (c0=1, r0).
|
||||
// Each c step (+= 64) advances both by 2 dim-tiles worth of fp16.
|
||||
@@ -174,7 +174,7 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
|
||||
const size_t tb_step = 2 * tile_stride_elms;
|
||||
|
||||
if (pv_in1) {
|
||||
for (int c = 0; c < head_dim; c += 64) {
|
||||
for (uint32_t c = 0; c < head_dim; c += 64) {
|
||||
HVX_Vector v0 = *pv_in0++;
|
||||
HVX_Vector v1 = *pv_in1++;
|
||||
HVX_VectorPair vp = Q6_W_vshuff_VVR(v1, v0, -2);
|
||||
@@ -185,7 +185,7 @@ static inline void hmx_interleave_cols_to_tiles(__fp16 * restrict tiles_out,
|
||||
}
|
||||
} else {
|
||||
const HVX_Vector vzero = Q6_V_vzero();
|
||||
for (int c = 0; c < head_dim; c += 64) {
|
||||
for (uint32_t c = 0; c < head_dim; c += 64) {
|
||||
HVX_Vector v0 = *pv_in0++;
|
||||
HVX_VectorPair vp = Q6_W_vshuff_VVR(vzero, v0, -2);
|
||||
((HVX_Vector *) tb0)[r1_half] = Q6_V_lo_W(vp);
|
||||
|
||||
@@ -60,6 +60,7 @@ enum htp_op_code {
|
||||
HTP_OP_MUL_MAT_ID,
|
||||
HTP_OP_MUL_MAT_QKV,
|
||||
HTP_OP_MUL_MAT_FFN,
|
||||
HTP_OP_MUL_MAT_ADD,
|
||||
HTP_OP_RMS_NORM,
|
||||
HTP_OP_RMS_NORM_MUL,
|
||||
HTP_OP_UNARY_SILU,
|
||||
@@ -175,6 +176,11 @@ enum htp_trace_event_id {
|
||||
HTP_TRACE_EVT_HVX_W_DEQUANT = 23,
|
||||
HTP_TRACE_EVT_HVX_W_PREP = 24,
|
||||
HTP_TRACE_EVT_HVX_O_PROC = 25,
|
||||
HTP_TRACE_EVT_HVX_FA_QK = 26,
|
||||
HTP_TRACE_EVT_HVX_FA_SFM = 27,
|
||||
HTP_TRACE_EVT_HVX_FA_Q_PREP = 28,
|
||||
HTP_TRACE_EVT_HVX_FA_K_PREP = 29,
|
||||
HTP_TRACE_EVT_HVX_FA_V_PREP = 30,
|
||||
|
||||
HTP_TRACE_EVT_HMX_COMP = 40,
|
||||
};
|
||||
|
||||
@@ -134,16 +134,7 @@ static inline HVX_Vector hvx_vec_f32_to_f16_shuff(HVX_Vector v0, HVX_Vector v1)
|
||||
}
|
||||
|
||||
static inline HVX_Vector hvx_vec_f32_to_f16(HVX_Vector v0, HVX_Vector v1) {
|
||||
HVX_Vector v = Q6_Vh_vdeal_Vh(hvx_vec_f32_to_f16_shuff(v0, v1));
|
||||
|
||||
#if __HVX_ARCH__ < 79
|
||||
// replace NaNs with -INF, older arches produce NaNs for (-INF + 0.0)
|
||||
const HVX_Vector neg_inf = hvx_vec_splat_f16(-INFINITY);
|
||||
HVX_VectorPred nan = hvx_vec_is_nan_f16(v);
|
||||
v = Q6_V_vmux_QVV(nan, neg_inf, v);
|
||||
#endif
|
||||
|
||||
return v;
|
||||
return Q6_Vh_vdeal_Vh(hvx_vec_f32_to_f16_shuff(v0, v1));
|
||||
}
|
||||
|
||||
#if __HVX_ARCH__ >= 79
|
||||
@@ -170,8 +161,6 @@ static inline HVX_VectorPair hvx_vec_f16_to_f32(HVX_Vector v) {
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
|
||||
static inline HVX_Vector hvx_vec_i16_from_hf_rnd_sat(HVX_Vector vin) {
|
||||
// This looks complicated.
|
||||
// Ideally should just be Q6_Vh_equals_Vhf(vin)
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
#define EXP_COEFF_0 (0x3F000000) // 0.5 = 1/(2!)
|
||||
#define EXP_LOGN2 (0x3F317218) // ln(2) = 0.6931471805
|
||||
#define EXP_LOG2E (0x3FB8AA3B) // log2(e) = 1/ln(2) = 1.4426950408
|
||||
#define EXP_LOG2E_F 1.44269504f
|
||||
#define EXP_ONE (0x3f800000) // 1.0
|
||||
#define EXP_RANGE_R (0x42B17218) // ln(FLT_MAX) approx = 88.7228
|
||||
#define EXP_RANGE_L (0xC2B00000) // -88.0 (approx log(FLT_MIN))
|
||||
@@ -213,4 +214,42 @@ static inline void hvx_exp_f32(uint8_t * restrict dst, const uint8_t * restrict
|
||||
}
|
||||
}
|
||||
|
||||
static inline HVX_Vector hvx_vec_exp2_f16(HVX_Vector x_v) {
|
||||
const HVX_Vector zero_v = Q6_V_vzero();
|
||||
const HVX_Vector half_hf_v = Q6_Vh_vsplat_R(0x3800); // fp16 0.5
|
||||
|
||||
// Clamp input to prevent integer underflow in FP16-to-INT16 conversion
|
||||
const HVX_Vector v_clamp_min = hvx_vec_splat_f16(-24.0f);
|
||||
x_v = Q6_Vhf_vmax_VhfVhf(v_clamp_min, x_v);
|
||||
|
||||
// k = round_toward_neg_inf(x); f = (float)k; frac = x - f
|
||||
HVX_Vector x_minus_half = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(x_v, half_hf_v));
|
||||
HVX_Vector k_v = Q6_Vh_equals_Vhf(x_minus_half); // truncate to int16
|
||||
HVX_Vector f_v = Q6_Vhf_equals_Vh(k_v); // back to fp16
|
||||
|
||||
HVX_Vector x_qf16 = Q6_Vqf16_vsub_VhfVhf(x_v, f_v); // fractional part in qf16
|
||||
|
||||
// Horner: y = ((((E5*x + E4)*x + E3)*x + E2)*x + E1)*x + E0
|
||||
HVX_Vector y = Q6_Vqf16_vmpy_Vqf16Vqf16(Q6_Vh_vsplat_R(0x5082), x_qf16); // E5*x
|
||||
y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x157d)); // + E4
|
||||
y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
|
||||
y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x20ed)); // + E3
|
||||
y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
|
||||
y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x2b1b)); // + E2
|
||||
y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
|
||||
y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x33b0)); // + E1
|
||||
y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16);
|
||||
y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x398c)); // + E0
|
||||
y = Q6_Vqf16_vmpy_Vqf16Vqf16(y, x_qf16); // y = y * x
|
||||
y = Q6_Vqf16_vadd_Vqf16Vhf(y, Q6_Vh_vsplat_R(0x3c00)); // + 1.0
|
||||
|
||||
// Combine polynomial (mantissa) with integer part (exponent): result = y * 2^k
|
||||
y = Q6_Vhf_equals_Vqf16(y);
|
||||
HVX_Vector y_exp = Q6_Vuh_vlsr_VuhR(Q6_Vh_vasl_VhR(y, 1), 11);
|
||||
y_exp = Q6_Vh_vadd_VhVh(k_v, y_exp);
|
||||
HVX_VectorPred q_underflow = Q6_Q_vcmp_gt_VhVh(zero_v, y_exp);
|
||||
y = Q6_Vh_vaslacc_VhVhR(y, k_v, 10);
|
||||
return Q6_V_vmux_QVV(q_underflow, zero_v, y);
|
||||
}
|
||||
|
||||
#endif /* HVX_EXP_H */
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
#ifndef HVX_FA_KERNELS_H
|
||||
#define HVX_FA_KERNELS_H
|
||||
|
||||
#include <assert.h>
|
||||
#include <math.h>
|
||||
#include "hvx-utils.h"
|
||||
|
||||
// Little inner kernels for HVX
|
||||
|
||||
#if __HVX_ARCH__ < 79
|
||||
#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
|
||||
#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b))
|
||||
#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
|
||||
#else
|
||||
#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
|
||||
#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b)
|
||||
#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
|
||||
#endif
|
||||
|
||||
// This is a bit of a hack because the compiler is struggling to properly inline
|
||||
// the default hvx_vec_f32_to_f16 with output into the local array.
|
||||
static __attribute__((unused)) __attribute__((noinline)) void hvx_vec_f32_to_f16_a(void *ptr, HVX_Vector v0, HVX_Vector v1)
|
||||
{
|
||||
*(HVX_Vector *) ptr = hvx_vec_f32_to_f16(v0, v1);
|
||||
}
|
||||
|
||||
// Dot product of two F16 vectors, accumulating to float
|
||||
static inline void hvx_dot_f16_f16_aa(float * restrict r, const void * restrict x, const void * restrict y, unsigned int n, float s) {
|
||||
const HVX_Vector * restrict vx = (const HVX_Vector * restrict) x; // fp16
|
||||
const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
HVX_VectorPair rsum_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, vx[i], vy[i]);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
|
||||
HVX_Vector x_hf = Q6_V_vand_QV(bmask, vx[i]);
|
||||
|
||||
rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
|
||||
}
|
||||
|
||||
HVX_Vector rsum = HVX_OP_ADD_F32(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p));
|
||||
rsum = HVX_OP_MUL_F32(hvx_vec_splat_f32(s), hvx_vec_reduce_sum_f32(rsum));
|
||||
hvx_vec_store_u(r, 4, rsum);
|
||||
}
|
||||
|
||||
static inline HVX_Vector hvx_dot_f16_f16_aa_rx4(const void * restrict y,
|
||||
const uint8_t * restrict x,
|
||||
const size_t stride_x,
|
||||
const size_t nvec,
|
||||
const size_t nloe) {
|
||||
const HVX_Vector * restrict vx0 = (const HVX_Vector * restrict) x; // fp16
|
||||
const HVX_Vector * restrict vx1 = (const HVX_Vector * restrict) (x + stride_x); // fp16
|
||||
const HVX_Vector * restrict vx2 = (const HVX_Vector * restrict) (x + stride_x * 2); // fp16
|
||||
const HVX_Vector * restrict vx3 = (const HVX_Vector * restrict) (x + stride_x * 3); // fp16
|
||||
const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
|
||||
|
||||
HVX_VectorPair rsum0_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
|
||||
HVX_VectorPair rsum1_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
|
||||
HVX_VectorPair rsum2_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
|
||||
HVX_VectorPair rsum3_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector y_hf = vy[i];
|
||||
HVX_Vector x0_hf = vx0[i];
|
||||
HVX_Vector x1_hf = vx1[i];
|
||||
HVX_Vector x2_hf = vx2[i];
|
||||
HVX_Vector x3_hf = vx3[i];
|
||||
|
||||
rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
|
||||
rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
|
||||
rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
|
||||
rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
// Load x (fp16) and zero-out unused elements
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
|
||||
HVX_Vector x0_hf = Q6_V_vand_QV(bmask, vx0[i]);
|
||||
HVX_Vector x1_hf = Q6_V_vand_QV(bmask, vx1[i]);
|
||||
HVX_Vector x2_hf = Q6_V_vand_QV(bmask, vx2[i]);
|
||||
HVX_Vector x3_hf = Q6_V_vand_QV(bmask, vx3[i]);
|
||||
|
||||
rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
|
||||
rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
|
||||
rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
|
||||
rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
|
||||
}
|
||||
|
||||
HVX_Vector rsum0 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p));
|
||||
HVX_Vector rsum1 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p));
|
||||
HVX_Vector rsum2 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum2_p), Q6_V_hi_W(rsum2_p));
|
||||
HVX_Vector rsum3 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum3_p), Q6_V_hi_W(rsum3_p));
|
||||
|
||||
HVX_Vector_x4 rsum0123 = { .v = { rsum0, rsum1, rsum2, rsum3 } };
|
||||
return hvx_vec_reduce_sum_f32x4(rsum0123);
|
||||
}
|
||||
|
||||
static inline HVX_Vector hvx_dot_f16_f16_aa_rx32(const void * restrict y,
|
||||
const uint8_t * restrict x,
|
||||
const size_t stride_x,
|
||||
const size_t n,
|
||||
float s) {
|
||||
|
||||
const size_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
const size_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
HVX_Vector sums = Q6_V_vzero();
|
||||
const size_t stride_x_4 = stride_x * 4;
|
||||
for (uint32_t j = 0; j < VLEN_FP32; j += 4) {
|
||||
HVX_Vector sums_x4 = hvx_dot_f16_f16_aa_rx4(y, x, stride_x, nvec, nloe);
|
||||
HVX_VectorPred pred = Q6_Q_vsetq_R(j * SIZEOF_FP32);
|
||||
sums = Q6_V_vmux_QVV(pred, sums, sums_x4);
|
||||
x += stride_x_4;
|
||||
}
|
||||
|
||||
return HVX_OP_MUL_F32(hvx_vec_splat_f32(s), sums);
|
||||
}
|
||||
|
||||
// MAD: y (F32) += x (F16) * s (F16)
|
||||
static inline void hvx_mad_f32_f16_aa(float * restrict y, const void * restrict x, const __fp16 * restrict s, uint32_t n) {
|
||||
const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
|
||||
|
||||
HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
|
||||
HVX_Vector * restrict vy = (HVX_Vector *) y;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
HVX_Vector S0 = hvx_vec_splat_f16(*s);
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; ++i) {
|
||||
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPair xy_p = vy_p[i];
|
||||
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
|
||||
|
||||
HVX_Vector xy = Q6_V_lo_W(xy_p);
|
||||
i = 2 * i; // index for vy
|
||||
|
||||
if (nloe >= VLEN_FP32) {
|
||||
vy[i] = xy;
|
||||
nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
hvx_vec_store_a(&vy[i], nloe * 4, xy);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MAD: y (F32) += x0 (F16) * s0 (F16) + x1 (F16) * s1 (F16)
|
||||
static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restrict x0, const void * restrict x1,
|
||||
const __fp16 * restrict s0, const __fp16 * restrict s1, uint32_t n) {
|
||||
const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
|
||||
const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
|
||||
|
||||
HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
|
||||
HVX_Vector * restrict vy = (HVX_Vector *) y;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
HVX_Vector S0 = hvx_vec_splat_f16(*s0);
|
||||
HVX_Vector S1 = hvx_vec_splat_f16(*s1);
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; ++i) {
|
||||
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
|
||||
vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPair xy_p = vy_p[i];
|
||||
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
|
||||
xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
|
||||
|
||||
HVX_Vector xy = Q6_V_lo_W(xy_p);
|
||||
i = 2 * i; // index for vy
|
||||
|
||||
if (nloe >= VLEN_FP32) {
|
||||
vy[i] = xy;
|
||||
nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
hvx_vec_store_a(&vy[i], nloe * 4, xy);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) {
|
||||
assert((size_t) dst % 128 == 0);
|
||||
assert((size_t) src % 128 == 0);
|
||||
|
||||
const HVX_Vector * restrict vsrc = (const HVX_Vector * restrict) src;
|
||||
HVX_Vector * restrict vdst = (HVX_Vector * restrict) dst;
|
||||
|
||||
const uint32_t nvec = n / VLEN_FP32;
|
||||
const uint32_t nloe = n % VLEN_FP32;
|
||||
|
||||
uint32_t i = 0;
|
||||
#pragma unroll(4)
|
||||
for (; i < nvec; ++i) {
|
||||
vdst[i] = HVX_OP_MUL_F32(vsrc[i], vs);
|
||||
}
|
||||
if (nloe) {
|
||||
hvx_vec_store_a(&vdst[i], nloe * sizeof(float), HVX_OP_MUL_F32(vsrc[i], vs));
|
||||
}
|
||||
}
|
||||
|
||||
#endif /* HVX_FA_KERNELS_H */
|
||||
@@ -256,7 +256,7 @@ static inline void quantize_f16_f16_flat_kernel(
|
||||
|
||||
// Dot kernels that consume flat (non-tiled) activations
|
||||
|
||||
static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -312,10 +312,14 @@ static void flat_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const v
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -397,11 +401,19 @@ static void flat_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -464,10 +476,14 @@ static void flat_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const v
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -561,11 +577,19 @@ static void flat_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -620,10 +644,14 @@ static void flat_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const v
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -704,11 +732,19 @@ static void flat_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -765,10 +801,14 @@ static void flat_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -851,11 +891,19 @@ static void flat_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -921,10 +969,14 @@ static void flat_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const
|
||||
|
||||
v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -1019,6 +1071,441 @@ static void flat_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
|
||||
v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
#if __HVX_ARCH__ < 79
|
||||
#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
|
||||
#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
|
||||
#else
|
||||
#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
|
||||
#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
|
||||
#endif
|
||||
|
||||
static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
|
||||
const HVX_Vector * restrict x = (const HVX_Vector *) vx;
|
||||
const HVX_Vector * restrict y = (const HVX_Vector *) vy;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP32; // leftover elements
|
||||
|
||||
HVX_Vector rsum = Q6_V_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector prod = HVX_OP_MUL_F32(x[i], y[i]);
|
||||
rsum = HVX_OP_ADD_F32(rsum, prod);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector x_sf = Q6_V_vand_QV(bmask, x[i]);
|
||||
HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
|
||||
HVX_Vector prod = HVX_OP_MUL_F32(x_sf, y_sf);
|
||||
rsum = HVX_OP_ADD_F32(rsum, prod);
|
||||
}
|
||||
|
||||
*s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum));
|
||||
}
|
||||
|
||||
static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0,
|
||||
const void * restrict vx0, const void * restrict vx1,
|
||||
const void * restrict vy0) {
|
||||
const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
|
||||
const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
|
||||
const HVX_Vector * restrict y = (const HVX_Vector *) vy0;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP32;
|
||||
uint32_t nloe = n % VLEN_FP32;
|
||||
|
||||
HVX_Vector rsum0 = Q6_V_vzero();
|
||||
HVX_Vector rsum1 = Q6_V_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector y_sf = y[i];
|
||||
HVX_Vector prod0 = HVX_OP_MUL_F32(x0[i], y_sf);
|
||||
HVX_Vector prod1 = HVX_OP_MUL_F32(x1[i], y_sf);
|
||||
rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
|
||||
rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
HVX_Vector y_sf = Q6_V_vand_QV(bmask, y[i]);
|
||||
HVX_Vector x0_sf = Q6_V_vand_QV(bmask, x0[i]);
|
||||
HVX_Vector x1_sf = Q6_V_vand_QV(bmask, x1[i]);
|
||||
HVX_Vector prod0 = HVX_OP_MUL_F32(x0_sf, y_sf);
|
||||
HVX_Vector prod1 = HVX_OP_MUL_F32(x1_sf, y_sf);
|
||||
rsum0 = HVX_OP_ADD_F32(rsum0, prod0);
|
||||
rsum1 = HVX_OP_ADD_F32(rsum1, prod1);
|
||||
}
|
||||
|
||||
HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
|
||||
hvx_vec_store_u(s0, 8, rsum);
|
||||
}
|
||||
|
||||
static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
|
||||
const void * restrict vx0, const void * restrict vx1,
|
||||
const void * restrict vy0, const void * restrict vy1) {
|
||||
const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
|
||||
const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
|
||||
const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
|
||||
const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP32;
|
||||
uint32_t nloe = n % VLEN_FP32;
|
||||
|
||||
HVX_Vector r0_c0_sum = Q6_V_vzero();
|
||||
HVX_Vector r0_c1_sum = Q6_V_vzero();
|
||||
HVX_Vector r1_c0_sum = Q6_V_vzero();
|
||||
HVX_Vector r1_c1_sum = Q6_V_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector r0_sf = x0[i];
|
||||
HVX_Vector r1_sf = x1[i];
|
||||
HVX_Vector c0_sf = y0[i];
|
||||
HVX_Vector c1_sf = y1[i];
|
||||
|
||||
r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
|
||||
r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
|
||||
r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
|
||||
r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
|
||||
HVX_Vector r0_sf = Q6_V_vand_QV(bmask, x0[i]);
|
||||
HVX_Vector r1_sf = Q6_V_vand_QV(bmask, x1[i]);
|
||||
HVX_Vector c0_sf = Q6_V_vand_QV(bmask, y0[i]);
|
||||
HVX_Vector c1_sf = Q6_V_vand_QV(bmask, y1[i]);
|
||||
|
||||
r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(r0_sf, c0_sf));
|
||||
r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(r0_sf, c1_sf));
|
||||
r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(r1_sf, c0_sf));
|
||||
r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(r1_sf, c1_sf));
|
||||
}
|
||||
|
||||
// Reduce and store results
|
||||
HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
|
||||
HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
|
||||
|
||||
hvx_vec_store_u(s0, 8, r0_r1_c0_sum);
|
||||
hvx_vec_store_u(s1, 8, r0_r1_c1_sum);
|
||||
}
|
||||
|
||||
static inline void vec_dot_f32_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
|
||||
const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
|
||||
const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP32; // num full fp32 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP32; // leftover elements
|
||||
|
||||
HVX_Vector rsum = Q6_V_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector x_sf = vx[i];
|
||||
HVX_Vector y_sf = vy[i];
|
||||
|
||||
rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_Vector x_sf = vx[i];
|
||||
HVX_Vector y_sf = vy[i];
|
||||
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 4);
|
||||
x_sf = Q6_V_vand_QV(bmask, x_sf);
|
||||
y_sf = Q6_V_vand_QV(bmask, y_sf);
|
||||
|
||||
rsum = HVX_OP_ADD_F32(rsum, HVX_OP_MUL_F32(x_sf, y_sf));
|
||||
}
|
||||
|
||||
rsum = hvx_vec_reduce_sum_f32(rsum);
|
||||
hvx_vec_store_u(&s[0], 4, rsum);
|
||||
}
|
||||
|
||||
#undef HVX_OP_ADD_F32
|
||||
#undef HVX_OP_MUL_F32
|
||||
|
||||
static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
|
||||
const HVX_Vector * restrict x = (const HVX_Vector *) vx;
|
||||
const HVX_Vector * restrict y = (const HVX_Vector *) vy;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
HVX_VectorPair rsum_p = Q6_W_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x[i], y[i]);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
|
||||
HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
|
||||
rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
|
||||
}
|
||||
|
||||
HVX_Vector rsum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p)));
|
||||
hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum));
|
||||
}
|
||||
|
||||
static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0,
|
||||
const void * restrict vx0, const void * restrict vx1,
|
||||
const void * restrict vy0) {
|
||||
const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
|
||||
const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
|
||||
const HVX_Vector * restrict y = (const HVX_Vector *) vy0;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16;
|
||||
uint32_t nloe = n % VLEN_FP16;
|
||||
|
||||
HVX_VectorPair rsum0_p = Q6_W_vzero();
|
||||
HVX_VectorPair rsum1_p = Q6_W_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector y_hf = y[i];
|
||||
rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0[i], y_hf);
|
||||
rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1[i], y_hf);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
|
||||
HVX_Vector x0_hf = Q6_V_vand_QV(bmask, x0[i]);
|
||||
HVX_Vector x1_hf = Q6_V_vand_QV(bmask, x1[i]);
|
||||
rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
|
||||
rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
|
||||
}
|
||||
|
||||
HVX_Vector rsum0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p)));
|
||||
HVX_Vector rsum1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p)));
|
||||
HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1);
|
||||
hvx_vec_store_u(s0, 8, rsum);
|
||||
}
|
||||
|
||||
static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1,
|
||||
const void * restrict vx0, const void * restrict vx1,
|
||||
const void * restrict vy0, const void * restrict vy1) {
|
||||
const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0;
|
||||
const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1;
|
||||
const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0;
|
||||
const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16;
|
||||
uint32_t nloe = n % VLEN_FP16;
|
||||
|
||||
// Row sums (sf) - 4 accumulators for 2x2 tile
|
||||
HVX_VectorPair r0_c0_sum_p = Q6_W_vzero();
|
||||
HVX_VectorPair r0_c1_sum_p = Q6_W_vzero();
|
||||
HVX_VectorPair r1_c0_sum_p = Q6_W_vzero();
|
||||
HVX_VectorPair r1_c1_sum_p = Q6_W_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_Vector r0_hf = x0[i];
|
||||
HVX_Vector r1_hf = x1[i];
|
||||
HVX_Vector c0_hf = y0[i];
|
||||
HVX_Vector c1_hf = y1[i];
|
||||
|
||||
// Compute 4 dot products: r0xc0, r0xc1, r1xc0, r1xc1
|
||||
r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
|
||||
r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
|
||||
r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
|
||||
r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
|
||||
HVX_Vector r0_hf = Q6_V_vand_QV(bmask, x0[i]);
|
||||
HVX_Vector r1_hf = Q6_V_vand_QV(bmask, x1[i]);
|
||||
HVX_Vector c0_hf = Q6_V_vand_QV(bmask, y0[i]);
|
||||
HVX_Vector c1_hf = Q6_V_vand_QV(bmask, y1[i]);
|
||||
|
||||
r0_c0_sum_p = hvx_vec_mpyacc_f32_f16(r0_c0_sum_p, r0_hf, c0_hf);
|
||||
r0_c1_sum_p = hvx_vec_mpyacc_f32_f16(r0_c1_sum_p, r0_hf, c1_hf);
|
||||
r1_c0_sum_p = hvx_vec_mpyacc_f32_f16(r1_c0_sum_p, r1_hf, c0_hf);
|
||||
r1_c1_sum_p = hvx_vec_mpyacc_f32_f16(r1_c1_sum_p, r1_hf, c1_hf);
|
||||
}
|
||||
|
||||
HVX_Vector r0_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c0_sum_p), Q6_V_hi_W(r0_c0_sum_p)));
|
||||
HVX_Vector r0_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r0_c1_sum_p), Q6_V_hi_W(r0_c1_sum_p)));
|
||||
HVX_Vector r1_c0_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c0_sum_p), Q6_V_hi_W(r1_c0_sum_p)));
|
||||
HVX_Vector r1_c1_sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(r1_c1_sum_p), Q6_V_hi_W(r1_c1_sum_p)));
|
||||
|
||||
// Reduce and store results
|
||||
HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum);
|
||||
HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum);
|
||||
|
||||
hvx_vec_store_u(&s0[0], 8, r0_r1_c0_sum); // row0,col0 row1,col0
|
||||
hvx_vec_store_u(&s1[0], 8, r0_r1_c1_sum); // row0,col1 row1,col1
|
||||
}
|
||||
|
||||
static inline void vec_dot_f16_f16_uu_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) {
|
||||
const HVX_UVector * restrict x = (const HVX_UVector *) vx;
|
||||
const HVX_UVector * restrict y = (const HVX_UVector *) vy;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
HVX_Vector rsum = Q6_V_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(4)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x[i], y[i]);
|
||||
rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
HVX_Vector x_hf = Q6_V_vand_QV(bmask, x[i]);
|
||||
HVX_Vector y_hf = Q6_V_vand_QV(bmask, y[i]);
|
||||
|
||||
HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
|
||||
rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
|
||||
}
|
||||
|
||||
rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
|
||||
hvx_vec_store_u(&s[0], 4, rsum);
|
||||
}
|
||||
|
||||
static inline void vec_dot_f16_f32_uu_1x1(const uint32_t n, float * restrict s, const void * restrict x, const void * restrict y) {
|
||||
const HVX_UVector * restrict vx = (const HVX_UVector * restrict) x;
|
||||
const HVX_UVector * restrict vy = (const HVX_UVector * restrict) y;
|
||||
|
||||
uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
|
||||
uint32_t nloe = n % VLEN_FP16; // leftover elements
|
||||
|
||||
const HVX_Vector zero = Q6_V_vzero();
|
||||
|
||||
HVX_Vector rsum = Q6_V_vzero();
|
||||
|
||||
uint32_t i = 0;
|
||||
|
||||
#pragma unroll(2)
|
||||
for (i = 0; i < nvec; i++) {
|
||||
// Load y (fp32) and convert into fp16
|
||||
HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements
|
||||
HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements
|
||||
HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
|
||||
|
||||
// Load x (fp16)
|
||||
HVX_Vector x_hf = vx[i];
|
||||
|
||||
HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
|
||||
|
||||
rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
|
||||
}
|
||||
|
||||
if (nloe) {
|
||||
// Load y (fp32) and convert into fp16
|
||||
HVX_Vector y0_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+0], zero); // 32 elements
|
||||
HVX_Vector y1_qf = Q6_Vqf32_vsub_VsfVsf(vy[i*2+1], zero); // 32 elements
|
||||
HVX_Vector y_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(y1_qf, y0_qf)));
|
||||
|
||||
// Load x (fp16)
|
||||
HVX_Vector x_hf = vx[i];
|
||||
|
||||
// Zero-out unused elements
|
||||
// Note that we need to clear both x and y because they may contain NANs
|
||||
HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
|
||||
x_hf = Q6_V_vand_QV(bmask, x_hf);
|
||||
y_hf = Q6_V_vand_QV(bmask, y_hf);
|
||||
|
||||
HVX_VectorPair xy_qf = Q6_Wqf32_vmpy_VhfVhf(x_hf, y_hf);
|
||||
|
||||
rsum = Q6_Vqf32_vadd_Vqf32Vqf32(rsum, Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(xy_qf), Q6_V_hi_W(xy_qf)));
|
||||
}
|
||||
|
||||
// Convert into fp32 and reduce
|
||||
rsum = hvx_vec_reduce_sum_f32(Q6_Vsf_equals_Vqf32(rsum));
|
||||
hvx_vec_store_u(&s[0], 4, rsum);
|
||||
}
|
||||
|
||||
static inline void hvx_tensor_add_f32_grid(
|
||||
const struct htp_tensor * restrict dst,
|
||||
const struct htp_tensor * restrict src2,
|
||||
uint32_t start_row,
|
||||
uint32_t end_row,
|
||||
uint32_t start_col,
|
||||
uint32_t end_col,
|
||||
const struct fastdiv_values * div_ne11_12,
|
||||
const struct fastdiv_values * div_ne11
|
||||
) {
|
||||
if (start_row >= end_row || start_col >= end_col) return;
|
||||
const uint32_t nb1 = dst->nb[1]; // row stride in bytes
|
||||
|
||||
const uint32_t ne11 = dst->ne[1];
|
||||
const uint32_t ne12 = dst->ne[2];
|
||||
const uint32_t ne11_12 = ne11 * ne12;
|
||||
|
||||
const bool is_broadcast1 = (src2->ne[1] == 1);
|
||||
const bool is_broadcast2 = (src2->ne[2] == 1);
|
||||
const bool is_broadcast3 = (src2->ne[3] == 1);
|
||||
|
||||
for (uint32_t r = start_row; r < end_row; r++) {
|
||||
float * dst_row = (float *) ((uint8_t *) dst->data + r * nb1);
|
||||
|
||||
uint32_t i13 = fastdiv(r, div_ne11_12);
|
||||
uint32_t i12 = fastdiv(r - i13 * ne11_12, div_ne11);
|
||||
uint32_t i11 = r - i13 * ne11_12 - i12 * ne11;
|
||||
|
||||
uint32_t i23 = is_broadcast3 ? 0 : i13;
|
||||
uint32_t i22 = is_broadcast2 ? 0 : i12;
|
||||
uint32_t i21 = is_broadcast1 ? 0 : i11;
|
||||
|
||||
const float * src2_row = (const float *) ((const uint8_t *) src2->data +
|
||||
i21 * src2->nb[1] + i22 * src2->nb[2] + i23 * src2->nb[3]);
|
||||
|
||||
float * dst_ptr = &dst_row[start_col];
|
||||
const float * src2_ptr = &src2_row[start_col];
|
||||
int remaining = end_col - start_col;
|
||||
while (remaining >= 32) {
|
||||
HVX_Vector v_out = hvx_vmemu(dst_ptr);
|
||||
HVX_Vector v_z = hvx_vmemu(src2_ptr);
|
||||
hvx_vmemu(dst_ptr) = hvx_vec_add_f32_f32(v_out, v_z);
|
||||
dst_ptr += 32;
|
||||
src2_ptr += 32;
|
||||
remaining -= 32;
|
||||
}
|
||||
if (remaining > 0) {
|
||||
HVX_Vector v_out = hvx_vmemu(dst_ptr);
|
||||
HVX_Vector v_z = hvx_vmemu(src2_ptr);
|
||||
hvx_vec_store_u(dst_ptr, remaining * sizeof(float), hvx_vec_add_f32_f32(v_out, v_z));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -378,7 +378,7 @@ static inline HVX_VectorPair accum_q8_0_32x2(
|
||||
return Q6_W_vcombine_VV(v_sum1, v_sum0);
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -401,10 +401,14 @@ static void tiled_vec_dot_q4_0_32x1(const uint32_t n, float * restrict s, const
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -484,11 +488,19 @@ static void tiled_vec_dot_q4_0_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -519,10 +531,14 @@ static void tiled_vec_dot_q4_1_32x1(const uint32_t n, float * restrict s, const
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -637,11 +653,19 @@ static void tiled_vec_dot_q4_1_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -663,10 +687,14 @@ static void tiled_vec_dot_q8_0_32x1(const uint32_t n, float * restrict s, const
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -745,11 +773,19 @@ static void tiled_vec_dot_q8_0_32x2(const uint32_t n, float * restrict s0, float
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -773,10 +809,14 @@ static void tiled_vec_dot_iq4nl_32x1(const uint32_t n, float * restrict s, const
|
||||
v_sum_float = hvx_vec_add_f32_f32(v_sum_float, v_sum_scaled);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -857,11 +897,19 @@ static void tiled_vec_dot_iq4nl_32x2(const uint32_t n, float * restrict s0, floa
|
||||
v_sum_float_c1 = hvx_vec_add_f32_f32(v_sum_float_c1, v_sum_scaled_c1);
|
||||
}
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y_q = vy;
|
||||
|
||||
@@ -896,10 +944,14 @@ static void tiled_vec_dot_mxfp4_32x1(const uint32_t n, float * restrict s, const
|
||||
|
||||
v_sum_float = hvx_vec_mul_f32_f32(v_sum_float, hvx_vec_splat_f32(0.5f));
|
||||
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
if (sz) {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float, hvx_vmemu(sz)));
|
||||
} else {
|
||||
hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum_float);
|
||||
}
|
||||
}
|
||||
|
||||
static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows) {
|
||||
static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) {
|
||||
const uint8_t * restrict tile_ptr = vx;
|
||||
const uint8_t * restrict y0_q = vy0;
|
||||
const uint8_t * restrict y1_q = vy1;
|
||||
@@ -1013,8 +1065,16 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa
|
||||
v_sum_float_c0 = hvx_vec_mul_f32_f32(v_sum_float_c0, hvx_vec_splat_f32(0.5f));
|
||||
v_sum_float_c1 = hvx_vec_mul_f32_f32(v_sum_float_c1, hvx_vec_splat_f32(0.5f));
|
||||
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
if (sz0) {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c0, hvx_vmemu(sz0)));
|
||||
} else {
|
||||
hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum_float_c0);
|
||||
}
|
||||
if (sz1) {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum_float_c1, hvx_vmemu(sz1)));
|
||||
} else {
|
||||
hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum_float_c1);
|
||||
}
|
||||
}
|
||||
|
||||
static inline void quantize_f32_q8_0_tiled_kernel(
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
#include "hvx-base.h"
|
||||
#include "hvx-inverse.h"
|
||||
#include "hvx-exp.h"
|
||||
|
||||
#define FAST_SIGMOID_LOG2F (0x3fb8aa3b) // 1.442695022
|
||||
#define FAST_SIGMOID_C1 (0x3d009076) // 0.03138777
|
||||
@@ -139,4 +140,42 @@ static inline void hvx_tanh_f32_aa(uint8_t * restrict dst, const uint8_t * restr
|
||||
hvx_tanh_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
|
||||
}
|
||||
|
||||
static inline HVX_Vector hvx_vec_fast_sigmoid_f16(HVX_Vector x_v) {
|
||||
const HVX_Vector v_one = hvx_vec_splat_f16(1.0f);
|
||||
const HVX_Vector v_neg_log2e = hvx_vec_splat_f16(-EXP_LOG2E_F);
|
||||
const HVX_Vector em_mask = Q6_Vh_vsplat_R(0x7FFF);
|
||||
|
||||
// Compute absolute value of x_v
|
||||
HVX_Vector abs_x = Q6_V_vand_VV(x_v, em_mask);
|
||||
|
||||
// Compute u = -abs_x * log2(e) <= 0.
|
||||
HVX_Vector u = hvx_vec_mul_f16_f16(abs_x, v_neg_log2e);
|
||||
|
||||
// Clamp input to prevent underflow in exp2
|
||||
const HVX_Vector v_clamp_min = hvx_vec_splat_f16(-24.0f);
|
||||
u = Q6_Vhf_vmax_VhfVhf(v_clamp_min, u);
|
||||
|
||||
HVX_Vector exp_val = hvx_vec_exp2_f16(u);
|
||||
HVX_Vector denom = hvx_vec_add_f16_f16(v_one, exp_val);
|
||||
HVX_Vector sig_abs = hvx_vec_inverse_f16(denom);
|
||||
|
||||
// check if x_v < 0 (using integer comparison on absolute value)
|
||||
HVX_VectorPred is_neg = Q6_Q_vcmp_gt_VhVh(abs_x, x_v);
|
||||
|
||||
// If x_v < 0, return 1.0f - sig_abs
|
||||
HVX_Vector sig_neg = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vsub_VhfVhf(v_one, sig_abs));
|
||||
return Q6_V_vmux_QVV(is_neg, sig_neg, sig_abs);
|
||||
}
|
||||
|
||||
static inline HVX_Vector hvx_vec_tanh_f16(HVX_Vector x) {
|
||||
// tanh(x) = 2 * sigmoid(2x) - 1
|
||||
const HVX_Vector v_two = hvx_vec_splat_f16(2.0f);
|
||||
|
||||
HVX_Vector x2 = hvx_vec_mul_f16_f16(x, v_two);
|
||||
HVX_Vector sig2x = hvx_vec_fast_sigmoid_f16(x2);
|
||||
|
||||
const HVX_Vector v_neg_one = hvx_vec_splat_f16(-1.0f);
|
||||
return hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(sig2x, v_two), v_neg_one);
|
||||
}
|
||||
|
||||
#endif /* HVX_SIGMOID_H */
|
||||
|
||||
@@ -575,6 +575,7 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) {
|
||||
static int execute_op(struct htp_ops_context * octx) {
|
||||
switch (octx->op) {
|
||||
case HTP_OP_MUL_MAT:
|
||||
case HTP_OP_MUL_MAT_ADD:
|
||||
return op_matmul(octx);
|
||||
|
||||
case HTP_OP_MUL_MAT_ID:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -392,56 +392,49 @@ static inline size_t htp_mm_hvx_get_vtcm_sizes(
|
||||
case HTP_MM_KERNEL_HVX_QUANT_ROW: {
|
||||
size_t q_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10);
|
||||
|
||||
vtcm_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 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);
|
||||
|
||||
// src0 spad is also used in dynamic quantizer to store padded src1 rows
|
||||
size_t src1_row_size_padded = htp_mm_round_up(q_src1_row_size, QK_Q8_0_TILED * sizeof(float));
|
||||
if (vtcm_src0_size < src1_row_size_padded) {
|
||||
vtcm_src0_size = src1_row_size_padded;
|
||||
}
|
||||
|
||||
vtcm_src0_size = vtcm_src0_size * n_threads;
|
||||
vtcm_dst_size = vtcm_dst_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);
|
||||
if (repacked_vtcm_size < src1_row_size_padded) {
|
||||
repacked_vtcm_size = src1_row_size_padded;
|
||||
}
|
||||
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;
|
||||
}
|
||||
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_dst_size = dst_nrows > 0 ? htp_mm_round_up(dst_row_size, 128) : 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);
|
||||
|
||||
size_t src1_row_size_padded = htp_mm_round_up(q_src1_row_size, 256);
|
||||
if (vtcm_src0_size < src1_row_size_padded) {
|
||||
vtcm_src0_size = src1_row_size_padded;
|
||||
}
|
||||
|
||||
vtcm_src0_size = vtcm_src0_size * n_threads;
|
||||
vtcm_dst_size = vtcm_dst_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);
|
||||
if (repacked_vtcm_size < src1_row_size_padded) {
|
||||
repacked_vtcm_size = src1_row_size_padded;
|
||||
}
|
||||
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;
|
||||
}
|
||||
default:
|
||||
@@ -463,7 +456,8 @@ static inline size_t htp_mm_hvx_id_get_vtcm_sizes(
|
||||
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_src1_size_out,
|
||||
size_t * vtcm_dst_size_out
|
||||
) {
|
||||
const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 ||
|
||||
wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL ||
|
||||
@@ -476,29 +470,22 @@ static inline size_t htp_mm_hvx_id_get_vtcm_sizes(
|
||||
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);
|
||||
|
||||
// src0 spad also holds temporary transposed src1 columns during dynamic quantization.
|
||||
const size_t src1_row_size_padded = htp_mm_round_up(src1_row_size, QK_Q8_0_TILED * sizeof(float));
|
||||
if (src0_sz_per_thread < src1_row_size_padded) {
|
||||
src0_sz_per_thread = src1_row_size_padded;
|
||||
}
|
||||
|
||||
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);
|
||||
if (repacked_vtcm_size < src1_row_size_padded) {
|
||||
repacked_vtcm_size = src1_row_size_padded;
|
||||
}
|
||||
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;
|
||||
return vtcm_src0_size + src1_sz + vtcm_dst_size;
|
||||
}
|
||||
|
||||
#ifdef __cplusplus
|
||||
|
||||
Reference in New Issue
Block a user