58 lines
2.2 KiB
C++
58 lines
2.2 KiB
C++
#include <riscv_vector.h>
|
|
|
|
extern "C" {
|
|
|
|
extern void MNNExp(float* dst, const float* src, float* offset, size_t dataSize);
|
|
|
|
void MNNSiLu(float* dst, const float* src, size_t dataSize) {
|
|
float offset[4] = {-1.0f, 0.0f, 0.0f, 0.0f};
|
|
MNNExp(dst, src, offset, dataSize);
|
|
size_t n = dataSize;
|
|
float* dstPtr = dst;
|
|
const float* srcPtr = src;
|
|
while (n > 0) {
|
|
size_t vl = __riscv_vsetvl_e32m4(n);
|
|
vfloat32m4_t vExp = __riscv_vle32_v_f32m4(dstPtr, vl);
|
|
vfloat32m4_t vSrc = __riscv_vle32_v_f32m4(srcPtr, vl);
|
|
vfloat32m4_t vDenom = __riscv_vfadd_vf_f32m4(vExp, 1.0f, vl);
|
|
vfloat32m4_t vRcp = __riscv_vfrec7_v_f32m4(vDenom, vl);
|
|
// Newton-Raphson x2: rcp *= (2 - denom * rcp)
|
|
vfloat32m4_t vErr = __riscv_vfmul_vv_f32m4(vDenom, vRcp, vl);
|
|
vErr = __riscv_vfrsub_vf_f32m4(vErr, 2.0f, vl);
|
|
vRcp = __riscv_vfmul_vv_f32m4(vRcp, vErr, vl);
|
|
vErr = __riscv_vfmul_vv_f32m4(vDenom, vRcp, vl);
|
|
vErr = __riscv_vfrsub_vf_f32m4(vErr, 2.0f, vl);
|
|
vRcp = __riscv_vfmul_vv_f32m4(vRcp, vErr, vl);
|
|
vfloat32m4_t vResult = __riscv_vfmul_vv_f32m4(vSrc, vRcp, vl);
|
|
__riscv_vse32_v_f32m4(dstPtr, vResult, vl);
|
|
dstPtr += vl;
|
|
srcPtr += vl;
|
|
n -= vl;
|
|
}
|
|
}
|
|
|
|
void MNNSiLuLowp(float* dst, const float* src, size_t dataSize) {
|
|
float offset[4] = {-1.0f, 0.0f, 0.0f, 0.0f};
|
|
MNNExp(dst, src, offset, dataSize);
|
|
size_t n = dataSize;
|
|
float* dstPtr = dst;
|
|
const float* srcPtr = src;
|
|
while (n > 0) {
|
|
size_t vl = __riscv_vsetvl_e32m4(n);
|
|
vfloat32m4_t vExp = __riscv_vle32_v_f32m4(dstPtr, vl);
|
|
vfloat32m4_t vSrc = __riscv_vle32_v_f32m4(srcPtr, vl);
|
|
vfloat32m4_t vDenom = __riscv_vfadd_vf_f32m4(vExp, 1.0f, vl);
|
|
vfloat32m4_t vRcp = __riscv_vfrec7_v_f32m4(vDenom, vl);
|
|
vfloat32m4_t vErr = __riscv_vfmul_vv_f32m4(vDenom, vRcp, vl);
|
|
vErr = __riscv_vfrsub_vf_f32m4(vErr, 2.0f, vl);
|
|
vRcp = __riscv_vfmul_vv_f32m4(vRcp, vErr, vl);
|
|
vfloat32m4_t vResult = __riscv_vfmul_vv_f32m4(vSrc, vRcp, vl);
|
|
__riscv_vse32_v_f32m4(dstPtr, vResult, vl);
|
|
dstPtr += vl;
|
|
srcPtr += vl;
|
|
n -= vl;
|
|
}
|
|
}
|
|
|
|
} // extern "C"
|