#if defined(__ANDROID__) || defined(__aarch64__)
#include <math.h>
#include <float.h>
#include "Arm82Functions.hpp"
#include "Arm82OptFunc.hpp"
#include "Arm82WinogradOptFunc.hpp"
#include "Arm82Vec.hpp"
#include "Arm82Binary.hpp"
#include "Arm82Unary.hpp"
#include "Arm82Relu.hpp"
#include "backend/cpu/compute/CommonOptFunction.h"
#include "backend/cpu/CPUPool.hpp"
#include "backend/cpu/CPURuntime.hpp"

#define FLOAT FLOAT16
#define PACK 8
using Vec = MNN::Math::Vec<FLOAT16, 8>;

#include "backend/cpu/GridSampler.hpp"

#if defined(MNN_USE_NEON)
#include <arm_neon.h>
#endif

extern "C" {
// (UP_DIV(l,8), e, 8) -> (UP_DIV(e,eP), l, eP)
void Arm82MNNPackForMatMul_A(float* destOrigin, float const** sourceGroup, const int32_t* info, const int32_t* el);
// void MNNPackTransposeInt16C8(int16_t* dst, const int16_t* src, size_t area, size_t depth, int32_t* areaOffset);

// C(UP_DIV(h,8), e, h8) = B(UP_DIV(h,hP), l, hP) * A(l, eP), hP = 24
// parameter: [aStride, l, h, cStride, bExtraStride]
// aStride in parameter is deprecated (useless), but for code clean, just retain it
void MNNPackedMatMulFP16(float* C, const float* A, const float* B, const size_t* parameter, const float* postParameters, const float* bias, const float* k, const float* b);

// C(UP_DIV(h,8), e, h8) = B(UP_DIV(h,hP), l, hP) * A(l, e), hP = 24, e >= 1
// parameter: [aStride, l, h, cStride, bExtraStride]
void MNNPackedMatMulRemainFP16(float* C, const float* A, const float* B, size_t eSize, const size_t* parameter, const float* postParameters, const float* bias, const float* k, const float* b);


#ifdef __aarch64__
#ifdef MNN_LOW_MEMORY
void MNNAbsMaxFP16_Pack8(const float* source, float* absmax, size_t src_depth_quad, size_t realSize, int pack);
void MNNAbsMaxFP16_Pack4(const float* source, float* absmax, size_t src_depth_quad, size_t realSize, int pack);
void MNNQuantScaleFP16(float* sum, float* absmax, float* quant_scale, float* dequant_scale, size_t thread, size_t batch);
void MNNDynamicQuantFP16_Pack8(const float* src, int8_t* dst, const float* scale, size_t src_depth_quad, size_t realSize, const float* bias, size_t pack);
void MNNDynamicQuantFP16_Pack4(const float* src, int8_t* dst, const float* scale, size_t src_depth_quad, size_t realSize, const float* bias, size_t pack);
void MNNGeneralIm2col_Arm82(float* destOrigin, float const** sourceGroup, const int32_t* info, const int32_t* el, int32_t LP, int32_t pack);
void MNNGeneralIm2col_Arm86(float* destOrigin, float const** sourceGroup, const int32_t* info, const int32_t* el, int32_t LP, int32_t pack);
#ifdef MNN_SME2
void MNNGeneralIm2col_Fp16Sme2(float* destOrigin, float const** sourceGroup, const int32_t* info, const int32_t* el, int32_t LP, int32_t pack);
#endif
void MNNLocalMinMaxFP16_Pack4(float* dstMin, float* dstMax, const float* source, size_t blockNum, size_t blockLU, size_t EP, size_t LP, size_t loadDstBuffer);
void MNNLocalMinMaxFP16_Pack8(float* dstMin, float* dstMax, const float* source, size_t blockNum, size_t blockLU, size_t EP, size_t LP, size_t loadDstBuffer);
#endif // MNN_LOW_MEMORY
void CountMinMaxValue_FP16(float* source, float* minVal, float* maxVal, size_t sizeQuad);
#ifdef MNN_SME2
void MNNPackedMatMulRemainFP16_SME2(float* C, const float* A, const float* B, size_t eSize, const size_t* parameter, const float* postParameters, const float* bias, const float* k, const float* b);
#endif
#endif

#if defined(__aarch64__)

void MNNDepthwiseConvFastKernelFP16(float* dst, const float* src, const float* weight, size_t width, size_t src_w_setup,
                                    size_t fw, size_t fh, size_t dilateX_step, size_t dilateY_step, size_t height,
                                    size_t srcHStep, size_t dstHStep, const float* bias, const float* parameters);
#endif

void MNNConvRunForLineDepthwiseFP16(float* dst, const float* src, const float* weight, size_t width, size_t src_w_setup,
                                size_t fw, size_t fh, size_t dilateX_step, size_t dilateY_step, size_t height, size_t srcHStep, size_t dstHStep);

// LinearAttention fp16 kernels
void MNNRankOneUpdateFp16(float* S, const float* k, const float* delta, size_t dk, size_t dv);
void MNNDualMatVecFp16(const float* S, const float* k, const float* q, float* out_k, float* out_q, size_t dk, size_t dv);
void MNNDecayRankOneUpdateFp16(float* S, const float* k, const float* delta, float decay, size_t dk, size_t dv);
}

#if defined(__aarch64__) && defined(MNN_USE_NEON)
// ──────────────────────────────────────────────────────────────────────────
// MNNFusedGatedDeltaFp16 — FP16 specialization of the fused gated-delta-rule
// kernel. See documentation in CommonOptFunction.h for the math.
//
// Pointers are typed `float*` for ABI uniformity with the FP32 version (the
// dispatch table holds a single function pointer signature), but the
// underlying memory is fp16. d_v is processed in chunks of 32 (= 4 v.8h
// registers per accumulator) to keep all per-chunk state in registers,
// which matters at d_v=128 where holding both out_k and out_q for the
// whole row would otherwise exceed the 32 NEON v register budget.
// ──────────────────────────────────────────────────────────────────────────
static void MNNFusedGatedDeltaFp16(float* S_, const float* k_, const float* q_, const float* v_, float* out_,
                                   float decay, float beta, float kq, size_t dk, size_t dv) {
    auto S = reinterpret_cast<__fp16*>(S_);
    auto k = reinterpret_cast<const __fp16*>(k_);
    auto q = reinterpret_cast<const __fp16*>(q_);
    auto vIn = reinterpret_cast<const __fp16*>(v_);
    auto out = reinterpret_cast<__fp16*>(out_);

    const __fp16 decayH = static_cast<__fp16>(decay);
    const __fp16 betaH = static_cast<__fp16>(beta);
    const __fp16 kqH = static_cast<__fp16>(kq);
    const float16x8_t vDecay = vdupq_n_f16(decayH);
    const float16x8_t vBeta = vdupq_n_f16(betaH);
    const float16x8_t vKq = vdupq_n_f16(kqH);

    const size_t kChunk = 32;
    size_t j = 0;
    for (; j + kChunk <= dv; j += kChunk) {
        // ── Pass 1: out_k = S^T @ k, out_q = S^T @ q for this column chunk ──
        float16x8_t ok0 = vdupq_n_f16((__fp16)0), ok1 = vdupq_n_f16((__fp16)0), ok2 = vdupq_n_f16((__fp16)0),
                    ok3 = vdupq_n_f16((__fp16)0);
        float16x8_t oq0 = vdupq_n_f16((__fp16)0), oq1 = vdupq_n_f16((__fp16)0), oq2 = vdupq_n_f16((__fp16)0),
                    oq3 = vdupq_n_f16((__fp16)0);
        size_t i = 0;
        // Unroll i by 8: load 8 k & q scalars at once, then use fma-by-lane
        // to amortize the scalar broadcast across 8 row iterations.
        for (; i + 8 <= dk; i += 8) {
            float16x8_t kVec = vld1q_f16(k + i);
            float16x8_t qVec = vld1q_f16(q + i);
#define LANE_STEP(lane)                                \
    {                                                  \
        const __fp16* row = S + (i + (lane)) * dv + j; \
        float16x8_t s0 = vld1q_f16(row);               \
        float16x8_t s1 = vld1q_f16(row + 8);           \
        float16x8_t s2 = vld1q_f16(row + 16);          \
        float16x8_t s3 = vld1q_f16(row + 24);          \
        ok0 = vfmaq_laneq_f16(ok0, s0, kVec, (lane));  \
        ok1 = vfmaq_laneq_f16(ok1, s1, kVec, (lane));  \
        ok2 = vfmaq_laneq_f16(ok2, s2, kVec, (lane));  \
        ok3 = vfmaq_laneq_f16(ok3, s3, kVec, (lane));  \
        oq0 = vfmaq_laneq_f16(oq0, s0, qVec, (lane));  \
        oq1 = vfmaq_laneq_f16(oq1, s1, qVec, (lane));  \
        oq2 = vfmaq_laneq_f16(oq2, s2, qVec, (lane));  \
        oq3 = vfmaq_laneq_f16(oq3, s3, qVec, (lane));  \
    }
            LANE_STEP(0);
            LANE_STEP(1);
            LANE_STEP(2);
            LANE_STEP(3);
            LANE_STEP(4);
            LANE_STEP(5);
            LANE_STEP(6);
            LANE_STEP(7);
#undef LANE_STEP
        }
        // Tail rows (dk % 8) — fall back to the broadcast form.
        for (; i < dk; ++i) {
            const __fp16* row = S + i * dv + j;
            float16x8_t s0 = vld1q_f16(row);
            float16x8_t s1 = vld1q_f16(row + 8);
            float16x8_t s2 = vld1q_f16(row + 16);
            float16x8_t s3 = vld1q_f16(row + 24);
            __fp16 ki = k[i];
            __fp16 qi = q[i];
            ok0 = vfmaq_n_f16(ok0, s0, ki);
            ok1 = vfmaq_n_f16(ok1, s1, ki);
            ok2 = vfmaq_n_f16(ok2, s2, ki);
            ok3 = vfmaq_n_f16(ok3, s3, ki);
            oq0 = vfmaq_n_f16(oq0, s0, qi);
            oq1 = vfmaq_n_f16(oq1, s1, qi);
            oq2 = vfmaq_n_f16(oq2, s2, qi);
            oq3 = vfmaq_n_f16(oq3, s3, qi);
        }

        // ── Inline analytic correction (regs only) ──
        float16x8_t v0 = vld1q_f16(vIn + j);
        float16x8_t v1 = vld1q_f16(vIn + j + 8);
        float16x8_t v2 = vld1q_f16(vIn + j + 16);
        float16x8_t v3 = vld1q_f16(vIn + j + 24);
        // delta = beta * (v - decay * out_k)
        float16x8_t d0 = vmulq_f16(vBeta, vsubq_f16(v0, vmulq_f16(vDecay, ok0)));
        float16x8_t d1 = vmulq_f16(vBeta, vsubq_f16(v1, vmulq_f16(vDecay, ok1)));
        float16x8_t d2 = vmulq_f16(vBeta, vsubq_f16(v2, vmulq_f16(vDecay, ok2)));
        float16x8_t d3 = vmulq_f16(vBeta, vsubq_f16(v3, vmulq_f16(vDecay, ok3)));
        // out = decay * out_q + kq * delta
        float16x8_t o0 = vfmaq_f16(vmulq_f16(vDecay, oq0), vKq, d0);
        float16x8_t o1 = vfmaq_f16(vmulq_f16(vDecay, oq1), vKq, d1);
        float16x8_t o2 = vfmaq_f16(vmulq_f16(vDecay, oq2), vKq, d2);
        float16x8_t o3 = vfmaq_f16(vmulq_f16(vDecay, oq3), vKq, d3);
        vst1q_f16(out + j, o0);
        vst1q_f16(out + j + 8, o1);
        vst1q_f16(out + j + 16, o2);
        vst1q_f16(out + j + 24, o3);

        // ── Pass 2: S = decay * S + k ⊗ delta (delta still in regs) ──
        size_t i2 = 0;
        for (; i2 + 8 <= dk; i2 += 8) {
            float16x8_t kVec = vld1q_f16(k + i2);
#define ROW_UPDATE(lane)                                                           \
    {                                                                              \
        __fp16* row = S + (i2 + (lane)) * dv + j;                                  \
        float16x8_t s0 = vld1q_f16(row);                                           \
        float16x8_t s1 = vld1q_f16(row + 8);                                       \
        float16x8_t s2 = vld1q_f16(row + 16);                                      \
        float16x8_t s3 = vld1q_f16(row + 24);                                      \
        float16x8_t r0 = vfmaq_laneq_f16(vmulq_f16(vDecay, s0), d0, kVec, (lane)); \
        float16x8_t r1 = vfmaq_laneq_f16(vmulq_f16(vDecay, s1), d1, kVec, (lane)); \
        float16x8_t r2 = vfmaq_laneq_f16(vmulq_f16(vDecay, s2), d2, kVec, (lane)); \
        float16x8_t r3 = vfmaq_laneq_f16(vmulq_f16(vDecay, s3), d3, kVec, (lane)); \
        vst1q_f16(row, r0);                                                        \
        vst1q_f16(row + 8, r1);                                                    \
        vst1q_f16(row + 16, r2);                                                   \
        vst1q_f16(row + 24, r3);                                                   \
    }
            ROW_UPDATE(0);
            ROW_UPDATE(1);
            ROW_UPDATE(2);
            ROW_UPDATE(3);
            ROW_UPDATE(4);
            ROW_UPDATE(5);
            ROW_UPDATE(6);
            ROW_UPDATE(7);
#undef ROW_UPDATE
        }
        for (; i2 < dk; ++i2) {
            __fp16* row = S + i2 * dv + j;
            float16x8_t s0 = vld1q_f16(row);
            float16x8_t s1 = vld1q_f16(row + 8);
            float16x8_t s2 = vld1q_f16(row + 16);
            float16x8_t s3 = vld1q_f16(row + 24);
            __fp16 ki = k[i2];
            float16x8_t r0 = vfmaq_n_f16(vmulq_f16(vDecay, s0), d0, ki);
            float16x8_t r1 = vfmaq_n_f16(vmulq_f16(vDecay, s1), d1, ki);
            float16x8_t r2 = vfmaq_n_f16(vmulq_f16(vDecay, s2), d2, ki);
            float16x8_t r3 = vfmaq_n_f16(vmulq_f16(vDecay, s3), d3, ki);
            vst1q_f16(row, r0);
            vst1q_f16(row + 8, r1);
            vst1q_f16(row + 16, r2);
            vst1q_f16(row + 24, r3);
        }
    }

    // ── Tail (chunks of 8) ──
    for (; j + 8 <= dv; j += 8) {
        float16x8_t ok = vdupq_n_f16((__fp16)0);
        float16x8_t oq = vdupq_n_f16((__fp16)0);
        for (size_t i = 0; i < dk; ++i) {
            float16x8_t s = vld1q_f16(S + i * dv + j);
            ok = vfmaq_n_f16(ok, s, k[i]);
            oq = vfmaq_n_f16(oq, s, q[i]);
        }
        float16x8_t vv = vld1q_f16(vIn + j);
        float16x8_t d = vmulq_f16(vBeta, vsubq_f16(vv, vmulq_f16(vDecay, ok)));
        float16x8_t o = vfmaq_f16(vmulq_f16(vDecay, oq), vKq, d);
        vst1q_f16(out + j, o);
        for (size_t i = 0; i < dk; ++i) {
            float16x8_t s = vld1q_f16(S + i * dv + j);
            float16x8_t r = vfmaq_n_f16(vmulq_f16(vDecay, s), d, k[i]);
            vst1q_f16(S + i * dv + j, r);
        }
    }

    // ── Scalar tail (defensive; d_v < multiple of 8 not used in current models) ──
    for (; j < dv; ++j) {
        float ok = 0.0f, oq = 0.0f;
        for (size_t i = 0; i < dk; ++i) {
            float s = (float)S[i * dv + j];
            ok += s * (float)k[i];
            oq += s * (float)q[i];
        }
        float delta_j = beta * ((float)vIn[j] - decay * ok);
        out[j] = (__fp16)(decay * oq + kq * delta_j);
        for (size_t i = 0; i < dk; ++i) {
            float s = (float)S[i * dv + j];
            S[i * dv + j] = (__fp16)(decay * s + (float)k[i] * delta_j);
        }
    }
}
#endif // __aarch64__ && MNN_USE_NEON

