From 81b128b5e97284f576ed8aaf18404980cb45fe34 Mon Sep 17 00:00:00 2001 From: Cheng Date: Thu, 18 Aug 2022 09:54:39 +0800 Subject: [PATCH] Qlinearsoftmax take FLOAT lookup-table (#12574) * [loopuptable] float-type * typed y-scale * round to nearest even --- .../cpu/quantization/qlinear_softmax.cc | 73 ++++++++++--------- .../cpu/quantization/qlinear_softmax.h | 14 +++- 2 files changed, 49 insertions(+), 38 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc index 89cbe521a5..7816510b34 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc @@ -3,6 +3,7 @@ #include "contrib_ops/cpu/quantization/qlinear_softmax.h" +#include #include #include #include @@ -23,25 +24,30 @@ constexpr int OPSET13 = 13; namespace { -void QlinearBuildLookupTableUint32(gsl::span table, +void QlinearBuildLookupTableUint32(gsl::span table, const float x_scale, size_t reduce_len, bool is_signed) { - const double qscale = - fmin(static_cast(UINT32_MAX) / static_cast(reduce_len), static_cast(0x7fffff)); + // make sure sum(exp(x)) < max() + double bit_shift = + log(std::numeric_limits::max() / reduce_len); + double reserve_bit = std::is_same_v ? 5 : 3; + bit_shift = std::max(0.0, bit_shift - reserve_bit) / x_scale; + for (int32_t i = 0; i < 256; i++) { - double scaled_exp_xi = qscale * exp(static_cast(i - 255) * static_cast(x_scale)); - // we can't get the real max value of input tensor here, so we just assume 255. + double scaled_exp_xi = exp((i - 255 + bit_shift) * static_cast(x_scale)); + // we can't get the real max value of input tensor here, so we just assume 255-bit_shift. // in the function of `QlinearSoftmaxCPU`, - // all numbers will have a shift (255-max_value) if its max value is not 255 + // all numbers will have a shift (255-bit_shift-max_value) if its max value is not 255 // // if is_signed index = [1 2 3 ......126 127 -128 -127 ..... -3 -2 -1] // else [0 1 2 3 4 ..... 256] uint8_t index = static_cast(is_signed ? i - 128 : i); - table[index] = static_cast(lrint(scaled_exp_xi)); + table[index] = static_cast((scaled_exp_xi)); } } -void BuildLookupTableIfFixed(const OpKernelInfo& info, std::vector& fixed_lookup_table, +void BuildLookupTableIfFixed(const OpKernelInfo& info, + std::vector& fixed_lookup_table, size_t reduce_len, bool is_signed) { const Tensor* tensor_x_scale = nullptr; @@ -103,9 +109,9 @@ Status QLinearSoftmax::Compute(OpKernelContext* ctx) const { return Status::OK(); } concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); - const size_t D = opset_ < OPSET13 ? X_shape.SizeFromDimension(axis_): X_shape[axis_]; - uint32_t tmp_lookup_table[256]; - gsl::span lookup_table = GetLookupTable(ctx, tmp_lookup_table, D); + const size_t D = opset_ < OPSET13 ? X_shape.SizeFromDimension(axis_) : X_shape[axis_]; + EXP_OUT_DTYPE tmp_lookup_table[256]; + gsl::span lookup_table = GetLookupTable(ctx, tmp_lookup_table, D); if (opset_ < OPSET13) { return ComputeInternal(ctx, *X, *Y, lookup_table, axis_, thread_pool); @@ -119,8 +125,8 @@ common::Status QlinearSoftmaxCPU(size_t N, size_t D, const T* x_data, T* y_data, - const uint32_t* lookup_table, - uint32_t y_scale, + const QLinearSoftmax::EXP_OUT_DTYPE* lookup_table, + QLinearSoftmax::EXP_OUT_DTYPE y_scale, T yzp, onnxruntime::concurrency::ThreadPool* thread_pool); @@ -129,8 +135,8 @@ common::Status QlinearSoftmaxCPU(size_t N, size_t D, const uint8_t* x_data, uint8_t* y_data, - const uint32_t* lookup_table, - uint32_t y_scale, + const QLinearSoftmax::EXP_OUT_DTYPE* lookup_table, + QLinearSoftmax::EXP_OUT_DTYPE y_scale, uint8_t yzp, onnxruntime::concurrency::ThreadPool* thread_pool) { using onnxruntime::TensorOpCost; @@ -156,11 +162,11 @@ common::Status QlinearSoftmaxCPU(size_t N, // 1 3 5 ... 10 // after the shift ---> // 235 237 239 .. 255 - const uint32_t* shifted_lookuptable = lookup_table + 255 - xmax; + const QLinearSoftmax::EXP_OUT_DTYPE* shifted_lookuptable = lookup_table + 255 - xmax; size_t elements_n = D; // reduceSumUin8ToUint32: need speedup // vsum = \sum_i{e^x_i} - uint32_t vsum = 0; + QLinearSoftmax::EXP_OUT_DTYPE vsum = 0; const uint8_t* x_t_cur = x_t; do { const size_t vx = *x_t_cur++; @@ -172,12 +178,11 @@ common::Status QlinearSoftmaxCPU(size_t N, elements_n = D; x_t_cur = x_t; // elementwise div, y_i=\frac{x_i}{vsum} - const uint32_t vrounding = (vsum >> 1); do { const size_t vx = *x_t_cur++; - const uint32_t vt = shifted_lookuptable[vx]; + const QLinearSoftmax::EXP_OUT_DTYPE vt = shifted_lookuptable[vx]; // simulate round function, and re-quant to uint8 - const uint32_t vq = ((vt * c_y_scale) + vrounding) / vsum + c_y_zp; + const uint32_t vq = static_cast(std::nearbyintf(((vt * c_y_scale)) / vsum)) + c_y_zp; const uint8_t vy = vq > 255 ? static_cast(255) : static_cast(vq); *y_t++ = vy; } while (--elements_n != 0); @@ -193,8 +198,8 @@ common::Status QlinearSoftmaxCPU(size_t N, size_t D, const int8_t* x_data, int8_t* y_data, - const uint32_t* lookup_table, - uint32_t y_scale, + const QLinearSoftmax::EXP_OUT_DTYPE* lookup_table, + QLinearSoftmax::EXP_OUT_DTYPE y_scale, int8_t yzp, onnxruntime::concurrency::ThreadPool* thread_pool) { using onnxruntime::TensorOpCost; @@ -215,10 +220,10 @@ common::Status QlinearSoftmaxCPU(size_t N, // reduceMaxInt8 int8_t xmax = *std::max_element(x_t, x_t + D); const size_t adjustment = 127 - xmax; - const uint32_t* shifted_lookuptable = lookup_table; + const QLinearSoftmax::EXP_OUT_DTYPE* shifted_lookuptable = lookup_table; size_t elements_n = D; // reduceSumUin8ToUint32: need speedup - uint32_t vsum = 0; + QLinearSoftmax::EXP_OUT_DTYPE vsum = 0; const int8_t* x_t_cur = x_t; do { const size_t vx = uint8_t(adjustment + (*x_t_cur++)); @@ -230,12 +235,11 @@ common::Status QlinearSoftmaxCPU(size_t N, elements_n = D; x_t_cur = x_t; // elementwise div - const uint32_t vrounding = (vsum >> 1); do { const size_t vx = uint8_t(adjustment + (*x_t_cur++)); - const uint32_t vt = shifted_lookuptable[vx]; + const QLinearSoftmax::EXP_OUT_DTYPE vt = shifted_lookuptable[vx]; // simulate round function, and re-quant to Int8 - const uint32_t vq = ((vt * c_y_scale) + vrounding) / vsum + c_y_zp; + const int32_t vq = static_cast(std::nearbyintf(((vt * c_y_scale)) / vsum)) + c_y_zp; const int8_t vy = static_cast(vq) > 255 ? static_cast(255) : static_cast(vq); *y_t++ = vy; } while (--elements_n != 0); @@ -246,10 +250,11 @@ common::Status QlinearSoftmaxCPU(size_t N, return Status::OK(); } -gsl::span QLinearSoftmax::GetLookupTable(OpKernelContext* context, - gsl::span lookup_table_span, - size_t reduce_len) const { - gsl::span lookup_table = fixed_lookup_table_; +gsl::span QLinearSoftmax::GetLookupTable( + OpKernelContext* context, + gsl::span lookup_table_span, + size_t reduce_len) const { + gsl::span lookup_table = fixed_lookup_table_; if (fixed_lookup_table_.size() == 0) { lookup_table = lookup_table_span; const float X_scale = *(context->Input(1)->Data()); @@ -260,11 +265,11 @@ gsl::span QLinearSoftmax::GetLookupTable(OpKernelContext* contex // opset-12 and below Status QLinearSoftmax::ComputeInternal(OpKernelContext* context, const Tensor& input, Tensor& output, - gsl::span lookup_table, int axis, + gsl::span lookup_table, int axis, concurrency::ThreadPool* thread_pool) const { const auto* Y_scale_tensor = context->Input(3); const auto* Y_zp_tensor = context->Input(4); - const auto Y_scale = gsl::narrow_cast(1.0F / (*(Y_scale_tensor->Data()))); + const QLinearSoftmax::EXP_OUT_DTYPE Y_scale = std::floor(1.0F / (*(Y_scale_tensor->Data()))); const auto& X_shape = input.Shape(); const size_t N = X_shape.SizeToDimension(axis); const size_t D = X_shape.SizeFromDimension(axis); @@ -286,7 +291,7 @@ Status QLinearSoftmax::ComputeInternal(OpKernelContext* context, const Tensor& i // opset-13 and above Status QLinearSoftmax::ComputeImplOpset13(OpKernelContext* context, const Tensor& input, Tensor& output, - gsl::span lookup_table, + gsl::span lookup_table, concurrency::ThreadPool* thread_pool) const { const auto& X_shape = input.Shape(); size_t rank = X_shape.NumDimensions(); diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.h b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.h index a90083cd40..08785f05a6 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.h +++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.h @@ -12,19 +12,25 @@ namespace contrib { class QLinearSoftmax final : public OpKernel { public: + using EXP_OUT_DTYPE = float; // or uint32_t if uint32_t is preferred. QLinearSoftmax(const OpKernelInfo& info); Status Compute(OpKernelContext* context) const override; private: - gsl::span GetLookupTable(OpKernelContext* context, gsl::span lookup_table_span, size_t reduce_len) const; + gsl::span GetLookupTable( + OpKernelContext* context, + gsl::span lookup_table_span, size_t reduce_len) const; - Status ComputeInternal(OpKernelContext* context, const Tensor& input, Tensor& output, gsl::span lookup_table, int axis, concurrency::ThreadPool* thread_pool) const; + Status ComputeInternal(OpKernelContext* context, const Tensor& input, + Tensor& output, gsl::span lookup_table, + int axis, concurrency::ThreadPool* thread_pool) const; Status ComputeImplOpset13(OpKernelContext* context, const Tensor& input, Tensor& output, - gsl::span lookup_table, concurrency::ThreadPool* thread_pool) const; + gsl::span lookup_table, + concurrency::ThreadPool* thread_pool) const; private: - std::vector fixed_lookup_table_; + std::vector fixed_lookup_table_; int axis_ = -1; int opset_ = 1; bool is_signed_{false};