#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;
    }
}