namespace MNN {

#define FP16_SME2_MATMUL_EP 16
#define FP16_SME2_MATMUL_LP 2
#define FP16_SME2_MATMUL_HP 64

static void Sme2MNNGetMatMulPackMode(int* eP, int *lP, int* hP) {
    *hP = FP16_SME2_MATMUL_HP;
    *eP = FP16_SME2_MATMUL_EP;
    *lP = FP16_SME2_MATMUL_LP;
}

static void MNNMatrixAddFP16(FLOAT16* C, const FLOAT16* A, const FLOAT16* B, size_t widthC8, size_t cStride, size_t aStride, size_t bStride, size_t height) {
    for (int y = 0; y < height; ++y) {
        auto a = A + aStride * y, b = B + bStride * y;
        auto c = C + cStride * y;
        for (int x = 0; x < widthC8; ++x) {
            vst1q_f16(c + x * 8, vaddq_f16(vld1q_f16(a + x * 8), vld1q_f16(b + x * 8)));
        }
    }
}
static void MNNMatrixSubFP16(FLOAT16* C, const FLOAT16* A, const FLOAT16* B, size_t widthC8, size_t cStride, size_t aStride, size_t bStride, size_t height) {
    for (int y = 0; y < height; ++y) {
        auto a = A + aStride * y, b = B + bStride * y;
        auto c = C + cStride * y;
        for (int x = 0; x < widthC8; ++x) {
            vst1q_f16(c + x * 8, vsubq_f16(vld1q_f16(a + x * 8), vld1q_f16(b + x * 8)));
        }
    }
}
#if defined(__aarch64__)
static void ARM82CountMinMaxValue(float* source, float* minVal, float* maxVal, size_t size) {
    if (size % 8 == 0) {
        CountMinMaxValue_FP16(source, minVal, maxVal, size / 8);
    } else {
        auto remain = size - 8 * (size / 8);
        auto max_ = ((__fp16*)source)[0];
        auto min_ = max_;
        if (size >= 8) {
            CountMinMaxValue_FP16(source, minVal, maxVal, size / 8);
            max_ = ((__fp16*)maxVal)[0];
            min_ = ((__fp16*)minVal)[0];
        }
        auto srcPtr = reinterpret_cast<__fp16*>(source);
        while (remain) {
            max_ = ALIMAX(srcPtr[0], max_);
            min_ = ALIMIN(srcPtr[0], min_);
            srcPtr += 1;
            remain--;
        }
        reinterpret_cast<__fp16*>(minVal)[0] = min_;
        reinterpret_cast<__fp16*>(maxVal)[0] = max_;
    }
}
#ifdef MNN_SME2
//(float* C, const float* A, const float* B, const size_t* parameter, const float* postParameters, const float* bias)
static void MNNPackedMatMulFP16_SME2(float* C, const float* A, const float* B, const size_t* parameter, const float* postParameters, const float* bias, const float* k, const float* b) {
    MNNPackedMatMulRemainFP16_SME2(C, A, B, 16, parameter, postParameters, bias, k, b);
}
#endif
#else
static void ARM82CountMinMaxValue(float* source, float* minVal, float* maxVal, size_t size) {
    auto srcPtr = (FLOAT16*)source;
    auto minPtr = (FLOAT16*)minVal;
    auto maxPtr = (FLOAT16*)maxVal;
    auto max_ = srcPtr[0], min_ = srcPtr[0];
    for (int i = 1; i < size; ++i) {
        if (max_ < srcPtr[i]) {
            max_ = srcPtr[i];
        }
        if (min_ > srcPtr[i]) {
            min_ = srcPtr[i];
        }
    }
    minPtr[0] = min_;
    maxPtr[0] = max_;
}
#endif

static void Arm82MNNPackForMatMul_B(float* destC, const float* sourceC, size_t h, size_t kernelsize, size_t ic, bool transpose) {
    auto l = kernelsize * ic;
    auto dest = (int16_t*)destC;
    auto source = (int16_t*)sourceC;
    int ePack, lPack, hPack;
    Arm82MNNGetMatMulPackMode(&ePack, &lPack, &hPack);
    auto hP = (int)h / hPack;
    auto hR = (int)hP * hPack;
    if (hR != h) {
        ::memset(dest, 0, UP_DIV(h, hPack) * hPack * l * sizeof(FLOAT16));
    }
    if (!transpose) {
        for (int y = 0; y < hP; ++y) {
            auto destY = dest + y * hPack * l;
            auto sourceY = source + y * hPack;
            for (int x = 0; x < l; ++x) {
                ::memcpy(destY + hPack * x, sourceY + x * h, hPack * sizeof(FLOAT16));
            }
        }
        auto hRemain = h - hR;
        if (hRemain > 0) {
            auto destY = dest + hP * hPack * l;
            auto sourceY = source + hP * hPack;
            for (int x = 0; x < l; ++x) {
                ::memcpy(destY + hPack * x, sourceY + x * h, hRemain * sizeof(FLOAT16));
            }
        }
        return;
    }
    for (int y = 0; y < h; ++y) {
        for (int x = 0; x < l; ++x) {
            dest[(y / hPack * l + x) * hPack + y % hPack] = source[y * l + x];
        }
    }
}

static void Sme2MNNPackForMatMul_B(float* destC, const float* sourceC, size_t h, size_t kernelsize, size_t ic ,bool transpose) {
    auto dest = (int16_t*)destC;
    auto source = (int16_t*)sourceC;
    int LP = FP16_SME2_MATMUL_LP;
    int HP = FP16_SME2_MATMUL_HP;
    auto l = kernelsize * ic;
    memset(dest, 0, ROUND_UP(h, HP) * ROUND_UP(ic, LP) * kernelsize * sizeof(FLOAT16));
    auto stride0 = ROUND_UP(ic, LP) * kernelsize * HP;
    auto stride1 = HP *  ROUND_UP(ic, LP);
    auto stride2 = HP * LP;

    size_t srcStride0 = l; // [h,k2,ic]->[hu,k2,ic/lp,hp,lp]
    size_t srcStride1 = 1;
    if (!transpose) { // [k2,ic,h]->[hu,k2,ic/lp,hp,lp]
        srcStride0 = 1;
        srcStride1 = h;
    }
    for (int y = 0; y < h; ++y) {
        auto yHu = y / HP;
        auto yHp = y % HP;
        for (int k = 0; k < kernelsize; ++k) {
            for (int x = 0; x < ic; ++x) {
                auto xLu = x / LP;
                auto xLp = x % LP;
                dest[yHu * stride0 + k * stride1 + xLu * stride2 + yHp * LP + xLp] = source[y * srcStride0 + (x + k * ic) * srcStride1];
            }
        }
    }
}

static void MNNScaleAndAddBiasFP16(FLOAT16* dst, const FLOAT16* src, const FLOAT16* bias, const FLOAT16* alpha, size_t planeNumber,
                                   size_t biasNumber) {
    for (int z = 0; z < biasNumber; ++z) {
        FLOAT16* dstZ         = dst + planeNumber * 8 * z;
        const FLOAT16* srcZ   = src + planeNumber * 8 * z;
        auto biasZ = vld1q_f16(bias + 8 * z), alphaZ = vld1q_f16(alpha + 8 * z);
        for (int p = 0; p < planeNumber; ++p) {
            FLOAT16* dstX       = dstZ + 8 * p;
            const FLOAT16* srcX = srcZ + 8 * p;
            auto res = vaddq_f16(vmulq_f16(vld1q_f16(srcX), alphaZ), biasZ);
            vst1q_f16(dstX, res);
        }
    }
}

static void MNNGridSampleComputeCordFP16(FLOAT16* dst, const FLOAT16* src, size_t inH, size_t inW, size_t outH, size_t outW, size_t stride, bool alignCorners) {
    float16x8_t zero = vdupq_n_f16(0);
    float16x8_t one = vdupq_n_f16(1);
    float16x8_t half = vdupq_n_f16(0.5f);
    float16x8_t a = alignCorners ? one : zero;
    float16x8_t b = alignCorners ? zero : one;
    float16x8_t inW_sub_a = vsubq_f16(vdupq_n_f16(inW), a);
    float16x8_t inH_sub_a = vsubq_f16(vdupq_n_f16(inH), a);

    int area = outH * outW;
    int areaC8 = area / 8;
    int areaRemain = area - areaC8 * 8;
    for (int i = 0; i < areaC8; ++i) {
        auto cordH = vld2q_f16(src);
        // float16x8_t x = cordH.val[0];
        // float16x8_t y = cordH.val[1];
        cordH.val[0] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[0]), inW_sub_a), b));
        cordH.val[1] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[1]), inH_sub_a), b));
        vst2q_f16(dst, cordH);

        src += 16;
        dst += 16;
    }
    if (areaRemain == 0) {
        return;
    }

    // areaRemain
    FLOAT16 tempDst[16];
    ::memcpy(tempDst, src, areaRemain * 2 * sizeof(int16_t));
    auto cordH = vld2q_f16(tempDst);
    cordH.val[0] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[0]), inW_sub_a), b));
    cordH.val[1] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[1]), inH_sub_a), b));
    vst2q_f16(tempDst, cordH);
    ::memcpy(dst, tempDst, areaRemain * 2 * sizeof(int16_t));
}

static void MNNGridSampleComputeCord3DFp16(FLOAT* dst, const FLOAT* src, size_t inD, size_t inH, size_t inW, size_t outD, size_t outH, size_t outW, bool alignCorners) {
    float16x8_t zero = vdupq_n_f16(0);
    float16x8_t one = vdupq_n_f16(1);
    float16x8_t half = vdupq_n_f16(0.5f);
    float16x8_t a = alignCorners ? one : zero;
    float16x8_t b = alignCorners ? zero : one;
    float16x8_t inW_sub_a = vsubq_f16(vdupq_n_f16(inW), a);
    float16x8_t inH_sub_a = vsubq_f16(vdupq_n_f16(inH), a);
    float16x8_t inD_sub_a = vsubq_f16(vdupq_n_f16(inD), a);
    size_t area = outH * outW * outD;
    size_t areaC8 = area / 8;
    size_t areaRemain = area - areaC8 * 8;

    for (int i = 0; i < areaC8; ++i) {
        auto cordH = vld3q_f16(src);
        // float16x8_t x = cordH.val[0];
        // float16x8_t y = cordH.val[1];
        cordH.val[0] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[0]), inW_sub_a), b));
        cordH.val[1] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[1]), inH_sub_a), b));
        cordH.val[2] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[2]), inD_sub_a), b));
        vst3q_f16(dst, cordH);
        src += 24;
        dst += 24;
    }
    if (areaRemain == 0) {
        return;
    }

    // areaRemain
    FLOAT16 tempDst[24];
    ::memcpy(tempDst, src, areaRemain * 3 * sizeof(int16_t));
    auto cordH = vld3q_f16(tempDst);
    cordH.val[0] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[0]), inW_sub_a), b));
    cordH.val[1] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[1]), inH_sub_a), b));
    cordH.val[2] = vmulq_f16(half, vsubq_f16(vmulq_f16(vaddq_f16(one, cordH.val[2]), inD_sub_a), b));
    vst3q_f16(tempDst, cordH);
    ::memcpy(dst, tempDst, areaRemain * 3 * sizeof(int16_t));
}
static void MNNRoiPoolingMaxFP16(FLOAT16* dst, const FLOAT16* src, int hLen, int wLen, int iw) {
    Vec max = Vec(-65504.0f);
    for (int h = 0; h < hLen; h++, src += iw * 8) {
        for (int w = 0; w < wLen; w++) {
            Vec in = Vec::load(src + w * 8);
            max = Vec::max(max, in);
        }
    }
    Vec::save(dst, max);
}

static void MNNRoiAlignMaxFP16(FLOAT16* dst, const FLOAT16* src, const std::vector<std::vector<int>> &vecPos, const std::vector<std::vector<float>> &vecArea, int samplingRatioArea, int pooledHeight, int pooledWidth) {
    for (int h = 0; h < pooledHeight; ++h, dst += pooledWidth * 8) {
        int preCalcIdx = h * pooledWidth * samplingRatioArea;
        for (int w = 0; w < pooledWidth; ++w) {
            Vec res = Vec(-65504.0f);
            for (int i = 0; i < samplingRatioArea; ++i) {
                const std::vector<int>& pos    = vecPos[preCalcIdx];
                const std::vector<float>& area = vecArea[preCalcIdx];

                Vec val0 = Vec::load(src + pos[0] * 8);
                Vec val1 = Vec::load(src + pos[1] * 8);
                Vec val2 = Vec::load(src + pos[2] * 8);
                Vec val3 = Vec::load(src + pos[3] * 8);
                Vec mla  = val0 * area[0];
                mla       = Vec::fma(mla, val1, area[1]);
                mla       = Vec::fma(mla, val2, area[2]);
                mla       = Vec::fma(mla, val3, area[3]);
                res       = Vec::max(res, mla);
                preCalcIdx++;
            }
            Vec::save(dst + w * 8, res);
        }
    }
}

static void MNNRoiAlignAvgFP16(FLOAT16* dst, const FLOAT16* src, const std::vector<std::vector<int>> &vecPos, const std::vector<std::vector<float>> &vecArea, int samplingRatioArea, int pooledHeight, int pooledWidth) {
    float invSamplingCnt = 1.f / samplingRatioArea;
    for (int h = 0; h < pooledHeight; ++h, dst += pooledWidth * 8) {
        int preCalcIdx = h * pooledWidth * samplingRatioArea;
        for (int w = 0; w < pooledWidth; ++w) {
            Vec res = Vec(0.f);
            for (int i = 0; i < samplingRatioArea; ++i) {
                const std::vector<int>& pos    = vecPos[preCalcIdx];
                const std::vector<float>& area = vecArea[preCalcIdx];

                Vec val0 = Vec::load(src + pos[0] * 8);
                Vec val1 = Vec::load(src + pos[1] * 8);
                Vec val2 = Vec::load(src + pos[2] * 8);
                Vec val3 = Vec::load(src + pos[3] * 8);
                Vec mla  = val0 * area[0];
                mla       = Vec::fma(mla, val1, area[1]);
                mla       = Vec::fma(mla, val2, area[2]);
                mla       = Vec::fma(mla, val3, area[3]);
                res       += mla;
                preCalcIdx++;
            }
            res = res * invSamplingCnt;
            Vec::save(dst + w * 8, res);
        }
    }
}

static void MNNCopyC8WithStrideFP16(const FLOAT16* source, FLOAT16* dest, size_t srcStride, size_t dstStride, size_t count) {
    using Vec = MNN::Math::Vec<FLOAT16, 8>;
    for (int i = 0; i < count; ++i) {
        auto srcPtr = source + i * srcStride;
        auto dstPtr = dest + i * dstStride;
        Vec::save(dstPtr, Vec::load(srcPtr));
    }
}

static void MNNAddC8WithStrideFP16(const FLOAT16* source, FLOAT16* dest, size_t srcStride, size_t dstStride, size_t count) {
    for (int i = 0; i < count; ++i) {
        auto srcPtr = source + i * srcStride;
        auto dstPtr = dest + i * dstStride;
        auto value = Vec::load(dstPtr) + Vec::load(srcPtr);
        Vec::save(dstPtr, value);
    }
}

static void MNNAxByClampBroadcastC8FP16(float* CF, const float* AF, const float* BF, size_t width, size_t cStride, size_t aStride, size_t height, const float* parameters) {
    auto C = (FLOAT16*)CF;
    auto A = (FLOAT16*)AF;
    auto B = (FLOAT16*)BF;
    using Vec = MNN::Math::Vec<FLOAT16, 8>;
    auto minF = Vec(parameters[2]);
    auto maxF = Vec(parameters[3]);
    auto beta = Vec(parameters[1]);
    for (int y = 0; y < height; ++y) {
        auto a = A + aStride * y;
        auto b = B + 8 * y;
        auto bv = Vec::load(b);
        auto c = C + cStride * y;
        for (int x = 0; x < width; ++x) {
            auto av = Vec::load(a + 8 * x);
            auto cv = av + bv * beta;
            cv = Vec::min(cv, maxF);
            cv = Vec::max(cv, minF);
            Vec::save(c + 8 * x, cv);
        }
    }
}

void ARM82StrassenMerge(FLOAT16* c11, FLOAT16* c12, FLOAT16* c21, FLOAT16* c22, FLOAT16* xAddr,
                        size_t cStride, size_t eSub, size_t hSub) {
    const int pack = 8;
    for (int y = 0; y < hSub; ++y) {
        auto c11Y = c11 + y * cStride;
        auto c12Y = c12 + y * cStride;
        auto c22Y = c22 + y * cStride;
        auto c21Y = c21 + y * cStride;
        auto xY = xAddr + y * eSub * pack;
        for (int x = 0; x < eSub; ++x) {
            auto xv = vld1q_f16(xY + x * pack);
            auto c21v = vld1q_f16(c21Y + x * pack);
            auto c11v = vld1q_f16(c11Y + x * pack);
            auto c22v = vld1q_f16(c22Y + x * pack);
            auto c12v = vld1q_f16(c12Y + x * pack);
            c12v = c12v + xv;
            c21v = c12v + c21v;
            c12v = c22v + c12v;
            c22v = c22v + c21v;
            c12v = c11v + c12v;
            vst1q_f16(c12Y + x * pack, c12v);
            vst1q_f16(c22Y + x * pack, c22v);
            vst1q_f16(c21Y + x * pack, c21v);
        }
    }
}

void MNNUnpackTransposeInt16C8(int16_t* dst, const int16_t* src, size_t area, size_t depth, int32_t* areaOffset) {
    // [depth/8, srcAreaOffset, 8] -> [area, dstAreaOffset]
    int srcAreaOffset = areaOffset[0];
    int dstAreaOffset = areaOffset[1];
    int c      = (int)depth;
    int cDiv8  = c / 8;
    int cAlign = cDiv8 * 8;
    int areaDiv4 = area / 4;
    int areaAlign = areaDiv4 * 4;

    if (areaAlign > 0) {
        for (int ci = 0; ci < cDiv8; ++ci) {
            auto srcH = src + ci * 8 * srcAreaOffset;
            auto dstH = dst + ci * 8;
            for (int hi = 0; hi < areaAlign; hi+=4) {
                auto src0 = srcH + hi * 8;
                auto src1 = srcH + hi * 8 + 8;
                auto src2 = srcH + hi * 8 + 16;
                auto src3 = srcH + hi * 8 + 24;

                auto dst0 = dstH + hi * dstAreaOffset;
                auto dst1 = dstH + hi * dstAreaOffset + dstAreaOffset;
                auto dst2 = dstH + hi * dstAreaOffset + 2 * dstAreaOffset;
                auto dst3 = dstH + hi * dstAreaOffset + 3 * dstAreaOffset;
                vst1q_s16(dst0, vld1q_s16(src0));
                vst1q_s16(dst1, vld1q_s16(src1));
                vst1q_s16(dst2, vld1q_s16(src2));
                vst1q_s16(dst3, vld1q_s16(src3));
            }
        }
    }
    if (areaAlign < area) {
        for (int ci = 0; ci < cDiv8; ++ci) {
            auto srcH = src + 8 * ci * srcAreaOffset;
            auto dstH = dst + ci * 8;
            for (int hi = areaAlign; hi < area; ++hi) {
                auto src0 = srcH + hi * 8;
                auto dst0 = dstH + hi * dstAreaOffset;
                vst1q_s16(dst0, vld1q_s16(src0));
            }
        }
    }
    if (c == cAlign) {
        return;
    }

    int cReamin   = c - cAlign;
    auto srcAlign = src + srcAreaOffset * cAlign;
    auto dstAlign = dst + cAlign;

    for (int hi = 0; hi < area; ++hi) {
        auto srcHeight = srcAlign + hi * 8;
        auto dstHeight = dstAlign + hi * dstAreaOffset;

        for (int ci = 0; ci < cReamin; ++ci) {
            dstHeight[ci] = srcHeight[ci];
        }
    }
}

void MNNPackTransposeInt16C8(int16_t* dst, const int16_t* src, size_t area, size_t depth, int32_t* areaOffset) {
    if (depth == 8) {
        ::memcpy(dst, src, area * depth * sizeof(int16_t));
        return;
    }
    int dstAreaOffset = areaOffset[1];
    int c      = (int)depth;
    int cDiv4  = c / 8;
    int cAlign = cDiv4 * 8;
    int areaDiv4 = area / 4;
    int areaAlign = areaDiv4 * 4;
    if (areaAlign > 0) {
        for (int ci = 0; ci < cDiv4; ++ci) {
            auto srcH = src + ci * 8;
            auto dstH = dst + ci * dstAreaOffset * 8;
            for (int hi = 0; hi < areaAlign; hi+=4) {
                auto src0 = srcH + hi * c;
                auto src1 = srcH + hi * c + c;
                auto src2 = srcH + hi * c + 2 * c;
                auto src3 = srcH + hi * c + 3 * c;
                auto dst0 = dstH + hi * 8;
                auto dst1 = dstH + hi * 8 + 8;
                auto dst2 = dstH + hi * 8 + 16;
                auto dst3 = dstH + hi * 8 + 24;
                vst1q_s16(dst0, vld1q_s16(src0));
                vst1q_s16(dst1, vld1q_s16(src1));
                vst1q_s16(dst2, vld1q_s16(src2));
                vst1q_s16(dst3, vld1q_s16(src3));
            }
        }
    }
    if (areaAlign < area) {
        for (int ci = 0; ci < cDiv4; ++ci) {
            auto srcH = src + ci * 8;
            auto dstH = dst + ci * dstAreaOffset * 8;
            for (int hi = areaAlign; hi < area; ++hi) {
                auto src0 = srcH + hi * c;
                auto dst0 = dstH + hi * 8;
                vst1q_s16(dst0, vld1q_s16(src0));
            }
        }
    }

    if (cAlign == c) {
        return;
    }

    int cReamin   = c - cAlign;
    auto srcAlign = src + cAlign;
    auto dstAlign = dst + dstAreaOffset * cAlign;

    for (int hi = 0; hi < area; ++hi) {
        auto srcHeight = srcAlign + hi * c;
        auto dstHeight = dstAlign + hi * 8;
        for (int i = 0; i < 8; ++i) {
            dstHeight[i] = 0;
        }
        for (int ci = 0; ci < cReamin; ++ci) {
            dstHeight[ci] = srcHeight[ci];
        }
    }
}

