From 47551da99451a1f20b34bf83e06c6ccbd638fe13 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Thu, 13 Dec 2018 22:53:40 -0800 Subject: [PATCH] Optimize Tanh/Sigmoid activations (#162) * optimized tanh/sigmoid * fix /W4 warnings from alternate build environment * use MLAS for tanh/sigmoid * fix my broken C++ templates * add x86_64 files --- cmake/onnxruntime_mlas.cmake | 6 + onnxruntime/core/mlas/inc/mlas.h | 20 ++ .../mlas/lib/amd64/LogisticKernelFma3.asm | 207 ++++++++++++++++++ .../core/mlas/lib/amd64/SgemmKernelAvx.asm | 2 +- .../mlas/lib/amd64/SgemmKernelAvx512F.asm | 2 +- .../core/mlas/lib/amd64/SgemmKernelFma3.asm | 2 +- .../core/mlas/lib/amd64/SgemmKernelSse2.asm | 2 +- .../core/mlas/lib/amd64/TanhKernelFma3.asm | 199 +++++++++++++++++ onnxruntime/core/mlas/lib/amd64/cvtfp16a.asm | 2 +- onnxruntime/core/mlas/lib/amd64/mlasi.inc | 43 ++++ onnxruntime/core/mlas/lib/amd64/sgemma.asm | 2 +- onnxruntime/core/mlas/lib/convolve.cpp | 2 +- onnxruntime/core/mlas/lib/logistic.cpp | 180 +++++++++++++++ onnxruntime/core/mlas/lib/mlasi.h | 95 +++++++- onnxruntime/core/mlas/lib/platform.cpp | 5 + onnxruntime/core/mlas/lib/tanh.cpp | 178 +++++++++++++++ onnxruntime/core/mlas/lib/threading.cpp | 3 + .../core/mlas/lib/x86_64/LogisticKernelFma3.S | 151 +++++++++++++ .../core/mlas/lib/x86_64/TanhKernelFma3.S | 143 ++++++++++++ .../providers/cpu/activation/activations.cc | 19 ++ .../providers/cpu/activation/activations.h | 6 + 21 files changed, 1261 insertions(+), 8 deletions(-) create mode 100644 onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm create mode 100644 onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm create mode 100644 onnxruntime/core/mlas/lib/amd64/mlasi.inc create mode 100644 onnxruntime/core/mlas/lib/logistic.cpp create mode 100644 onnxruntime/core/mlas/lib/tanh.cpp create mode 100644 onnxruntime/core/mlas/lib/x86_64/LogisticKernelFma3.S create mode 100644 onnxruntime/core/mlas/lib/x86_64/TanhKernelFma3.S diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index c15e794673..3339221ebb 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -8,6 +8,8 @@ set(mlas_common_srcs ${ONNXRUNTIME_ROOT}/core/mlas/lib/convolve.cpp ${ONNXRUNTIME_ROOT}/core/mlas/lib/pooling.cpp ${ONNXRUNTIME_ROOT}/core/mlas/lib/bias.cpp + ${ONNXRUNTIME_ROOT}/core/mlas/lib/logistic.cpp + ${ONNXRUNTIME_ROOT}/core/mlas/lib/tanh.cpp ) if (MSVC) @@ -61,6 +63,8 @@ if (MSVC) ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/SgemmKernelAvx512F.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/sgemma.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/cvtfp16a.asm + ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/LogisticKernelFma3.asm + ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/TanhKernelFma3.asm ) endif() @@ -114,6 +118,8 @@ else() set(mlas_platform_srcs_avx2 ${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/SgemmKernelFma3.S + ${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/LogisticKernelFma3.S + ${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/TanhKernelFma3.S ) set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma") diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 738777183e..7317516736 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -169,6 +169,26 @@ MlasBiasAdd( size_t ldc ); +// +// Miscellaneous compute routines. +// + +void +MLASCALL +MlasComputeLogistic( + const float* Input, + float* Output, + size_t N + ); + +void +MLASCALL +MlasComputeTanh( + const float* Input, + float* Output, + size_t N + ); + // // Half-precision floating-point routines. // diff --git a/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm new file mode 100644 index 0000000000..28bc3507e3 --- /dev/null +++ b/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm @@ -0,0 +1,207 @@ +;++ +; +; Copyright (c) Microsoft Corporation. All rights reserved. +; +; Licensed under the MIT License. +; +; Module Name: +; +; LogisticKernelFma3.asm +; +; Abstract: +; +; This module implements a kernel for computing the logistic function for a +; buffer of elements. +; +; This implementation uses AVX fused multiply/add instructions. +; +;-- + + .xlist +INCLUDE mlasi.inc +INCLUDE SgemmKernelCommon.inc + .list + + EXTERN MlasMaskMoveAvx:NEAR + EXTERN MlasLogisticConstants:NEAR + +; +; Structure layout for the logistic constants block. +; + +LogisticConstants STRUCT + + LowerRange DWORD ? + UpperRange DWORD ? + alpha_9 DWORD ? + alpha_7 DWORD ? + alpha_5 DWORD ? + alpha_3 DWORD ? + alpha_1 DWORD ? + beta_10 DWORD ? + beta_8 DWORD ? + beta_6 DWORD ? + beta_4 DWORD ? + beta_2 DWORD ? + beta_0 DWORD ? + one_half DWORD ? + +LogisticConstants ENDS + +; +; Stack frame layout for the logistic kernel. +; + +LogisticKernelFrame STRUCT + + SavedXmm6 OWORD ? + SavedXmm7 OWORD ? + SavedXmm8 OWORD ? + SavedXmm9 OWORD ? + SavedXmm10 OWORD ? + SavedXmm11 OWORD ? + SavedXmm12 OWORD ? + SavedXmm13 OWORD ? + SavedXmm14 OWORD ? + SavedXmm15 OWORD ? + Padding0 QWORD ? + Padding1 QWORD ? + CountN QWORD ? + ReturnAddress QWORD ? + PreviousP1Home QWORD ? + PreviousP2Home QWORD ? + PreviousP3Home QWORD ? + PreviousP4Home QWORD ? + +LogisticKernelFrame ENDS + +;++ +; +; Routine Description: +; +; This routine implements the a vectorized kernel for the logistic function. +; +; Arguments: +; +; Input (rcx) - Supplies the input buffer. +; +; Output (rdx) - Supplies the output buffer. +; +; N (r8) - Supplies the number of elements to process. +; +; Return Value: +; +; None. +; +;-- + + NESTED_ENTRY MlasLogisticKernelFma3, _TEXT + + alloc_stack (LogisticKernelFrame.ReturnAddress) + + save_xmm128_avx xmm6,LogisticKernelFrame.SavedXmm6 + save_xmm128_avx xmm7,LogisticKernelFrame.SavedXmm7 + save_xmm128_avx xmm8,LogisticKernelFrame.SavedXmm8 + save_xmm128_avx xmm9,LogisticKernelFrame.SavedXmm9 + save_xmm128_avx xmm10,LogisticKernelFrame.SavedXmm10 + save_xmm128_avx xmm11,LogisticKernelFrame.SavedXmm11 + save_xmm128_avx xmm12,LogisticKernelFrame.SavedXmm12 + save_xmm128_avx xmm13,LogisticKernelFrame.SavedXmm13 + save_xmm128_avx xmm14,LogisticKernelFrame.SavedXmm14 + save_xmm128_avx xmm15,LogisticKernelFrame.SavedXmm15 + + END_PROLOGUE + + lea rax,MlasLogisticConstants + vbroadcastss ymm4,LogisticConstants.LowerRange[rax] + vbroadcastss ymm5,LogisticConstants.UpperRange[rax] + vbroadcastss ymm6,LogisticConstants.alpha_9[rax] + vbroadcastss ymm7,LogisticConstants.alpha_7[rax] + vbroadcastss ymm8,LogisticConstants.alpha_5[rax] + vbroadcastss ymm9,LogisticConstants.alpha_3[rax] + vbroadcastss ymm10,LogisticConstants.alpha_1[rax] + vbroadcastss ymm11,LogisticConstants.beta_10[rax] + vbroadcastss ymm12,LogisticConstants.beta_6[rax] + vbroadcastss ymm13,LogisticConstants.beta_4[rax] + vbroadcastss ymm14,LogisticConstants.beta_2[rax] + vbroadcastss ymm15,LogisticConstants.beta_0[rax] + + sub r8,8 + jb ProcessRemainingCount + +ComputeLogisticBy8Loop: + vmaxps ymm0,ymm4,YMMWORD PTR [rcx] ; clamp lower bound + vmovaps ymm2,ymm7 + vminps ymm0,ymm5,ymm0 ; clamp upper bound + vmulps ymm1,ymm0,ymm0 ; x2 + vbroadcastss ymm3,LogisticConstants.beta_8[rax] + vfmadd231ps ymm2,ymm1,ymm6 ; p = x2 * alpha_9 + alpha_7 + vfmadd213ps ymm2,ymm1,ymm8 ; p = x2 * p + alpha_5 + vfmadd213ps ymm2,ymm1,ymm9 ; p = x2 * p + alpha_3 + vfmadd213ps ymm2,ymm1,ymm10 ; p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm11 ; q = x2 * beta_10 + beta_8 + vfmadd213ps ymm3,ymm1,ymm12 ; q = x2 * q + beta_6 + vfmadd213ps ymm3,ymm1,ymm13 ; q = x2 * q + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 ; q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 ; q = x2 * q + beta_0 + vmulps ymm2,ymm0,ymm2 ; p = x * p + vbroadcastss ymm0,LogisticConstants.one_half[rax] + vdivps ymm2,ymm2,ymm3 + vxorps ymm3,ymm3,ymm3 + vaddps ymm0,ymm2,ymm0 ; logistic = p / q + 0.5 + vmaxps ymm0,ymm3,ymm0 ; clamp lower bound + add rcx,8*4 ; advance input by 8 elements + vmovups YMMWORD PTR [rdx],ymm0 + add rdx,8*4 ; advance output by 8 elements + sub r8,8 + jae ComputeLogisticBy8Loop + +ProcessRemainingCount: + add r8,8 ; correct for over-subtract above + jz ExitKernel + mov DWORD PTR LogisticKernelFrame.CountN[rsp],r8d + vbroadcastss ymm2,DWORD PTR LogisticKernelFrame.CountN[rsp] + vpcmpgtd ymm2,ymm2,YMMWORD PTR [MlasMaskMoveAvx] + vmaskmovps ymm0,ymm2,YMMWORD PTR [rcx] + vmaxps ymm0,ymm4,ymm0 ; clamp lower bound + vminps ymm0,ymm5,ymm0 ; clamp upper bound + vmulps ymm1,ymm0,ymm0 ; x2 + vbroadcastss ymm3,LogisticConstants.beta_8[rax] + vfmadd231ps ymm7,ymm1,ymm6 ; p = x2 * alpha_9 + alpha_7 + vfmadd213ps ymm7,ymm1,ymm8 ; p = x2 * p + alpha_5 + vfmadd213ps ymm7,ymm1,ymm9 ; p = x2 * p + alpha_3 + vfmadd213ps ymm7,ymm1,ymm10 ; p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm11 ; q = x2 * beta_10 + beta_8 + vfmadd213ps ymm3,ymm1,ymm12 ; q = x2 * q + beta_6 + vfmadd213ps ymm3,ymm1,ymm13 ; q = x2 * q + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 ; q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 ; q = x2 * q + beta_0 + vmulps ymm7,ymm0,ymm7 ; p = x * p + vbroadcastss ymm0,LogisticConstants.one_half[rax] + vdivps ymm7,ymm7,ymm3 + vxorps ymm3,ymm3,ymm3 + vaddps ymm0,ymm7,ymm0 ; logistic = p / q + 0.5 + vmaxps ymm0,ymm3,ymm0 ; clamp lower bound + vmaskmovps YMMWORD PTR [rdx],ymm2,ymm0 + +ExitKernel: + vzeroupper + vmovaps xmm6,LogisticKernelFrame.SavedXmm6[rsp] + vmovaps xmm7,LogisticKernelFrame.SavedXmm7[rsp] + vmovaps xmm8,LogisticKernelFrame.SavedXmm8[rsp] + vmovaps xmm9,LogisticKernelFrame.SavedXmm9[rsp] + vmovaps xmm10,LogisticKernelFrame.SavedXmm10[rsp] + vmovaps xmm11,LogisticKernelFrame.SavedXmm11[rsp] + vmovaps xmm12,LogisticKernelFrame.SavedXmm12[rsp] + vmovaps xmm13,LogisticKernelFrame.SavedXmm13[rsp] + vmovaps xmm14,LogisticKernelFrame.SavedXmm14[rsp] + vmovaps xmm15,LogisticKernelFrame.SavedXmm15[rsp] + add rsp,(LogisticKernelFrame.ReturnAddress) + + BEGIN_EPILOGUE + + ret + + NESTED_END MlasLogisticKernelFma3, _TEXT + + END diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm index 4c31cc8c73..2cf23e0897 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm @@ -18,7 +18,7 @@ ;-- .xlist -INCLUDE macamd64.inc +INCLUDE mlasi.inc INCLUDE SgemmKernelCommon.inc .list diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm index 85ef906a4a..1df94ce606 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm @@ -18,7 +18,7 @@ ;-- .xlist -INCLUDE macamd64.inc +INCLUDE mlasi.inc INCLUDE SgemmKernelCommon.inc .list diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm index 895709ed31..2343cc463a 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm @@ -18,7 +18,7 @@ ;-- .xlist -INCLUDE macamd64.inc +INCLUDE mlasi.inc INCLUDE SgemmKernelCommon.inc .list diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelSse2.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelSse2.asm index c9e53fce39..fb3b3dfb7e 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelSse2.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelSse2.asm @@ -18,7 +18,7 @@ ;-- .xlist -INCLUDE macamd64.inc +INCLUDE mlasi.inc INCLUDE SgemmKernelCommon.inc .list diff --git a/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm new file mode 100644 index 0000000000..f8c423238c --- /dev/null +++ b/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm @@ -0,0 +1,199 @@ +;++ +; +; Copyright (c) Microsoft Corporation. All rights reserved. +; +; Licensed under the MIT License. +; +; Module Name: +; +; TanhKernelFma3.asm +; +; Abstract: +; +; This module implements a kernel for computing the hyperbolic tangent +; function for a buffer of elements. +; +; This implementation uses AVX fused multiply/add instructions. +; +;-- + + .xlist +INCLUDE mlasi.inc +INCLUDE SgemmKernelCommon.inc + .list + + EXTERN MlasMaskMoveAvx:NEAR + EXTERN MlasTanhConstants:NEAR + +; +; Structure layout for the tanh constants block. +; + +TanhConstants STRUCT + + LowerRange DWORD ? + UpperRange DWORD ? + alpha_13 DWORD ? + alpha_11 DWORD ? + alpha_9 DWORD ? + alpha_7 DWORD ? + alpha_5 DWORD ? + alpha_3 DWORD ? + alpha_1 DWORD ? + beta_6 DWORD ? + beta_4 DWORD ? + beta_2 DWORD ? + beta_0 DWORD ? + +TanhConstants ENDS + +; +; Stack frame layout for the tanh kernel. +; + +TanhKernelFrame STRUCT + + SavedXmm6 OWORD ? + SavedXmm7 OWORD ? + SavedXmm8 OWORD ? + SavedXmm9 OWORD ? + SavedXmm10 OWORD ? + SavedXmm11 OWORD ? + SavedXmm12 OWORD ? + SavedXmm13 OWORD ? + SavedXmm14 OWORD ? + SavedXmm15 OWORD ? + Padding0 QWORD ? + Padding1 QWORD ? + CountN QWORD ? + ReturnAddress QWORD ? + PreviousP1Home QWORD ? + PreviousP2Home QWORD ? + PreviousP3Home QWORD ? + PreviousP4Home QWORD ? + +TanhKernelFrame ENDS + +;++ +; +; Routine Description: +; +; This routine implements a vectorized kernel for the hyperbolic tangent +; function. +; +; Arguments: +; +; Input (rcx) - Supplies the input buffer. +; +; Output (rdx) - Supplies the output buffer. +; +; N (r8) - Supplies the number of elements to process. +; +; Return Value: +; +; None. +; +;-- + + NESTED_ENTRY MlasTanhKernelFma3, _TEXT + + alloc_stack (TanhKernelFrame.ReturnAddress) + + save_xmm128_avx xmm6,TanhKernelFrame.SavedXmm6 + save_xmm128_avx xmm7,TanhKernelFrame.SavedXmm7 + save_xmm128_avx xmm8,TanhKernelFrame.SavedXmm8 + save_xmm128_avx xmm9,TanhKernelFrame.SavedXmm9 + save_xmm128_avx xmm10,TanhKernelFrame.SavedXmm10 + save_xmm128_avx xmm11,TanhKernelFrame.SavedXmm11 + save_xmm128_avx xmm12,TanhKernelFrame.SavedXmm12 + save_xmm128_avx xmm13,TanhKernelFrame.SavedXmm13 + save_xmm128_avx xmm14,TanhKernelFrame.SavedXmm14 + save_xmm128_avx xmm15,TanhKernelFrame.SavedXmm15 + + END_PROLOGUE + + lea rax,MlasTanhConstants + vbroadcastss ymm4,TanhConstants.LowerRange[rax] + vbroadcastss ymm5,TanhConstants.UpperRange[rax] + vbroadcastss ymm6,TanhConstants.alpha_13[rax] + vbroadcastss ymm7,TanhConstants.alpha_11[rax] + vbroadcastss ymm8,TanhConstants.alpha_9[rax] + vbroadcastss ymm9,TanhConstants.alpha_7[rax] + vbroadcastss ymm10,TanhConstants.alpha_5[rax] + vbroadcastss ymm11,TanhConstants.alpha_3[rax] + vbroadcastss ymm12,TanhConstants.alpha_1[rax] + vbroadcastss ymm13,TanhConstants.beta_6[rax] + vbroadcastss ymm14,TanhConstants.beta_2[rax] + vbroadcastss ymm15,TanhConstants.beta_0[rax] + + sub r8,8 + jb ProcessRemainingCount + +ComputeTanhBy8Loop: + vmaxps ymm0,ymm4,YMMWORD PTR [rcx] ; clamp lower bound + vmovaps ymm2,ymm7 + vminps ymm0,ymm5,ymm0 ; clamp upper bound + vmulps ymm1,ymm0,ymm0 ; x2 + vbroadcastss ymm3,TanhConstants.beta_4[rax] + vfmadd231ps ymm2,ymm1,ymm6 ; p = x2 * alpha_13 + alpha_11 + vfmadd213ps ymm2,ymm1,ymm8 ; p = x2 * p + alpha_9 + vfmadd213ps ymm2,ymm1,ymm9 ; p = x2 * p + alpha_7 + vfmadd213ps ymm2,ymm1,ymm10 ; p = x2 * p + alpha_5 + vfmadd213ps ymm2,ymm1,ymm11 ; p = x2 * p + alpha_3 + vfmadd213ps ymm2,ymm1,ymm12 ; p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm13 ; q = x2 * beta_6 + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 ; q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 ; q = x2 * q + beta_0 + vmulps ymm2,ymm0,ymm2 ; p = x * p + vdivps ymm0,ymm2,ymm3 ; tanh = p / q + add rcx,8*4 ; advance input by 8 elements + vmovups YMMWORD PTR [rdx],ymm0 + add rdx,8*4 ; advance output by 8 elements + sub r8,8 + jae ComputeTanhBy8Loop + +ProcessRemainingCount: + add r8,8 ; correct for over-subtract above + jz ExitKernel + mov DWORD PTR TanhKernelFrame.CountN[rsp],r8d + vbroadcastss ymm2,DWORD PTR TanhKernelFrame.CountN[rsp] + vpcmpgtd ymm2,ymm2,YMMWORD PTR [MlasMaskMoveAvx] + vmaskmovps ymm0,ymm2,YMMWORD PTR [rcx] + vmaxps ymm0,ymm4,ymm0 ; clamp lower bound + vminps ymm0,ymm5,ymm0 ; clamp upper bound + vmulps ymm1,ymm0,ymm0 ; x2 + vbroadcastss ymm3,TanhConstants.beta_4[rax] + vfmadd231ps ymm7,ymm1,ymm6 ; p = x2 * alpha_13 + alpha_11 + vfmadd213ps ymm7,ymm1,ymm8 ; p = x2 * p + alpha_9 + vfmadd213ps ymm7,ymm1,ymm9 ; p = x2 * p + alpha_7 + vfmadd213ps ymm7,ymm1,ymm10 ; p = x2 * p + alpha_5 + vfmadd213ps ymm7,ymm1,ymm11 ; p = x2 * p + alpha_3 + vfmadd213ps ymm7,ymm1,ymm12 ; p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm13 ; q = x2 * beta_6 + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 ; q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 ; q = x2 * q + beta_0 + vmulps ymm7,ymm0,ymm7 ; p = x * p + vdivps ymm0,ymm7,ymm3 ; tanh = p / q + vmaskmovps YMMWORD PTR [rdx],ymm2,ymm0 + +ExitKernel: + vzeroupper + vmovaps xmm6,TanhKernelFrame.SavedXmm6[rsp] + vmovaps xmm7,TanhKernelFrame.SavedXmm7[rsp] + vmovaps xmm8,TanhKernelFrame.SavedXmm8[rsp] + vmovaps xmm9,TanhKernelFrame.SavedXmm9[rsp] + vmovaps xmm10,TanhKernelFrame.SavedXmm10[rsp] + vmovaps xmm11,TanhKernelFrame.SavedXmm11[rsp] + vmovaps xmm12,TanhKernelFrame.SavedXmm12[rsp] + vmovaps xmm13,TanhKernelFrame.SavedXmm13[rsp] + vmovaps xmm14,TanhKernelFrame.SavedXmm14[rsp] + vmovaps xmm15,TanhKernelFrame.SavedXmm15[rsp] + add rsp,(TanhKernelFrame.ReturnAddress) + + BEGIN_EPILOGUE + + ret + + NESTED_END MlasTanhKernelFma3, _TEXT + + END diff --git a/onnxruntime/core/mlas/lib/amd64/cvtfp16a.asm b/onnxruntime/core/mlas/lib/amd64/cvtfp16a.asm index 03351ced05..50315146ca 100644 --- a/onnxruntime/core/mlas/lib/amd64/cvtfp16a.asm +++ b/onnxruntime/core/mlas/lib/amd64/cvtfp16a.asm @@ -15,7 +15,7 @@ ;-- .xlist -INCLUDE macamd64.inc +INCLUDE mlasi.inc .list .const diff --git a/onnxruntime/core/mlas/lib/amd64/mlasi.inc b/onnxruntime/core/mlas/lib/amd64/mlasi.inc new file mode 100644 index 0000000000..3db76dc257 --- /dev/null +++ b/onnxruntime/core/mlas/lib/amd64/mlasi.inc @@ -0,0 +1,43 @@ +;++ +; +; Copyright (c) Microsoft Corporation. All rights reserved. +; +; Licensed under the MIT License. +; +; Module Name: +; +; mlasi.inc +; +; Abstract: +; +; This module contains common kernel macros and structures for the Microsoft +; Machine Learning algebra subprogram library. +; +;-- + + .xlist +INCLUDE macamd64.inc + .list + +; +; save_xmm128_avx +; +; Macro Description: +; +; This macro uses AVX instructions to save a vector register as part of a +; function prologue as an alternative to save_xmm128. +; +; Arguments: +; +; Register - Supplies the vector register to be saved. +; +; Offset - Supplies the stack frame offset to store the contents of the +; vector register. +; + +save_xmm128_avx MACRO Register, Offset + + vmovaps Offset[rsp], Register + .savexmm128 Register, Offset + + ENDM diff --git a/onnxruntime/core/mlas/lib/amd64/sgemma.asm b/onnxruntime/core/mlas/lib/amd64/sgemma.asm index 8522c8bdde..80805d337e 100644 --- a/onnxruntime/core/mlas/lib/amd64/sgemma.asm +++ b/onnxruntime/core/mlas/lib/amd64/sgemma.asm @@ -16,7 +16,7 @@ ;-- .xlist -INCLUDE macamd64.inc +INCLUDE mlasi.inc .list ;++ diff --git a/onnxruntime/core/mlas/lib/convolve.cpp b/onnxruntime/core/mlas/lib/convolve.cpp index b933032011..53f48ccef7 100644 --- a/onnxruntime/core/mlas/lib/convolve.cpp +++ b/onnxruntime/core/mlas/lib/convolve.cpp @@ -902,7 +902,7 @@ Return Value: int32_t TargetThreadCount = MlasPlatform.GetMaximumThreadCount(); - if (TargetThreadCount >= BatchGroupCount) { + if (size_t(TargetThreadCount) >= BatchGroupCount) { TargetThreadCount = int32_t(BatchGroupCount); } diff --git a/onnxruntime/core/mlas/lib/logistic.cpp b/onnxruntime/core/mlas/lib/logistic.cpp new file mode 100644 index 0000000000..03061bb6ba --- /dev/null +++ b/onnxruntime/core/mlas/lib/logistic.cpp @@ -0,0 +1,180 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + logistic.cpp + +Abstract: + + This module implements routines to compute the logistic function. + + This implementation uses the same polynomial coefficients and algorithm as + found in Eigen. Our usage requires building platform specific versions of + the algorithm to target different instruction sets. The implementation below + targets the base instruction set (typically SSE2) while assembly + implementations target newer instruction sets (such as FMA3). + +--*/ + +#include "mlasi.h" + +// +// Bundles the floating point constants for use by kernels written in assembly. +// + +extern "C" const struct { + float LowerRange; + float UpperRange; + float alpha_9; + float alpha_7; + float alpha_5; + float alpha_3; + float alpha_1; + float beta_10; + float beta_8; + float beta_6; + float beta_4; + float beta_2; + float beta_0; + float one_half; +} MlasLogisticConstants = { + -18.0f, + 18.0f, + 4.37031012579801e-11f, + 1.15627324459942e-07f, + 6.08574864600143e-05f, + 8.51377133304701e-03f, + 2.48287947061529e-01f, + 6.10247389755681e-13f, + 5.76102136993427e-09f, + 6.29106785017040e-06f, + 1.70198817374094e-03f, + 1.16817656904453e-01f, + 9.93151921023180e-01f, + 0.5f, +}; + +void +MLASCALL +MlasLogisticKernel( + const float* Input, + float* Output, + size_t N + ) +/*++ + +Routine Description: + + This routine implements the generic kernel for the logistic function. + +Arguments: + + Input - Supplies the input buffer. + + Output - Supplies the output buffer. + + N - Supplies the number of elements to process. + +Return Value: + + None. + +--*/ +{ + while (N >= 4) { + + MLAS_FLOAT32X4 Value = MlasLoadFloat32x4(Input); + + Value = MlasMaximumFloat32x4(MlasBroadcastFloat32x4(MlasLogisticConstants.LowerRange), Value); + Value = MlasMinimumFloat32x4(MlasBroadcastFloat32x4(MlasLogisticConstants.UpperRange), Value); + + MLAS_FLOAT32X4 ValueSquared = MlasMultiplyFloat32x4(Value, Value); + + MLAS_FLOAT32X4 p; + p = MlasMultiplyAddFloat32x4(ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.alpha_9), + MlasBroadcastFloat32x4(MlasLogisticConstants.alpha_7)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.alpha_5)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.alpha_3)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.alpha_1)); + p = MlasMultiplyFloat32x4(p, Value); + + MLAS_FLOAT32X4 q; + q = MlasMultiplyAddFloat32x4(ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.beta_10), + MlasBroadcastFloat32x4(MlasLogisticConstants.beta_8)); + q = MlasMultiplyAddFloat32x4(q, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.beta_6)); + q = MlasMultiplyAddFloat32x4(q, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.beta_4)); + q = MlasMultiplyAddFloat32x4(q, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.beta_2)); + q = MlasMultiplyAddFloat32x4(q, ValueSquared, MlasBroadcastFloat32x4(MlasLogisticConstants.beta_0)); + + MlasStoreFloat32x4(Output, MlasAddFloat32x4(MlasDivideFloat32x4(p, q), MlasBroadcastFloat32x4(0.5f))); + + Input += 4; + Output += 4; + N -= 4; + } + + while (N > 0) { + + float Value = *Input++; + + Value = (std::min)(MlasLogisticConstants.UpperRange, (std::max)(MlasLogisticConstants.LowerRange, Value)); + + float ValueSquared = Value * Value; + + float p; + p = ValueSquared * MlasLogisticConstants.alpha_9 + MlasLogisticConstants.alpha_7; + p = p * ValueSquared + MlasLogisticConstants.alpha_5; + p = p * ValueSquared + MlasLogisticConstants.alpha_3; + p = p * ValueSquared + MlasLogisticConstants.alpha_1; + p = p * Value; + + float q; + q = ValueSquared * MlasLogisticConstants.beta_10 + MlasLogisticConstants.beta_8; + q = q * ValueSquared + MlasLogisticConstants.beta_6; + q = q * ValueSquared + MlasLogisticConstants.beta_4; + q = q * ValueSquared + MlasLogisticConstants.beta_2; + q = q * ValueSquared + MlasLogisticConstants.beta_0; + + *Output++ = (p / q) + 0.5f; + + N -= 1; + } +} + +void +MLASCALL +MlasComputeLogistic( + const float* Input, + float* Output, + size_t N + ) +/*++ + +Routine Description: + + This routine computes the logistic function. + +Arguments: + + Input - Supplies the input buffer. + + Output - Supplies the output buffer. + + N - Supplies the number of elements to process. + +Return Value: + + None. + +--*/ +{ +#if defined(MLAS_TARGET_AMD64) + MlasPlatform.LogisticKernelRoutine(Input, Output, N); +#else + MlasLogisticKernel(Input, Output, N); +#endif +} diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 52f8ea3b85..bbdbe56637 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -113,7 +113,7 @@ Abstract: #define MLAS_SGEMM_STRIDEN_THREAD_ALIGN 16 // -// Define the prototypes of the SGEMM platform optimized routines. +// Define the prototypes of the platform optimized routines. // typedef @@ -156,6 +156,26 @@ void typedef MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE* PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE; +typedef +void +(MLASCALL MLAS_LOGISTIC_KERNEL_ROUTINE)( + const float* Input, + float* Output, + size_t N + ); + +typedef MLAS_LOGISTIC_KERNEL_ROUTINE* PMLAS_LOGISTIC_KERNEL_ROUTINE; + +typedef +void +(MLASCALL MLAS_TANH_KERNEL_ROUTINE)( + const float* Input, + float* Output, + size_t N + ); + +typedef MLAS_TANH_KERNEL_ROUTINE* PMLAS_TANH_KERNEL_ROUTINE; + extern "C" { MLAS_SGEMM_KERNEL_ROUTINE MlasSgemmKernelZero; @@ -183,6 +203,13 @@ extern "C" { MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Avx; #endif + MLAS_TANH_KERNEL_ROUTINE MlasLogisticKernel; + MLAS_TANH_KERNEL_ROUTINE MlasTanhKernel; +#if defined(MLAS_TARGET_AMD64) + MLAS_TANH_KERNEL_ROUTINE MlasLogisticKernelFma3; + MLAS_TANH_KERNEL_ROUTINE MlasTanhKernelFma3; +#endif + } // @@ -242,6 +269,8 @@ struct MLAS_PLATFORM { PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1Routine; PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1TransposeBRoutine; PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE TransposePackB16x4Routine; + PMLAS_LOGISTIC_KERNEL_ROUTINE LogisticKernelRoutine; + PMLAS_TANH_KERNEL_ROUTINE TanhKernelRoutine; #endif #if defined(MLAS_USE_WIN32_THREADPOOL) @@ -289,6 +318,8 @@ MlasExecuteThreaded( // Define the missing ARM64 NEON intrinsic macros from arm64_neon.h that enable // cross-compiler support. // +// Also define additional standard NEON intrinsics using the MSVC aliases. +// #if defined(_M_ARM64) #ifndef vmaxvq_f32 @@ -296,6 +327,15 @@ MlasExecuteThreaded( #endif #endif +#if defined(_M_ARM) || defined(_M_ARM64) || defined(_M_HYBRID_X86_ARM64) +#ifndef vcgezq_f32 +#define vcgezq_f32(src) vcgeq_z_f32_ex(src) +#endif +#ifndef vclezq_f32 +#define vclezq_f32(src) vcleq_z_f32_ex(src) +#endif +#endif + // // Cross-platform wrappers for vector intrinsics. // @@ -308,6 +348,15 @@ MlasExecuteThreaded( #define MLAS_NEON64_INTRINSICS #elif defined(MLAS_TARGET_AMD64_IX86) #define MLAS_SSE2_INTRINSICS +#if defined(__AVX__) +#define MLAS_AVX_INTRINSICS +#endif +#if defined(__AVX2__) +#define MLAS_AVX2_INTRINSICS +#endif +#if defined(__FMA__) || (defined(_MSC_VER) && defined(__AVX2__)) +#define MLAS_FMA3_INTRINSICS +#endif #else #error Unsupported architecture. #endif @@ -387,6 +436,18 @@ MlasStoreLaneFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector) #endif } +template +inline +float +MlasExtractLaneFloat32x4(MLAS_FLOAT32X4 Vector) +{ +#if defined(MLAS_NEON_INTRINSICS) + return vgetq_lane_f32(Vector, Lane); +#elif defined(MLAS_SSE2_INTRINSICS) + return _mm_cvtss_f32(_mm_shuffle_ps(Vector, Vector, _MM_SHUFFLE(Lane, Lane, Lane, Lane))); +#endif +} + #if defined(MLAS_SSE2_INTRINSICS) template<> @@ -397,6 +458,14 @@ MlasStoreLaneFloat32x4<0>(float* Buffer, MLAS_FLOAT32X4 Vector) _mm_store_ss(Buffer, Vector); } +template<> +inline +float +MlasExtractLaneFloat32x4<0>(MLAS_FLOAT32X4 Vector) +{ + return _mm_cvtss_f32(Vector); +} + #endif inline @@ -410,6 +479,17 @@ MlasBroadcastFloat32x4(float Value) #endif } +inline +MLAS_FLOAT32X4 +MlasBroadcastFloat32x4(const float* Value) +{ +#if defined(MLAS_NEON_INTRINSICS) + return vld1q_dup_f32(Value); +#elif defined(MLAS_SSE2_INTRINSICS) + return _mm_load_ps1(Value); +#endif +} + inline MLAS_FLOAT32X4 MlasAddFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) @@ -443,6 +523,19 @@ MlasMultiplyFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) #endif } +inline +MLAS_FLOAT32X4 +MlasMultiplyAddFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2, MLAS_FLOAT32X4 Vector3) +{ +#if defined(MLAS_NEON_INTRINSICS) + return vmlaq_f32(Vector3, Vector1, Vector2); +#elif defined(MLAS_FMA3_INTRINSICS) + return _mm_fmadd_ps(Vector1, Vector2, Vector3); +#elif defined(MLAS_SSE2_INTRINSICS) + return _mm_add_ps(_mm_mul_ps(Vector1, Vector2), Vector3); +#endif +} + inline MLAS_FLOAT32X4 MlasDivideFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 136517d72a..8d6c6e1c3f 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -53,6 +53,8 @@ Return Value: this->KernelAddRoutine = MlasSgemmKernelAddSse; #if defined(MLAS_TARGET_AMD64) this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Sse; + this->LogisticKernelRoutine = MlasLogisticKernel; + this->TanhKernelRoutine = MlasTanhKernel; #endif // @@ -109,6 +111,9 @@ Return Value: this->KernelAddRoutine = MlasSgemmKernelAddFma3; } + this->LogisticKernelRoutine = MlasLogisticKernelFma3; + this->TanhKernelRoutine = MlasTanhKernelFma3; + } else { this->KernelZeroRoutine = MlasSgemmKernelZeroAvx; diff --git a/onnxruntime/core/mlas/lib/tanh.cpp b/onnxruntime/core/mlas/lib/tanh.cpp new file mode 100644 index 0000000000..430afdf60e --- /dev/null +++ b/onnxruntime/core/mlas/lib/tanh.cpp @@ -0,0 +1,178 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + tanh.cpp + +Abstract: + + This module implements routines to compute the hyperbolic tangent function. + + This implementation uses the same polynomial coefficients and algorithm as + found in Eigen. Our usage requires building platform specific versions of + the algorithm to target different instruction sets. The implementation below + targets the base instruction set (typically SSE2) while assembly + implementations target newer instruction sets (such as FMA3). + +--*/ + +#include "mlasi.h" + +// +// Bundles the floating point constants for use by kernels written in assembly. +// + +extern "C" const struct { + float LowerRange; + float UpperRange; + float alpha_13; + float alpha_11; + float alpha_9; + float alpha_7; + float alpha_5; + float alpha_3; + float alpha_1; + float beta_6; + float beta_4; + float beta_2; + float beta_0; +} MlasTanhConstants = { + -9.0f, + 9.0f, + -2.76076847742355e-16f, + 2.00018790482477e-13f, + -8.60467152213735e-11f, + 5.12229709037114e-08f, + 1.48572235717979e-05f, + 6.37261928875436e-04f, + 4.89352455891786e-03f, + 1.19825839466702e-06f, + 1.18534705686654e-04f, + 2.26843463243900e-03f, + 4.89352518554385e-03f, +}; + +void +MLASCALL +MlasTanhKernel( + const float* Input, + float* Output, + size_t N + ) +/*++ + +Routine Description: + + This routine implements the generic kernel for the hyperbolic tangent function. + +Arguments: + + Input - Supplies the input buffer. + + Output - Supplies the output buffer. + + N - Supplies the number of elements to process. + +Return Value: + + None. + +--*/ +{ + while (N >= 4) { + + MLAS_FLOAT32X4 Value = MlasLoadFloat32x4(Input); + + Value = MlasMaximumFloat32x4(MlasBroadcastFloat32x4(MlasTanhConstants.LowerRange), Value); + Value = MlasMinimumFloat32x4(MlasBroadcastFloat32x4(MlasTanhConstants.UpperRange), Value); + + MLAS_FLOAT32X4 ValueSquared = MlasMultiplyFloat32x4(Value, Value); + + MLAS_FLOAT32X4 p; + p = MlasMultiplyAddFloat32x4(ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.alpha_13), + MlasBroadcastFloat32x4(MlasTanhConstants.alpha_11)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.alpha_9)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.alpha_7)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.alpha_5)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.alpha_3)); + p = MlasMultiplyAddFloat32x4(p, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.alpha_1)); + p = MlasMultiplyFloat32x4(p, Value); + + MLAS_FLOAT32X4 q; + q = MlasMultiplyAddFloat32x4(ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.beta_6), + MlasBroadcastFloat32x4(MlasTanhConstants.beta_4)); + q = MlasMultiplyAddFloat32x4(q, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.beta_2)); + q = MlasMultiplyAddFloat32x4(q, ValueSquared, MlasBroadcastFloat32x4(MlasTanhConstants.beta_0)); + + MlasStoreFloat32x4(Output, MlasDivideFloat32x4(p, q)); + + Input += 4; + Output += 4; + N -= 4; + } + + while (N > 0) { + + float Value = *Input++; + + Value = (std::min)(MlasTanhConstants.UpperRange, (std::max)(MlasTanhConstants.LowerRange, Value)); + + float ValueSquared = Value * Value; + + float p; + p = ValueSquared * MlasTanhConstants.alpha_13 + MlasTanhConstants.alpha_11; + p = p * ValueSquared + MlasTanhConstants.alpha_9; + p = p * ValueSquared + MlasTanhConstants.alpha_7; + p = p * ValueSquared + MlasTanhConstants.alpha_5; + p = p * ValueSquared + MlasTanhConstants.alpha_3; + p = p * ValueSquared + MlasTanhConstants.alpha_1; + p = p * Value; + + float q; + q = ValueSquared * MlasTanhConstants.beta_6 + MlasTanhConstants.beta_4; + q = q * ValueSquared + MlasTanhConstants.beta_2; + q = q * ValueSquared + MlasTanhConstants.beta_0; + + *Output++ = (p / q); + + N -= 1; + } +} + +void +MLASCALL +MlasComputeTanh( + const float* Input, + float* Output, + size_t N + ) +/*++ + +Routine Description: + + This routine computes the hyperbolic tangent function. + +Arguments: + + Input - Supplies the input buffer. + + Output - Supplies the output buffer. + + N - Supplies the number of elements to process. + +Return Value: + + None. + +--*/ +{ +#if defined(MLAS_TARGET_AMD64) + MlasPlatform.TanhKernelRoutine(Input, Output, N); +#else + MlasTanhKernel(Input, Output, N); +#endif +} diff --git a/onnxruntime/core/mlas/lib/threading.cpp b/onnxruntime/core/mlas/lib/threading.cpp index b2c27ce697..8ccadf1441 100644 --- a/onnxruntime/core/mlas/lib/threading.cpp +++ b/onnxruntime/core/mlas/lib/threading.cpp @@ -57,6 +57,9 @@ Return Value: --*/ { + MLAS_UNREFERENCED_PARAMETER(Instance); + MLAS_UNREFERENCED_PARAMETER(WorkObject); + MLAS_THREADED_WORK_BLOCK* WorkBlock = (MLAS_THREADED_WORK_BLOCK*)Context; LONG Index = InterlockedIncrement(&WorkBlock->Counter) - 1; diff --git a/onnxruntime/core/mlas/lib/x86_64/LogisticKernelFma3.S b/onnxruntime/core/mlas/lib/x86_64/LogisticKernelFma3.S new file mode 100644 index 0000000000..243b355398 --- /dev/null +++ b/onnxruntime/core/mlas/lib/x86_64/LogisticKernelFma3.S @@ -0,0 +1,151 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + LogisticKernelFma3.s + +Abstract: + + This module implements a kernel for computing the logistic function for a + buffer of elements. + + This implementation uses AVX fused multiply/add instructions. + +--*/ + +#include "asmmacro.h" + + .intel_syntax noprefix + + .text + +// +// Structure layout for the logistic constants block. +// + + .equ LogisticConstants_LowerRange, 0 + .equ LogisticConstants_UpperRange, 4 + .equ LogisticConstants_alpha_9, 8 + .equ LogisticConstants_alpha_7, 12 + .equ LogisticConstants_alpha_5, 16 + .equ LogisticConstants_alpha_3, 20 + .equ LogisticConstants_alpha_1, 24 + .equ LogisticConstants_beta_10, 28 + .equ LogisticConstants_beta_8, 32 + .equ LogisticConstants_beta_6, 36 + .equ LogisticConstants_beta_4, 40 + .equ LogisticConstants_beta_2, 44 + .equ LogisticConstants_beta_0, 48 + .equ LogisticConstants_one_half, 52 + +// +// Stack frame layout for the logistic kernel. +// + + .equ LogisticKernelFrame_CountN, -8 + .equ LogisticKernelFrame_ReturnAddress, 0 + +/*++ + +Routine Description: + + This routine implements a vectorized kernel for the logistic function. + +Arguments: + + Input (rdi) - Supplies the input buffer. + + Output (rsi) - Supplies the output buffer. + + N (rdx) - Supplies the number of elements to process. + +Return Value: + + None. + +--*/ + + .globl C_UNDERSCORE(MlasLogisticKernelFma3) +C_UNDERSCORE(MlasLogisticKernelFma3): + + lea rax,C_UNDERSCORE(MlasLogisticConstants)[rip] + vbroadcastss ymm4,LogisticConstants_LowerRange[rax] + vbroadcastss ymm5,LogisticConstants_UpperRange[rax] + vbroadcastss ymm6,LogisticConstants_alpha_9[rax] + vbroadcastss ymm7,LogisticConstants_alpha_7[rax] + vbroadcastss ymm8,LogisticConstants_alpha_5[rax] + vbroadcastss ymm9,LogisticConstants_alpha_3[rax] + vbroadcastss ymm10,LogisticConstants_alpha_1[rax] + vbroadcastss ymm11,LogisticConstants_beta_10[rax] + vbroadcastss ymm12,LogisticConstants_beta_6[rax] + vbroadcastss ymm13,LogisticConstants_beta_4[rax] + vbroadcastss ymm14,LogisticConstants_beta_2[rax] + vbroadcastss ymm15,LogisticConstants_beta_0[rax] + + sub rdx,8 + jb .LProcessRemainingCount + +.LComputeLogisticBy8Loop: + vmaxps ymm0,ymm4,YMMWORD PTR [rdi] # clamp lower bound + vmovaps ymm2,ymm7 + vminps ymm0,ymm5,ymm0 # clamp upper bound + vmulps ymm1,ymm0,ymm0 # x2 + vbroadcastss ymm3,LogisticConstants_beta_8[rax] + vfmadd231ps ymm2,ymm1,ymm6 # p = x2 * alpha_9 + alpha_7 + vfmadd213ps ymm2,ymm1,ymm8 # p = x2 * p + alpha_5 + vfmadd213ps ymm2,ymm1,ymm9 # p = x2 * p + alpha_3 + vfmadd213ps ymm2,ymm1,ymm10 # p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm11 # q = x2 * beta_10 + beta_8 + vfmadd213ps ymm3,ymm1,ymm12 # q = x2 * q + beta_6 + vfmadd213ps ymm3,ymm1,ymm13 # q = x2 * q + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 # q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 # q = x2 * q + beta_0 + vmulps ymm2,ymm0,ymm2 # p = x * p + vbroadcastss ymm0,LogisticConstants_one_half[rax] + vdivps ymm2,ymm2,ymm3 + vxorps ymm3,ymm3,ymm3 + vaddps ymm0,ymm2,ymm0 # logistic = p / q + 0.5 + vmaxps ymm0,ymm3,ymm0 # clamp lower bound + add rdi,8*4 # advance input by 8 elements + vmovups YMMWORD PTR [rsi],ymm0 + add rsi,8*4 # advance output by 8 elements + sub rdx,8 + jae .LComputeLogisticBy8Loop + +.LProcessRemainingCount: + add rdx,8 # correct for over-subtract above + jz .LExitKernel + mov DWORD PTR LogisticKernelFrame_CountN[rsp],edx + vbroadcastss ymm2,DWORD PTR LogisticKernelFrame_CountN[rsp] + vpcmpgtd ymm2,ymm2,YMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip] + vmaskmovps ymm0,ymm2,YMMWORD PTR [rdi] + vmaxps ymm0,ymm4,ymm0 # clamp lower bound + vminps ymm0,ymm5,ymm0 # clamp upper bound + vmulps ymm1,ymm0,ymm0 # x2 + vbroadcastss ymm3,LogisticConstants_beta_8[rax] + vfmadd231ps ymm7,ymm1,ymm6 # p = x2 * alpha_9 + alpha_7 + vfmadd213ps ymm7,ymm1,ymm8 # p = x2 * p + alpha_5 + vfmadd213ps ymm7,ymm1,ymm9 # p = x2 * p + alpha_3 + vfmadd213ps ymm7,ymm1,ymm10 # p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm11 # q = x2 * beta_10 + beta_8 + vfmadd213ps ymm3,ymm1,ymm12 # q = x2 * q + beta_6 + vfmadd213ps ymm3,ymm1,ymm13 # q = x2 * q + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 # q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 # q = x2 * q + beta_0 + vmulps ymm7,ymm0,ymm7 # p = x * p + vbroadcastss ymm0,LogisticConstants_one_half[rax] + vdivps ymm7,ymm7,ymm3 + vxorps ymm3,ymm3,ymm3 + vaddps ymm0,ymm7,ymm0 # logistic = p / q + 0.5 + vmaxps ymm0,ymm3,ymm0 # clamp lower bound + vmaskmovps YMMWORD PTR [rsi],ymm2,ymm0 + +.LExitKernel: + vzeroupper + ret + + .end diff --git a/onnxruntime/core/mlas/lib/x86_64/TanhKernelFma3.S b/onnxruntime/core/mlas/lib/x86_64/TanhKernelFma3.S new file mode 100644 index 0000000000..dd5584648d --- /dev/null +++ b/onnxruntime/core/mlas/lib/x86_64/TanhKernelFma3.S @@ -0,0 +1,143 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + TanhKernelFma3.s + +Abstract: + + This module implements a kernel for computing the hyperbolic tangent + function for a buffer of elements. + + This implementation uses AVX fused multiply/add instructions. + +--*/ + +#include "asmmacro.h" + + .intel_syntax noprefix + + .text + +// +// Structure layout for the tanh constants block. +// + + .equ TanhConstants_LowerRange, 0 + .equ TanhConstants_UpperRange, 4 + .equ TanhConstants_alpha_13, 8 + .equ TanhConstants_alpha_11, 12 + .equ TanhConstants_alpha_9, 16 + .equ TanhConstants_alpha_7, 20 + .equ TanhConstants_alpha_5, 24 + .equ TanhConstants_alpha_3, 28 + .equ TanhConstants_alpha_1, 32 + .equ TanhConstants_beta_6, 36 + .equ TanhConstants_beta_4, 40 + .equ TanhConstants_beta_2, 44 + .equ TanhConstants_beta_0, 48 + +// +// Stack frame layout for the tanh kernel. +// + + .equ TanhKernelFrame_CountN, -8 + .equ TanhKernelFrame_ReturnAddress, 0 + +/*++ + +Routine Description: + + This routine implements a vectorized kernel for the hyperbolic tangent + function. + +Arguments: + + Input (rdi) - Supplies the input buffer. + + Output (rsi) - Supplies the output buffer. + + N (rdx) - Supplies the number of elements to process. + +Return Value: + + None. + +--*/ + + .globl C_UNDERSCORE(MlasTanhKernelFma3) +C_UNDERSCORE(MlasTanhKernelFma3): + + lea rax,C_UNDERSCORE(MlasTanhConstants)[rip] + vbroadcastss ymm4,TanhConstants_LowerRange[rax] + vbroadcastss ymm5,TanhConstants_UpperRange[rax] + vbroadcastss ymm6,TanhConstants_alpha_13[rax] + vbroadcastss ymm7,TanhConstants_alpha_11[rax] + vbroadcastss ymm8,TanhConstants_alpha_9[rax] + vbroadcastss ymm9,TanhConstants_alpha_7[rax] + vbroadcastss ymm10,TanhConstants_alpha_5[rax] + vbroadcastss ymm11,TanhConstants_alpha_3[rax] + vbroadcastss ymm12,TanhConstants_alpha_1[rax] + vbroadcastss ymm13,TanhConstants_beta_6[rax] + vbroadcastss ymm14,TanhConstants_beta_2[rax] + vbroadcastss ymm15,TanhConstants_beta_0[rax] + + sub rdx,8 + jb .LProcessRemainingCount + +.LComputeTanhBy8Loop: + vmaxps ymm0,ymm4,YMMWORD PTR [rdi] # clamp lower bound + vmovaps ymm2,ymm7 + vminps ymm0,ymm5,ymm0 # clamp upper bound + vmulps ymm1,ymm0,ymm0 # x2 + vbroadcastss ymm3,TanhConstants_beta_4[rax] + vfmadd231ps ymm2,ymm1,ymm6 # p = x2 * alpha_13 + alpha_11 + vfmadd213ps ymm2,ymm1,ymm8 # p = x2 * p + alpha_9 + vfmadd213ps ymm2,ymm1,ymm9 # p = x2 * p + alpha_7 + vfmadd213ps ymm2,ymm1,ymm10 # p = x2 * p + alpha_5 + vfmadd213ps ymm2,ymm1,ymm11 # p = x2 * p + alpha_3 + vfmadd213ps ymm2,ymm1,ymm12 # p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm13 # q = x2 * beta_6 + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 # q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 # q = x2 * q + beta_0 + vmulps ymm2,ymm0,ymm2 # p = x * p + vdivps ymm0,ymm2,ymm3 # tanh = p / q + add rdi,8*4 # advance input by 8 elements + vmovups YMMWORD PTR [rsi],ymm0 + add rsi,8*4 # advance output by 8 elements + sub rdx,8 + jae .LComputeTanhBy8Loop + +.LProcessRemainingCount: + add rdx,8 # correct for over-subtract above + jz .LExitKernel + mov DWORD PTR TanhKernelFrame_CountN[rsp],edx + vbroadcastss ymm2,DWORD PTR TanhKernelFrame_CountN[rsp] + vpcmpgtd ymm2,ymm2,YMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip] + vmaskmovps ymm0,ymm2,YMMWORD PTR [rdi] + vmaxps ymm0,ymm4,ymm0 # clamp lower bound + vminps ymm0,ymm5,ymm0 # clamp upper bound + vmulps ymm1,ymm0,ymm0 # x2 + vbroadcastss ymm3,TanhConstants_beta_4[rax] + vfmadd231ps ymm7,ymm1,ymm6 # p = x2 * alpha_13 + alpha_11 + vfmadd213ps ymm7,ymm1,ymm8 # p = x2 * p + alpha_9 + vfmadd213ps ymm7,ymm1,ymm9 # p = x2 * p + alpha_7 + vfmadd213ps ymm7,ymm1,ymm10 # p = x2 * p + alpha_5 + vfmadd213ps ymm7,ymm1,ymm11 # p = x2 * p + alpha_3 + vfmadd213ps ymm7,ymm1,ymm12 # p = x2 * p + alpha_1 + vfmadd231ps ymm3,ymm1,ymm13 # q = x2 * beta_6 + beta_4 + vfmadd213ps ymm3,ymm1,ymm14 # q = x2 * q + beta_2 + vfmadd213ps ymm3,ymm1,ymm15 # q = x2 * q + beta_0 + vmulps ymm7,ymm0,ymm7 # p = x * p + vdivps ymm0,ymm7,ymm3 # tanh = p / q + vmaskmovps YMMWORD PTR [rsi],ymm2,ymm0 + +.LExitKernel: + vzeroupper + ret + + .end diff --git a/onnxruntime/core/providers/cpu/activation/activations.cc b/onnxruntime/core/providers/cpu/activation/activations.cc index 7d381a9edc..145831ec99 100644 --- a/onnxruntime/core/providers/cpu/activation/activations.cc +++ b/onnxruntime/core/providers/cpu/activation/activations.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cpu/activation/activations.h" +#include "core/mlas/inc/mlas.h" namespace onnxruntime { @@ -29,4 +30,22 @@ REGISTER_UNARY_ELEMENTWISE_KERNEL(Softsign, 1); REGISTER_UNARY_ELEMENTWISE_KERNEL(Tanh, 6); REGISTER_UNARY_ELEMENTWISE_KERNEL(ThresholdedRelu, 1); +template <> +Status Sigmoid::Compute(OpKernelContext* context) const { + const Tensor* X = context->Input(0); + const auto& x_shape = X->Shape(); + Tensor* Y = context->Output(0, x_shape); + MlasComputeLogistic(X->template Data(), Y->template MutableData(), x_shape.Size()); + return Status::OK(); +} + +template <> +Status Tanh::Compute(OpKernelContext* context) const { + const Tensor* X = context->Input(0); + const auto& x_shape = X->Shape(); + Tensor* Y = context->Output(0, x_shape); + MlasComputeTanh(X->template Data(), Y->template MutableData(), x_shape.Size()); + return Status::OK(); +} + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/activation/activations.h b/onnxruntime/core/providers/cpu/activation/activations.h index 7b5e213d8c..42c434a2e2 100644 --- a/onnxruntime/core/providers/cpu/activation/activations.h +++ b/onnxruntime/core/providers/cpu/activation/activations.h @@ -157,6 +157,9 @@ class Sigmoid final : public OpKernel { } }; +template <> +Status Sigmoid::Compute(OpKernelContext* context) const; + template class Softsign final : public OpKernel { public: @@ -184,6 +187,9 @@ class Tanh final : public OpKernel { } }; +template <> +Status Tanh::Compute(OpKernelContext* context) const; + template class ThresholdedRelu final : public OpKernel { public: