From ccb8963705205eef1f7447be7964dce008c7b997 Mon Sep 17 00:00:00 2001 From: tensor-tang Date: Fri, 16 Nov 2018 16:54:48 +0000 Subject: [PATCH] refine exp jitcode with all size test=develop --- paddle/fluid/operators/math/jit_code.cc | 223 +++-------------------- paddle/fluid/operators/math/jit_code.h | 132 +++++++++++++- paddle/fluid/operators/math/jit_kernel.h | 1 + 3 files changed, 153 insertions(+), 203 deletions(-) diff --git a/paddle/fluid/operators/math/jit_code.cc b/paddle/fluid/operators/math/jit_code.cc index 2a10cd78216..fd18256b0c9 100644 --- a/paddle/fluid/operators/math/jit_code.cc +++ b/paddle/fluid/operators/math/jit_code.cc @@ -13,8 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/fluid/operators/math/jit_code.h" -#include "paddle/fluid/operators/math/jit_kernel.h" -#include "paddle/fluid/platform/cpu_info.h" +#include "paddle/fluid/operators/math/jit_kernel.h" // TODO(TJ): remove me namespace paddle { namespace operators { @@ -111,60 +110,26 @@ void VXXJitCode::generate() { ret(); } -#define ALIGN32 __attribute__((aligned(32))) -#define EXP_HIG 88.3762626647949f -#define EXP_LOW -88.3762626647949f -#define CEPHES_LOG2EF 1.44269504088896341 -#define CEPHES_EXP_C1 0.693359375 -#define CEPHES_EXP_C2 -2.12194440e-4 -#define CEPHES_EXP_P0 1.9875691500E-4 -#define CEPHES_EXP_P1 1.3981999507E-3 -#define CEPHES_EXP_P2 8.3334519073E-3 -#define CEPHES_EXP_P3 4.1665795894E-2 -#define CEPHES_EXP_P4 1.6666665459E-1 -#define CEPHES_EXP_P5 5.0000001201E-1 +const float exp_float_consts[] ALIGN32 = {REPEAT_8TIMES(1.f), + REPEAT_8TIMES(2.f), + REPEAT_8TIMES(0.5f), + REPEAT_8TIMES(EXP_HIG), + REPEAT_8TIMES(EXP_LOW), + REPEAT_8TIMES(CEPHES_LOG2EF), + REPEAT_8TIMES(CEPHES_EXP_C1), + REPEAT_8TIMES(CEPHES_EXP_C2), + REPEAT_8TIMES(CEPHES_EXP_P0), + REPEAT_8TIMES(CEPHES_EXP_P1), + REPEAT_8TIMES(CEPHES_EXP_P2), + REPEAT_8TIMES(CEPHES_EXP_P3), + REPEAT_8TIMES(CEPHES_EXP_P4), + REPEAT_8TIMES(CEPHES_EXP_P5), + REPEAT_8TIMES(EXP_MAX_INPUT), + REPEAT_8TIMES(SIGMOID_THRESHOLD_MAX), + REPEAT_8TIMES(SIGMOID_THRESHOLD_MIN)}; -#define REPEAT_8TIMES(val) val, val, val, val, val, val, val, val - -#define OFFSET_EXP_ONE 0 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_TWO 1 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_0P5 2 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_HIG 3 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_LOW 4 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_LOG2EF 5 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_C1 6 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_C2 7 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_P0 8 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_P1 9 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_P2 10 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_P3 11 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_P4 12 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_P5 13 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_EXP_MAX_INPUT 14 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_SIGMOID_MAX 15 * YMM_FLOAT_BLOCK * sizeof(float) -#define OFFSET_SIGMOID_MIN 16 * YMM_FLOAT_BLOCK * sizeof(float) - -static const float exp_float_consts[] ALIGN32 = { - REPEAT_8TIMES(1.f), - REPEAT_8TIMES(2.f), - REPEAT_8TIMES(0.5f), - REPEAT_8TIMES(EXP_HIG), - REPEAT_8TIMES(EXP_LOW), - REPEAT_8TIMES(CEPHES_LOG2EF), - REPEAT_8TIMES(CEPHES_EXP_C1), - REPEAT_8TIMES(CEPHES_EXP_C2), - REPEAT_8TIMES(CEPHES_EXP_P0), - REPEAT_8TIMES(CEPHES_EXP_P1), - REPEAT_8TIMES(CEPHES_EXP_P2), - REPEAT_8TIMES(CEPHES_EXP_P3), - REPEAT_8TIMES(CEPHES_EXP_P4), - REPEAT_8TIMES(CEPHES_EXP_P5), - REPEAT_8TIMES(EXP_MAX_INPUT), - REPEAT_8TIMES(SIGMOID_THRESHOLD_MAX), - REPEAT_8TIMES(SIGMOID_THRESHOLD_MIN)}; - -static const int exp_int_0x7f[] ALIGN32 = {REPEAT_8TIMES(0x7f)}; -static int g_tmp_mem[16] ALIGN32 = {0}; +const int exp_int_0x7f[] ALIGN32 = {REPEAT_8TIMES(0x7f)}; +int g_tmp_mem[16] ALIGN32 = {0}; bool VActJitCode::init(int d, operand_type type) { bool ok = MayIUse(avx); @@ -177,146 +142,6 @@ bool VActJitCode::init(int d, operand_type type) { } } -void VActJitCode::exp_ymm(ymm_t& ymm_dst, ymm_t& ymm_src, int fx_idx, - int fy_idx, int mask_idx, int tmp_idx) { - assert(ymm_src.getIdx() != ymm_dst.getIdx()); // TODO(TJ): use enfore - // check all idx can not equal - ymm_t ymm_fx = ymm_t(fx_idx); - ymm_t ymm_fy = ymm_t(fy_idx); - ymm_t ymm_mask = ymm_t(mask_idx); - ymm_t ymm_tmp = ymm_t(tmp_idx); - reg64_t reg_ptr_global = rax; - push(reg_ptr_global); - mov(reg_ptr_global, reinterpret_cast(exp_float_consts)); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_HIG]); - vminps(ymm_src, ymm_src, ymm_tmp); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_LOW]); - vmaxps(ymm_src, ymm_src, ymm_tmp); - // express exp(x) as exp(g + n*log(2)) - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_LOG2EF]); - vmulps(ymm_fx, ymm_src, ymm_tmp); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_0P5]); - vaddps(ymm_fx, ymm_fx, ymm_tmp); - vroundps(ymm_fy, ymm_fx, 0x01); - // if greater, substract 1 - vcmpgtps(ymm_mask, ymm_fy, ymm_fx); - vmovaps(ymm_tmp, ptr[reg_ptr_global]); - vandps(ymm_mask, ymm_mask, ymm_tmp); - vsubps(ymm_fx, ymm_fy, ymm_mask); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_C1]); - vmulps(ymm_fy, ymm_fx, ymm_tmp); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_C2]); - ymm_t ymm_z = ymm_t(ymm_mask.getIdx()); - vmulps(ymm_z, ymm_fx, ymm_tmp); - vsubps(ymm_src, ymm_src, ymm_fy); - vsubps(ymm_src, ymm_src, ymm_z); - vmulps(ymm_z, ymm_src, ymm_src); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_P0]); - vmulps(ymm_dst, ymm_src, ymm_tmp); - for (size_t i = OFFSET_EXP_P1; i < OFFSET_EXP_P5; - i += (YMM_FLOAT_BLOCK * sizeof(float))) { - vmovaps(ymm_tmp, ptr[reg_ptr_global + i]); // P1~P4 - vaddps(ymm_dst, ymm_dst, ymm_tmp); - vmulps(ymm_dst, ymm_dst, ymm_src); - } - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_P5]); - vaddps(ymm_dst, ymm_dst, ymm_tmp); - vmulps(ymm_dst, ymm_dst, ymm_z); - vaddps(ymm_dst, ymm_dst, ymm_src); - vmovaps(ymm_tmp, ptr[reg_ptr_global]); - vaddps(ymm_dst, ymm_dst, ymm_tmp); - // build 2^n - ymm_t ymm_int = ymm_fx; - vcvttps2dq(ymm_int, ymm_fx); - mov(reg_ptr_global, reinterpret_cast(exp_int_0x7f)); - vmovdqa(ymm_tmp, ptr[reg_ptr_global]); - if (MayIUse(avx2)) { - vpaddd(ymm_int, ymm_int, ymm_tmp); - vpslld(ymm_int, ymm_int, 23); - } else if (MayIUse(avx)) { - xmm_t xtmp1 = xmm_t(ymm_int.getIdx()); - xmm_t xtmp2 = xmm_t(ymm_tmp.getIdx()); - reg64_t reg_ptr_tmp = reg_ptr_global; - mov(reg_ptr_tmp, reinterpret_cast(g_tmp_mem)); - vmovdqa(ptr[reg_ptr_tmp], ymm_int); - vmovdqa(ptr[reg_ptr_tmp + YMM_FLOAT_BLOCK * sizeof(float)], ymm_tmp); - vpaddd(xtmp1, xtmp1, xtmp2); - vpslld(xtmp1, xtmp1, 23); - vmovdqa(ptr[reg_ptr_tmp], xtmp1); - // next 128bits - vmovdqa(xtmp1, ptr[reg_ptr_tmp + 4 /*xmm float block*/ * sizeof(float)]); - vmovdqa(xtmp2, - ptr[reg_ptr_tmp + - (YMM_FLOAT_BLOCK + 4 /*xmm float block*/) * sizeof(float)]); - vpaddd(xtmp1, xtmp1, xtmp2); - vpslld(xtmp1, xtmp1, 23); - vmovdqa(ptr[reg_ptr_tmp + 4 /*xmm float block*/ * sizeof(float)], xtmp1); - // load out - vmovdqa(ymm_int, ptr[reg_ptr_tmp]); - } - vmulps(ymm_dst, ymm_dst, ymm_int); - pop(reg_ptr_global); -} - -void VActJitCode::exp_xmm(xmm_t& ymm_dst, xmm_t& ymm_src, int fx_idx, - int fy_idx, int mask_idx, int tmp_idx) { - assert(ymm_src.getIdx() != ymm_dst.getIdx()); // TODO(TJ): use enfore - // check all idx can not equal - xmm_t ymm_fx = xmm_t(fx_idx); - xmm_t ymm_fy = xmm_t(fy_idx); - xmm_t ymm_mask = xmm_t(mask_idx); - xmm_t ymm_tmp = xmm_t(tmp_idx); - reg64_t reg_ptr_global = rax; - push(reg_ptr_global); - mov(reg_ptr_global, reinterpret_cast(exp_float_consts)); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_HIG]); - vminps(ymm_src, ymm_src, ymm_tmp); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_LOW]); - vmaxps(ymm_src, ymm_src, ymm_tmp); - // express exp(x) as exp(g + n*log(2)) - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_LOG2EF]); - vmulps(ymm_fx, ymm_src, ymm_tmp); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_0P5]); - vaddps(ymm_fx, ymm_fx, ymm_tmp); - vroundps(ymm_fy, ymm_fx, 0x01); - // if greater, substract 1 - vcmpgtps(ymm_mask, ymm_fy, ymm_fx); - vmovaps(ymm_tmp, ptr[reg_ptr_global]); - vandps(ymm_mask, ymm_mask, ymm_tmp); - vsubps(ymm_fx, ymm_fy, ymm_mask); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_C1]); - vmulps(ymm_fy, ymm_fx, ymm_tmp); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_C2]); - xmm_t ymm_z = xmm_t(ymm_mask.getIdx()); - vmulps(ymm_z, ymm_fx, ymm_tmp); - vsubps(ymm_src, ymm_src, ymm_fy); - vsubps(ymm_src, ymm_src, ymm_z); - vmulps(ymm_z, ymm_src, ymm_src); - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_P0]); - vmulps(ymm_dst, ymm_src, ymm_tmp); - for (size_t i = OFFSET_EXP_P1; i < OFFSET_EXP_P5; - i += (YMM_FLOAT_BLOCK * sizeof(float))) { - vmovaps(ymm_tmp, ptr[reg_ptr_global + i]); // P1~P4 - vaddps(ymm_dst, ymm_dst, ymm_tmp); - vmulps(ymm_dst, ymm_dst, ymm_src); - } - vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_P5]); - vaddps(ymm_dst, ymm_dst, ymm_tmp); - vmulps(ymm_dst, ymm_dst, ymm_z); - vaddps(ymm_dst, ymm_dst, ymm_src); - vmovaps(ymm_tmp, ptr[reg_ptr_global]); - vaddps(ymm_dst, ymm_dst, ymm_tmp); - // build 2^n - xmm_t ymm_int = ymm_fx; - vcvttps2dq(ymm_int, ymm_fx); - mov(reg_ptr_global, reinterpret_cast(exp_int_0x7f)); - vmovdqa(ymm_tmp, ptr[reg_ptr_global]); - vpaddd(ymm_int, ymm_int, ymm_tmp); - vpslld(ymm_int, ymm_int, 23); - vmulps(ymm_dst, ymm_dst, ymm_int); - pop(reg_ptr_global); -} - void VActJitCode::sigmoid_ymm(ymm_t& ymm_dst, ymm_t& ymm_src, int fx_idx, int fy_idx, int mask_idx, int tmp_idx) { // y = 1 / (1 + e^-x) @@ -330,7 +155,7 @@ void VActJitCode::sigmoid_ymm(ymm_t& ymm_dst, ymm_t& ymm_src, int fx_idx, vmaxps(ymm_src, ymm_src, ymm_tmp); vxorps(ymm_tmp, ymm_tmp, ymm_tmp); vsubps(ymm_src, ymm_tmp, ymm_src); - exp_ymm(ymm_dst, ymm_src, fx_idx, fy_idx, mask_idx, tmp_idx); + exp_jmm(ymm_dst, ymm_src, fx_idx, fy_idx, mask_idx, tmp_idx); vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_ONE]); vaddps(ymm_dst, ymm_dst, ymm_tmp); vdivps(ymm_dst, ymm_tmp, ymm_dst); @@ -349,7 +174,7 @@ void VActJitCode::tanh_ymm(ymm_t& ymm_dst, ymm_t& ymm_src, int fx_idx, vxorps(ymm_zero, ymm_zero, ymm_zero); vsubps(ymm_tmp, ymm_zero, ymm_tmp); vmulps(ymm_src, ymm_src, ymm_tmp); - exp_ymm(ymm_dst, ymm_src, fx_idx, fy_idx, mask_idx, tmp_idx); + exp_jmm(ymm_dst, ymm_src, fx_idx, fy_idx, mask_idx, tmp_idx); vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_ONE]); vaddps(ymm_dst, ymm_dst, ymm_tmp); vmovaps(ymm_tmp, ptr[reg_ptr_global + OFFSET_EXP_TWO]); @@ -373,7 +198,7 @@ void VActJitCode::generate() { relu_jmm(ymm_dst, ymm_src, ymm_zero); break; case operand_type::exp: - exp_ymm(ymm_dst, ymm_src, 2, 3, 4, 5); + exp_jmm(ymm_dst, ymm_src, 2, 3, 4, 5); break; case operand_type::sigmoid: sigmoid_ymm(ymm_dst, ymm_src, 2, 3, 4, 5); @@ -409,7 +234,7 @@ void VActJitCode::generate() { relu_jmm(xmm_dst, xmm_src, xmm_zero); break; case operand_type::exp: - exp_xmm(xmm_dst, xmm_src, 2, 3, 4, 5); + exp_jmm(xmm_dst, xmm_src, 2, 3, 4, 5); break; default: break; diff --git a/paddle/fluid/operators/math/jit_code.h b/paddle/fluid/operators/math/jit_code.h index 6adeebca7ca..534398f4a42 100644 --- a/paddle/fluid/operators/math/jit_code.h +++ b/paddle/fluid/operators/math/jit_code.h @@ -16,6 +16,8 @@ limitations under the License. */ #include #include "paddle/fluid/operators/math/jit_gen.h" +#include "paddle/fluid/platform/cpu_info.h" + namespace paddle { namespace operators { namespace math { @@ -40,6 +42,51 @@ typedef enum { identity } operand_type; +extern const float exp_float_consts[]; +extern const int exp_int_0x7f[]; +extern int g_tmp_mem[]; + +// TODO(TJ): move these to some proper place +#define SIGMOID_THRESHOLD_MIN -40.0 +#define SIGMOID_THRESHOLD_MAX 13.0 +#define EXP_MAX_INPUT 40.0 +#define XMM_FLOAT_BLOCK 4 +#define YMM_FLOAT_BLOCK 8 +#define ZMM_FLOAT_BLOCK 16 + +#define ALIGN32 __attribute__((aligned(32))) +#define EXP_HIG 88.3762626647949f +#define EXP_LOW -88.3762626647949f +#define CEPHES_LOG2EF 1.44269504088896341 +#define CEPHES_EXP_C1 0.693359375 +#define CEPHES_EXP_C2 -2.12194440e-4 +#define CEPHES_EXP_P0 1.9875691500E-4 +#define CEPHES_EXP_P1 1.3981999507E-3 +#define CEPHES_EXP_P2 8.3334519073E-3 +#define CEPHES_EXP_P3 4.1665795894E-2 +#define CEPHES_EXP_P4 1.6666665459E-1 +#define CEPHES_EXP_P5 5.0000001201E-1 + +#define REPEAT_8TIMES(val) val, val, val, val, val, val, val, val + +#define OFFSET_EXP_ONE 0 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_TWO 1 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_0P5 2 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_HIG 3 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_LOW 4 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_LOG2EF 5 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_C1 6 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_C2 7 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_P0 8 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_P1 9 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_P2 10 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_P3 11 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_P4 12 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_P5 13 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_EXP_MAX_INPUT 14 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_SIGMOID_MAX 15 * YMM_FLOAT_BLOCK * sizeof(float) +#define OFFSET_SIGMOID_MIN 16 * YMM_FLOAT_BLOCK * sizeof(float) + // function: vec = Operand(vec(or scalar), vec(or scalar)) (maybe with relu) class VXXJitCode : public JitCode { public: @@ -134,10 +181,87 @@ class VActJitCode : public JitCode { } // compute exp with ymm, xmm - void exp_ymm(const Xbyak::Ymm& dst, const Xbyak::Ymm& src, int fx_idx = 2, - int fy_idx = 3, int mask_idx = 4, int tmp_idx = 5); - void exp_xmm(const Xbyak::Xmm& dst, const Xbyak::Xmm& src, int fx_idx = 2, - int fy_idx = 3, int mask_idx = 4, int tmp_idx = 5); + template + void exp_jmm(JMM& dst, JMM& src, int fx_idx = 2, int fy_idx = 3, // NOLINT + int mask_idx = 4, int tmp_idx = 5) { + using namespace platform::jit; // NOLINT + assert(src.getIdx() != dst.getIdx()); // TODO(TJ): use enfore + // check all idx can not equal + JMM jmm_fx = JMM(fx_idx); + JMM jmm_fy = JMM(fy_idx); + JMM jmm_mask = JMM(mask_idx); + JMM jmm_tmp = JMM(tmp_idx); + reg64_t reg_ptr_global = rax; + push(reg_ptr_global); + mov(reg_ptr_global, reinterpret_cast(exp_float_consts)); + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_HIG]); + vminps(src, src, jmm_tmp); + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_LOW]); + vmaxps(src, src, jmm_tmp); + // express exp(x) as exp(g + n*log(2)) + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_LOG2EF]); + vmulps(jmm_fx, src, jmm_tmp); + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_0P5]); + vaddps(jmm_fx, jmm_fx, jmm_tmp); + vroundps(jmm_fy, jmm_fx, 0x01); + // if greater, substract 1 + vcmpgtps(jmm_mask, jmm_fy, jmm_fx); + vmovaps(jmm_tmp, ptr[reg_ptr_global]); + vandps(jmm_mask, jmm_mask, jmm_tmp); + vsubps(jmm_fx, jmm_fy, jmm_mask); + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_C1]); + vmulps(jmm_fy, jmm_fx, jmm_tmp); + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_C2]); + JMM ymm_z = JMM(jmm_mask.getIdx()); + vmulps(ymm_z, jmm_fx, jmm_tmp); + vsubps(src, src, jmm_fy); + vsubps(src, src, ymm_z); + vmulps(ymm_z, src, src); + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_P0]); + vmulps(dst, src, jmm_tmp); + for (size_t i = OFFSET_EXP_P1; i < OFFSET_EXP_P5; + i += (YMM_FLOAT_BLOCK * sizeof(float))) { + vmovaps(jmm_tmp, ptr[reg_ptr_global + i]); // P1~P4 + vaddps(dst, dst, jmm_tmp); + vmulps(dst, dst, src); + } + vmovaps(jmm_tmp, ptr[reg_ptr_global + OFFSET_EXP_P5]); + vaddps(dst, dst, jmm_tmp); + vmulps(dst, dst, ymm_z); + vaddps(dst, dst, src); + vmovaps(jmm_tmp, ptr[reg_ptr_global]); + vaddps(dst, dst, jmm_tmp); + // build 2^n + JMM ymm_int = jmm_fx; + vcvttps2dq(ymm_int, jmm_fx); + mov(reg_ptr_global, reinterpret_cast(exp_int_0x7f)); + vmovdqa(jmm_tmp, ptr[reg_ptr_global]); + if (MayIUse(avx2) || std::is_same::value) { + vpaddd(ymm_int, ymm_int, jmm_tmp); + vpslld(ymm_int, ymm_int, 23); + } else if (MayIUse(avx)) { + xmm_t xtmp1 = xmm_t(ymm_int.getIdx()); + xmm_t xtmp2 = xmm_t(jmm_tmp.getIdx()); + reg64_t reg_ptr_tmp = reg_ptr_global; + mov(reg_ptr_tmp, reinterpret_cast(g_tmp_mem)); + vmovdqa(ptr[reg_ptr_tmp], ymm_int); + vmovdqa(ptr[reg_ptr_tmp + YMM_FLOAT_BLOCK * sizeof(float)], jmm_tmp); + vpaddd(xtmp1, xtmp1, xtmp2); + vpslld(xtmp1, xtmp1, 23); + vmovdqa(ptr[reg_ptr_tmp], xtmp1); + // next 128bits + vmovdqa(xtmp1, ptr[reg_ptr_tmp + XMM_FLOAT_BLOCK * sizeof(float)]); + vmovdqa(xtmp2, ptr[reg_ptr_tmp + + (YMM_FLOAT_BLOCK + XMM_FLOAT_BLOCK) * sizeof(float)]); + vpaddd(xtmp1, xtmp1, xtmp2); + vpslld(xtmp1, xtmp1, 23); + vmovdqa(ptr[reg_ptr_tmp + XMM_FLOAT_BLOCK * sizeof(float)], xtmp1); + // load out + vmovdqa(ymm_int, ptr[reg_ptr_tmp]); + } + vmulps(dst, dst, ymm_int); + pop(reg_ptr_global); + } // compute sigmoid with ymm void sigmoid_ymm(const Xbyak::Ymm& dst, const Xbyak::Ymm& src, int fx_idx = 2, diff --git a/paddle/fluid/operators/math/jit_kernel.h b/paddle/fluid/operators/math/jit_kernel.h index 4d8d3cd79a1..117baaee2b9 100644 --- a/paddle/fluid/operators/math/jit_kernel.h +++ b/paddle/fluid/operators/math/jit_kernel.h @@ -26,6 +26,7 @@ namespace operators { namespace math { namespace jitkernel { +// TODO(TJ): move these to some proper place #define SIGMOID_THRESHOLD_MIN -40.0 #define SIGMOID_THRESHOLD_MAX 13.0 #define EXP_MAX_INPUT 40.0 -- GitLab