static void _MNNDeconvRunForUnitDepthWise(const FLOAT16* dst, FLOAT16* src, const FLOAT16* weight, size_t fw, size_t fh,
                                          size_t weight_y_step, size_t dilateX_step, size_t dilateY_step) {
    int fx, fy;
    auto src_z          = src;
    auto weight_z = weight;
    Vec dstV           = Vec::load(dst);
    for (fy = 0; fy < fh; ++fy) {
        auto src_y          = src_z + fy * dilateY_step;
        auto weight_y = weight_z + fy * weight_y_step;
        for (fx = 0; fx < fw; ++fx) {
            Vec weight_x = Vec::load(weight_y + 8 * fx);
            Vec src_x    = Vec::load(src_y + fx * dilateX_step);
            Vec::save(src_y + fx * dilateX_step, src_x + weight_x * dstV);
        }
    }
}
static void _MNNDeconvRunForLineDepthwise(const FLOAT16* dst, FLOAT16* src, const FLOAT16* weight, size_t width, size_t src_w_setup,
                                          size_t fw, size_t fh, size_t dilateX_step, size_t dilateY_step) {
    int dx;
    for (dx = 0; dx < width; ++dx) {
        auto dst_x = dst + dx * 8;
        auto src_dx      = src + src_w_setup * dx;
        _MNNDeconvRunForUnitDepthWise(dst_x, src_dx, weight, fw, fh, fw * 8, dilateX_step, dilateY_step);
    }
}

static void _MNNComputeMatMulForH_1_FP16(const float* AF, const float* BF, float* CF, const float* biasPtrF, const MatMulParam* param, size_t tId) {
    auto A = (const FLOAT16*)AF;
    auto B = (const FLOAT16*)BF;
    auto C = (FLOAT16*)CF;
    auto biasPtr = (const FLOAT16*)biasPtrF;
    int e = param->e;
    int l = param->l;
    int numberThread = param->numberThread;
    float biasValue = 0.0f;
    if (nullptr != biasPtr) {
        biasValue = biasPtr[0];
    }
    if (param->ATranspose) {
        auto eC4 = e / 8;
        auto eR = e % 8;
        for (int y=tId; y<eC4; y+=numberThread) {
            Vec sumValue = Vec(biasValue);
            auto srcY = A + y * 8;
            for (int x=0; x<l; ++x) {
                sumValue = sumValue + Vec::load(srcY + x * e) * Vec(B[x]);
            }
            Vec::save(C + 8 * y, sumValue);
        }
        if (0 == tId && eR > 0) {
            Vec sumValue = Vec(biasValue);
            auto srcY = A + eC4 * 8;
            FLOAT16 AR[8];
            for (int x=0; x<l; ++x) {
                ::memcpy(AR, srcY + x * e, eR * sizeof(int16_t));
                sumValue = sumValue + Vec::load(AR) * Vec(B[x]);
            }
            FLOAT16 CR[8];
            Vec::save(CR, sumValue);
            ::memcpy(C + 8 * eC4, CR, eR * sizeof(int16_t));
        }
        return;
    }
    auto lC4 = l / 8;
    auto lR = l % 8;
    for (int y=tId; y<e; y+=numberThread) {
        Vec sumValue = Vec(biasValue);
        auto srcY = A + y * l;
        for (int x=0; x<lC4; ++x) {
            sumValue = sumValue + Vec::load(srcY + 8 * x) * Vec::load(B + 8 * x);
        }
        if (lR > 0) {
            FLOAT16 AR[8] = {0, 0, 0, 0, 0, 0, 0, 0};
            FLOAT16 BR[8] = {0, 0, 0, 0, 0, 0, 0, 0};
            ::memcpy(AR, srcY + lC4 * 8, lR * sizeof(int16_t));
            ::memcpy(BR, B + 8 * lC4, lR * sizeof(int16_t));
            sumValue = sumValue + Vec::load(AR) * Vec::load(BR);
        }
        float sumSingle = sumValue[0] + sumValue[1] + sumValue[2] + sumValue[3] + sumValue[4] + sumValue[5] + sumValue[6] + sumValue[7];
        C[y] = sumSingle;
    }
}

static void _MNNComputeMatMulForE_1_FP16(const float* AF, const float* BF, float* CF, const float* biasPtrF, const MatMulParam* param, size_t tId) {
    auto l = param->l;
    auto h = param->h;
    auto numberThread = param->numberThread;
    auto lC4 = l / 8;
    auto lR = l % 8;
    auto A = (const FLOAT16*)AF;
    auto B = (const FLOAT16*)BF;
    auto C = (FLOAT16*)CF;
    auto biasPtr = (const FLOAT16*)biasPtrF;
    if (param->BTranspose) {
        for (int y=tId; y<h; y+=numberThread) {
            Vec sumValue = Vec(0.0f);
            auto by = B + y * l;
            for (int x=0; x<lC4; ++x) {
                sumValue = Vec::fma(sumValue, Vec::load(A + x * 8), Vec::load(by + x * 8));
            }
            if (lR > 0) {
                FLOAT16 AR[8] = {0, 0, 0, 0, 0, 0, 0, 0};
                FLOAT16 BR[8] = {0, 0, 0, 0, 0, 0, 0, 0};
                ::memcpy(AR, A + lC4 * 8, lR * sizeof(int16_t));
                ::memcpy(BR, by + 8 * lC4, lR * sizeof(int16_t));
                sumValue = sumValue + Vec::load(AR) * Vec::load(BR);
            }
            float sumRemain = sumValue[0] + sumValue[1] + sumValue[2] + sumValue[3] + sumValue[4] + sumValue[5] + sumValue[6] + sumValue[7];
            if (nullptr != biasPtr) {
                sumRemain += biasPtr[y];
            }
            C[y] = sumRemain;
        }
    } else {
        auto hC4 = h / 8;
        auto hR = h % 8;
        auto hC16 = hC4 / 4;
        auto hC4R = hC4 % 4;
        for (int y=tId; y<hC16; y+=numberThread) {
            auto biasP = biasPtr + 8 * 4 * y;
            auto bs = B + 8 * 4 * y;
            Vec s0 = Vec(0.0f);
            Vec s1 = Vec(0.0f);
            Vec s2 = Vec(0.0f);
            Vec s3 = Vec(0.0f);
            if (biasPtr != nullptr) {
                s0 = Vec::load(biasP + 8 * 0);
                s1 = Vec::load(biasP + 8 * 1);
                s2 = Vec::load(biasP + 8 * 2);
                s3 = Vec::load(biasP + 8 * 3);
            }
            auto srcY = A + y * l * 8 * 4;
            for (int x=0; x<l; ++x) {
                auto a = Vec(A[x]);
                s0 = Vec::fma(s0, a, Vec::load(bs + h * x + 0 * 8));
                s1 = Vec::fma(s1, a, Vec::load(bs + h * x + 1 * 8));
                s2 = Vec::fma(s2, a, Vec::load(bs + h * x + 2 * 8));
                s3 = Vec::fma(s3, a, Vec::load(bs + h * x + 3 * 8));
            }
            Vec::save(C + 4 * 8 * y + 8 * 0, s0);
            Vec::save(C + 4 * 8 * y + 8 * 1, s1);
            Vec::save(C + 4 * 8 * y + 8 * 2, s2);
            Vec::save(C + 4 * 8 * y + 8 * 3, s3);
        }

        for (int y=hC16*4+tId; y<hC4; y+=numberThread) {
            auto bs = B + 8 * y;
            Vec sumValue = Vec(0.0f);
            if (biasPtr != nullptr) {
                sumValue = Vec::load(biasPtr + 8 * y);
            }
            auto srcY = A + y * l * 8;
            for (int x=0; x<l; ++x) {
                sumValue = Vec::fma(sumValue, Vec(A[x]), Vec::load(bs + h * x));
            }
            Vec::save(C + 8 * y, sumValue);
        }
        if (tId == 0 && hR > 0) {
            auto bs = B + 8 * hC4;
            Vec sumValue = Vec(0.0f);
            if (biasPtr != nullptr) {
                FLOAT16 biasTemp[8];
                ::memcpy(biasTemp, biasPtr + 8 * hC4, hR * sizeof(int16_t));
                sumValue = Vec::load(biasTemp);
            }
            auto srcY = A + 8 * hC4 * l;
            FLOAT16 bTemp[8];
            for (int x=0; x<l; ++x) {
                ::memcpy(bTemp, bs + h * x, hR * sizeof(int16_t));
                sumValue = sumValue + Vec(A[x]) * Vec::load(bTemp);
            }
            FLOAT16 cTemp[8];
            Vec::save(cTemp, sumValue);
            ::memcpy(C + 8 * hC4, cTemp, hR * sizeof(int16_t));
        }
    }
}

template<int EP, int LP>
static void _Arm82MNNPackC4ForMatMul_A(int8_t* destOrigin, int8_t const** sourceGroup, const int32_t* info, const int32_t* el) {
    const int pack = 8;
    int number = info[0];
    int eReal = info[1];
    int xStride = info[3];
    int xS4 = xStride * pack / sizeof(int32_t);
    int PUNIT = pack / LP;
    int FLOATPACK = pack / sizeof(int32_t);
    int eOutsideStride = info[2] / sizeof(int32_t);
    int eDest = EP;
    int realDstCount = info[4];
    for (int n=0; n<number; ++n) {
        int e = el[4 * n + 0];
        int l = el[4 * n + 1];
        int eOffset = el[4 * n + 2];
        int lOffset = el[4 * n + 3];
        int eC = eOffset / EP;
        int eR = eOffset % EP;
        int eS = eDest - eR;
        bool lastBag = false;
        int eOutsideStride4LastBag = eOutsideStride;
        if (realDstCount % EP > 0) {
            int jobsE = realDstCount - eOffset - e;
            if (jobsE == 0 || (jobsE < (realDstCount % EP))) {
                lastBag = true;
            }
        }
        auto source = (int32_t*)sourceGroup[n];
        auto dest = (int32_t*)(destOrigin + eC * info[2] + eR * LP + lOffset * EP);
        //printf("e=%d, l=%d, eOffset=%d, lOffset=%d, eDest=%d\n", e, l, eOffset, lOffset, eDest);
        l = l / 4; // Use float instead of int8 * 4
        if (lastBag && e + eR < EP) {
            int elast = ALIMAX(eR + e, realDstCount % EP);
            dest = (int32_t*)(destOrigin + lOffset * elast + eC * info[2] + eR * LP);
        }
        int offsetLC = lOffset / 4;
        for (int x = 0; x < l; ++x) {
            int eRemain = e;
            auto xR                  = x % PUNIT;
            auto xC                  = x / PUNIT;
            auto d = dest;
            auto s = source + xC * eReal * FLOATPACK + xR;
            if (eR > 0) {
                int eStep = ALIMIN(eRemain, eS);
                for (int yi=0; yi<eStep; ++yi) {
                    d[yi] = s[yi * xS4];
                }
                eRemain-=eStep;
                if (!lastBag ||eRemain >= EP) {
                    d += (eOutsideStride - eR);
                } else {
                    int eFill = ALIMAX(eRemain, realDstCount % EP); // maybe padding>0
                    eOutsideStride4LastBag = eOutsideStride - (EP * 4 * offsetLC / sizeof(float));
                    d += (eOutsideStride4LastBag - eR + offsetLC * eFill);
                }
                s += eS * xS4;
            }
            while (eRemain > 0) {
                int eStep = ALIMIN(eDest, eRemain);
                for (int yi=0; yi<eStep; ++yi) {
                    d[yi] = s[yi * xS4];
                }
                eRemain-=eStep;
                if (!lastBag || eRemain >= EP) {
                    d+= eOutsideStride;
                } else {
                    int eFill = ALIMAX(eRemain, realDstCount % EP); // maybe padding>0
                    eOutsideStride4LastBag = eOutsideStride - (EP * 4 * offsetLC / sizeof(float));
                    d+= (eOutsideStride4LastBag + offsetLC * eFill);
                }
                s+= eStep * xS4;
            }
            if (lastBag && e + eR < EP) {
                int efill = ALIMAX(e + eR, realDstCount % EP);
                dest += efill;
            } else {
                dest += eDest;
            }
            offsetLC++;
        }
    }
}

template<int EP, int HP>
static void _ArmBasicMNNPackC4ForMatMul_A_L8(int8_t* destOrigin, int8_t const** sourceGroup, const int32_t* info, const int32_t* el) {
    int number = info[0];
    int eReal = info[1];
    int eDest = EP;
    int offset = info[3];
    const int LP = 8;
    int eOutsideStride = info[2] / sizeof(int64_t);
    int realDstCount = info[4];
    for (int n=0; n<number; ++n) {
        int e = el[4 * n + 0];
        int l = el[4 * n + 1];
        int eOffset = el[4 * n + 2];
        int lOffset = el[4 * n + 3];
        int eC = eOffset / EP;
        int eR = eOffset % EP;
        int eS = eDest - eR;
        bool lastBag = false;
        int eOutsideStride4LastBag = eOutsideStride;
        int eres = realDstCount - eOffset;
        if (realDstCount % EP > 0) {
            int jobsE = realDstCount - eOffset - e;
            if (jobsE == 0 || (jobsE < (realDstCount % EP))) {
                lastBag = true;
            }
        }
        auto dest = (int64_t*)(destOrigin + lOffset * eDest + eC * info[2] + eR * LP);
        auto source = (int64_t*)sourceGroup[n];
        int lRemain = l / LP;
        if (lastBag && e + eR < EP) {
            int elast = ALIMIN(ALIMAX(eR + e, realDstCount % EP), EP);
            dest = (int64_t*)(destOrigin + lOffset * elast + eC * info[2] + eR * LP);
        }
        int offsetLC = lOffset / LP;
        for (int x = 0; x < lRemain; ++x) {
            int eRemain = e;
            auto d = dest;
            auto s = source;
            if (1 == offset) {
                if (eR > 0) {
                    int eStep = ALIMIN(eRemain, eS);
                    ::memcpy(d, s, eStep * sizeof(int64_t));
                    eRemain-=eStep;
                    if (!lastBag ||eRemain >= EP) {
                        d += (eOutsideStride - eR);
                    } else {
                        int eFill = ALIMAX(eRemain, realDstCount % EP); // maybe padding>0
                        eOutsideStride4LastBag = eOutsideStride - (EP * offsetLC);
                        d += (eOutsideStride4LastBag - eR + offsetLC * eFill);
                    }
                    s += (eS * offset);
                }
                while (eRemain > 0) {
                    int eStep = ALIMIN(eDest, eRemain);
                    ::memcpy(d, s, eStep * sizeof(int64_t));
                    eRemain-=eStep;
                    if (!lastBag || eRemain >= EP) {
                        d+= eOutsideStride;
                    } else {
                        int eFill = ALIMAX(eRemain, realDstCount % EP); // maybe padding>0
                        eOutsideStride4LastBag = eOutsideStride - (EP * offsetLC);
                        d+= (eOutsideStride4LastBag + offsetLC * eFill);
                    }
                    s+= (eStep * offset);
                }
            } else {
                if (eR > 0) {
                    int eStep = ALIMIN(eRemain, eS);
                    for (int yi=0; yi<eStep; ++yi) {
                        d[yi] = s[yi * offset];
                    }
                    eRemain-=eStep;
                    if (!lastBag ||eRemain >= EP) {
                        d += (eOutsideStride - eR);
                    } else {
                        int eFill = ALIMAX(eRemain, realDstCount % EP); // maybe padding>0
                        eOutsideStride4LastBag = eOutsideStride - (EP * offsetLC);
                        d += (eOutsideStride4LastBag - eR + offsetLC * eFill);
                    }
                    s += eS * offset;
                }
                while (eRemain > 0) {
                    int eStep = ALIMIN(eDest, eRemain);
                    for (int yi=0; yi<eStep; ++yi) {
                        d[yi] = s[yi * offset];
                    }
                    eRemain-=eStep;
                    if (!lastBag || eRemain >= EP) {
                        d+= eOutsideStride;
                    } else {
                        int eFill = ALIMAX(eRemain, realDstCount % EP); // maybe padding>0
                        eOutsideStride4LastBag = eOutsideStride - (EP * offsetLC);
                        d+= (eOutsideStride4LastBag + offsetLC * eFill);
                    }
                    s+= eStep * offset;
                }
            }
            source += eReal;
            if (lastBag && e + eR < EP ) { // eR=0;eR>0
                int efill = ALIMAX(e + eR, realDstCount % EP);
                dest += efill;
            } else {
                dest += eDest;
            }
            offsetLC++;
        }
    }
}

inline void transpose_4x4_f32(float32x4_t& r0, float32x4_t& r1, float32x4_t& r2, float32x4_t& r3) {
    // Stage 1: Transpose 2x2 blocks of float32 elements between pairs of adjacent rows.
    float32x4x2_t temp0 = vtrnq_f32(r0, r1);
    float32x4x2_t temp1 = vtrnq_f32(r2, r3);

    // Intermediate state:
    // temp0.val[0] = [A0, B0, A2, B2]
    // temp0.val[1] = [A1, B1, A3, B3]
    // temp1.val[0] = [C0, D0, C2, D2]
    // temp1.val[1] = [C1, D1, C3, D3]

    // Stage 2: Manually swap the 64-bit blocks to finalize the transpose.
    // This correctly simulates the non-existent 64-bit transpose/zip.
    float64x2_t i0_f64 = vreinterpretq_f64_f32(temp0.val[0]);
    float64x2_t i1_f64 = vreinterpretq_f64_f32(temp0.val[1]);
    float64x2_t i2_f64 = vreinterpretq_f64_f32(temp1.val[0]);
    float64x2_t i3_f64 = vreinterpretq_f64_f32(temp1.val[1]);

    // Combine the low 64 bits of i0 and i2 to form the first part of the result.
    float32x4_t t0 = vreinterpretq_f32_f64(vcombine_f64(vget_low_f64(i0_f64), vget_low_f64(i2_f64)));
    // Combine the low 64 bits of i1 and i3 for the second part.
    float32x4_t t1 = vreinterpretq_f32_f64(vcombine_f64(vget_low_f64(i1_f64), vget_low_f64(i3_f64)));
    // Combine the high 64 bits of i0 and i2 for the third part.
    float32x4_t t2 = vreinterpretq_f32_f64(vcombine_f64(vget_high_f64(i0_f64), vget_high_f64(i2_f64)));
    // Combine the high 64 bits of i1 and i3 for the final part.
    float32x4_t t3 = vreinterpretq_f32_f64(vcombine_f64(vget_high_f64(i1_f64), vget_high_f64(i3_f64)));

    r0 = t0;
    r1 = t1;
    r2 = t2;
    r3 = t3;
}

