diff --git a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc index de1798e54874f..ff52d9d0444a2 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc @@ -3,8 +3,10 @@ #include "contrib_ops/cpu/quantization/qlinear_softmax.h" +#include #include #include +#include #include #include @@ -216,6 +218,8 @@ common::Status QlinearSoftmaxCPU(size_t N, [x_data, y_data, D, y_scale, yzp, &lookup_table](std::ptrdiff_t first, std::ptrdiff_t last) { const auto c_y_scale = y_scale; const auto c_y_zp = yzp; + constexpr int32_t low = std::numeric_limits::lowest(); + constexpr int32_t high = std::numeric_limits::max(); const int8_t* x_t = x_data + first * D; int8_t* y_t = y_data + first * D; @@ -243,7 +247,7 @@ common::Status QlinearSoftmaxCPU(size_t N, const QLinearSoftmax::EXP_OUT_DTYPE vt = shifted_lookuptable[vx]; // simulate round function, and re-quant to Int8 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); + const int8_t vy = static_cast(std::clamp(vq, low, high)); *y_t++ = vy; } while (--elements_n != 0); x_t = x_t_cur; diff --git a/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc b/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc index 4980e9cfd14fd..9186acf6e348a 100644 --- a/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc +++ b/onnxruntime/test/contrib_ops/qlinear_lookup_table_test.cc @@ -289,5 +289,40 @@ TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_v12) { run_test(false); } +TEST(QLinearLookupTableBasedOperatorTests, QLinearSoftmax_Int8_Saturate) { + auto run_test = [](bool add_shape_to_input) { + OpTester test("QLinearSoftmax", 1, onnxruntime::kMSDomain); + test.AddAttribute("axis", -1); + test.AddAttribute("opset", 13); + float X_scale = 0.1f; + + int8_t X_zero_point = 0; + // Quantizer currently uses this Y_scale and Y_zero_point by default. + // See onnxruntime/python/tools/quantization/operators/softmax.py + float Y_scale = 1.0f / 256.0f; + int8_t Y_zero_point = -128; + + std::vector dims = {1, 4}; + auto x_in = std::vector{127, -128, -128, -128}; + // First element expected to saturate to 127 (would otherwise + // mathematically be 128) + auto y_out = std::vector{127, -128, -128, -128}; + + test.AddShapeToTensorData(add_shape_to_input); + test.AddInput("X", dims, x_in); + test.AddInput("X_scale", {}, {X_scale}); + test.AddInput("X_zero_point", {}, {X_zero_point}); + test.AddInput("Y_scale", {}, {Y_scale}); + test.AddInput("Y_zero_point", {}, {Y_zero_point}); + test.AddOutput("Y", dims, y_out); + auto origin_round_mode = std::fegetround(); + std::fesetround(FE_TONEAREST); + test.Run(); + std::fesetround(origin_round_mode); + }; + run_test(true); + run_test(false); +} + } // namespace test } // namespace onnxruntime