Max Krasnyansky
0a50d9909a
* hex-l2: use dirty ranges for flushing * hex-l2: simplify range based flush logic * hex-l2: optimize dirty range scans * hex-hvx: support for reduce_max_i32 * hex-mm: optimize fused MUL_MAT+ADD to use vtcm for bias when it fits * hex-mmid: optimize mmid row-mapping generation * hex-mmid: optimize mmid row-mapping generation * hex-mmid: optimize mmid row-mapping generation (round2) * hmx-mm: optimize output proc by tiling (col-chunking) * hex-fa: start the next q dmas a bit earlier * hex-fa: prefetch Q even earlier * hvx-fa: optimize softmax to keep things in hvx registers * hex-fa: hoist const register init in softmax loop * hmx-fa: kick off next-qkv DMAs before o-proc * hmx-fa: hoist various checks out of the inner loop * hmx-fa: adjust the cost model to better balance softmax work across hvx threads * hmx-fa: overlap diag rescale build with last HMX task * hmx-fa: optimize idx update in output proc * hmx-fa: unroll the softmax loops for improved perf * hmx-fa: overlap qk-dot with softmax, double-buffer p and s tiles * hex-trace: double the default number of trace entries * hex-trace: add trace events for opbatch and buffer mgmt * hex-trace: overhaul tracing to simplify runtime event handling and support opbatch stats * hex-trace: replace ascii timeline diagram with pipeline bubbles detector * hex-trace: handle missing start/stop events * hex-dma: always log stop/start trace events even for dummy dmas * hex-scripts: fix flake warnings
337 lines
14 KiB
C
337 lines
14 KiB
C
#ifndef HVX_REDUCE_H
|
|
#define HVX_REDUCE_H
|
|
|
|
#include <math.h>
|
|
#include <stdbool.h>
|
|
#include <stdint.h>
|
|
#include <assert.h>
|
|
|
|
#include "hex-utils.h"
|
|
#include "hvx-base.h"
|
|
#include "hvx-types.h"
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_n_i32(HVX_Vector in, unsigned int n) {
|
|
unsigned int total = n * 4; // total vec nbytes
|
|
unsigned int width = 4; // int32
|
|
|
|
HVX_Vector sum = in, sum_t;
|
|
while (width < total) {
|
|
sum_t = Q6_V_vror_VR(sum, width); // rotate right
|
|
sum = Q6_Vw_vadd_VwVw(sum_t, sum); // elementwise sum
|
|
width = width << 1;
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_i32(HVX_Vector in) {
|
|
return hvx_vec_reduce_sum_n_i32(in, 32);
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_n_qf32(HVX_Vector in, unsigned int n) {
|
|
unsigned int total = n * 4; // total vec nbytes
|
|
unsigned int width = 4; // fp32 nbytes
|
|
|
|
HVX_Vector sum = in, sum_t;
|
|
while (width < total) {
|
|
sum_t = Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum), width); // rotate right
|
|
sum = Q6_Vqf32_vadd_Vqf32Vsf(sum, sum_t); // elementwise sum
|
|
width = width << 1;
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_qf32(HVX_Vector in) {
|
|
return hvx_vec_reduce_sum_n_qf32(in, 32);
|
|
}
|
|
|
|
#if __HVX_ARCH__ > 75
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) {
|
|
HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4);
|
|
HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4);
|
|
HVX_Vector sum_sf01 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01));
|
|
HVX_Vector sum_sf23 = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23));
|
|
|
|
HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(sum_sf23, sum_sf01, 8);
|
|
HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123));
|
|
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2));
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4));
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8));
|
|
return sum_sf;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) {
|
|
HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4);
|
|
HVX_Vector sum_sf = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump));
|
|
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2));
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4));
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8));
|
|
sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 16));
|
|
return sum_sf;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) {
|
|
unsigned int total = n * 4; // total vec nbytes
|
|
unsigned int width = 4; // fp32 nbytes
|
|
|
|
HVX_Vector sum = in, sum_t;
|
|
while (width < total) {
|
|
sum_t = Q6_V_vror_VR(sum, width); // rotate right
|
|
sum = Q6_Vsf_vadd_VsfVsf(sum, sum_t); // elementwise sum
|
|
width = width << 1;
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
#else
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) {
|
|
HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4);
|
|
HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4);
|
|
HVX_Vector sum_qf01 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01));
|
|
HVX_Vector sum_qf23 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23));
|
|
|
|
HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(Q6_Vsf_equals_Vqf32(sum_qf23), Q6_Vsf_equals_Vqf32(sum_qf01), 8);
|
|
HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123));
|
|
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2));
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4));
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8));
|
|
return Q6_Vsf_equals_Vqf32(sum_qf);
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) {
|
|
HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4);
|
|
HVX_Vector sum_qf = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump));
|
|
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2));
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4));
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8));
|
|
sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 16));
|
|
return Q6_Vsf_equals_Vqf32(sum_qf);
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) {
|
|
unsigned int total = n * 4; // total vec nbytes
|
|
unsigned int width = 4; // fp32 nbytes
|
|
|
|
HVX_Vector sum = in, sum_t;
|
|
while (width < total) {
|
|
sum_t = Q6_V_vror_VR(sum, width); // rotate right
|
|
sum = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(sum, sum_t)); // elementwise sum
|
|
width = width << 1;
|
|
}
|
|
return sum;
|
|
}
|
|
|
|
#endif
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_sum_f32(HVX_Vector in) {
|
|
return hvx_vec_reduce_sum_n_f32(in, 32);
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_max_f16(HVX_Vector in) {
|
|
unsigned total = 128; // total vec nbytes
|
|
unsigned width = 2; // fp16 nbytes
|
|
|
|
HVX_Vector _max = in, _max_t;
|
|
while (width < total) {
|
|
_max_t = Q6_V_vror_VR(_max, width); // rotate right
|
|
_max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max
|
|
width = width << 1;
|
|
}
|
|
|
|
return _max;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_max2_f16(HVX_Vector in, HVX_Vector _max) {
|
|
unsigned total = 128; // total vec nbytes
|
|
unsigned width = 2; // fp32 nbytes
|
|
|
|
HVX_Vector _max_t;
|
|
|
|
_max = Q6_Vhf_vmax_VhfVhf(in, _max);
|
|
while (width < total) {
|
|
_max_t = Q6_V_vror_VR(_max, width); // rotate right
|
|
_max = Q6_Vhf_vmax_VhfVhf(_max_t, _max); // elementwise max
|
|
width = width << 1;
|
|
}
|
|
|
|
return _max;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_max_f32(HVX_Vector in) {
|
|
unsigned total = 128; // total vec nbytes
|
|
unsigned width = 4; // fp32 nbytes
|
|
|
|
HVX_Vector _max = in, _max_t;
|
|
while (width < total) {
|
|
_max_t = Q6_V_vror_VR(_max, width); // rotate right
|
|
_max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max
|
|
width = width << 1;
|
|
}
|
|
|
|
return _max;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_max2_f32(HVX_Vector in, HVX_Vector _max) {
|
|
unsigned total = 128; // total vec nbytes
|
|
unsigned width = 4; // fp32 nbytes
|
|
|
|
HVX_Vector _max_t;
|
|
|
|
_max = Q6_Vsf_vmax_VsfVsf(in, _max);
|
|
while (width < total) {
|
|
_max_t = Q6_V_vror_VR(_max, width); // rotate right
|
|
_max = Q6_Vsf_vmax_VsfVsf(_max_t, _max); // elementwise max
|
|
width = width << 1;
|
|
}
|
|
|
|
return _max;
|
|
}
|
|
|
|
#define hvx_reduce_loop_body(src_type, init_vec, pad_vec, vec_op, reduce_op, scalar_reduce) \
|
|
do { \
|
|
src_type * restrict vsrc = (src_type *) src; \
|
|
HVX_Vector acc = init_vec; \
|
|
\
|
|
const uint32_t elem_size = sizeof(float); \
|
|
const uint32_t epv = 128 / elem_size; \
|
|
const uint32_t nvec = num_elems / epv; \
|
|
const uint32_t nloe = num_elems % epv; \
|
|
\
|
|
uint32_t i = 0; \
|
|
_Pragma("unroll(4)") \
|
|
for (; i < nvec; i++) { \
|
|
acc = vec_op(acc, vsrc[i]); \
|
|
} \
|
|
if (nloe) { \
|
|
const float * srcf = (const float *) src + i * epv; \
|
|
HVX_Vector in = *(HVX_UVector *) srcf; \
|
|
HVX_Vector temp = Q6_V_valign_VVR(in, pad_vec, nloe * elem_size); \
|
|
acc = vec_op(acc, temp); \
|
|
} \
|
|
HVX_Vector v = reduce_op(acc); \
|
|
return scalar_reduce(v); \
|
|
} while(0)
|
|
|
|
#define HVX_REDUCE_MAX_OP(acc, val) Q6_Vsf_vmax_VsfVsf(acc, val)
|
|
#define HVX_REDUCE_SUM_OP(acc, val) Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(acc), val)
|
|
#define HVX_SUM_SQ_OP(acc, val) Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(val, val))
|
|
#define HVX_REDUCE_MAX_SCALAR(v) hvx_vec_get_f32(v)
|
|
#define HVX_REDUCE_SUM_SCALAR(v) hvx_vec_get_f32(Q6_Vsf_equals_Vqf32(v))
|
|
|
|
// Max variants
|
|
|
|
static inline float hvx_reduce_max_f32_a(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]);
|
|
assert((unsigned long) src % 128 == 0);
|
|
hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR);
|
|
}
|
|
|
|
static inline float hvx_reduce_max_f32_u(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]);
|
|
hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR);
|
|
}
|
|
|
|
static inline float hvx_reduce_max_f32(const uint8_t * restrict src, const int num_elems) {
|
|
if (hex_is_aligned((void *) src, 128)) {
|
|
return hvx_reduce_max_f32_a(src, num_elems);
|
|
} else {
|
|
return hvx_reduce_max_f32_u(src, num_elems);
|
|
}
|
|
}
|
|
|
|
// Sum variants
|
|
|
|
static inline float hvx_reduce_sum_f32_a(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = Q6_V_vsplat_R(0);
|
|
assert((unsigned long) src % 128 == 0);
|
|
hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
|
|
}
|
|
|
|
static inline float hvx_reduce_sum_f32_u(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = Q6_V_vsplat_R(0);
|
|
hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
|
|
}
|
|
|
|
static inline float hvx_reduce_sum_f32(const uint8_t * restrict src, const int num_elems) {
|
|
if (hex_is_aligned((void *) src, 128)) {
|
|
return hvx_reduce_sum_f32_a(src, num_elems);
|
|
} else {
|
|
return hvx_reduce_sum_f32_u(src, num_elems);
|
|
}
|
|
}
|
|
|
|
// Sum of squares variants
|
|
|
|
static inline float hvx_sum_of_squares_f32_a(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = Q6_V_vsplat_R(0);
|
|
assert((uintptr_t) src % 128 == 0);
|
|
hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
|
|
}
|
|
|
|
static inline float hvx_sum_of_squares_f32_u(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = Q6_V_vsplat_R(0);
|
|
hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
|
|
}
|
|
|
|
static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const int num_elems) {
|
|
if (hex_is_aligned((void *) src, 128)) {
|
|
return hvx_sum_of_squares_f32_a(src, num_elems);
|
|
} else {
|
|
return hvx_sum_of_squares_f32_u(src, num_elems);
|
|
}
|
|
}
|
|
|
|
// Signed 32-bit Integer Max variants
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_max_n_i32(HVX_Vector in, unsigned int n) {
|
|
unsigned int total = n * 4; // total vec nbytes
|
|
unsigned int width = 4; // int32 nbytes
|
|
|
|
HVX_Vector max_val = in, max_t;
|
|
while (width < total) {
|
|
max_t = Q6_V_vror_VR(max_val, width); // rotate right
|
|
max_val = Q6_Vw_vmax_VwVw(max_t, max_val); // elementwise signed max
|
|
width = width << 1;
|
|
}
|
|
return max_val;
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_reduce_max_i32(HVX_Vector in) {
|
|
return hvx_vec_reduce_max_n_i32(in, 32);
|
|
}
|
|
|
|
static inline int32_t hvx_reduce_max_i32_a(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
|
|
HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
|
|
assert((uintptr_t) src % 128 == 0);
|
|
hvx_reduce_loop_body(HVX_Vector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
|
|
}
|
|
|
|
static inline int32_t hvx_reduce_max_i32_u(const uint8_t * restrict src, const int num_elems) {
|
|
HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
|
|
HVX_Vector pad_vec = Q6_V_vsplat_R(0x80000000);
|
|
hvx_reduce_loop_body(HVX_UVector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
|
|
}
|
|
|
|
static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int num_elems) {
|
|
if (hex_is_aligned((void *) src, 128)) {
|
|
return hvx_reduce_max_i32_a(src, num_elems);
|
|
} else {
|
|
return hvx_reduce_max_i32_u(src, num_elems);
|
|
}
|
|
}
|
|
|
|
#undef hvx_reduce_loop_body
|
|
#undef HVX_REDUCE_MAX_OP
|
|
#undef HVX_REDUCE_SUM_OP
|
|
#undef HVX_REDUCE_MAX_SCALAR
|
|
#undef HVX_REDUCE_SUM_SCALAR
|
|
#undef HVX_SUM_SQ_OP
|
|
|
|
#endif /* HVX_REDUCE_H */
|