static void Sme2MNNPackC4ForMatMul_A_FP16(float* destOrigin, float const** sourceGroup, const int32_t* info, const int32_t* el) {
    const int lP = FP16_SME2_MATMUL_LP;
    const int pack = 8;
    int number = info[0];
    int eReal = info[1];
    int eDest = info[2];
    int offset = info[3];

    float32x4_t v0, v1, v2, v3, v4, v5, v6, v7;

    for (int n = 0; n < number; ++n) {
        int e = el[4 * n + 0];
        int l = el[4 * n + 1];
        int eOffset = el[4 * n + 2];
        int lOffset = el[4 * n + 3];

        auto destBase = (FLOAT16*)destOrigin + lOffset * eDest + eOffset * lP;
        auto sourceBase = (const FLOAT16*)(sourceGroup[n]);

        const int eTile = 8;
        const int lTile = 8;

        const int eMain = e / eTile;
        const int lMain = l / lTile;

        const size_t srcRowStride = (size_t)pack * offset;
        const size_t srcColBlockStride = (size_t)eReal * pack;
        const size_t dstColBlockStride = (size_t)eDest * lP;

        for (int y0 = 0; y0 < eMain; ++y0) {
            const int yBase = y0 * eTile;
            for (int x0 = 0; x0 < lMain; ++x0) {
                const int xBase = x0 * lTile;

                const auto srcBlockBase = sourceBase + yBase * srcRowStride + x0 * srcColBlockStride;

                v0 = vld1q_f32((const float*)(srcBlockBase + 0 * srcRowStride));
                v1 = vld1q_f32((const float*)(srcBlockBase + 1 * srcRowStride));
                v2 = vld1q_f32((const float*)(srcBlockBase + 2 * srcRowStride));
                v3 = vld1q_f32((const float*)(srcBlockBase + 3 * srcRowStride));
                v4 = vld1q_f32((const float*)(srcBlockBase + 4 * srcRowStride));
                v5 = vld1q_f32((const float*)(srcBlockBase + 5 * srcRowStride));
                v6 = vld1q_f32((const float *)(srcBlockBase + 6 * srcRowStride));
                v7 = vld1q_f32((const float *)(srcBlockBase + 7 * srcRowStride));

                transpose_4x4_f32(v0, v1, v2, v3);
                transpose_4x4_f32(v4, v5, v6, v7);

                float* addr0 = (float*)(destBase + yBase * lP + (xBase / lP) * dstColBlockStride);
                float* addr1= (float*)(destBase + yBase * lP + (xBase / lP + 1) * dstColBlockStride);
                float* addr2= (float*)(destBase + yBase * lP + (xBase / lP + 2) * dstColBlockStride);
                float* addr3= (float*)(destBase + yBase * lP + (xBase / lP + 3) * dstColBlockStride);

                vst1q_f32(addr0, v0);
                vst1q_f32(addr0 + 4, v4);
                vst1q_f32(addr1, v1);
                vst1q_f32(addr1 + 4, v5);
                vst1q_f32(addr2, v2);
                vst1q_f32(addr2 + 4, v6);
                vst1q_f32(addr3, v3);
                vst1q_f32(addr3 + 4, v7);
            }
        }

        const int eHandled = eMain * eTile;
        const int lHandled = lMain * lTile;

        // Process remaining rows
        for (int y = eHandled; y < e; ++y) {
            int yR = y % eDest;
            for (int x = 0; x < l; ++x) {
                int xR = x % pack;
                int xC = x / pack;
                destBase[(x / lP) * dstColBlockStride + yR * lP + (x % lP)] = sourceBase[xC * srcColBlockStride + y * srcRowStride + xR];
            }
        }

        // Process remaining columns for the already handled rows
        for (int y = 0; y < eHandled; ++y) {
            int yR = y % eDest;
            for (int x = lHandled; x < l; ++x) {
                int xR = x % pack;
                int xC = x / pack;
                destBase[(x / lP) * dstColBlockStride + yR * lP + (x % lP)] = sourceBase[xC * srcColBlockStride + y * srcRowStride + xR];
            }
        }
    }
}

#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
static void MNNAttenPackAndScaleSingleHead(float* dst, const float* srcHeadBase, size_t srcRowStride, const float* scale, const int32_t* units, size_t seqLen, size_t headDim) {
    const int32_t eP = units[0];
    const int32_t lP = units[1];

    if (lP != 1 && lP != 2) {
        MNN_ERROR("This function only supports lP=1 or 2\n");
        return;
    }

    const float scaleVal = scale[0];
    const float16x8_t vScale = vdupq_n_f16(scaleVal);

    const size_t packedHeadDim = UP_DIV(headDim, lP);
    const size_t dstStrideDOuter = (size_t)eP * lP;
    const size_t dstStrideSOuter = packedHeadDim * dstStrideDOuter;

    for (int s = 0; s < seqLen; ++s) {
        const int sOuter = s / eP;
        const int sInner = s % eP;
        const FLOAT16* srcRowPtr = (FLOAT16*)srcHeadBase + s * srcRowStride;
        FLOAT16* dstBasePtr = (FLOAT16*)dst + sOuter * dstStrideSOuter + sInner * lP;

        if (lP == 1) {
            size_t d = 0;
            for (; d + 7 < headDim; d += 8) {
                float16x8_t sVec = vld1q_f16(srcRowPtr + d);
                sVec = vmulq_f16(sVec, vScale);

                dstBasePtr[(d + 0) * dstStrideDOuter] = sVec[0];
                dstBasePtr[(d + 1) * dstStrideDOuter] = sVec[1];
                dstBasePtr[(d + 2) * dstStrideDOuter] = sVec[2];
                dstBasePtr[(d + 3) * dstStrideDOuter] = sVec[3];
                dstBasePtr[(d + 4) * dstStrideDOuter] = sVec[4];
                dstBasePtr[(d + 5) * dstStrideDOuter] = sVec[5];
                dstBasePtr[(d + 6) * dstStrideDOuter] = sVec[6];
                dstBasePtr[(d + 7) * dstStrideDOuter] = sVec[7];
            }
            for (; d < headDim; ++d) {
                dstBasePtr[d * dstStrideDOuter] = srcRowPtr[d] * scaleVal;
            }
        } else { // lP == 2
            const FLOAT16* srcDPtr = srcRowPtr;
            FLOAT16* dstDPtr = dstBasePtr;
            size_t dRealSize = headDim;

            while (dRealSize >= 16) {
                float16x8_t s0 = vld1q_f16(srcDPtr);
                float16x8_t s1 = vld1q_f16(srcDPtr + 8);
                s0 = vmulq_f16(s0, vScale);
                s1 = vmulq_f16(s1, vScale);

                float16x4_t lowS0_f16 = vget_low_f16(s0);   // {s0, s1, s2, s3}
                float16x4_t highS0_f16 = vget_high_f16(s0); // {s4, s5, s6, s7}
                uint32x2_t lowS0_u32 = vreinterpret_u32_f16(lowS0_f16);
                uint32x2_t highS0_u32 = vreinterpret_u32_f16(highS0_f16);

                *((uint32_t*)(dstDPtr + 0 * dstStrideDOuter)) = vget_lane_u32(lowS0_u32, 0); // Store pair {s0, s1}
                *((uint32_t*)(dstDPtr + 1 * dstStrideDOuter)) = vget_lane_u32(lowS0_u32, 1); // Store pair {s2, s3}
                *((uint32_t*)(dstDPtr + 2 * dstStrideDOuter)) = vget_lane_u32(highS0_u32, 0); // Store pair {s4, s5}
                *((uint32_t*)(dstDPtr + 3 * dstStrideDOuter)) = vget_lane_u32(highS0_u32, 1); // Store pair {s6, s7}

                float16x4_t lowS1_f16 = vget_low_f16(s1);   // {s8, s9, s10, s11}
                float16x4_t highS1_f16 = vget_high_f16(s1); // {s12, s13, s14, s15}
                uint32x2_t lowS1_u32 = vreinterpret_u32_f16(lowS1_f16);
                uint32x2_t highS1_u32 = vreinterpret_u32_f16(highS1_f16);

                *((uint32_t*)(dstDPtr + 4 * dstStrideDOuter)) = vget_lane_u32(lowS1_u32, 0);
                *((uint32_t*)(dstDPtr + 5 * dstStrideDOuter)) = vget_lane_u32(lowS1_u32, 1);
                *((uint32_t*)(dstDPtr + 6 * dstStrideDOuter)) = vget_lane_u32(highS1_u32, 0);
                *((uint32_t*)(dstDPtr + 7 * dstStrideDOuter)) = vget_lane_u32(highS1_u32, 1);

                dRealSize -= 16;
                srcDPtr += 16;
                dstDPtr += 8 * dstStrideDOuter;
            }
            // Remainder loop with padding
            while (dRealSize > 0) {
                if (dRealSize >= 2) {
                    dstDPtr[0] = srcDPtr[0] * scaleVal;
                    dstDPtr[1] = srcDPtr[1] * scaleVal;

                    dRealSize -= 2;
                    srcDPtr += 2;
                    dstDPtr += dstStrideDOuter;
                } else { // dRealSize == 1
                    dstDPtr[0] = srcDPtr[0] * scaleVal;
                    dstDPtr[1] = (FLOAT16)0.0f; // Pad with zero
                    dRealSize = 0;
                }
            }
        }
    }
}

static void MNNFlashAttentionUpdateBlockOutput(float* dst, float* src, float* scale, float* normalizeScale, int depthQuad, int plane, int pack, int idx, int kvBlocks, int size, int bytes, int seqStart) {
    auto dstPtr = (float16_t*)dst;
    auto srcPtr = (float16_t*)src;
    const auto stride0 = plane * pack;

    if (idx == 0) {
        memcpy(dst, src, size * bytes);
    } else {
        for (int j = 0; j < depthQuad; ++j) {
            const auto baseOffset = j * stride0;
            int i = seqStart;

            for (; i + 4 < plane; i += 4) {

                auto pdst0 = dstPtr + baseOffset + (i + 0) * pack;
                auto psrc0 = srcPtr + baseOffset + (i + 0) * pack;
                auto pdst1 = dstPtr + baseOffset + (i + 1) * pack;
                auto psrc1 = srcPtr + baseOffset + (i + 1) * pack;
                auto pdst2 = dstPtr + baseOffset + (i + 2) * pack;
                auto psrc2 = srcPtr + baseOffset + (i + 2) * pack;
                auto pdst3 = dstPtr + baseOffset + (i + 3) * pack;
                auto psrc3 = srcPtr + baseOffset + (i + 3) * pack;

                float16x8_t src0 = vld1q_f16(psrc0);
                float16x8_t dst0 = vld1q_f16(pdst0);
                float16x8_t src1 = vld1q_f16(psrc1);
                float16x8_t dst1 = vld1q_f16(pdst1);
                float16x8_t src2 = vld1q_f16(psrc2);
                float16x8_t dst2 = vld1q_f16(pdst2);
                float16x8_t src3 = vld1q_f16(psrc3);
                float16x8_t dst3 = vld1q_f16(pdst3);

                float32x4_t svec0 = vdupq_n_f32(scale[i + 0]);
                float32x4_t svec1 = vdupq_n_f32(scale[i + 1]);
                float32x4_t svec2 = vdupq_n_f32(scale[i + 2]);
                float32x4_t svec3 = vdupq_n_f32(scale[i + 3]);


                float32x4_t res00 = vfmaq_f32(vcvt_f32_f16(vget_low_f16(src0)),  vcvt_f32_f16(vget_low_f16(dst0)),  svec0);
                float32x4_t res10 = vfmaq_f32(vcvt_f32_f16(vget_high_f16(src0)), vcvt_f32_f16(vget_high_f16(dst0)), svec0);

                float32x4_t res01 = vfmaq_f32(vcvt_f32_f16(vget_low_f16(src1)),  vcvt_f32_f16(vget_low_f16(dst1)),  svec1);
                float32x4_t res11 = vfmaq_f32(vcvt_f32_f16(vget_high_f16(src1)), vcvt_f32_f16(vget_high_f16(dst1)), svec1);

                float32x4_t res02 = vfmaq_f32(vcvt_f32_f16(vget_low_f16(src2)),  vcvt_f32_f16(vget_low_f16(dst2)),  svec2);
                float32x4_t res12 = vfmaq_f32(vcvt_f32_f16(vget_high_f16(src2)), vcvt_f32_f16(vget_high_f16(dst2)), svec2);

                float32x4_t res03 = vfmaq_f32(vcvt_f32_f16(vget_low_f16(src3)),  vcvt_f32_f16(vget_low_f16(dst3)),  svec3);
                float32x4_t res13 = vfmaq_f32(vcvt_f32_f16(vget_high_f16(src3)), vcvt_f32_f16(vget_high_f16(dst3)), svec3);

                vst1q_f16(pdst0, vcombine_f16(vcvt_f16_f32(res00), vcvt_f16_f32(res10)));
                vst1q_f16(pdst1, vcombine_f16(vcvt_f16_f32(res01), vcvt_f16_f32(res11)));
                vst1q_f16(pdst2, vcombine_f16(vcvt_f16_f32(res02), vcvt_f16_f32(res12)));
                vst1q_f16(pdst3, vcombine_f16(vcvt_f16_f32(res03), vcvt_f16_f32(res13)));
            }

            for (; i < plane; ++i) {
                auto pdst = dstPtr + baseOffset + i * pack;
                auto psrc = srcPtr + baseOffset + i * pack;

                float16x8_t srcF16 = vld1q_f16(psrc);
                float16x8_t dstF16 = vld1q_f16(pdst);
                float32x4_t svec = vdupq_n_f32(scale[i]);

                float32x4_t s0 = vcvt_f32_f16(vget_low_f16(srcF16));
                float32x4_t s1 = vcvt_f32_f16(vget_high_f16(srcF16));
                float32x4_t d0 = vcvt_f32_f16(vget_low_f16(dstF16));
                float32x4_t d1 = vcvt_f32_f16(vget_high_f16(dstF16));

                float32x4_t res0 = vfmaq_f32(s0, d0, svec);
                float32x4_t res1 = vfmaq_f32(s1, d1, svec);

                vst1q_f16(pdst, vcombine_f16(vcvt_f16_f32(res0), vcvt_f16_f32(res1)));
            }
        }
    }

    if (idx == kvBlocks - 1) {
        for (int j = 0; j < depthQuad; ++j) {
            const auto baseOffset = j * stride0;
            int i = 0;
            const int plane4 = plane - (plane % 4);
            for (; i < plane4; i += 4) {
                auto pdst0 = dstPtr + baseOffset + (i + 0) * pack;
                auto pdst1 = dstPtr + baseOffset + (i + 1) * pack;
                auto pdst2 = dstPtr + baseOffset + (i + 2) * pack;
                auto pdst3 = dstPtr + baseOffset + (i + 3) * pack;

                float16x8_t dst0 = vld1q_f16(pdst0);
                float16x8_t dst1 = vld1q_f16(pdst1);
                float16x8_t dst2 = vld1q_f16(pdst2);
                float16x8_t dst3 = vld1q_f16(pdst3);

                float32x4_t ns0 = vdupq_n_f32(1.0f / normalizeScale[i + 0]);
                float32x4_t ns1 = vdupq_n_f32(1.0f / normalizeScale[i + 1]);
                float32x4_t ns2 = vdupq_n_f32(1.0f / normalizeScale[i + 2]);
                float32x4_t ns3 = vdupq_n_f32(1.0f / normalizeScale[i + 3]);

                float32x4_t d00 = vmulq_f32(vcvt_f32_f16(vget_low_f16(dst0)),  ns0);
                float32x4_t d10 = vmulq_f32(vcvt_f32_f16(vget_high_f16(dst0)), ns0);
                float32x4_t d01 = vmulq_f32(vcvt_f32_f16(vget_low_f16(dst1)),  ns1);
                float32x4_t d11 = vmulq_f32(vcvt_f32_f16(vget_high_f16(dst1)), ns1);
                float32x4_t d02 = vmulq_f32(vcvt_f32_f16(vget_low_f16(dst2)),  ns2);
                float32x4_t d12 = vmulq_f32(vcvt_f32_f16(vget_high_f16(dst2)), ns2);
                float32x4_t d03 = vmulq_f32(vcvt_f32_f16(vget_low_f16(dst3)),  ns3);
                float32x4_t d13 = vmulq_f32(vcvt_f32_f16(vget_high_f16(dst3)), ns3);

                vst1q_f16(pdst0, vcombine_f16(vcvt_f16_f32(d00), vcvt_f16_f32(d10)));
                vst1q_f16(pdst1, vcombine_f16(vcvt_f16_f32(d01), vcvt_f16_f32(d11)));
                vst1q_f16(pdst2, vcombine_f16(vcvt_f16_f32(d02), vcvt_f16_f32(d12)));
                vst1q_f16(pdst3, vcombine_f16(vcvt_f16_f32(d03), vcvt_f16_f32(d13)));
            }

            for (; i < plane; ++i) {
                auto pdst = dstPtr + baseOffset + i * pack;
                float32x4_t nsvec = vdupq_n_f32(1.0f / normalizeScale[i]);

                float16x8_t dstF16 = vld1q_f16(pdst);
                float32x4_t d0 = vcvt_f32_f16(vget_low_f16(dstF16));
                float32x4_t d1 = vcvt_f32_f16(vget_high_f16(dstF16));

                d0 = vmulq_f32(d0, nsvec);
                d1 = vmulq_f32(d1, nsvec);

                vst1q_f16(pdst, vcombine_f16(vcvt_f16_f32(d0), vcvt_f16_f32(d1)));
            }
        }
    }
}

