saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Implement QuantizeLinear with avx512 (#6260)

Browse source 
* Implement QuantizeLinear with AVX512
This commit is contained in:
Yufeng Li 2021-02-01 14:33:44 -08:00 committed by GitHub
parent 5b69cbe80e
commit 7264a067a9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 372 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
cmake/onnxruntime_mlas.cmake
View file

@ -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 <immintrin.h>
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()

2
onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm
View file

@ -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

170
onnxruntime/core/mlas/lib/intrinsics/avx512/quantize_avx512f.cpp Normal file
View file

@ -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 <typename OutputType>
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<OutputType>::min();
constexpr int32_t MaximumValue = std::numeric_limits<OutputType>::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<uint8_t>(Input, Output, N, Scale, ZeroPoint);
}
void
MLASCALL
MlasQuantizeLinearS8KernelAvx512F(
const float* Input,
int8_t* Output,
size_t N,
float Scale,
int8_t ZeroPoint
)
{
MlasQuantizeLinearAvx512F<int8_t>(Input, Output, N, Scale, ZeroPoint);
}

30
onnxruntime/core/mlas/lib/mlasi.h
View file

@ -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

9
onnxruntime/core/mlas/lib/platform.cpp
View file

@ -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).

69
onnxruntime/core/mlas/lib/quantize.cpp
View file

@ -133,7 +133,7 @@ MlasQuantizeLinearPackBytes<int8_t>(
template<typename OutputType>
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<int8_t>(Input, Output, N, Scale, ZeroPoint);
}
void
MLASCALL
MlasQuantizeLinearU8Kernel(
const float* Input,
uint8_t* Output,
size_t N,
float Scale,
uint8_t ZeroPoint
)
{
MlasQuantizeLinearKernel<uint8_t>(Input, Output, N, Scale, ZeroPoint);
}
template<>
void
MLASCALL
MlasQuantizeLinear<int8_t>(
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<uint8_t>(
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<uint8_t>(
float Scale,
uint8_t ZeroPoint
);
#endif
#if defined(MLAS_SSE2_INTRINSICS)

78
onnxruntime/test/mlas/unittest.cpp
View file

@ -3033,6 +3033,80 @@ public:
}
};
template<typename OutputType>
class MlasQuantizeLinearTest : public MlasTestBase
{
private:
MatrixGuardBuffer<float> BufferInput;
MatrixGuardBuffer<OutputType> BufferOutput;
MatrixGuardBuffer<OutputType> 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<OutputType>::min()));
FloatValue = std::min(FloatValue, float(std::numeric_limits<OutputType>::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<unsigned>(N));
std::uniform_real_distribution<float> min_gen(-10.f, -10e-3f);
float MinimumValue = min_gen(generator);
std::uniform_real_distribution<float> max_gen(10e-3f, 10.f);
float MaximumValue = max_gen(generator);
float Scale = (MaximumValue - MinimumValue) / 512.f;
std::uniform_int_distribution<int32_t> zp_distribution(std::numeric_limits<OutputType>::min(), std::numeric_limits<OutputType>::max());
OutputType ZeroPoint = static_cast<OutputType>(zp_distribution(generator));
std::uniform_real_distribution<float> 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<MlasQLinearGlobalAveragePoolU8Test>()->ExecuteShort();
printf("MlasQuantizeLinear tests.\n");
onnxruntime::make_unique<MlasQuantizeLinearTest<int8_t>>()->ExecuteShort();
onnxruntime::make_unique<MlasQuantizeLinearTest<uint8_t>>()->ExecuteShort();
printf("Done.\n");
return 0;