1220 lines
57 KiB
C
1220 lines
57 KiB
C
#pragma clang diagnostic ignored "-Wunused-variable"
|
|
#pragma clang diagnostic ignored "-Wunused-function"
|
|
#pragma clang diagnostic ignored "-Wunused-but-set-variable"
|
|
|
|
#include <HAP_farf.h>
|
|
|
|
#include <math.h>
|
|
#include <string.h>
|
|
|
|
#include "dma-queue.h"
|
|
#include "hvx-utils.h"
|
|
|
|
#define GGML_COMMON_DECL_C
|
|
#include "ggml-common.h"
|
|
#include "htp-ctx.h"
|
|
#include "htp-ops.h"
|
|
#include "hex-common.h"
|
|
#include "htp-tensor.h"
|
|
#include "htp-vtcm.h"
|
|
#include "glu-ops.h"
|
|
|
|
#define htp_glu_preamble \
|
|
const struct htp_tensor * src0 = gctx->octx->src[0]; \
|
|
const struct htp_tensor * src1 = gctx->octx->src[1]; \
|
|
const struct htp_tensor * dst = gctx->octx->dst; \
|
|
\
|
|
const uint32_t ne00 = src0->ne[0]; \
|
|
const uint32_t ne01 = src0->ne[1]; \
|
|
const uint32_t ne02 = src0->ne[2]; \
|
|
const uint32_t ne03 = src0->ne[3]; \
|
|
\
|
|
const uint32_t nb00 = src0->nb[0]; \
|
|
const uint32_t nb01 = src0->nb[1]; \
|
|
const uint32_t nb02 = src0->nb[2]; \
|
|
const uint32_t nb03 = src0->nb[3]; \
|
|
\
|
|
const uint32_t ne10 = src1 ? src1->ne[0] : 0; \
|
|
const uint32_t ne11 = src1 ? src1->ne[1] : 0; \
|
|
const uint32_t ne12 = src1 ? src1->ne[2] : 0; \
|
|
const uint32_t ne13 = src1 ? src1->ne[3] : 0; \
|
|
\
|
|
const uint32_t nb10 = src1 ? src1->nb[0] : 0; \
|
|
const uint32_t nb11 = src1 ? src1->nb[1] : 0; \
|
|
const uint32_t nb12 = src1 ? src1->nb[2] : 0; \
|
|
const uint32_t nb13 = src1 ? src1->nb[3] : 0; \
|
|
\
|
|
const uint32_t ne0 = dst->ne[0]; \
|
|
const uint32_t ne1 = dst->ne[1]; \
|
|
const uint32_t ne2 = dst->ne[2]; \
|
|
const uint32_t ne3 = dst->ne[3]; \
|
|
\
|
|
const uint32_t nb0 = dst->nb[0]; \
|
|
const uint32_t nb1 = dst->nb[1]; \
|
|
const uint32_t nb2 = dst->nb[2]; \
|
|
const uint32_t nb3 = dst->nb[3];
|
|
|
|
struct htp_glu_context;
|
|
|
|
typedef void (*glu_compute_f32_fn_t)(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx);
|
|
|
|
typedef void (*glu_compute_f16_fn_t)(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx);
|
|
|
|
struct htp_glu_context {
|
|
struct htp_ops_context * octx;
|
|
|
|
void * compute;
|
|
const char * op_str;
|
|
|
|
// Precomputed values
|
|
dma_addr_t data_src0;
|
|
dma_addr_t data_src1;
|
|
dma_addr_t data_dst;
|
|
|
|
size_t src0_row_size;
|
|
size_t src1_row_size;
|
|
size_t dst_row_size;
|
|
|
|
size_t src0_row_stride;
|
|
size_t src1_row_stride;
|
|
|
|
size_t src0_row_size_aligned;
|
|
size_t src1_row_size_aligned;
|
|
size_t dst_row_size_aligned;
|
|
|
|
size_t src0_spad_half_size;
|
|
size_t src1_spad_half_size;
|
|
size_t dst_spad_half_size;
|
|
|
|
uint32_t block;
|
|
uint32_t src0_nrows;
|
|
uint32_t src0_nrows_per_thread;
|
|
uint32_t row_start;
|
|
int nc;
|
|
|
|
uint8_t * vtcm_src0;
|
|
uint8_t * vtcm_src1;
|
|
uint8_t * vtcm_dst;
|
|
|
|
size_t vtcm_src0_size_per_thread;
|
|
size_t vtcm_src1_size_per_thread;
|
|
size_t vtcm_dst_size_per_thread;
|
|
};
|
|
|
|
#define htp_glu_op_preamble \
|
|
const size_t src0_row_size_aligned = gctx->src0_row_size_aligned; \
|
|
const size_t src1_row_size_aligned = gctx->src1_row_size_aligned; \
|
|
const size_t dst_row_size_aligned = gctx->dst_row_size_aligned; \
|
|
const int nc = gctx->nc;
|
|
|
|
// swiglu(x) = x1 * sigmoid(x0)
|
|
static void swiglu_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_sigmoid_f32_aa(dst_ptr, src0_ptr, nc);
|
|
hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
// out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f)
|
|
static void swiglu_oai_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
const float alpha = ((const float *) (gctx->octx->op_params))[2];
|
|
const float limit = ((const float *) (gctx->octx->op_params))[3];
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
// x (src0_ptr) = std::min(src0_p[k], limit);
|
|
hvx_min_scalar_f32((uint8_t *) src0_ptr, src0_ptr, limit, nc);
|
|
// y1 (src1_ptr) = std::clamp(src1_p[k], -limit, limit);
|
|
hvx_clamp_scalar_f32((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
|
|
// y (src1_ptr) = y1 + 1.f
|
|
hvx_add_scalar_f32((uint8_t *) src1_ptr, src1_ptr, 1.0, nc);
|
|
// x1 (dst_ptr) = alpha * x
|
|
hvx_mul_scalar_f32(dst_ptr, src0_ptr, alpha, nc);
|
|
// x2 (dst_ptr) = sigmoid(x1) = 1/(1+exp(-x1))
|
|
hvx_sigmoid_f32_aa(dst_ptr, dst_ptr, nc);
|
|
// out = x * sigmoid(alpha * x) * (y + 1.f)
|
|
hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static void swiglu_clamp_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
const float limit = ((const float *) (gctx->octx->op_params))[3];
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_min_scalar_f32((uint8_t *) src0_ptr, src0_ptr, limit, nc);
|
|
hvx_clamp_scalar_f32((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
|
|
hvx_sigmoid_f32_aa(dst_ptr, src0_ptr, nc);
|
|
hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static const float GELU_COEF_A = 0.044715f;
|
|
static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
|
|
|
|
static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
|
|
v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
|
|
v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
|
|
|
|
HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
|
|
HVX_Vector x = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
|
|
HVX_Vector xx = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
|
|
|
|
HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
|
|
v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
|
|
|
|
HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
|
|
v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
|
|
v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
|
|
|
|
HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
|
|
v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
|
|
|
|
HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
|
|
HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
|
|
|
|
// Newton-Raphson with 2 iterations
|
|
HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
|
|
HVX_Vector i_sf = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
|
|
HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf(
|
|
i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
|
|
r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
|
|
r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
|
|
HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
|
|
|
|
res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
|
|
|
|
return Q6_Vsf_equals_Vqf32(res);
|
|
}
|
|
|
|
static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
|
|
HVX_Vector one,
|
|
HVX_Vector max_exp,
|
|
HVX_Vector min_exp) {
|
|
const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
|
|
const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
|
|
|
|
HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
|
|
out = Q6_V_vmux_QVV(pred_max, out, one);
|
|
return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
|
|
}
|
|
|
|
static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(float);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
const float GELU_COEF_A_TIMES_SQRT = GELU_COEF_A * SQRT_2_OVER_PI;
|
|
|
|
const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT);
|
|
const HVX_Vector v_sqrt_2_pi = hvx_vec_splat_f32(SQRT_2_OVER_PI);
|
|
const HVX_Vector v_one = hvx_vec_splat_f32(1.0f);
|
|
const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
|
|
const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
|
|
|
|
uint32_t i = 0;
|
|
|
|
_Pragma("unroll(4)")
|
|
for (; i < nvec; i++) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
|
|
HVX_Vector x2 = hvx_vec_mul_f32_f32(x, x);
|
|
HVX_Vector coef = hvx_vec_mul_f32_f32(x2, v_coef_a_times_sqrt);
|
|
coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
|
|
HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
|
|
|
|
// y2 = 2 * inner = inner + inner
|
|
HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);
|
|
|
|
// Fast sigmoid approximation (2 iterations)
|
|
HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);
|
|
|
|
HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
|
|
vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g);
|
|
}
|
|
|
|
if (nloe) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
|
|
HVX_Vector x2 = hvx_vec_mul_f32_f32(x, x);
|
|
HVX_Vector coef = hvx_vec_mul_f32_f32(x2, v_coef_a_times_sqrt);
|
|
coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
|
|
HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
|
|
|
|
HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);
|
|
|
|
HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);
|
|
|
|
HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
|
|
HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g);
|
|
hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
|
|
}
|
|
}
|
|
|
|
static inline void hvx_geglu_quick_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(float);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
const HVX_Vector v_scale = hvx_vec_splat_f32(1.702f);
|
|
const HVX_Vector v_one = hvx_vec_splat_f32(1.0f);
|
|
const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
|
|
const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
|
|
|
|
uint32_t i = 0;
|
|
|
|
_Pragma("unroll(4)")
|
|
for (; i < nvec; i++) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
HVX_Vector scaled_x = hvx_vec_mul_f32_f32(x, v_scale);
|
|
HVX_Vector sigmoid_x = hvx_vec_fast_sigmoid_f32_guard_2it(scaled_x, v_one, v_max_exp, v_min_exp);
|
|
vdst[i] = hvx_vec_mul_f32_f32(hvx_vec_mul_f32_f32(x, sigmoid_x), g);
|
|
}
|
|
|
|
if (nloe) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
HVX_Vector scaled_x = hvx_vec_mul_f32_f32(x, v_scale);
|
|
HVX_Vector sigmoid_x = hvx_vec_fast_sigmoid_f32_guard_2it(scaled_x, v_one, v_max_exp, v_min_exp);
|
|
HVX_Vector result = hvx_vec_mul_f32_f32(hvx_vec_mul_f32_f32(x, sigmoid_x), g);
|
|
hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), result);
|
|
}
|
|
}
|
|
|
|
static inline void hvx_geglu_erf_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(float);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
uint32_t i = 0;
|
|
|
|
_Pragma("unroll(4)")
|
|
for (; i < nvec; i++) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
vdst[i] = hvx_vec_mul_f32_f32(hvx_vec_gelu_erf_f32(x), g);
|
|
}
|
|
|
|
if (nloe) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
HVX_Vector result = hvx_vec_mul_f32_f32(hvx_vec_gelu_erf_f32(x), g);
|
|
hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), result);
|
|
}
|
|
}
|
|
|
|
static inline void hvx_geglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(_Float16);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
const HVX_Vector v_coef_a = hvx_vec_splat_f16((_Float16) GELU_COEF_A);
|
|
const HVX_Vector v_sqrt_2_pi = hvx_vec_splat_f16((_Float16) SQRT_2_OVER_PI);
|
|
const HVX_Vector v_one = hvx_vec_splat_f16((_Float16) 1.0f);
|
|
const HVX_Vector v_half = hvx_vec_splat_f16((_Float16) 0.5f);
|
|
|
|
uint32_t i = 0;
|
|
|
|
_Pragma("unroll(4)")
|
|
for (; i < nvec; i++) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
|
|
// gelu(x) = 0.5*x*(1 + tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x)))
|
|
HVX_Vector x2 = hvx_vec_mul_f16_f16(x, x);
|
|
HVX_Vector coef = hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(x2, v_coef_a), v_one);
|
|
HVX_Vector inner = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, coef), v_sqrt_2_pi);
|
|
|
|
HVX_Vector t = hvx_vec_add_f16_f16(hvx_vec_tanh_f16(inner), v_one);
|
|
|
|
HVX_Vector gelu_x = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, t), v_half);
|
|
vdst[i] = hvx_vec_mul_f16_f16(gelu_x, g);
|
|
}
|
|
|
|
if (nloe) {
|
|
HVX_Vector x = vsrc0[i];
|
|
HVX_Vector g = vsrc1[i];
|
|
|
|
HVX_Vector x2 = hvx_vec_mul_f16_f16(x, x);
|
|
HVX_Vector coef = hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(x2, v_coef_a), v_one);
|
|
HVX_Vector inner = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, coef), v_sqrt_2_pi);
|
|
|
|
HVX_Vector t = hvx_vec_add_f16_f16(hvx_vec_tanh_f16(inner), v_one);
|
|
|
|
HVX_Vector gelu_x = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, t), v_half);
|
|
HVX_Vector res = hvx_vec_mul_f16_f16(gelu_x, g);
|
|
hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(_Float16), res);
|
|
}
|
|
}
|
|
|
|
// geglu(x, g) = gelu(x) * g
|
|
static void geglu_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_geglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
// geglu_quick(x, g) = x * sigmoid(1.702 * x) * g
|
|
static void geglu_quick_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_geglu_quick_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
// geglu_erf(x, g) = gelu_erf(x) * g
|
|
static void geglu_erf_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_geglu_erf_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static inline void hvx_reglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(float);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
const HVX_Vector v_zero = Q6_V_vzero();
|
|
|
|
uint32_t i = 0;
|
|
|
|
_Pragma("unroll(4)")
|
|
for (; i < nvec; i++) {
|
|
HVX_Vector x = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
|
|
vdst[i] = hvx_vec_mul_f32_f32(x, vsrc1[i]);
|
|
}
|
|
|
|
if (nloe) {
|
|
HVX_Vector x = Q6_Vsf_vmax_VsfVsf(vsrc0[i], v_zero);
|
|
HVX_Vector res = hvx_vec_mul_f32_f32(x, vsrc1[i]);
|
|
hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
|
|
}
|
|
}
|
|
|
|
// reglu(x, g) = relu(x) * g
|
|
static void reglu_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_reglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static __attribute__((noinline)) void hvx_sigglu_f32_vec(HVX_Vector * restrict vdst,
|
|
const HVX_Vector * restrict vsrc0,
|
|
const HVX_Vector * restrict vsrc1,
|
|
uint32_t nvec) {
|
|
const HVX_Vector v_one = hvx_vec_splat_f32(1.0f);
|
|
const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
|
|
const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
|
|
|
|
for (uint32_t i = 0; i < nvec; i++) {
|
|
HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[i], v_one, v_max_exp, v_min_exp);
|
|
vdst[i] = hvx_vec_mul_f32_f32(sig, vsrc1[i]);
|
|
}
|
|
}
|
|
|
|
static __attribute__((noinline)) void hvx_sigglu_f32_tail(HVX_Vector * restrict vdst,
|
|
const HVX_Vector * restrict vsrc0,
|
|
const HVX_Vector * restrict vsrc1,
|
|
uint32_t nloe) {
|
|
const HVX_Vector v_one = hvx_vec_splat_f32(1.0f);
|
|
const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f);
|
|
const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f);
|
|
|
|
HVX_Vector sig = hvx_vec_fast_sigmoid_f32_guard(vsrc0[0], v_one, v_max_exp, v_min_exp);
|
|
HVX_Vector res = hvx_vec_mul_f32_f32(sig, vsrc1[0]);
|
|
hvx_vec_store_a((void *) vdst, nloe * sizeof(float), res);
|
|
}
|
|
|
|
static inline void hvx_sigglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(float);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
if (nvec) {
|
|
hvx_sigglu_f32_vec(vdst, vsrc0, vsrc1, nvec);
|
|
}
|
|
if (nloe) {
|
|
hvx_sigglu_f32_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
|
|
}
|
|
}
|
|
|
|
// sigglu(x, g) = sigmoid(x) * g
|
|
static void sigglu_f32(const float * restrict src0,
|
|
const float * restrict src1,
|
|
float * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_sigglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
// swiglu(x) = x1 * sigmoid(x0)
|
|
static void swiglu_f16(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_sigmoid_f16_aa(dst_ptr, src0_ptr, nc);
|
|
hvx_mul_mul_f16_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
// out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f)
|
|
static void swiglu_oai_f16(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
const _Float16 alpha = (_Float16) ((const float *) (gctx->octx->op_params))[2];
|
|
const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
// x (src0_ptr) = std::min(src0_p[k], limit);
|
|
hvx_min_scalar_f16((uint8_t *) src0_ptr, src0_ptr, limit, nc);
|
|
// y1 (src1_ptr) = std::clamp(src1_p[k], -limit, limit);
|
|
hvx_clamp_scalar_f16((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
|
|
// y (src1_ptr) = y1 + 1.f
|
|
hvx_add_scalar_f16((uint8_t *) src1_ptr, src1_ptr, (_Float16) 1.0f, nc);
|
|
// x1 (dst_ptr) = alpha * x
|
|
hvx_mul_scalar_f16(dst_ptr, src0_ptr, alpha, nc);
|
|
// x2 (dst_ptr) = sigmoid(x1) = 1/(1+exp(-x1))
|
|
hvx_sigmoid_f16_aa(dst_ptr, dst_ptr, nc);
|
|
// out = x * sigmoid(alpha * x) * (y + 1.f)
|
|
hvx_mul_mul_f16_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static void swiglu_clamp_f16(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
const _Float16 limit = (_Float16) ((const float *) (gctx->octx->op_params))[3];
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_min_scalar_f16((uint8_t *) src0_ptr, src0_ptr, limit, nc);
|
|
hvx_clamp_scalar_f16((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
|
|
hvx_sigmoid_f16_aa(dst_ptr, src0_ptr, nc);
|
|
hvx_mul_mul_f16_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
// geglu(x, g) = gelu(x) * g, gelu via tanh approximation
|
|
static void geglu_f16(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_geglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static inline void hvx_reglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(_Float16);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
const HVX_Vector v_zero = Q6_V_vzero();
|
|
|
|
uint32_t i = 0;
|
|
|
|
_Pragma("unroll(4)")
|
|
for (; i < nvec; i++) {
|
|
HVX_Vector x = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
|
|
vdst[i] = hvx_vec_mul_f16_f16(x, vsrc1[i]);
|
|
}
|
|
|
|
if (nloe) {
|
|
HVX_Vector x = Q6_Vhf_vmax_VhfVhf(vsrc0[i], v_zero);
|
|
HVX_Vector res = hvx_vec_mul_f16_f16(x, vsrc1[i]);
|
|
hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(_Float16), res);
|
|
}
|
|
}
|
|
|
|
// reglu(x, g) = relu(x) * g
|
|
static void reglu_f16(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_reglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
static __attribute__((noinline)) void hvx_sigglu_f16_vec(HVX_Vector * restrict vdst,
|
|
const HVX_Vector * restrict vsrc0,
|
|
const HVX_Vector * restrict vsrc1,
|
|
uint32_t nvec) {
|
|
for (uint32_t i = 0; i < nvec; i++) {
|
|
HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[i]);
|
|
vdst[i] = hvx_vec_mul_f16_f16(sig, vsrc1[i]);
|
|
}
|
|
}
|
|
|
|
static __attribute__((noinline)) void hvx_sigglu_f16_tail(HVX_Vector * restrict vdst,
|
|
const HVX_Vector * restrict vsrc0,
|
|
const HVX_Vector * restrict vsrc1,
|
|
uint32_t nloe) {
|
|
HVX_Vector sig = hvx_vec_fast_sigmoid_f16(vsrc0[0]);
|
|
HVX_Vector res = hvx_vec_mul_f16_f16(sig, vsrc1[0]);
|
|
hvx_vec_store_a((void *) vdst, nloe * sizeof(_Float16), res);
|
|
}
|
|
|
|
static inline void hvx_sigglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
|
|
assert((unsigned long) dst % 128 == 0);
|
|
assert((unsigned long) src0 % 128 == 0);
|
|
assert((unsigned long) src1 % 128 == 0);
|
|
|
|
HVX_Vector * restrict vdst = (HVX_Vector *) dst;
|
|
const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
|
|
const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
|
|
|
|
const uint32_t epv = 128 / sizeof(_Float16);
|
|
const uint32_t nvec = n / epv;
|
|
const uint32_t nloe = n % epv;
|
|
|
|
if (nvec) {
|
|
hvx_sigglu_f16_vec(vdst, vsrc0, vsrc1, nvec);
|
|
}
|
|
if (nloe) {
|
|
hvx_sigglu_f16_tail(vdst + nvec, vsrc0 + nvec, vsrc1 + nvec, nloe);
|
|
}
|
|
}
|
|
|
|
// sigglu(x, g) = sigmoid(x) * g
|
|
static void sigglu_f16(const _Float16 * restrict src0,
|
|
const _Float16 * restrict src1,
|
|
_Float16 * restrict dst,
|
|
const uint32_t num_rows,
|
|
const struct htp_glu_context * gctx) {
|
|
htp_glu_op_preamble;
|
|
|
|
for (uint32_t ib = 0; ib < num_rows; ib++) {
|
|
const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
|
|
const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
|
|
uint8_t * restrict dst_ptr = (uint8_t *) dst + (ib * dst_row_size_aligned);
|
|
|
|
hvx_sigglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
|
|
}
|
|
}
|
|
|
|
#define DEFINE_GLU_PER_THREAD(TYPE, SUFFIX) \
|
|
static void glu_##SUFFIX##_per_thread(unsigned int nth, unsigned int ith, void * data) { \
|
|
struct htp_glu_context * gctx = (struct htp_glu_context *) data; \
|
|
htp_glu_preamble; \
|
|
\
|
|
struct htp_thread_trace * tr = gctx->octx->ctx ? &gctx->octx->ctx->trace[ith] : NULL; \
|
|
\
|
|
size_t src0_row_size = gctx->src0_row_size; \
|
|
size_t src1_row_size = gctx->src1_row_size; \
|
|
size_t dst_row_size = gctx->dst_row_size; \
|
|
\
|
|
size_t src0_row_stride = gctx->src0_row_stride; \
|
|
size_t src1_row_stride = gctx->src1_row_stride; \
|
|
\
|
|
const uint32_t src0_nrows = gctx->src0_nrows; \
|
|
const uint32_t src0_nrows_per_thread = gctx->src0_nrows_per_thread; \
|
|
\
|
|
const uint32_t src0_start_row = gctx->row_start + src0_nrows_per_thread * ith; \
|
|
const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, gctx->row_start + src0_nrows); \
|
|
\
|
|
/* no work for this thread */ \
|
|
if (src0_start_row >= src0_end_row) { \
|
|
return; \
|
|
} \
|
|
\
|
|
const dma_addr_t data_src0 = gctx->data_src0; \
|
|
const dma_addr_t data_src1 = gctx->data_src1; \
|
|
const dma_addr_t data_dst = gctx->data_dst; \
|
|
\
|
|
const size_t src0_row_size_aligned = gctx->src0_row_size_aligned; \
|
|
const size_t src1_row_size_aligned = gctx->src1_row_size_aligned; \
|
|
const size_t dst_row_size_aligned = gctx->dst_row_size_aligned; \
|
|
\
|
|
uint8_t * restrict src0_spad_data = gctx->vtcm_src0 + (ith * gctx->vtcm_src0_size_per_thread); \
|
|
uint8_t * restrict src1_spad_data = gctx->vtcm_src1 + (ith * gctx->vtcm_src1_size_per_thread); \
|
|
uint8_t * restrict dst_spad_data = gctx->vtcm_dst + (ith * gctx->vtcm_dst_size_per_thread); \
|
|
\
|
|
size_t src0_spad_half_size = gctx->src0_spad_half_size; \
|
|
size_t src1_spad_half_size = gctx->src1_spad_half_size; \
|
|
size_t dst_spad_half_size = gctx->dst_spad_half_size; \
|
|
\
|
|
const int BLOCK = gctx->block; \
|
|
\
|
|
dma_queue * dma_q = gctx->octx->ctx->dma[ith]; \
|
|
glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) gctx->compute; \
|
|
\
|
|
for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \
|
|
const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); \
|
|
\
|
|
/* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */ \
|
|
dma_queue_push(dma_q, \
|
|
dma_make_data(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)), \
|
|
dst_row_size, dst_row_size_aligned, dst_row_size, 0); \
|
|
\
|
|
dma_queue_push(dma_q, \
|
|
dma_make_data(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)), \
|
|
src0_row_size_aligned, src0_row_stride, src0_row_size, block_size); \
|
|
\
|
|
dma_queue_push(dma_q, \
|
|
dma_make_data(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)), \
|
|
src1_row_size_aligned, src1_row_stride, src1_row_size, block_size); \
|
|
} \
|
|
\
|
|
for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) { \
|
|
const uint32_t block_size = MIN(BLOCK, src0_end_row - ir); \
|
|
\
|
|
TYPE * dst_spad = (TYPE *) dma_queue_pop(dma_q).src; \
|
|
TYPE * src0_spad = (TYPE *) dma_queue_pop(dma_q).dst; \
|
|
TYPE * src1_spad = (TYPE *) dma_queue_pop(dma_q).dst; \
|
|
\
|
|
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir); \
|
|
compute(src0_spad, src1_spad, dst_spad, block_size, gctx); \
|
|
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir); \
|
|
\
|
|
dma_queue_push(dma_q, dma_make_data(data_dst + (ir * dst_row_size), dst_spad), \
|
|
dst_row_size, dst_row_size_aligned, dst_row_size, block_size); \
|
|
\
|
|
/* prefetch N+2 loop iteration if any */ \
|
|
const uint32_t pref_block = (ir + BLOCK * 2); \
|
|
if (pref_block < src0_end_row) { \
|
|
const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block); \
|
|
dma_queue_push(dma_q, dma_make_data(src0_spad, data_src0 + (pref_block * src0_row_stride)), \
|
|
src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size); \
|
|
dma_queue_push(dma_q, dma_make_data(src1_spad, data_src1 + (pref_block * src1_row_stride)), \
|
|
src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size); \
|
|
} \
|
|
} \
|
|
\
|
|
dma_queue_flush(dma_q); \
|
|
}
|
|
|
|
DEFINE_GLU_PER_THREAD(float, f32)
|
|
DEFINE_GLU_PER_THREAD(_Float16, f16)
|
|
|
|
static int execute_op_glu_f32(struct htp_ops_context * octx) {
|
|
const struct htp_tensor * src0 = octx->src[0];
|
|
const struct htp_tensor * src1 = octx->src[1];
|
|
const struct htp_tensor * dst = octx->dst;
|
|
|
|
glu_compute_f32_fn_t compute_fn = NULL;
|
|
const char * op_type = NULL;
|
|
|
|
switch (octx->op) {
|
|
case HTP_OP_GLU_SWIGLU:
|
|
compute_fn = swiglu_f32;
|
|
op_type = "swiglu-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_SWIGLU_OAI:
|
|
compute_fn = swiglu_oai_f32;
|
|
op_type = "swiglu-oai-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_SWIGLU_CLAMP:
|
|
compute_fn = swiglu_clamp_f32;
|
|
op_type = "swiglu-clamp-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_GEGLU:
|
|
compute_fn = geglu_f32;
|
|
op_type = "geglu-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_GEGLU_QUICK:
|
|
compute_fn = geglu_quick_f32;
|
|
op_type = "geglu-quick-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_GEGLU_ERF:
|
|
compute_fn = geglu_erf_f32;
|
|
op_type = "geglu-erf-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_REGLU:
|
|
compute_fn = reglu_f32;
|
|
op_type = "reglu-f32";
|
|
break;
|
|
|
|
case HTP_OP_GLU_SIGMOID:
|
|
compute_fn = sigglu_f32;
|
|
op_type = "sigglu-f32";
|
|
break;
|
|
|
|
default:
|
|
FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
|
|
return HTP_STATUS_NO_SUPPORT;
|
|
}
|
|
|
|
const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
|
|
|
|
if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
|
|
return HTP_STATUS_INVAL_PARAMS;
|
|
}
|
|
|
|
const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
|
|
const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;
|
|
|
|
uint32_t row_start = 0;
|
|
uint32_t nrows = src0_nrows;
|
|
|
|
if (octx->ctx->mdev.count > 1) {
|
|
uint32_t rows_per_chunk = 0;
|
|
htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
|
|
const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
|
|
row_start = range.start;
|
|
nrows = range.count;
|
|
}
|
|
|
|
if (nrows == 0) {
|
|
return HTP_STATUS_OK;
|
|
}
|
|
|
|
const uint32_t n_threads = octx->n_threads;
|
|
|
|
// row_size = bytes of useful data per row (what the kernel touches / what DMA copies).
|
|
// row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
|
|
const size_t nc_bytes = dst_row_size;
|
|
const size_t src0_row_size = nc_bytes;
|
|
const size_t src1_row_size = nc_bytes;
|
|
const size_t src0_row_stride = src0->nb[1];
|
|
const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
|
|
|
|
struct htp_glu_vtcm_layout L;
|
|
htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);
|
|
|
|
const size_t src0_row_size_aligned = L.row_size_aligned;
|
|
const size_t src1_row_size_aligned = L.row_size_aligned;
|
|
const size_t dst_row_size_aligned = L.row_size_aligned;
|
|
|
|
// Make sure the reserved vtcm size is sufficient
|
|
if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
|
|
FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
|
|
(src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
|
|
return HTP_STATUS_VTCM_TOO_SMALL;
|
|
}
|
|
|
|
if (src1) {
|
|
FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n",
|
|
op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
|
|
src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], L.src0_bytes_per_thread * n_threads,
|
|
L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
|
|
} else {
|
|
FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n", op_type,
|
|
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
|
|
L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
|
|
}
|
|
|
|
// Prepare context
|
|
struct htp_glu_context gctx;
|
|
gctx.octx = octx;
|
|
gctx.compute = (void *) compute_fn;
|
|
gctx.op_str = op_type;
|
|
|
|
gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
|
|
|
|
gctx.src0_row_size = src0_row_size;
|
|
gctx.src1_row_size = src1_row_size;
|
|
gctx.dst_row_size = dst_row_size;
|
|
|
|
gctx.src0_row_size_aligned = src0_row_size_aligned;
|
|
gctx.src1_row_size_aligned = src1_row_size_aligned;
|
|
gctx.dst_row_size_aligned = dst_row_size_aligned;
|
|
|
|
gctx.src0_row_stride = src0_row_stride;
|
|
gctx.src1_row_stride = src1_row_stride;
|
|
|
|
uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
|
|
gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
|
|
gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
|
|
gctx.vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
|
|
|
|
gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
|
|
gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
|
|
gctx.vtcm_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
|
|
|
|
gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
|
|
gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
|
|
gctx.dst_spad_half_size = gctx.vtcm_dst_size_per_thread >> 1;
|
|
|
|
gctx.block = L.block;
|
|
gctx.src0_nrows = nrows;
|
|
gctx.row_start = row_start;
|
|
|
|
gctx.nc = dst->ne[0];
|
|
|
|
// Addresses and GLU logic
|
|
dma_addr_t data_src0 = src0->data;
|
|
dma_addr_t data_src1 = src1 ? src1->data : 0;
|
|
|
|
if (!src1) {
|
|
const int32_t swapped = octx->op_params[1];
|
|
data_src1 = data_src0;
|
|
gctx.src1_row_size = gctx.src0_row_size;
|
|
|
|
size_t nc_in_bytes = gctx.nc * SIZEOF_FP32;
|
|
if (swapped) {
|
|
data_src0 += nc_in_bytes;
|
|
} else {
|
|
data_src1 += nc_in_bytes;
|
|
}
|
|
}
|
|
|
|
gctx.data_src0 = data_src0;
|
|
gctx.data_src1 = data_src1;
|
|
gctx.data_dst = dst->data;
|
|
|
|
work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f32_per_thread, &gctx, n_threads);
|
|
return HTP_STATUS_OK;
|
|
}
|
|
|
|
static int execute_op_glu_f16(struct htp_ops_context * octx) {
|
|
const struct htp_tensor * src0 = octx->src[0];
|
|
const struct htp_tensor * src1 = octx->src[1];
|
|
const struct htp_tensor * dst = octx->dst;
|
|
|
|
glu_compute_f16_fn_t compute_fn = NULL;
|
|
const char * op_type = NULL;
|
|
|
|
switch (octx->op) {
|
|
case HTP_OP_GLU_SWIGLU:
|
|
compute_fn = swiglu_f16;
|
|
op_type = "swiglu-f16";
|
|
break;
|
|
|
|
case HTP_OP_GLU_SWIGLU_OAI:
|
|
compute_fn = swiglu_oai_f16;
|
|
op_type = "swiglu-oai-f16";
|
|
break;
|
|
|
|
case HTP_OP_GLU_SWIGLU_CLAMP:
|
|
compute_fn = swiglu_clamp_f16;
|
|
op_type = "swiglu-clamp-f16";
|
|
break;
|
|
|
|
case HTP_OP_GLU_GEGLU:
|
|
compute_fn = geglu_f16;
|
|
op_type = "geglu-f16";
|
|
break;
|
|
|
|
case HTP_OP_GLU_REGLU:
|
|
compute_fn = reglu_f16;
|
|
op_type = "reglu-f16";
|
|
break;
|
|
|
|
case HTP_OP_GLU_SIGMOID:
|
|
compute_fn = sigglu_f16;
|
|
op_type = "sigglu-f16";
|
|
break;
|
|
|
|
default:
|
|
FARF(ERROR, "Unsupported GLU Op %u\n", octx->op);
|
|
return HTP_STATUS_NO_SUPPORT;
|
|
}
|
|
|
|
const struct htp_glu_kernel_params * kparams = (const struct htp_glu_kernel_params *) octx->kernel_params;
|
|
|
|
if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
|
|
return HTP_STATUS_INVAL_PARAMS;
|
|
}
|
|
|
|
const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
|
|
const size_t dst_row_size = dst->ne[0] * SIZEOF_FP16;
|
|
|
|
uint32_t row_start = 0;
|
|
uint32_t nrows = src0_nrows;
|
|
|
|
if (octx->ctx->mdev.count > 1) {
|
|
uint32_t rows_per_chunk = 0;
|
|
htp_tensor_mdev_rows_per_chunk(dst, sizeof(_Float16), (uint32_t) dst_row_size, &rows_per_chunk);
|
|
const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
|
|
row_start = range.start;
|
|
nrows = range.count;
|
|
}
|
|
|
|
if (nrows == 0) {
|
|
return HTP_STATUS_OK;
|
|
}
|
|
|
|
const uint32_t n_threads = octx->n_threads;
|
|
|
|
// row_size = bytes of useful data per row (what the kernel touches / what DMA copies).
|
|
// row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
|
|
const size_t nc_bytes = dst_row_size;
|
|
const size_t src0_row_size = nc_bytes;
|
|
const size_t src1_row_size = nc_bytes;
|
|
const size_t src0_row_stride = src0->nb[1];
|
|
const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
|
|
|
|
struct htp_glu_vtcm_layout L;
|
|
htp_glu_vtcm_layout_build(&L, dst_row_size, n_threads, kparams->vtcm_row_per_thread);
|
|
|
|
const size_t src0_row_size_aligned = L.row_size_aligned;
|
|
const size_t src1_row_size_aligned = L.row_size_aligned;
|
|
const size_t dst_row_size_aligned = L.row_size_aligned;
|
|
|
|
// Make sure the reserved vtcm size is sufficient
|
|
if (kparams->vtcm_row_per_thread == 0 || octx->ctx->vtcm_size < kparams->vtcm_size) {
|
|
FARF(ERROR, "glu-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
|
|
(src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
|
|
return HTP_STATUS_VTCM_TOO_SMALL;
|
|
}
|
|
|
|
if (src1) {
|
|
FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n",
|
|
op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
|
|
src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], L.src0_bytes_per_thread * n_threads,
|
|
L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
|
|
} else {
|
|
FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n", op_type,
|
|
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
|
|
L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
|
|
}
|
|
|
|
// Prepare context
|
|
struct htp_glu_context gctx;
|
|
gctx.octx = octx;
|
|
gctx.compute = (void *) compute_fn;
|
|
gctx.op_str = op_type;
|
|
|
|
gctx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
|
|
|
|
gctx.src0_row_size = src0_row_size;
|
|
gctx.src1_row_size = src1_row_size;
|
|
gctx.dst_row_size = dst_row_size;
|
|
|
|
gctx.src0_row_size_aligned = src0_row_size_aligned;
|
|
gctx.src1_row_size_aligned = src1_row_size_aligned;
|
|
gctx.dst_row_size_aligned = dst_row_size_aligned;
|
|
|
|
gctx.src0_row_stride = src0_row_stride;
|
|
gctx.src1_row_stride = src1_row_stride;
|
|
|
|
uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
|
|
gctx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
|
|
gctx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
|
|
gctx.vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
|
|
|
|
gctx.vtcm_src0_size_per_thread = (uint32_t) L.src0_bytes_per_thread;
|
|
gctx.vtcm_src1_size_per_thread = (uint32_t) L.src1_bytes_per_thread;
|
|
gctx.vtcm_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
|
|
|
|
gctx.src0_spad_half_size = gctx.vtcm_src0_size_per_thread >> 1;
|
|
gctx.src1_spad_half_size = gctx.vtcm_src1_size_per_thread >> 1;
|
|
gctx.dst_spad_half_size = gctx.vtcm_dst_size_per_thread >> 1;
|
|
|
|
gctx.block = L.block;
|
|
gctx.src0_nrows = nrows;
|
|
gctx.row_start = row_start;
|
|
|
|
gctx.nc = dst->ne[0];
|
|
|
|
// Addresses and GLU logic
|
|
dma_addr_t data_src0 = src0->data;
|
|
dma_addr_t data_src1 = src1 ? src1->data : 0;
|
|
|
|
if (!src1) {
|
|
const int32_t swapped = octx->op_params[1];
|
|
data_src1 = data_src0;
|
|
gctx.src1_row_size = gctx.src0_row_size;
|
|
|
|
size_t nc_in_bytes = gctx.nc * SIZEOF_FP16;
|
|
if (swapped) {
|
|
data_src0 += nc_in_bytes;
|
|
} else {
|
|
data_src1 += nc_in_bytes;
|
|
}
|
|
}
|
|
|
|
gctx.data_src0 = data_src0;
|
|
gctx.data_src1 = data_src1;
|
|
gctx.data_dst = dst->data;
|
|
|
|
work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &gctx, n_threads);
|
|
return HTP_STATUS_OK;
|
|
}
|
|
|
|
|
|
int op_glu(struct htp_ops_context * octx) {
|
|
switch (octx->src[0]->type) {
|
|
case HTP_TYPE_F32:
|
|
return execute_op_glu_f32(octx);
|
|
|
|
case HTP_TYPE_F16:
|
|
return execute_op_glu_f16(octx);
|
|
|
|
default:
|
|
return HTP_STATUS_NO_SUPPORT;
|
|
}
|
|
}
|