static void MNNAttenUnpackAndConvertFp16(float* dst, float* src, size_t depth, size_t planesize, int pack) {
    // src: (UP_DIV(depth, pack), planesize, pack), float16
    // dst: (planesize, depth), float32
    // pack=8

    if (planesize == 1) {
        MNNDequantizeFP16((int16_t*)src, dst, depth);
        return; // no need to convert
    }
    const auto depthDiv8 = UP_DIV(depth, pack);
    const auto srcStep = pack * planesize;
    const auto dstStep = depth;

    auto remainDepth = depth % pack;
    auto depthQuad = depthDiv8;
    if (remainDepth > 0) {
        depthQuad -= 1; // last quad is not full
    }

    for (int i = 0; i < depthQuad; ++i) {
        auto realsize = planesize;
        auto srcPtr = (FLOAT16*)src + i * srcStep;
        auto dstPtr = (float*)dst + i * pack;
        while (realsize >= 8) {
            float16x8_t s0_f16 = vld1q_f16(srcPtr + 0 * pack);
            float16x8_t s1_f16 = vld1q_f16(srcPtr + 1 * pack);
            float16x8_t s2_f16 = vld1q_f16(srcPtr + 2 * pack);
            float16x8_t s3_f16 = vld1q_f16(srcPtr + 3 * pack);
            float16x8_t s4_f16 = vld1q_f16(srcPtr + 4 * pack);
            float16x8_t s5_f16 = vld1q_f16(srcPtr + 5 * pack);
            float16x8_t s6_f16 = vld1q_f16(srcPtr + 6 * pack);
            float16x8_t s7_f16 = vld1q_f16(srcPtr + 7 * pack);

            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            float32x4_t d01_f32 = vcvt_f32_f16(vget_high_f16(s0_f16));
            float32x4_t d10_f32 = vcvt_f32_f16(vget_low_f16(s1_f16));
            float32x4_t d11_f32 = vcvt_f32_f16(vget_high_f16(s1_f16));
            float32x4_t d20_f32 = vcvt_f32_f16(vget_low_f16(s2_f16));
            float32x4_t d21_f32 = vcvt_f32_f16(vget_high_f16(s2_f16));
            float32x4_t d30_f32 = vcvt_f32_f16(vget_low_f16(s3_f16));
            float32x4_t d31_f32 = vcvt_f32_f16(vget_high_f16(s3_f16));
            float32x4_t d40_f32 = vcvt_f32_f16(vget_low_f16(s4_f16));
            float32x4_t d41_f32 = vcvt_f32_f16(vget_high_f16(s4_f16));
            float32x4_t d50_f32 = vcvt_f32_f16(vget_low_f16(s5_f16));
            float32x4_t d51_f32 = vcvt_f32_f16(vget_high_f16(s5_f16));
            float32x4_t d60_f32 = vcvt_f32_f16(vget_low_f16(s6_f16));
            float32x4_t d61_f32 = vcvt_f32_f16(vget_high_f16(s6_f16));
            float32x4_t d70_f32 = vcvt_f32_f16(vget_low_f16(s7_f16));
            float32x4_t d71_f32 = vcvt_f32_f16(vget_high_f16(s7_f16));

            vst1q_f32(dstPtr + 0 * dstStep, d00_f32); vst1q_f32(dstPtr + 0 * dstStep + 4, d01_f32);
            vst1q_f32(dstPtr + 1 * dstStep, d10_f32); vst1q_f32(dstPtr + 1 * dstStep + 4, d11_f32);
            vst1q_f32(dstPtr + 2 * dstStep, d20_f32); vst1q_f32(dstPtr + 2 * dstStep + 4, d21_f32);
            vst1q_f32(dstPtr + 3 * dstStep, d30_f32); vst1q_f32(dstPtr + 3 * dstStep + 4, d31_f32);
            vst1q_f32(dstPtr + 4 * dstStep, d40_f32); vst1q_f32(dstPtr + 4 * dstStep + 4, d41_f32);
            vst1q_f32(dstPtr + 5 * dstStep, d50_f32); vst1q_f32(dstPtr + 5 * dstStep + 4, d51_f32);
            vst1q_f32(dstPtr + 6 * dstStep, d60_f32); vst1q_f32(dstPtr + 6 * dstStep + 4, d61_f32);
            vst1q_f32(dstPtr + 7 * dstStep, d70_f32); vst1q_f32(dstPtr + 7 * dstStep + 4, d71_f32);

            srcPtr += 8 * pack;
            dstPtr += 8 * dstStep;
            realsize -= 8;
        }
        if (realsize >= 4) {
            float16x8_t s0_f16 = vld1q_f16(srcPtr + 0 * pack);
            float16x8_t s1_f16 = vld1q_f16(srcPtr + 1 * pack);
            float16x8_t s2_f16 = vld1q_f16(srcPtr + 2 * pack);
            float16x8_t s3_f16 = vld1q_f16(srcPtr + 3 * pack);

            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            float32x4_t d01_f32 = vcvt_f32_f16(vget_high_f16(s0_f16));
            float32x4_t d10_f32 = vcvt_f32_f16(vget_low_f16(s1_f16));
            float32x4_t d11_f32 = vcvt_f32_f16(vget_high_f16(s1_f16));
            float32x4_t d20_f32 = vcvt_f32_f16(vget_low_f16(s2_f16));
            float32x4_t d21_f32 = vcvt_f32_f16(vget_high_f16(s2_f16));
            float32x4_t d30_f32 = vcvt_f32_f16(vget_low_f16(s3_f16));
            float32x4_t d31_f32 = vcvt_f32_f16(vget_high_f16(s3_f16));

            vst1q_f32(dstPtr + 0 * dstStep, d00_f32); vst1q_f32(dstPtr + 0 * dstStep + 4, d01_f32);
            vst1q_f32(dstPtr + 1 * dstStep, d10_f32); vst1q_f32(dstPtr + 1 * dstStep + 4, d11_f32);
            vst1q_f32(dstPtr + 2 * dstStep, d20_f32); vst1q_f32(dstPtr + 2 * dstStep + 4, d21_f32);
            vst1q_f32(dstPtr + 3 * dstStep, d30_f32); vst1q_f32(dstPtr + 3 * dstStep + 4, d31_f32);

            srcPtr += 4 * pack;
            dstPtr += 4 * dstStep;
            realsize -= 4;
        }
        while (realsize > 0) {
            auto s0_fp16 = vld1q_f16(srcPtr);
            auto s00_fp32 = vcvt_f32_f16(vget_low_f16(s0_fp16));
            auto s01_fp32 = vcvt_f32_f16(vget_high_f16(s0_fp16));
            vst1q_f32(dstPtr, s00_fp32);
            vst1q_f32(dstPtr + 4, s01_fp32);
            srcPtr += pack;
            dstPtr += dstStep;
            realsize--;
        }
    }

    // process remain depth < 8
    if (remainDepth >= 4) {
        auto realsize = planesize;
        auto srcPtr = (FLOAT16*)src + (depthDiv8 - 1) * srcStep;
        auto dstPtr = (float*)dst + (depthDiv8 - 1) * pack;
        auto extraDepth = remainDepth - 4;

        float tmp0[4];
        float tmp1[4];
        float tmp2[4];
        float tmp3[4];
        float tmp4[4];
        float tmp5[4];
        float tmp6[4];
        float tmp7[4];

        while (realsize >= 8) {
            float16x8_t s0_f16 = vld1q_f16(srcPtr + 0 * pack);
            float16x8_t s1_f16 = vld1q_f16(srcPtr + 1 * pack);
            float16x8_t s2_f16 = vld1q_f16(srcPtr + 2 * pack);
            float16x8_t s3_f16 = vld1q_f16(srcPtr + 3 * pack);
            float16x8_t s4_f16 = vld1q_f16(srcPtr + 4 * pack);
            float16x8_t s5_f16 = vld1q_f16(srcPtr + 5 * pack);
            float16x8_t s6_f16 = vld1q_f16(srcPtr + 6 * pack);
            float16x8_t s7_f16 = vld1q_f16(srcPtr + 7 * pack);


            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            float32x4_t d01_f32 = vcvt_f32_f16(vget_high_f16(s0_f16));
            float32x4_t d10_f32 = vcvt_f32_f16(vget_low_f16(s1_f16));
            float32x4_t d11_f32 = vcvt_f32_f16(vget_high_f16(s1_f16));
            float32x4_t d20_f32 = vcvt_f32_f16(vget_low_f16(s2_f16));
            float32x4_t d21_f32 = vcvt_f32_f16(vget_high_f16(s2_f16));
            float32x4_t d30_f32 = vcvt_f32_f16(vget_low_f16(s3_f16));
            float32x4_t d31_f32 = vcvt_f32_f16(vget_high_f16(s3_f16));
            float32x4_t d40_f32 = vcvt_f32_f16(vget_low_f16(s4_f16));
            float32x4_t d41_f32 = vcvt_f32_f16(vget_high_f16(s4_f16));
            float32x4_t d50_f32 = vcvt_f32_f16(vget_low_f16(s5_f16));
            float32x4_t d51_f32 = vcvt_f32_f16(vget_high_f16(s5_f16));
            float32x4_t d60_f32 = vcvt_f32_f16(vget_low_f16(s6_f16));
            float32x4_t d61_f32 = vcvt_f32_f16(vget_high_f16(s6_f16));
            float32x4_t d70_f32 = vcvt_f32_f16(vget_low_f16(s7_f16));
            float32x4_t d71_f32 = vcvt_f32_f16(vget_high_f16(s7_f16));

            vst1q_f32(dstPtr + 0 * dstStep, d00_f32); vst1q_f32(tmp0, d01_f32);
            vst1q_f32(dstPtr + 1 * dstStep, d10_f32); vst1q_f32(tmp1, d11_f32);
            vst1q_f32(dstPtr + 2 * dstStep, d20_f32); vst1q_f32(tmp2, d21_f32);
            vst1q_f32(dstPtr + 3 * dstStep, d30_f32); vst1q_f32(tmp3, d31_f32);
            vst1q_f32(dstPtr + 4 * dstStep, d40_f32); vst1q_f32(tmp4, d41_f32);
            vst1q_f32(dstPtr + 5 * dstStep, d50_f32); vst1q_f32(tmp5, d51_f32);
            vst1q_f32(dstPtr + 6 * dstStep, d60_f32); vst1q_f32(tmp6, d61_f32);
            vst1q_f32(dstPtr + 7 * dstStep, d70_f32); vst1q_f32(tmp7, d71_f32);

            memcpy(dstPtr + 0 * dstStep + 4, tmp0, sizeof(float) * extraDepth);
            memcpy(dstPtr + 1 * dstStep + 4, tmp1, sizeof(float) * extraDepth);
            memcpy(dstPtr + 2 * dstStep + 4, tmp2, sizeof(float) * extraDepth);
            memcpy(dstPtr + 3 * dstStep + 4, tmp3, sizeof(float) * extraDepth);
            memcpy(dstPtr + 4 * dstStep + 4, tmp4, sizeof(float) * extraDepth);
            memcpy(dstPtr + 5 * dstStep + 4, tmp5, sizeof(float) * extraDepth);
            memcpy(dstPtr + 6 * dstStep + 4, tmp6, sizeof(float) * extraDepth);
            memcpy(dstPtr + 7 * dstStep + 4, tmp7, sizeof(float) * extraDepth);

            srcPtr += 8 * pack;
            dstPtr += 8 * dstStep;
            realsize -= 8;
        }
        if (realsize >= 4) {
            float16x8_t s0_f16 = vld1q_f16(srcPtr + 0 * pack);
            float16x8_t s1_f16 = vld1q_f16(srcPtr + 1 * pack);
            float16x8_t s2_f16 = vld1q_f16(srcPtr + 2 * pack);
            float16x8_t s3_f16 = vld1q_f16(srcPtr + 3 * pack);

            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            float32x4_t d01_f32 = vcvt_f32_f16(vget_high_f16(s0_f16));
            float32x4_t d10_f32 = vcvt_f32_f16(vget_low_f16(s1_f16));
            float32x4_t d11_f32 = vcvt_f32_f16(vget_high_f16(s1_f16));
            float32x4_t d20_f32 = vcvt_f32_f16(vget_low_f16(s2_f16));
            float32x4_t d21_f32 = vcvt_f32_f16(vget_high_f16(s2_f16));
            float32x4_t d30_f32 = vcvt_f32_f16(vget_low_f16(s3_f16));
            float32x4_t d31_f32 = vcvt_f32_f16(vget_high_f16(s3_f16));

            vst1q_f32(dstPtr + 0 * dstStep, d00_f32); vst1q_f32(tmp0, d01_f32);
            vst1q_f32(dstPtr + 1 * dstStep, d10_f32); vst1q_f32(tmp1, d11_f32);
            vst1q_f32(dstPtr + 2 * dstStep, d20_f32); vst1q_f32(tmp2, d21_f32);
            vst1q_f32(dstPtr + 3 * dstStep, d30_f32); vst1q_f32(tmp3, d31_f32);

            memcpy(dstPtr + 0 * dstStep + 4, tmp0, sizeof(float) * extraDepth);
            memcpy(dstPtr + 1 * dstStep + 4, tmp1, sizeof(float) * extraDepth);
            memcpy(dstPtr + 2 * dstStep + 4, tmp2, sizeof(float) * extraDepth);
            memcpy(dstPtr + 3 * dstStep + 4, tmp3, sizeof(float) * extraDepth);

            srcPtr += 4 * pack;
            dstPtr += 4 * dstStep;
            realsize -= 4;
        }
        while (realsize > 0) {
            auto s0_fp16 = vld1q_f16(srcPtr);
            auto d00_fp32 = vcvt_f32_f16(vget_low_f16(s0_fp16));
            auto d01_fp32 = vcvt_f32_f16(vget_high_f16(s0_fp16));
            vst1q_f32(dstPtr, d00_fp32);
            vst1q_f32(tmp0, d01_fp32);
            memcpy(dstPtr + 4, tmp0, sizeof(float) * extraDepth);
            srcPtr += pack;
            dstPtr += dstStep;
            realsize--;
        }
    }

    if (remainDepth > 0 && remainDepth < 4) {
        auto realsize = planesize;
        auto srcPtr = (FLOAT16*)src + (depthDiv8 - 1) * srcStep;
        auto dstPtr = (float*)dst + (depthDiv8 - 1) * pack;

        float tmp0[4];
        float tmp1[4];
        float tmp2[4];
        float tmp3[4];
        float tmp4[4];
        float tmp5[4];
        float tmp6[4];
        float tmp7[4];

        while (realsize >= 8) {
            float16x8_t s0_f16 = vld1q_f16(srcPtr + 0 * pack);
            float16x8_t s1_f16 = vld1q_f16(srcPtr + 1 * pack);
            float16x8_t s2_f16 = vld1q_f16(srcPtr + 2 * pack);
            float16x8_t s3_f16 = vld1q_f16(srcPtr + 3 * pack);
            float16x8_t s4_f16 = vld1q_f16(srcPtr + 4 * pack);
            float16x8_t s5_f16 = vld1q_f16(srcPtr + 5 * pack);
            float16x8_t s6_f16 = vld1q_f16(srcPtr + 6 * pack);
            float16x8_t s7_f16 = vld1q_f16(srcPtr + 7 * pack);

            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            float32x4_t d10_f32 = vcvt_f32_f16(vget_low_f16(s1_f16));
            float32x4_t d20_f32 = vcvt_f32_f16(vget_low_f16(s2_f16));
            float32x4_t d30_f32 = vcvt_f32_f16(vget_low_f16(s3_f16));
            float32x4_t d40_f32 = vcvt_f32_f16(vget_low_f16(s4_f16));
            float32x4_t d50_f32 = vcvt_f32_f16(vget_low_f16(s5_f16));
            float32x4_t d60_f32 = vcvt_f32_f16(vget_low_f16(s6_f16));
            float32x4_t d70_f32 = vcvt_f32_f16(vget_low_f16(s7_f16));

            vst1q_f32(tmp0, d00_f32);
            vst1q_f32(tmp1, d10_f32);
            vst1q_f32(tmp2, d20_f32);
            vst1q_f32(tmp3, d30_f32);
            vst1q_f32(tmp4, d40_f32);
            vst1q_f32(tmp5, d50_f32);
            vst1q_f32(tmp6, d60_f32);
            vst1q_f32(tmp7, d70_f32);

            memcpy(dstPtr + 0 * dstStep, tmp0, sizeof(float) * remainDepth);
            memcpy(dstPtr + 1 * dstStep, tmp1, sizeof(float) * remainDepth);
            memcpy(dstPtr + 2 * dstStep, tmp2, sizeof(float) * remainDepth);
            memcpy(dstPtr + 3 * dstStep, tmp3, sizeof(float) * remainDepth);
            memcpy(dstPtr + 4 * dstStep, tmp4, sizeof(float) * remainDepth);
            memcpy(dstPtr + 5 * dstStep, tmp5, sizeof(float) * remainDepth);
            memcpy(dstPtr + 6 * dstStep, tmp6, sizeof(float) * remainDepth);
            memcpy(dstPtr + 7 * dstStep, tmp7, sizeof(float) * remainDepth);

            srcPtr += 8 * pack;
            dstPtr += 8 * dstStep;
            realsize -= 8;
        }
        if (realsize >= 4) {
            float16x8_t s0_f16 = vld1q_f16(srcPtr + 0 * pack);
            float16x8_t s1_f16 = vld1q_f16(srcPtr + 1 * pack);
            float16x8_t s2_f16 = vld1q_f16(srcPtr + 2 * pack);
            float16x8_t s3_f16 = vld1q_f16(srcPtr + 3 * pack);

            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            float32x4_t d10_f32 = vcvt_f32_f16(vget_low_f16(s1_f16));
            float32x4_t d20_f32 = vcvt_f32_f16(vget_low_f16(s2_f16));
            float32x4_t d30_f32 = vcvt_f32_f16(vget_low_f16(s3_f16));

            vst1q_f32(tmp0, d00_f32);
            vst1q_f32(tmp1, d10_f32);
            vst1q_f32(tmp2, d20_f32);
            vst1q_f32(tmp3, d30_f32);

            memcpy(dstPtr + 0 * dstStep, tmp0, sizeof(float) * remainDepth);
            memcpy(dstPtr + 1 * dstStep, tmp1, sizeof(float) * remainDepth);
            memcpy(dstPtr + 2 * dstStep, tmp2, sizeof(float) * remainDepth);
            memcpy(dstPtr + 3 * dstStep, tmp3, sizeof(float) * remainDepth);

            srcPtr += 4 * pack;
            dstPtr += 4 * dstStep;
            realsize -= 4;
        }
        while (realsize > 0) {
            auto s0_f16 = vld1q_f16(srcPtr);
            float32x4_t d00_f32 = vcvt_f32_f16(vget_low_f16(s0_f16));
            vst1q_f32(tmp0, d00_f32);
            memcpy(dstPtr + 0 * dstStep, tmp0, sizeof(float) * remainDepth);
            srcPtr += pack;
            dstPtr += dstStep;
            realsize--;
        }
    }
}

static void MNNQuantAttentionKeyFP16(int8_t* dst, const float* source, float* sumKeyPtr, float* maxKeyPtr, int32_t* params) {
    int32_t kvNumHead = params[0];
    int32_t seqLen = params[1];
    int32_t headDim = params[2];
    int32_t blockNum = params[3];

    int32_t lP = params[5];
    int32_t hP = params[6];
    int32_t pastLength = params[7];
    int32_t kvHeadIdx = params[8];

    auto blockHeadDim = UP_DIV(headDim, blockNum);
    auto weightStride1 = ROUND_UP(blockHeadDim, lP) * hP;
    auto weightStride2 = lP * hP;
    auto packedWeightStride1 = weightStride1 + 2 * sizeof(float) * hP;

    auto sourceFp16 = (FLOAT16*)source;
    auto maxKeyFp16 = (FLOAT16*)maxKeyPtr;
    int8_t tempBuffer[8];
    float32x4_t neg128Vec = vdupq_n_f32(-128.0f);

    // Get max: [1, headDim]
    if (seqLen > 1) {
        for (int s = 0; s < seqLen; ++s) {
            const FLOAT16* keySrc = sourceFp16 + s * kvNumHead * headDim + kvHeadIdx * headDim;
            int d = 0;
            for (; d <= headDim - 16; d += 16) {
                float16x8_t maxVec0 = vld1q_f16(maxKeyFp16 + d);
                float16x8_t maxVec1 = vld1q_f16(maxKeyFp16 + d + 8);
                float16x8_t srcVec0 = vld1q_f16(keySrc + d);
                float16x8_t srcVec1 = vld1q_f16(keySrc + d + 8);
                maxVec0 = vmaxq_f16(maxVec0, srcVec0);
                maxVec1 = vmaxq_f16(maxVec1, srcVec1);
                vst1q_f16(maxKeyFp16 + d, maxVec0);
                vst1q_f16(maxKeyFp16 + d + 8, maxVec1);
            }
            for (; d <= headDim - 8; d += 8) {
                float16x8_t maxVec = vld1q_f16(maxKeyFp16 + d);
                float16x8_t srcVec = vld1q_f16(keySrc + d);
                maxVec = vmaxq_f16(maxVec, srcVec);
                vst1q_f16(maxKeyFp16 + d, maxVec);
            }
            for (; d < headDim; ++d) {
                maxKeyFp16[d] = ALIMAX(maxKeyFp16[d], keySrc[d]);
            }
        }
    }

    // Quant fp16
    for (int s = 0; s < seqLen; s++) {
        const FLOAT16* keySrc = sourceFp16 + s * kvNumHead * headDim + kvHeadIdx * headDim;

        float16x8_t minVec = vdupq_n_f16(keySrc[0]);
        float16x8_t maxVec = vdupq_n_f16(keySrc[0]);

        int d = 0;
        for (; d <= headDim - 8; d += 8) {
            float16x8_t srcVec = vld1q_f16(keySrc + d);
            float16x8_t maxKeyVec = vld1q_f16(maxKeyFp16 + d);
            float16x8_t keyDataF16 = vsubq_f16(srcVec, maxKeyVec);

            minVec = vminq_f16(minVec, keyDataF16);
            maxVec = vmaxq_f16(maxVec, keyDataF16);

            float32x4_t keyDataF32Low = vcvt_f32_f16(vget_low_f16(keyDataF16));
            float32x4_t keyDataF32High = vcvt_f32_f16(vget_high_f16(keyDataF16));
        }

        FLOAT16 minKey = vminvq_f16(minVec);
        FLOAT16 maxKey = vmaxvq_f16(maxVec);

        for (; d < headDim; ++d) {
            auto keydata = keySrc[d] - maxKeyFp16[d];
            minKey = ALIMIN(minKey, keydata);
            maxKey = ALIMAX(maxKey, keydata);
        }

        int outIndex = (pastLength + s) / hP;
        int inIndex  = (pastLength + s) % hP;

        float range = (float)maxKey - (float)minKey;
        float quantScaleVal = 0;
        float biasVal = minKey + 128.0f * range / 255.0;
        if (range <= 1e-6f) {
            quantScaleVal = 0.f;
        } else {
            quantScaleVal = 255.0f / range;
        }

        for (int k = 0; k < blockNum; ++k) {
            int8_t* weightDstBase = dst + outIndex * blockNum * packedWeightStride1 + k * packedWeightStride1;
            float* scaleDst = (float*)(weightDstBase + weightStride1);
            float* biasDst = scaleDst + hP;

            scaleDst[inIndex] = range / 255.f;
            biasDst[inIndex] = biasVal;

            float32x4_t scaleVecFp32 = vdupq_n_f32(quantScaleVal);
            float32x4_t negMinKeyVecF32 = vdupq_n_f32(-(float)minKey);

            const FLOAT16* currentKeyBlock = keySrc + k * blockHeadDim;
            const FLOAT16* currentMaxBlock = maxKeyFp16 + k * blockHeadDim;

            int32x4_t sumInt32_0 = vdupq_n_s32(0);
            int32x4_t sumInt32_1 = vdupq_n_s32(0);
            int headDimIdx = 0;
            for (; headDimIdx <= blockHeadDim - 8; headDimIdx += 8) {
                float16x8_t srcVecFp16 = vld1q_f16(currentKeyBlock + headDimIdx);
                float16x8_t maxVecFp16 = vld1q_f16(currentMaxBlock + headDimIdx);

                float16x8_t keyDataF16 = vsubq_f16(srcVecFp16, maxVecFp16);

                float32x4_t keyDataLowFp32 = vcvt_f32_f16(vget_low_f16(keyDataF16));
                float32x4_t keyDataHighFp32 = vcvt_f32_f16(vget_high_f16(keyDataF16));

                keyDataLowFp32 = vaddq_f32(keyDataLowFp32, negMinKeyVecF32);
                keyDataHighFp32 = vaddq_f32(keyDataHighFp32, negMinKeyVecF32);

                keyDataLowFp32 = vmulq_f32(keyDataLowFp32, scaleVecFp32);
                keyDataHighFp32 = vmulq_f32(keyDataHighFp32, scaleVecFp32);

                keyDataLowFp32 = vaddq_f32(keyDataLowFp32, neg128Vec);
                keyDataHighFp32 = vaddq_f32(keyDataHighFp32, neg128Vec);

                int32x4_t keyDataLowInt32 = vcvtaq_s32_f32(keyDataLowFp32);
                int32x4_t keyDataHighInt32 = vcvtaq_s32_f32(keyDataHighFp32);

                int16x4_t s16Low = vmovn_s32(keyDataLowInt32);
                int16x4_t s16High = vmovn_s32(keyDataHighInt32);

                int16x8_t s16Combined = vcombine_s16(s16Low, s16High);

                // sum
                sumInt32_0 = vaddq_s32(sumInt32_0, keyDataLowInt32);
                sumInt32_1 = vaddq_s32(sumInt32_1, keyDataHighInt32);

                int8x8_t s8Vec = vqmovn_s16(s16Combined);

                if (lP == 8) {
                    int i = headDimIdx / lP;
                    int8_t* dstPtr = weightDstBase + i * weightStride2 + inIndex * lP;
                    vst1_s8(dstPtr, s8Vec);
                } else if (lP == 4) {
                    vst1_s8(tempBuffer, s8Vec);
                    int iLow = headDimIdx / lP;
                    int iHigh = (headDimIdx + 4) / lP;

                    int8_t* dstPtrLow = weightDstBase + iLow * weightStride2 + inIndex * lP;
                    int8_t* dstPtrHigh = weightDstBase + iHigh * weightStride2 + inIndex * lP;

                    std::memcpy(dstPtrLow, tempBuffer, 4);
                    std::memcpy(dstPtrHigh, tempBuffer + 4, 4);
                } else {
                    vst1_s8(tempBuffer, s8Vec);
                    for (int nk = 0; nk < 8; ++nk) {
                        int headDimCurr = headDimIdx + nk;
                        int i = headDimCurr / lP;
                        int j = headDimCurr % lP;
                        weightDstBase[i * weightStride2 + inIndex * lP + j] = tempBuffer[nk];
                    }
                }

            }

            int32_t sumInt32 = vaddvq_s32(sumInt32_0) + vaddvq_s32(sumInt32_1);

            for (; headDimIdx < blockHeadDim; ++headDimIdx) {
                int i = headDimIdx / lP;
                int j = headDimIdx % lP;
                float keyVal = (float)currentKeyBlock[headDimIdx] - (float)currentMaxBlock[headDimIdx];
                float quantVal = (keyVal - minKey) * quantScaleVal - 128.0f;
                int32_t roundedVal = static_cast<int32_t>(roundf(quantVal));
                int8_t finalVal = static_cast<int8_t>(std::max(-128, std::min(127, roundedVal)));
                weightDstBase[i * weightStride2 + inIndex * lP + j] = finalVal;
                sumInt32 += finalVal;
            }

            // store sum
            sumKeyPtr[outIndex * hP + inIndex] = sumInt32 * range / 255.f + (minKey * (float)(blockHeadDim) + 128.0f * range * (float)(blockHeadDim) / 255.0);
        }
    }
}

static void MNNQuantAttentionValueFP16(int8_t* dst, const float* source, float* valueSum, int32_t* params) {
    // float   value src : [kvSeq,kvNumHead,headDim]
    // int8_t  value dest: [updiv(maxLength,flashAttentionBlockKv), updiv(headDim,hp),updiv(flashAttentionBlockKv,lp),hp,lp]
    // float   value sum: [updiv(maxLength,flashAttentionBlockKv), roundup(headDim,hp)]
    int32_t kvNumHead = params[0];
    int32_t seqLen = params[1];
    int32_t headDim = params[2];
    int32_t blockNum = params[3];
    int32_t maxLength = params[4];

    int32_t lP = params[5];
    int32_t hP = params[6];
    int32_t pastLength = params[7];
    int32_t kvHeadIdx = params[8];

    int32_t flashAttentionBlockKv = params[9];

    auto blockKvseq = UP_DIV(seqLen + pastLength, blockNum);
    auto weightStride2 = lP * hP;
    auto weightStride1 = UP_DIV(flashAttentionBlockKv, lP) * weightStride2;

    auto packedStride1 = (int)(weightStride1 + 2 * hP * sizeof(float));
    auto packedStride0 = UP_DIV(headDim, hP) * packedStride1;

    auto srcStride0 = kvNumHead * headDim;

    auto sourceFp16 = (FLOAT16*)source;

    // quant scale & bias
    if (pastLength == 0) {
        for (int d = 0; d < headDim; ++d) {
            float* scalePtr = (float*)(dst + (d / hP) * packedStride1 + weightStride1) + (d % hP);
            float* biasPtr = scalePtr + hP;

            // find min,max
            float dMax = sourceFp16[d + kvHeadIdx * headDim];
            float dMin = dMax;
            for (int s = 0; s < seqLen; ++s) {
                float data = sourceFp16[s * srcStride0 + d + kvHeadIdx * headDim];
                dMax = ALIMAX(dMax, data);
                dMin = ALIMIN(dMin, data);
            }

            // scale & bias
            float range = dMax - dMin;
            if (range < 1e-6) {
                scalePtr[0] = 0.f;
                biasPtr[0] = dMax;
            } else {
                float scale = range / 255.f;
                float bias  = range / 255.f * 128.f + dMin;
                scalePtr[0] = scale;
                biasPtr[0] = bias;
            }
        }
    }

    // copy the scale&bias to each blockKv
    //                                    pastLength == 0: First time prefill
    // (seqLen + pastLength) % flashAttentionBlockKv == 0: Open a new blockKv
    if (pastLength == 0 || (pastLength % flashAttentionBlockKv) == 0) {
        int32_t d0 = UP_DIV(maxLength, flashAttentionBlockKv);
        int32_t d1 = UP_DIV(headDim, hP);
        for (int k = 0; k < d0; ++k) {
            for (int r = 0; r < d1; ++r) {
                float* scalePtr = (float*)(dst + k * packedStride0 + r * packedStride1 + weightStride1);
                float* biasPtr  = scalePtr + hP;
                memcpy(scalePtr, dst + r * packedStride1 + weightStride1, hP * sizeof(float));
                memcpy(biasPtr, dst + r * packedStride1 + weightStride1 + hP * sizeof(float), hP * sizeof(float));
            }
        }
    }

    std::vector<float> qScales(headDim);
    std::vector<float> qBiases(headDim);
    std::vector<float> deqScales(headDim);
    std::vector<float> deqBiases(headDim);
    int8_t tmpQ[8];

    for (int d = 0; d < headDim; ++d) {
        float* scaleBase = (float*)(dst + (d / hP) * packedStride1 + weightStride1) + (d % hP);
        float* biasBase = scaleBase + hP;

        float s_val = scaleBase[0];
        float b_val = biasBase[0];

        deqScales[d] = s_val;
        deqBiases[d] = b_val;

        bool is_small = s_val < 1e-6f;
        qScales[d] = is_small ? 0.0f : (1.0f / s_val);
        qBiases[d] = is_small ? 0.0f : (-b_val / s_val);
    }

    const __fp16* srcBasePtr = sourceFp16 + kvHeadIdx * headDim;

    const int32_t sumStride = ROUND_UP(headDim, hP);

    for (int s = 0; s < seqLen; ++s) {
        int kvSeqIndx = s + pastLength;

        int blkIdx = kvSeqIndx / flashAttentionBlockKv;
        int blkRem = kvSeqIndx % flashAttentionBlockKv;

        int idxInnerCommon = blkIdx * packedStride0 + (blkRem / lP) * weightStride2 + (blkRem % lP);

        float* curSumRow = valueSum + blkIdx * sumStride;

        const __fp16* srcRow = srcBasePtr + s * srcStride0;

        int d = 0;
        for (; d <= headDim - 8; d += 8) {
            // --- Load Source ---
            float16x8_t vSrc16 = vld1q_f16(srcRow + d);
            float32x4_t vSrc0 = vcvt_f32_f16(vget_low_f16(vSrc16));
            float32x4_t vSrc1 = vcvt_high_f32_f16(vSrc16);

            // --- Load Quant Params ---
            float32x4_t vQs0 = vld1q_f32(&qScales[d]);
            float32x4_t vQb0 = vld1q_f32(&qBiases[d]);
            float32x4_t vQs1 = vld1q_f32(&qScales[d + 4]);
            float32x4_t vQb1 = vld1q_f32(&qBiases[d + 4]);

            // --- Quantize: x * qs + qb ---
            float32x4_t vRes0 = vaddq_f32(vmulq_f32(vSrc0, vQs0), vQb0);
            float32x4_t vRes1 = vaddq_f32(vmulq_f32(vSrc1, vQs1), vQb1);

            // --- Round & Saturate ---
            int32x4_t vInt32_0 = vcvtaq_s32_f32(vRes0);
            int32x4_t vInt32_1 = vcvtaq_s32_f32(vRes1);

            int16x8_t vInt16 = vcombine_s16(vqmovn_s32(vInt32_0), vqmovn_s32(vInt32_1));
            int8x8_t vInt8 = vqmovn_s16(vInt16); // Clamp to [-128, 127]

            vst1_s8(tmpQ, vInt8);
            for (int k = 0; k < 8; ++k) {
                int cur_d = d + k;
                int dstOffset = (cur_d / hP) * packedStride1 + idxInnerCommon + (cur_d % hP) * lP;
                dst[dstOffset] = tmpQ[k];
            }

            int16x8_t vXq16 = vmovl_s8(vInt8);
            float32x4_t vXqF0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(vXq16)));
            float32x4_t vXqF1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(vXq16)));

            float32x4_t vDs0 = vld1q_f32(&deqScales[d]);
            float32x4_t vDb0 = vld1q_f32(&deqBiases[d]);
            float32x4_t vDs1 = vld1q_f32(&deqScales[d + 4]);
            float32x4_t vDb1 = vld1q_f32(&deqBiases[d + 4]);

            // Dequant
            float32x4_t vDeq0 = vaddq_f32(vmulq_f32(vXqF0, vDs0), vDb0);
            float32x4_t vDeq1 = vaddq_f32(vmulq_f32(vXqF1, vDs1), vDb1);

            float* sumPtr = curSumRow + d;
            vst1q_f32(sumPtr, vaddq_f32(vld1q_f32(sumPtr), vDeq0));
            vst1q_f32(sumPtr + 4, vaddq_f32(vld1q_f32(sumPtr + 4), vDeq1));
        }

        for (; d < headDim; ++d) {
            float xf = (float)srcRow[d];

            float val_f = xf * qScales[d] + qBiases[d];
            int32_t val_i = (int32_t)roundf(val_f);
            if (val_i > 127) val_i = 127;
            if (val_i < -128) val_i = -128;
            int8_t xq = (int8_t)val_i;

            int dstOffset = (d / hP) * packedStride1 + idxInnerCommon + (d % hP) * lP;
            dst[dstOffset] = xq;

            curSumRow[d] += ((float)xq * deqScales[d] + deqBiases[d]);
        }
    }

/*
    // Quant fp16
    for (int d = 0; d < headDim; ++d) {
        // dst address
        int idxBase = (d / hP) * packedStride1 + (d % hP) * lP;
        int8_t*   dstBase = dst + idxBase;
        float*  scaleBase = (float*)(dst + (d / hP) * packedStride1 + weightStride1) + (d % hP);
        float*   biasBase = scaleBase + hP;
        float*   sumBase = valueSum + (d / hP) * hP + (d % hP);

        float qscale = scaleBase[0] < 1e-6 ? 0 : 1.0f / scaleBase[0];
        float qbias = scaleBase[0] < 1e-6 ? 0 : (-biasBase[0] / scaleBase[0]);
        // quant
        for (int s = 0; s < seqLen; ++s) {
            int kvSeqIndx = s + pastLength;
            int idxInner = (kvSeqIndx / flashAttentionBlockKv) * packedStride0 + (kvSeqIndx % flashAttentionBlockKv) / lP * weightStride2 + (kvSeqIndx % flashAttentionBlockKv) % lP;
            float xf = sourceFp16[s * srcStride0 + d + kvHeadIdx * headDim];
            int8_t xq = ALIMAX(ALIMIN(127, static_cast<int32_t>(roundf(xf * qscale + qbias))), -128);
            dstBase[idxInner] = xq;

            // sum
            int idxSum = (kvSeqIndx / flashAttentionBlockKv) * ROUND_UP(headDim, hP);
            sumBase[idxSum] += ((float)xq * scaleBase[0] + biasBase[0]);
        }
    }
        */
}

#endif // MNN_SUPPORT_TRANSFORMER_FUSE

#ifdef MNN_LOW_MEMORY
void MNNAbsMaxFP16(const float* source, float* absmax, size_t src_depth_quad, size_t realSize, int pack) {
    if (pack == 4) {
        MNNAbsMaxFP16_Pack4(source, absmax, src_depth_quad, realSize, pack);
        return;
    }
    if (pack == 8) {
        MNNAbsMaxFP16_Pack8(source, absmax, src_depth_quad, realSize, pack);
        return;
    }
    // source: (src_depth_quad, realSize, pack)
    auto srcStep = pack * realSize;
    auto srcPtr = (FLOAT16*)source;
    auto dstPtr = (FLOAT16*)absmax;
    for (int i = 0; i < realSize; ++i) {
        FLOAT16 absmaxVal = 0; // absmaxVal>=0
        for (int c = 0; c < src_depth_quad; ++c) {
            auto src = srcPtr + c * srcStep + i * pack;
            for (int k = 0; k < pack; ++k) {
                if (std::abs(src[k]) > absmaxVal) {
                    absmaxVal = std::abs(src[k]);
                }
            }
        }
        dstPtr[i] = absmaxVal;
    }
    return;
}

static void MNNDynamicQuantFP16(const float* src, int8_t* dst, const float* scale, size_t src_depth_quad, size_t realSize, int pack, const float* bias = nullptr) {
    if (pack == 8) {
        MNNDynamicQuantFP16_Pack8(src, dst, scale, src_depth_quad,realSize, nullptr, pack);
        return;
    }
    if (pack == 4) {
        MNNDynamicQuantFP16_Pack4(src, dst, scale, src_depth_quad,realSize, nullptr, pack);
        return;
    }
    int8_t* dstPtr = dst;
    auto srcPtr = (FLOAT16*)src;

    for (int i = 0; i < realSize; ++i) {
        auto scaleVal = static_cast<FLOAT16>(scale[i]);
        for (int c = 0; c < src_depth_quad; ++c) {
            auto srcZ = srcPtr + c * pack * realSize + i * pack;
            auto dstZ = dstPtr + c * pack * realSize + i * pack;
            for (int k = 0; k < pack; ++k) {
                int val = (int)roundf(srcZ[k] * scaleVal);
                dstZ[k] = val;
            }
        }
    }
    return;
}

