[CPU] Fix output saturation in int8 QLinearSoftmax

onnxruntime/contrib_ops/cpu/quantization/qlinear_softmax.cc

@@ -3,8 +3,10 @@
 
 #include "contrib_ops/cpu/quantization/qlinear_softmax.h"
 
+#include <algorithm>
 #include <cmath>
 #include <cstdint>
+#include <limits>
 #include <type_traits>
 #include <utility>
 
@@ -216,6 +218,8 @@ common::Status QlinearSoftmaxCPU<int8_t>(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<int8_t>::lowest();
+        constexpr int32_t high = std::numeric_limits<int8_t>::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<int8_t>(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<int32_t>(std::nearbyintf(((vt * c_y_scale)) / vsum)) + c_y_zp;
-            const int8_t vy = static_cast<int32_t>(vq) > 255 ? static_cast<int8_t>(255) : static_cast<int8_t>(vq);
+            const int8_t vy = static_cast<int8_t>(std::clamp(vq, low, high));
             *y_t++ = vy;
           } while (--elements_n != 0);
           x_t = x_t_cur;

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<int64_t>("axis", -1);
+    test.AddAttribute<int64_t>("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<int64_t> dims = {1, 4};
+    auto x_in = std::vector<int8_t>{127, -128, -128, -128};
+    // First element expected to saturate to 127 (would otherwise
+    // mathematically be 128)
+    auto y_out = std::vector<int8_t>{127, -128, -128, -128};
+
+    test.AddShapeToTensorData(add_shape_to_input);
+    test.AddInput<int8_t>("X", dims, x_in);
+    test.AddInput<float>("X_scale", {}, {X_scale});
+    test.AddInput<int8_t>("X_zero_point", {}, {X_zero_point});
+    test.AddInput<float>("Y_scale", {}, {Y_scale});
+    test.AddInput<int8_t>("Y_zero_point", {}, {Y_zero_point});
+    test.AddOutput<int8_t>("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