diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 0e247e1386..46b429f21c 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -78,6 +78,7 @@ if(MSVC) ${ONNXRUNTIME_ROOT}/core/mlas/lib/dgemm.cpp ${mlas_platform_srcs_avx} ${mlas_platform_srcs_avx2} + ${ONNXRUNTIME_ROOT}/core/mlas/lib/intrinsics/avx512/quantize_avx512f.cpp ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemvU8S8KernelAvx2.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8S8KernelAvx512Core.asm @@ -286,6 +287,24 @@ else() ${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/SpoolKernelAvx512F.S ${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/TransKernelAvx512F.S ) + + check_cxx_source_compiles(" + #include + int main() { + __m512 zeros = _mm512_set1_ps(0.f); + (void)zeros; + return 0; + }" + COMPILES_AVX512F_INTRINSICS + ) + if(COMPILES_AVX512F_INTRINSICS) + set(mlas_platform_srcs_avx512f + ${ONNXRUNTIME_ROOT}/core/mlas/lib/intrinsics/avx512/quantize_avx512f.cpp + ${mlas_platform_srcs_avx512f} + ) + else() + set_source_files_properties(${mlas_common_srcs} PROPERTIES COMPILE_FLAGS "-DMLAS_AVX512F_INTRINSICS_UNSUPPORTED") + endif() if(HAS_AVX512F) set_source_files_properties(${mlas_platform_srcs_avx512f} PROPERTIES COMPILE_FLAGS "-mavx512f") endif() diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm index 9aec9a195e..f785833d98 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm @@ -166,7 +166,7 @@ ProcessNextRowM4: vpxor xmm1,xmm1,xmm1 vpxor xmm2,xmm2,xmm2 vpxor xmm3,xmm3,xmm3 - lea r13,[r8+r8*2] ; compute ldb * 3 + lea r13,[r8+r8*2] ; compute lda * 3 lea rax,[r11+r11*2] ; compute output stride * 3 mov rdx,rsi mov rcx,rdi diff --git a/onnxruntime/core/mlas/lib/intrinsics/avx512/quantize_avx512f.cpp b/onnxruntime/core/mlas/lib/intrinsics/avx512/quantize_avx512f.cpp new file mode 100644 index 0000000000..47d2dce45c --- /dev/null +++ b/onnxruntime/core/mlas/lib/intrinsics/avx512/quantize_avx512f.cpp @@ -0,0 +1,170 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + quantize_avx512f.cpp + +Abstract: + + This module implements routines to quantize buffers with AVX512F instructions. + + For quantization formula as specified in the ONNX operator documentation is: + + Output = Saturate(RoundToEven(Input / Scale) + ZeroPoint) + +--*/ + +#include "mlasi.h" + +#ifndef _MM_K0_REG16 +#define _MM_K0_REG16 0xffff +#endif + +// +// QuantizeLinear implementation using AVX512 intrinsics. +// + +template +void +MLASCALL +MlasQuantizeLinearAvx512F( + const float* Input, + OutputType* Output, + size_t N, + float Scale, + OutputType ZeroPoint + ) +/*++ + +Routine Description: + + This routine quantizes the input buffer using the supplied quantization + parameters with AVX512 instructions. + +Arguments: + + Input - Supplies the input buffer. + + Output - Supplies the output buffer. + + N - Supplies the number of elements to process. + + Scale - Supplies the quantization scale. + + ZeroPoint - Supplies the quantization zero point value. + +Return Value: + + None. + +--*/ +{ + constexpr int32_t MinimumValue = std::numeric_limits::min(); + constexpr int32_t MaximumValue = std::numeric_limits::max(); + + auto ScaleVector = _mm512_set1_ps(Scale); + auto MinimumValueVector = _mm512_set1_ps(float(MinimumValue - ZeroPoint)); + auto MaximumValueVector = _mm512_set1_ps(float(MaximumValue - ZeroPoint)); + auto ZeroPointVector = _mm512_set1_epi32(ZeroPoint); + + while (N >= 64) { + + auto FloatVector0 = _mm512_loadu_ps(Input); + auto FloatVector1 = _mm512_loadu_ps(Input + 16); + auto FloatVector2 = _mm512_loadu_ps(Input + 32); + auto FloatVector3 = _mm512_loadu_ps(Input + 48); + + FloatVector0 = _mm512_div_ps(FloatVector0, ScaleVector); + FloatVector1 = _mm512_div_ps(FloatVector1, ScaleVector); + FloatVector2 = _mm512_div_ps(FloatVector2, ScaleVector); + FloatVector3 = _mm512_div_ps(FloatVector3, ScaleVector); + + FloatVector0 = _mm512_max_ps(FloatVector0, MinimumValueVector); + FloatVector1 = _mm512_max_ps(FloatVector1, MinimumValueVector); + FloatVector2 = _mm512_max_ps(FloatVector2, MinimumValueVector); + FloatVector3 = _mm512_max_ps(FloatVector3, MinimumValueVector); + + FloatVector0 = _mm512_min_ps(FloatVector0, MaximumValueVector); + FloatVector1 = _mm512_min_ps(FloatVector1, MaximumValueVector); + FloatVector2 = _mm512_min_ps(FloatVector2, MaximumValueVector); + FloatVector3 = _mm512_min_ps(FloatVector3, MaximumValueVector); + + auto IntegerVector0 = _mm512_cvtps_epi32(FloatVector0); + auto IntegerVector1 = _mm512_cvtps_epi32(FloatVector1); + auto IntegerVector2 = _mm512_cvtps_epi32(FloatVector2); + auto IntegerVector3 = _mm512_cvtps_epi32(FloatVector3); + + IntegerVector0 = _mm512_add_epi32(IntegerVector0, ZeroPointVector); + IntegerVector1 = _mm512_add_epi32(IntegerVector1, ZeroPointVector); + IntegerVector2 = _mm512_add_epi32(IntegerVector2, ZeroPointVector); + IntegerVector3 = _mm512_add_epi32(IntegerVector3, ZeroPointVector); + + _mm512_mask_cvtepi32_storeu_epi8(Output, _MM_K0_REG16, IntegerVector0); + _mm512_mask_cvtepi32_storeu_epi8(Output + 16, _MM_K0_REG16, IntegerVector1); + _mm512_mask_cvtepi32_storeu_epi8(Output + 32, _MM_K0_REG16, IntegerVector2); + _mm512_mask_cvtepi32_storeu_epi8(Output + 48, _MM_K0_REG16, IntegerVector3); + + Input += 64; + Output += 64; + N -= 64; + } + + while (N >= 16) { + auto FloatVector = _mm512_loadu_ps(Input); + FloatVector = _mm512_div_ps(FloatVector, ScaleVector); + FloatVector = _mm512_max_ps(FloatVector, MinimumValueVector); + FloatVector = _mm512_min_ps(FloatVector, MaximumValueVector); + + auto IntegerVector = _mm512_cvtps_epi32(FloatVector); + IntegerVector = _mm512_add_epi32(IntegerVector, ZeroPointVector); + + _mm512_mask_cvtepi32_storeu_epi8(Output, _MM_K0_REG16, IntegerVector); + + Input += 16; + Output += 16; + N -= 16; + } + + if (N > 0) { + __mmask16 mask = uint16_t((uint32_t(1) << N) - uint32_t(1)); + auto FloatVector = _mm512_maskz_loadu_ps(mask, Input); + FloatVector = _mm512_div_ps(FloatVector, ScaleVector); + FloatVector = _mm512_max_ps(FloatVector, MinimumValueVector); + FloatVector = _mm512_min_ps(FloatVector, MaximumValueVector); + + auto IntegerVector = _mm512_cvtps_epi32(FloatVector); + IntegerVector = _mm512_add_epi32(IntegerVector, ZeroPointVector); + + _mm512_mask_cvtepi32_storeu_epi8(Output, mask, IntegerVector); + } +} + +void +MLASCALL +MlasQuantizeLinearU8KernelAvx512F( + const float* Input, + uint8_t* Output, + size_t N, + float Scale, + uint8_t ZeroPoint + ) +{ + MlasQuantizeLinearAvx512F(Input, Output, N, Scale, ZeroPoint); +} + +void +MLASCALL +MlasQuantizeLinearS8KernelAvx512F( + const float* Input, + int8_t* Output, + size_t N, + float Scale, + int8_t ZeroPoint + ) +{ + MlasQuantizeLinearAvx512F(Input, Output, N, Scale, ZeroPoint); +} diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 95164f7352..da03f355b1 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -489,6 +489,30 @@ void typedef MLAS_QLINEAR_BINARY_OP_U8_KERNEL* PMLAS_QLINEAR_BINARY_OP_U8_KERNEL; +typedef +void +(MLASCALL MLAS_QUANTIZE_LINEAR_U8_KERNEL)( + const float* Input, + uint8_t* Output, + size_t N, + float Scale, + uint8_t ZeroPoint + ); + +typedef MLAS_QUANTIZE_LINEAR_U8_KERNEL* PMLAS_QUANTIZE_LINEAR_U8_KERNEL; + +typedef +void +(MLASCALL MLAS_QUANTIZE_LINEAR_S8_KERNEL)( + const float* Input, + int8_t* Output, + size_t N, + float Scale, + int8_t ZeroPoint + ); + +typedef MLAS_QUANTIZE_LINEAR_S8_KERNEL* PMLAS_QUANTIZE_LINEAR_S8_KERNEL; + extern "C" { #if defined(MLAS_TARGET_AMD64_IX86) @@ -581,6 +605,8 @@ extern "C" { MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32Kernel; MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8Kernel; MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8Kernel; + MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8Kernel; + MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8Kernel; #if defined(MLAS_TARGET_AMD64) MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernelFma3; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelFma3; @@ -593,6 +619,8 @@ extern "C" { MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32KernelAvx; MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8KernelAvx2; MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8KernelAvx2; + MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8KernelAvx512F; + MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8KernelAvx512F; #endif MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32Kernel; @@ -713,6 +741,8 @@ struct MLAS_PLATFORM { PMLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL ComputeLogSoftmaxOutputF32Kernel; PMLAS_REDUCE_MAXIMUM_FLOAT_KERNEL ReduceMaximumF32Kernel; PMLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL ReduceMinimumMaximumF32Kernel; + PMLAS_QUANTIZE_LINEAR_S8_KERNEL QuantizeLinearS8Kernel; + PMLAS_QUANTIZE_LINEAR_U8_KERNEL QuantizeLinearU8Kernel; uint32_t NchwcBlockSize; uint32_t PreferredBufferAlignment; #endif diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 42c1fe4d18..64f7134eac 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -68,7 +68,7 @@ inline uint64_t MlasReadExtendedControlRegister( unsigned int ext_ctrl_reg - ) +) { #if defined(_WIN32) return _xgetbv(ext_ctrl_reg); @@ -140,6 +140,8 @@ Return Value: this->ReduceMinimumMaximumF32Kernel = MlasReduceMinimumMaximumF32Kernel; this->QLinearAddS8Kernel = MlasQLinearAddS8Kernel; this->QLinearAddU8Kernel = MlasQLinearAddU8Kernel; + this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8Kernel; + this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8Kernel; this->NchwcBlockSize = 8; this->PreferredBufferAlignment = MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT; @@ -264,6 +266,11 @@ Return Value: this->NchwcBlockSize = 16; this->PreferredBufferAlignment = 64; +#if !defined(MLAS_AVX512F_INTRINSICS_UNSUPPORTED) + this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelAvx512F; + this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelAvx512F; +#endif + // // Check if the processor supports AVX512 core features // (AVX512BW/AVX512DQ/AVX512VL). diff --git a/onnxruntime/core/mlas/lib/quantize.cpp b/onnxruntime/core/mlas/lib/quantize.cpp index b24ffe0d69..facb060218 100644 --- a/onnxruntime/core/mlas/lib/quantize.cpp +++ b/onnxruntime/core/mlas/lib/quantize.cpp @@ -133,7 +133,7 @@ MlasQuantizeLinearPackBytes( template void MLASCALL -MlasQuantizeLinear( +MlasQuantizeLinearKernel( const float* Input, OutputType* Output, size_t N, @@ -210,6 +210,70 @@ Return Value: } } +void +MLASCALL +MlasQuantizeLinearS8Kernel( + const float* Input, + int8_t* Output, + size_t N, + float Scale, + int8_t ZeroPoint + ) +{ + MlasQuantizeLinearKernel(Input, Output, N, Scale, ZeroPoint); +} + +void +MLASCALL +MlasQuantizeLinearU8Kernel( + const float* Input, + uint8_t* Output, + size_t N, + float Scale, + uint8_t ZeroPoint +) +{ + MlasQuantizeLinearKernel(Input, Output, N, Scale, ZeroPoint); +} + +template<> +void +MLASCALL +MlasQuantizeLinear( + const float* Input, + int8_t* Output, + size_t N, + float Scale, + int8_t ZeroPoint + ) +{ +#if defined(MLAS_TARGET_AMD64) + MlasPlatform.QuantizeLinearS8Kernel( +#else + MlasQuantizeLinearS8Kernel( +#endif + Input, Output, N, Scale, ZeroPoint); +} + +template<> +void +MLASCALL +MlasQuantizeLinear( + const float* Input, + uint8_t* Output, + size_t N, + float Scale, + uint8_t ZeroPoint + ) +{ +#if defined(MLAS_TARGET_AMD64) + MlasPlatform.QuantizeLinearU8Kernel( +#else + MlasQuantizeLinearU8Kernel( +#endif + Input, Output, N, Scale, ZeroPoint); +} + #else // @@ -263,8 +327,6 @@ Return Value: } } -#endif - template void MLASCALL @@ -286,6 +348,7 @@ MlasQuantizeLinear( float Scale, uint8_t ZeroPoint ); +#endif #if defined(MLAS_SSE2_INTRINSICS) diff --git a/onnxruntime/test/mlas/unittest.cpp b/onnxruntime/test/mlas/unittest.cpp index 15606b5b2b..6fe67f5b18 100644 --- a/onnxruntime/test/mlas/unittest.cpp +++ b/onnxruntime/test/mlas/unittest.cpp @@ -3033,6 +3033,80 @@ public: } }; +template +class MlasQuantizeLinearTest : public MlasTestBase +{ +private: + MatrixGuardBuffer BufferInput; + MatrixGuardBuffer BufferOutput; + MatrixGuardBuffer BufferOutputReference; + + void + GenerateReference( + const float* Input, + OutputType* OutputReference, + size_t N, + float Scale, + OutputType ZeroPoint + ) + { + for (size_t n = 0; n < N; n++) { + float FloatValue = std::nearbyintf(Input[n] / Scale) + float(ZeroPoint); + FloatValue = std::max(FloatValue, float(std::numeric_limits::min())); + FloatValue = std::min(FloatValue, float(std::numeric_limits::max())); + OutputReference[n] = (OutputType)FloatValue; + } + } + + void + Test( + size_t N + ) + { + float* Input = BufferInput.GetBuffer(N); + OutputType* Output = BufferOutput.GetBuffer(N); + OutputType* OutputReference = BufferOutputReference.GetBuffer(N); + + std::default_random_engine generator(static_cast(N)); + + std::uniform_real_distribution min_gen(-10.f, -10e-3f); + float MinimumValue = min_gen(generator); + + std::uniform_real_distribution max_gen(10e-3f, 10.f); + float MaximumValue = max_gen(generator); + + float Scale = (MaximumValue - MinimumValue) / 512.f; + + std::uniform_int_distribution zp_distribution(std::numeric_limits::min(), std::numeric_limits::max()); + OutputType ZeroPoint = static_cast(zp_distribution(generator)); + + std::uniform_real_distribution distribution(MinimumValue, MaximumValue); + for (size_t n = 0; n < N; n++) { + Input[n] = distribution(generator); + } + + GenerateReference(Input, OutputReference, N, Scale, ZeroPoint); + MlasQuantizeLinear(Input, Output, N, Scale, ZeroPoint); + + for (size_t n = 0; n < N; n++) { + if (Output[n] != OutputReference[n]) { + printf("exp difference: size=%u, index=%u, output=%d, expected=%d\n", unsigned(N), unsigned(n), Output[n], OutputReference[n]); + } + } + } + +public: + void + ExecuteShort( + void + ) override + { + for (size_t n = 1; n <= 512; n++) { + Test(n); + } + } +}; + void RunThreadedTests( void @@ -3151,6 +3225,10 @@ main( printf("MlasGlobalAveragePool tests.\n"); onnxruntime::make_unique()->ExecuteShort(); + printf("MlasQuantizeLinear tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); + printf("Done.\n"); return 0;