static void MNNAsyQuantFunc_Arm82(int8_t* dst, const float* src, float* qscale, float* qbias, const size_t* info) {
    // input shape: [kernelsize, blockNum, blockLU, EP, LP]
    // qscale&qbias [blockNum, EP]
    auto blockNum = info[0];
    auto EP = info[1];        // real area for data
    auto LP = info[2];        // Innermost data layout, may come from backend's pack or gemmint8 units' SRC_UNIT
    auto DST_XUNIT = info[3]; // backend gemmint8 units
    auto SRC_UNIT = info[4];
    auto kernelsize = info[5];
    auto blockLU = info[6];
    auto stride0 = blockNum * blockLU * EP * LP;
    auto stride1 = blockLU * EP * LP;
    auto srcPtr = (FLOAT16*)src;
#ifdef __aarch64__
    if (LP == 4 || LP == 8) {
        for (int k = 0; k < kernelsize; ++k) {
            for (int i = 0; i < blockNum; ++i) {
                if (LP == 4) {
                    MNNDynamicQuantFP16_Pack4((float*)(srcPtr + k * stride0 + i * stride1), dst + k * stride0 + i * stride1, qscale + i * EP, blockLU, EP, qbias + i * EP, LP);
                }
                if (LP == 8) {
                    MNNDynamicQuantFP16_Pack8((float*)(srcPtr + k * stride0 + i * stride1), dst + k * stride0 + i * stride1, qscale + i * EP, blockLU, EP, qbias + i * EP, LP);
                }
            }
        }
        return;
    }
#endif
    for (int i = 0; i < EP; ++i) {
        for (int bk = 0; bk < blockNum; ++bk) {
            float quant_scale = qscale[i + bk * EP];
            float quant_bias  = qbias[i + bk * EP];
            for (int n = 0; n < kernelsize; ++n) {
                for (int k = 0; k < blockLU; ++k) {
                    for (int j = 0; j < LP; ++j) {
                        int dataIndx = n * stride0 + bk * stride1 + k * EP * LP + i * LP + j;
                        auto data_ = static_cast<float>(srcPtr[dataIndx]);
                        int qval = static_cast<int32_t>(roundf(data_ * quant_scale + quant_bias));
                        dst[dataIndx] = qval;
                    }
                }
            }
        }
    }
}

static void MNNAsyQuantInfo_FP16(float* scale, float* bias, float* qscale, float* qbias, float* dstMin, float* dstMax, float* src, const size_t* info) {
    auto blockNum = info[0];
    auto plane = info[1];        // real area for data
    auto innerSide = info[2];    // Innermost data layout, may come from backend's pack or gemmint8 units' SRC_UNIT
    auto DST_XUNIT = info[3];
    auto kernelsize = info[5];
    auto blockLU = info[6];
    auto stride0 = blockNum * blockLU * plane * innerSide;
    auto stride1 = blockLU * plane * innerSide;
    auto srcPtr = (FLOAT16*)src;

    // input shape: [kernelsize,blocknum,blocklu,DST_XUNIT,SRC_UNIT] or [ic/core->pack, plane, core->pack]
    // dequant scale/bias : [EU, blockNum, step]
    // quant scale/bias: [blockNum, plane]
    if (info[7] == 1) { // scale&bias:[1]
        FLOAT16 maxval, minval;
        ARM82CountMinMaxValue(src, (float*)(&minval), (float*)(&maxval) , kernelsize * stride0);
        if (info[8] == 1 && (maxval - minval) > 1e-7) {
            if (minval > 0.f) {
                minval = 0.f;
            } else if (maxval < 0.f){
                maxval = 0.f;
            }
        }
        auto range = maxval - minval;
        if (range <= 1e-7) {
            scale[0] = 1.f;
            qscale[0] = 1.f;
            qbias[0] = -maxval;
            bias[0] = maxval;
        } else {
            qscale[0] = 255.f / range;
            scale[0] = range / 255.f;
            qbias[0] = roundf(-minval * 255.f / range)- 128.f;
            bias[0] = -qbias[0] * scale[0];
        }
        return;
    }

#ifdef __aarch64__
    if (DST_XUNIT == 12 || DST_XUNIT == 16) { // Arm82/SME2, fp16: core->pack=8, SRC_UNIT=4
        // max,min shape: [blockNum, EP]
        if (innerSide == 4) {
            for (int i = 0; i < kernelsize; ++i) {
                MNNLocalMinMaxFP16_Pack4(dstMin, dstMax, (float*)(srcPtr + i * stride0), blockNum, blockLU, plane, innerSide, i);
            }
        }
        if (innerSide == 8) {
            for (int i = 0; i < kernelsize; ++i) {
                MNNLocalMinMaxFP16_Pack8(dstMin, dstMax, (float*)(srcPtr + i * stride0), blockNum, blockLU, plane, innerSide, i);
            }
        }
        // scale, bias
        if (DST_XUNIT == 12) {
            auto success = MNNAsyLocalQuantInfo_EP12_FP16(scale, bias, qscale, qbias, dstMin, dstMax, info);
            if (!success) {
                MNN_ERROR("Call error: MNNAsyLocalQuantInfo_EP12_FP16\n");
                return;
            }
            return;
        }
        if (DST_XUNIT == 16) {
            auto success = MNNAsyLocalQuantInfo_EP16_FP16(scale, bias, qscale, qbias, dstMin, dstMax, info);
            if (!success) {
                MNN_ERROR("Call error: MNNAsyLocalQuantInfo_EP16_FP16\n");
                return;
            }
            return;
        }
    }
    if (DST_XUNIT == 10 && innerSide == 8) { // Arm86, fp16: core->pack=8, SRC_UNIT=8
        // max,min shape: [blockNum, plane]
        for (int i = 0; i < kernelsize; ++i) {
            MNNLocalMinMaxFP16_Pack8(dstMin, dstMax, (float*)(srcPtr + i * stride0), blockNum, blockLU, plane, innerSide, i);
        }
        // scale, bias
        auto success = MNNAsyLocalQuantInfo_EP10_FP16(scale, bias, qscale, qbias, dstMin, dstMax, info);
        if (!success) {
            MNN_ERROR("Call error: MNNAsyLocalQuantInfo_EP10_FP16\n");
            return;
        }
        return;
    }
#else
    // aarch32
    // max,min shape: [blockNum, plane]
    auto minPtr = (FLOAT16*)dstMin;
    auto maxPtr = (FLOAT16*)dstMax;
    for (int i = 0; i < plane; ++i) {
        for (int bk = 0; bk < blockNum; ++bk) {
            auto idx0 = i * innerSide + bk * stride1;
            auto max_ = srcPtr[idx0];
            auto min_ = max_;
            for (int n = 0; n < kernelsize; ++n) {
                for (int k = 0; k < blockLU; ++k) {
                    for (int j = 0; j < innerSide; ++j) {
                        auto dataIndx = idx0 + n * stride0 + k * (plane * innerSide) + j;
                        auto data_ = srcPtr[dataIndx];
                        max_ = ALIMAX(max_, data_);
                        min_ = ALIMIN(min_, data_);
                    }
                }
            }
            auto sindx = i + bk * plane;
            minPtr[sindx] = min_;
            maxPtr[sindx] = max_;
        }
    }
    // scale, bias
    for (int i = 0; i < plane; ++i) {
        auto step = ALIMIN(DST_XUNIT, plane - (i / DST_XUNIT) * DST_XUNIT);
        auto sind0 = (i / DST_XUNIT) * DST_XUNIT * blockNum + (i % DST_XUNIT);
        for (int k = 0; k < blockNum; ++k) {
            auto sind = sind0 + k * step;
            auto qind = i + k * plane;
            auto max_ = (float)maxPtr[qind];
            auto min_ = (float)minPtr[qind];
            auto range = max_ - min_;
            if (fabs(range) < 1e-7) {
                qscale[qind] = 0.f;
                qbias[qind] = 0.f;
                scale[sind] = 0.f;
                bias[sind] = max_;
            } else {
                qscale[qind] = 255.f / range;
                qbias[qind] = -min_ * 255.f / range - 128.0f;
                scale[sind] = range / 255.f;
                bias[sind] = min_ + (128.f / 255.f) * range;

            }
        }
    }
#endif
}

#endif // MNN_LOW_MEMORY

#define EXP_APPROX_MIN_INPUT vdupq_n_f32(-88.0f)
#define EXP_APPROX_MAX_INPUT vdupq_n_f32(88.0f)
#define EXP_APPROX_LN2         vdupq_n_f32(0.69314718056f)  // ln(2)
#define EXP_APPROX_LN2_INV     vdupq_n_f32(1.44269504089f)   // 1/ln(2)
// Fourth-order polynomial approximation coefficients of exp(r):
// P(x) = c4*x^4 + c3*x^3 + c2*x^2 + c1*x + c0
#define EXP_APPROX_C4          vdupq_n_f32(0.0416624f)
#define EXP_APPROX_C3          vdupq_n_f32(0.166665f)
#define EXP_APPROX_C2          vdupq_n_f32(0.500000f)
#define EXP_APPROX_C1          vdupq_n_f32(1.0f)
#define EXP_APPROX_C0          vdupq_n_f32(1.0f)

#ifndef __aarch64__
static inline float32x4_t vrndaq_f32_compat(float32x4_t x) {
    float32x4_t sign = vbslq_f32(vdupq_n_u32(0x80000000), x, vdupq_n_f32(0.0f));
    return vcvtq_f32_s32(vcvtq_s32_f32(vaddq_f32(x, vbslq_f32(vcltq_f32(x, vdupq_n_f32(0.0f)), vdupq_n_f32(-0.5f), vdupq_n_f32(0.5f)))));
}
#endif

static inline float32x4_t expApprox(float32x4_t x) {
    x = vminq_f32(vmaxq_f32(x, EXP_APPROX_MIN_INPUT), EXP_APPROX_MAX_INPUT);

    float32x4_t k_float;
    float32x4_t r;
    float32x4_t exp_r;
#if defined(__aarch64__)
    k_float = vrndaq_f32(vmulq_f32(x, EXP_APPROX_LN2_INV));

    // r = x - k * ln(2)
    r = vfmsq_f32(x, k_float, EXP_APPROX_LN2);

    // P(r) = (c0 + c2*r^2 + c4*r^4) + r*(c1 + c3*r^2)
    float32x4_t r2 = vmulq_f32(r, r);
    float32x4_t p_odd = vfmaq_f32(EXP_APPROX_C1, EXP_APPROX_C3, r2);

    float32x4_t p_even = vfmaq_f32(EXP_APPROX_C0, EXP_APPROX_C2, r2);
    p_even = vfmaq_f32(p_even, EXP_APPROX_C4, vmulq_f32(r2, r2));
    exp_r = vfmaq_f32(p_even, p_odd, r);
#else

    k_float = vrndaq_f32_compat(vmulq_f32(x, EXP_APPROX_LN2_INV));


    r = vsubq_f32(x, vmulq_f32(k_float, EXP_APPROX_LN2));

    // 2. c0 + r*(c1 + r*(c2 + r*(c3 + r*c4)))
    exp_r = vmlaq_f32(EXP_APPROX_C3, EXP_APPROX_C4, r); // c3 + c4*r
    exp_r = vmlaq_f32(EXP_APPROX_C2, exp_r, r);         // c2 + r*(...)
    exp_r = vmlaq_f32(EXP_APPROX_C1, exp_r, r);         // c1 + r*(...)
    exp_r = vmlaq_f32(EXP_APPROX_C0, exp_r, r);         // c0 + r*(...)

#endif

    int32x4_t k_int = vcvtq_s32_f32(k_float);
    int32x4_t k_shifted = vshlq_n_s32(k_int, 23);
    return vreinterpretq_f32_s32(vaddq_s32(vreinterpretq_s32_f32(exp_r), k_shifted));
}
static void MNNSoftmaxFp16_Pack8(float* dest, const float* source, float* runningMax, float* runningSum, float* updateScale, int outside, int reduceSize, int kvSeqOffset, int validOffset, int pack, bool mask) {
    auto softmaxDst = (FLOAT16*)dest;
    auto softmaxSrc = (FLOAT16*)source;

    if (pack != 8) {
        MNN_ERROR("MNNSoftmaxFp16_Pack8 only support pack=8\n");
        return;
    }

    const int packUnit = 8;
    int reduceSizeOuter = UP_DIV(reduceSize, packUnit);
    int stride0 = outside * packUnit;

    // Loop Tiling: Unroll K by 16
    // 16 * 8 * 2 = 256 Bytes
    for (int k = 0; k < outside; k += 16) {
        int count = ALIMIN(16, outside - k);

        int validLens[16];
        bool isRowValid[16];

        for (int i = 0; i < count; ++i) {
            int currentK = k + i;
            if (mask && kvSeqOffset > currentK + validOffset) {
                isRowValid[i] = false;
                validLens[i] = 0;
                if (updateScale) updateScale[currentK] = 1.0f;
            } else {
                isRowValid[i] = true;
                validLens[i] = mask ? ALIMIN(reduceSize, currentK + (validOffset + 1) - kvSeqOffset) : reduceSize;
            }
        }

        float currentMax[16];
        for (int i = 0; i < count; ++i) {
            currentMax[i] = runningMax ? runningMax[k + i] : -65504.0f;
        }

        for (int j = 0; j < reduceSizeOuter; ++j) {
            auto blockSrcBase = softmaxSrc + j * stride0 + k * packUnit;

            for (int i = 0; i < count; ++i) {
                if (!isRowValid[i]) continue;

                int len = validLens[i];
                int blockStart = j * packUnit;
                if (blockStart >= len) continue;

                auto srcPtr = blockSrcBase + i * packUnit;
                int remain = len - blockStart;

                if (remain >= packUnit) {
                    float16x8_t val = vld1q_f16(srcPtr);
                    float maxInVec = vmaxvq_f16(val);
                    currentMax[i] = ALIMAX(currentMax[i], maxInVec);
                } else {
                    for (int p = 0; p < remain; ++p) {
                        currentMax[i] = ALIMAX(currentMax[i], (float)srcPtr[p]);
                    }
                }
            }
        }

        float currentSum[16] = {0.0f};
        float32x4_t vecSum0[16]; // Low part accumulator
        float32x4_t vecSum1[16]; // High part accumulator
        float32x4_t finalMaxVec[16];

        for (int i = 0; i < count; ++i) {
            vecSum0[i] = vdupq_n_f32(0.0f);
            vecSum1[i] = vdupq_n_f32(0.0f);
            finalMaxVec[i] = vdupq_n_f32(currentMax[i]);
        }

        for (int j = 0; j < reduceSizeOuter; ++j) {
            auto blockSrcBase = softmaxSrc + j * stride0 + k * packUnit;
            auto blockDstBase = softmaxDst + j * stride0 + k * packUnit;

            for (int i = 0; i < count; ++i) {
                if (!isRowValid[i]) {
                    memset(blockDstBase + i * packUnit, 0, packUnit * sizeof(__fp16));
                    continue;
                }

                int len = validLens[i];
                int blockStart = j * packUnit;
                if (blockStart >= len) {
                    memset(blockDstBase + i * packUnit, 0, packUnit * sizeof(__fp16));
                    continue;
                }

                auto srcPtr = blockSrcBase + i * packUnit;
                auto dstPtr = blockDstBase + i * packUnit;
                int remain = len - blockStart;

                if (remain >= packUnit) {
                    float16x8_t srcVal = vld1q_f16(srcPtr);

                    // F16 -> F32 expansion
                    float32x4_t low = vcvt_f32_f16(vget_low_f16(srcVal));
                    float32x4_t high = vcvt_f32_f16(vget_high_f16(srcVal));

                    // Subtract Max
                    low = vsubq_f32(low, finalMaxVec[i]);
                    high = vsubq_f32(high, finalMaxVec[i]);

                    // Exp
                    low = expApprox(low);
                    high = expApprox(high);

                    // Accumulate Sum
                    vecSum0[i] = vaddq_f32(vecSum0[i], low);
                    vecSum1[i] = vaddq_f32(vecSum1[i], high);

                    // Store Exp result temporarily
                    vst1q_f16(dstPtr, vcombine_f16(vcvt_f16_f32(low), vcvt_f16_f32(high)));
                } else {
                    // Handle Tail
                    for (int p = 0; p < remain; ++p) {
                        float val = expf((float)srcPtr[p] - currentMax[i]);
                        currentSum[i] += val;
                        dstPtr[p] = (__fp16)val;
                    }
                    memset(dstPtr + remain, 0, (packUnit - remain) * sizeof(__fp16));
                }
            }
        }

        // Horizontal reduction for sums
        for (int i = 0; i < count; ++i) {
            currentSum[i] += vaddvq_f32(vecSum0[i]) + vaddvq_f32(vecSum1[i]);
        }

        for (int i = 0; i < count; ++i) {
            int currentK = k + i;
            if (!isRowValid[i]) continue;

            float scale;
            if (runningMax && runningSum && updateScale) {
                // Incremental Softmax logic
                float oldMax = runningMax[currentK];
                float scaleForSum = expf(oldMax - currentMax[i]);
                runningSum[currentK] = runningSum[currentK] * scaleForSum + currentSum[i];
                runningMax[currentK] = currentMax[i];
                updateScale[currentK] = scaleForSum;
                continue;
            } else {
                // Standard Softmax logic
                if (runningMax && runningSum) {
                    currentSum[i] += runningSum[currentK] * expf(runningMax[currentK] - currentMax[i]);
                }
                scale = 1.0f / (currentSum[i] + 1e-20f);
            }

            float16x8_t scaleVec = vdupq_n_f16((__fp16)scale);

            // Normalize Pass
            for (int j = 0; j < reduceSizeOuter; ++j) {
                int len = validLens[i];
                int blockStart = j * packUnit;
                if (blockStart >= len) break;

                auto dstPtr = softmaxDst + j * stride0 + k * packUnit + i * packUnit;

                if (len - blockStart >= packUnit) {
                    float16x8_t val = vld1q_f16(dstPtr);
                    val = vmulq_f16(val, scaleVec);
                    vst1q_f16(dstPtr, val);
                } else {
                    int remain = len - blockStart;
                    for (int p = 0; p < remain; ++p) {
                        dstPtr[p] = (__fp16)((float)dstPtr[p] * scale);
                    }
                }
            }
        }
    }
}

static void MNNSoftmaxFp16_Pack1(float* dest, const float* source, float* runningMax, float* runningSum, float* updateScale, int outside, int reduceSize, int kvSeqOffset, int validOffset, bool mask) {
    auto softmaxDst = (FLOAT16*)dest;
    auto softmaxSrc = (FLOAT16*)source;

    for (int k = 0; k < outside; ++k) {
        int currentValidSize = reduceSize;
        bool isRowValid = true;

        if (mask) {
            if (kvSeqOffset > k + validOffset) {
                isRowValid = false;
                currentValidSize = 0;
                if (updateScale) updateScale[k] = 1.0f;
            } else {
                currentValidSize = ALIMIN(reduceSize, k + (validOffset + 1) - kvSeqOffset);
            }
        }

        if (!isRowValid || currentValidSize == 0) {
            memset(softmaxDst + k * reduceSize, 0, reduceSize * sizeof(__fp16));
            continue;
        }

        auto srcRow = softmaxSrc + k * reduceSize;
        auto dstRow = softmaxDst + k * reduceSize;

        float oldMax = runningMax ? runningMax[k] : -65504.0f;
        float16x8_t maxVec = vdupq_n_f16(-65504.0f);

        // Unroll 4 (32 elements per loop)
        int i = 0;
        for (; i <= currentValidSize - 32; i += 32) {
            float16x8_t v0 = vld1q_f16(srcRow + i + 0);
            float16x8_t v1 = vld1q_f16(srcRow + i + 8);
            float16x8_t v2 = vld1q_f16(srcRow + i + 16);
            float16x8_t v3 = vld1q_f16(srcRow + i + 24);

            maxVec = vmaxq_f16(maxVec, v0);
            maxVec = vmaxq_f16(maxVec, v1);
            maxVec = vmaxq_f16(maxVec, v2);
            maxVec = vmaxq_f16(maxVec, v3);
        }
        // Handle remaining blocks of 8
        for (; i <= currentValidSize - 8; i += 8) {
            maxVec = vmaxq_f16(maxVec, vld1q_f16(srcRow + i));
        }

        // Horizontal Max reduction
        float newMax = vmaxvq_f16(maxVec);

        // Handle remaining scalars (Tail)
        for (; i < currentValidSize; ++i) {
            newMax = ALIMAX(newMax, (float)srcRow[i]);
        }

        float finalMax = ALIMAX(oldMax, newMax);
        float32x4_t finalMaxVec = vdupq_n_f32(finalMax);

        float sum = 0.0f;
        float32x4_t sumVec = vdupq_n_f32(0.0f);

        i = 0;
        // Unroll 2 (16 elements). Exp is heavy, unroll 4 might cause register spilling.
        for (; i <= currentValidSize - 16; i += 16) {
            float16x8_t v0 = vld1q_f16(srcRow + i);
            float16x8_t v1 = vld1q_f16(srcRow + i + 8);

            // Process v0
            float32x4_t v0_lo = vcvt_f32_f16(vget_low_f16(v0));
            float32x4_t v0_hi = vcvt_f32_f16(vget_high_f16(v0));
            v0_lo = expApprox(vsubq_f32(v0_lo, finalMaxVec));
            v0_hi = expApprox(vsubq_f32(v0_hi, finalMaxVec));
            sumVec = vaddq_f32(sumVec, v0_lo);
            sumVec = vaddq_f32(sumVec, v0_hi);
            vst1q_f16(dstRow + i, vcombine_f16(vcvt_f16_f32(v0_lo), vcvt_f16_f32(v0_hi)));

            // Process v1
            float32x4_t v1_lo = vcvt_f32_f16(vget_low_f16(v1));
            float32x4_t v1_hi = vcvt_f32_f16(vget_high_f16(v1));
            v1_lo = expApprox(vsubq_f32(v1_lo, finalMaxVec));
            v1_hi = expApprox(vsubq_f32(v1_hi, finalMaxVec));
            sumVec = vaddq_f32(sumVec, v1_lo);
            sumVec = vaddq_f32(sumVec, v1_hi);
            vst1q_f16(dstRow + i + 8, vcombine_f16(vcvt_f16_f32(v1_lo), vcvt_f16_f32(v1_hi)));
        }

        // Handle remaining blocks of 8
        for (; i <= currentValidSize - 8; i += 8) {
            float16x8_t v = vld1q_f16(srcRow + i);
            float32x4_t v_lo = vcvt_f32_f16(vget_low_f16(v));
            float32x4_t v_hi = vcvt_f32_f16(vget_high_f16(v));

            v_lo = expApprox(vsubq_f32(v_lo, finalMaxVec));
            v_hi = expApprox(vsubq_f32(v_hi, finalMaxVec));

            sumVec = vaddq_f32(sumVec, v_lo);
            sumVec = vaddq_f32(sumVec, v_hi);

            vst1q_f16(dstRow + i, vcombine_f16(vcvt_f16_f32(v_lo), vcvt_f16_f32(v_hi)));
        }

        // Handle Tail scalars
        if (i < currentValidSize) {
             __fp16 tempDst[8];
            int remain = currentValidSize - i;
            auto sPtr = srcRow + i;
            for (int p = 0; p < remain; ++p) {
                float val = expf((float)sPtr[p] - finalMax);
                sum += val;
                tempDst[p] = (__fp16)val;
            }
            memcpy(dstRow + i, tempDst, remain * sizeof(__fp16));
            i += remain; // align i to currentValidSize
        }

        sum += vaddvq_f32(sumVec);

        // Fill remaining invalid part with 0
        if (currentValidSize < reduceSize) {
            memset(dstRow + currentValidSize, 0, (reduceSize - currentValidSize) * sizeof(__fp16));
        }

        if (runningMax && runningSum && updateScale) {
            float scaleForSum = expf(oldMax - finalMax);
            runningSum[k] = runningSum[k] * scaleForSum + sum;
            runningMax[k] = finalMax;
            updateScale[k] = scaleForSum;
        } else {
            if (runningMax && runningSum) {
                sum += runningSum[k] * expf(oldMax - finalMax);
            }
            float scale = 1.0f / (sum + 1e-20f);
            float16x8_t scaleVec = vdupq_n_f16((__fp16)scale);

            // Unroll 4 (32 elements) for throughput
            i = 0;
            for (; i <= currentValidSize - 32; i += 32) {
                float16x8_t v0 = vld1q_f16(dstRow + i);
                float16x8_t v1 = vld1q_f16(dstRow + i + 8);
                float16x8_t v2 = vld1q_f16(dstRow + i + 16);
                float16x8_t v3 = vld1q_f16(dstRow + i + 24);

                vst1q_f16(dstRow + i, vmulq_f16(v0, scaleVec));
                vst1q_f16(dstRow + i + 8, vmulq_f16(v1, scaleVec));
                vst1q_f16(dstRow + i + 16, vmulq_f16(v2, scaleVec));
                vst1q_f16(dstRow + i + 24, vmulq_f16(v3, scaleVec));
            }
            for (; i <= currentValidSize - 8; i += 8) {
                float16x8_t v = vld1q_f16(dstRow + i);
                vst1q_f16(dstRow + i, vmulq_f16(v, scaleVec));
            }
            for (; i < currentValidSize; ++i) {
                dstRow[i] = (__fp16)((float)dstRow[i] * scale);
            }
        }
    }
}


static void MNNSoftmaxFp16(float* dest, const float* source, float* runningMax, float* runningSum, float* updateScale, int outside, int reduceSize, int kvSeqOffset, int validOffset, int pack, bool mask) {
    // source shape: [reduceSizeOuter, outside, reduceSizeInner]
    // for C4, [up_div(reduceSize,8), outside,8] => reduceSizeOuter=up_div(reduceSize,8), reduceSizeInner=8
    // for C,  [outside, reduceSize]             => reduceSizeOuter=1, reduceSizeInner=reduceSize
    if (pack == 8) {
        MNNSoftmaxFp16_Pack8(dest, source, runningMax, runningSum, updateScale, outside, reduceSize, kvSeqOffset, validOffset, pack, mask);
        return;
    }
    if (pack == 1) {
        MNNSoftmaxFp16_Pack1(dest, source, runningMax, runningSum, updateScale, outside, reduceSize, kvSeqOffset, validOffset, mask);
        return;
    }
    MNN_ERROR("MNNSoftMaxFp16 not support pack!=8 and pack!=1\n");
    return;
}

static CoreFunctions* gInstance = nullptr;
static CoreInt8Functions* gArm82CoreInt8Functions = nullptr;

bool Arm82Functions::init() {
    using Vec = MNN::Math::Vec<FLOAT16, 8>;
    auto origin = MNNGetCoreFunctions();
#define FUNC_PTR_ASSIGN(dst, src) dst = (decltype(dst))(src)
    gInstance = new CoreFunctions;
    gArm82CoreInt8Functions = new CoreInt8Functions;
    *gArm82CoreInt8Functions = *MNNGetInt8CoreFunctions();
    gInstance->int8MatmulRelatedFunctions = origin->int8MatmulRelatedFunctions;
    {
        if (origin->supportSDot) {
            gArm82CoreInt8Functions->MNNPackC4Int8ForMatMul_A = _Arm82MNNPackC4ForMatMul_A<12, 4>;
            gInstance->arm82MatmulRelatedFunctions = origin->arm82MatmulRelatedFunctions;
            gInstance->arm82MatmulRelatedFunctions.MNNPackC4Int8ForMatMul_A = _Arm82MNNPackC4ForMatMul_A<12, 4>;
        }
        if (origin->supportI8mm) {
            gArm82CoreInt8Functions->MNNPackC4Int8ForMatMul_A = _ArmBasicMNNPackC4ForMatMul_A_L8<10, 8>;
            gInstance->supportI8mm = true;
        }
    }

    FUNC_PTR_ASSIGN(gInstance->MNNFp32ToLowp, MNNQuantizeFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNLowpToFp32, MNNDequantizeFP16);
    gInstance->bytes = 2;

    // Packed
    gInstance->pack = 8;
    FUNC_PTR_ASSIGN(gInstance->MNNPackCUnit, MNNPackC8FP16);
    FUNC_PTR_ASSIGN(gInstance->MNNUnpackCUnit, MNNUnPackC8FP16);
    FUNC_PTR_ASSIGN(gInstance->MNNPackCUnitTranspose, MNNPackTransposeInt16C8);
    FUNC_PTR_ASSIGN(gInstance->MNNUnpackCUnitTranspose, MNNUnpackTransposeInt16C8);
    FUNC_PTR_ASSIGN(gInstance->MNNConvRunForLineDepthwise, MNNConvRunForLineDepthwiseFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNAxByClampBroadcastUnit, MNNAxByClampBroadcastC8FP16);
    FUNC_PTR_ASSIGN(gInstance->MNNMatrixSub, MNNMatrixSubFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNMatrixAdd, MNNMatrixAddFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNStrassenMergeCFunction, ARM82StrassenMerge);
    gInstance->MNNReorderWeightInt4 = origin->MNNReorderWeightInt4;
    gInstance->MNNSumWeightInt8 = origin->MNNSumWeightInt8;
    gInstance->MNNSumWeightInt8SmeHp128 = origin->MNNSumWeightInt8SmeHp128;
    gInstance->penalty = 2.0f;
    FUNC_PTR_ASSIGN(gInstance->MNNScaleAndAddBias, MNNScaleAndAddBiasFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNGridSampleComputeCord, MNNGridSampleComputeCordFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNGridSampleInterp, MNNGridSampleInterp);
    FUNC_PTR_ASSIGN(gInstance->MNNGridSampleInterpGrad, MNNGridSampleInterpGrad);
    FUNC_PTR_ASSIGN(gInstance->MNNGridSampleComputeCord3D, MNNGridSampleComputeCord3DFp16);
    FUNC_PTR_ASSIGN(gInstance->MNNGridSampleInterp3D, MNNGridSampleInterp3D);
    FUNC_PTR_ASSIGN(gInstance->MNNRoiPoolingMax, MNNRoiPoolingMaxFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNRoiAlignMax, MNNRoiAlignMaxFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNRoiAlignAvg, MNNRoiAlignAvgFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNCopyC4WithStride, MNNCopyC8WithStrideFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNAddC4WithStride, MNNAddC8WithStrideFP16);

    // MatMul
    FUNC_PTR_ASSIGN(gInstance->MNNGetMatMulPackMode, Arm82MNNGetMatMulPackMode);
    FUNC_PTR_ASSIGN(gInstance->MNNPackedMatMul, MNNPackedMatMulFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNPackedMatMulRemain, MNNPackedMatMulRemainFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNPackC4ForMatMul_A, Arm82MNNPackForMatMul_A);
    FUNC_PTR_ASSIGN(gInstance->MNNPackForMatMul_B, Arm82MNNPackForMatMul_B);

    FUNC_PTR_ASSIGN(gInstance->MNNSoftmax, MNNSoftmaxFp16);
#if defined(__aarch64__)
    gInstance->supportFp16arith = origin->supportFp16arith;
    gInstance->supportSDot = origin->supportSDot;
    gInstance->supportI8mm = origin->supportI8mm;
    gInstance->supportSME2 = origin->supportSME2;
    gInstance->smeCoreNumber = origin->smeCoreNumber;
#ifdef MNN_LOW_MEMORY
    // Dynamic Qaunt Helper Functions
    FUNC_PTR_ASSIGN(gInstance->MNNAbsMax, MNNAbsMaxFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNQuantScale, MNNQuantScaleFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNDynamicQuant, MNNDynamicQuantFP16);
    FUNC_PTR_ASSIGN(gInstance->MNNAsyQuantFunc, MNNAsyQuantFunc_Arm82);
    FUNC_PTR_ASSIGN(gInstance->MNNAsyQuantInfo, MNNAsyQuantInfo_FP16); // return 'plane' min&max
    FUNC_PTR_ASSIGN(gInstance->MNNDynamicUpdateConvBiasScale, origin->MNNDynamicUpdateConvBiasScale);

    if (origin->supportSDot) {
        FUNC_PTR_ASSIGN(gInstance->MNNGeneralIm2Col, MNNGeneralIm2col_Arm82);
        gInstance->arm82MatmulRelatedFunctions.MNNGeneralIm2Col = MNNGeneralIm2col_Arm82;
    }
    if (origin->supportI8mm) {
        FUNC_PTR_ASSIGN(gInstance->MNNGeneralIm2Col, MNNGeneralIm2col_Arm86);
    }
#endif // MNN_LOW_MEMORY
    FUNC_PTR_ASSIGN(gInstance->MNNCountMaxMinValue, ARM82CountMinMaxValue); // return one min&max
    FUNC_PTR_ASSIGN(gInstance->MNNSumByAxisLForMatmul_A, origin->MNNSumByAxisLForMatmul_A);
    FUNC_PTR_ASSIGN(gInstance->MNNDepthwiseConvFastKernel, MNNDepthwiseConvFastKernelFP16);
#endif // __aarch64__

#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
    // Attention
    FUNC_PTR_ASSIGN(gInstance->MNNAttenPackAndScaleSingleHead, MNNAttenPackAndScaleSingleHead);
    FUNC_PTR_ASSIGN(gInstance->MNNFlashAttentionUpdateBlockOutput, MNNFlashAttentionUpdateBlockOutput);
    gInstance->MNNQuantAttentionKey = MNNQuantAttentionKeyFP16;
    gInstance->MNNQuantAttentionValue = MNNQuantAttentionValueFP16;

    // LinearAttention fp16 kernels
    FUNC_PTR_ASSIGN(gInstance->MNNRankOneUpdate, MNNRankOneUpdateFp16);
    FUNC_PTR_ASSIGN(gInstance->MNNDualMatVec, MNNDualMatVecFp16);
    FUNC_PTR_ASSIGN(gInstance->MNNDecayRankOneUpdate, MNNDecayRankOneUpdateFp16);
#if defined(__aarch64__) && defined(MNN_USE_NEON)
    // Fused kernel uses NEON intrinsics directly (not extern asm), so the
    // assignment must follow the same guard as the function body above.
    FUNC_PTR_ASSIGN(gInstance->MNNFusedGatedDelta, MNNFusedGatedDeltaFp16);
#endif
#endif // MNN_SUPPORT_TRANSFORMER_FUSE

    gInstance->MNNComputeMatMulForH_1 = _MNNComputeMatMulForH_1_FP16;
    gInstance->MNNComputeMatMulForE_1 = _MNNComputeMatMulForE_1_FP16;

    FUNC_PTR_ASSIGN(gInstance->chooseWinoSourceTransformPack, Arm82WinogradFunction::chooseWinoSourceTransformPack);
    FUNC_PTR_ASSIGN(gInstance->chooseWinoSourceUnrollTransform, Arm82WinogradFunction::chooseSourceUnrollTransform);
    FUNC_PTR_ASSIGN(gInstance->chooseWinoDestUnrollTransform, Arm82WinogradFunction::chooseWinoDestUnrollTransform);

    gInstance->MNNDeconvRunForLineDepthwise = (decltype(gInstance->MNNDeconvRunForLineDepthwise))_MNNDeconvRunForLineDepthwise;
    gInstance->MNNDeconvRunForUnitDepthWise = (decltype(gInstance->MNNDeconvRunForUnitDepthWise))_MNNDeconvRunForUnitDepthWise;

    // Binary and Unary
    gInstance->MNNSelectBinaryFunctionForFloat = Arm82BinaryFloat::select;
    gInstance->MNNSelectUnaryFunctionForFloat = Arm82Unary::select;

    // Relu with slope
    gInstance->MNNReluWithSlopeChannel = Arm82Relu::reluWithSlopeChannel;

    gInstance->MNNPoolingMax = (decltype(gInstance->MNNPoolingMax))(poolingMax<float16_t, Vec, 8, -65535>);
    gInstance->MNNPoolingAvg = (decltype(gInstance->MNNPoolingAvg))(poolingAvg<float16_t, Vec, 8>);

    {
        gInstance->int8MatmulRelatedFunctions.MNNPackC4Int8ForMatMul_A = gArm82CoreInt8Functions->MNNPackC4Int8ForMatMul_A;
        gInstance->int8MatmulRelatedFunctions.MNNGeneralIm2Col = gInstance->MNNGeneralIm2Col;
    }

#ifdef __aarch64__
#ifdef MNN_SME2
        if (origin->supportSME2) {
            gArm82CoreInt8Functions->MNNPackC4Int8ForMatMul_A = _Arm82MNNPackC4ForMatMul_A<16, 4>;

            FUNC_PTR_ASSIGN(gInstance->MNNPackedMatMul, MNNPackedMatMulFP16_SME2);
            FUNC_PTR_ASSIGN(gInstance->MNNPackedMatMulRemain, MNNPackedMatMulRemainFP16_SME2);
            FUNC_PTR_ASSIGN(gInstance->MNNGetMatMulPackMode, Sme2MNNGetMatMulPackMode);
            FUNC_PTR_ASSIGN(gInstance->MNNPackC4ForMatMul_A, Sme2MNNPackC4ForMatMul_A_FP16);
            FUNC_PTR_ASSIGN(gInstance->MNNPackForMatMul_B, Sme2MNNPackForMatMul_B);

#ifdef MNN_LOW_MEMORY
            FUNC_PTR_ASSIGN(gInstance->MNNGeneralIm2Col, MNNGeneralIm2col_Fp16Sme2);
#endif
        }
#endif // MNN_SME2
#endif // __aarch64__

    // Update the function pointers in the int8MatmulRelatedFunctions struct.
    gInstance->int8MatmulRelatedFunctions.MNNPackC4Int8ForMatMul_A = gArm82CoreInt8Functions->MNNPackC4Int8ForMatMul_A;
    gInstance->int8MatmulRelatedFunctions.MNNGeneralIm2Col = gInstance->MNNGeneralIm2Col;


    return true;
}

CoreFunctions* Arm82Functions::get() {
    return gInstance;
}
CoreInt8Functions* Arm82Functions::getInt8() {
    return gArm82CoreInt8Functions;
}
};
#endif