[Frontend] Add ONNX importer for QLinearSoftmax (#14425)

Add a missing importer for quantized softmax.

Author: Krzysztof Parzyszek

python/tvm/relay/frontend/onnx.py

@@ -5715,6 +5715,26 @@ def _impl_v10(cls, inputs, attr, params):
         return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=dtype)
 
 
+class QLinearSoftmax(OnnxOpConverter):
+    """Operator converter for QLinearSoftmax from Microsoft onnxruntime contrib opset."""
+
+    @classmethod
+    def _impl_v1(cls, inputs, attr, params):
+        axis = attr["axis"]
+
+        x = inputs[0]
+        x_scale = get_scalar(inputs[1], params)
+        x_zero_point = get_scalar(inputs[2], params, "int32")
+        y_scale = fold_constant(get_scalar(inputs[3], params))
+        y_zero_point = get_scalar(inputs[4], params, "int32")
+
+        dtype = infer_type(x).checked_type.dtype
+
+        x = _qnn.op.dequantize(x, x_scale, x_zero_point)
+        out = _op.nn.softmax(x, axis)
+        return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=dtype)
+
+
 class QLinearConcat(OnnxOpConverter):
     """Operator converter for QLinearConcat from Microsoft onnxruntime contrib opset."""
 
@@ -6812,6 +6832,7 @@ def _get_convert_map(opset):
         "QLinearMatMul": QLinearMatMul.get_converter(opset),
         "QLinearMul": QLinearMul.get_converter(opset),
         "QLinearSigmoid": QLinearSigmoid.get_converter(opset),
+        "QLinearSoftmax": QLinearSoftmax.get_converter(opset),
         "ConvInteger": ConvInteger.get_converter(opset),
         "QLinearAveragePool": QLinearAveragePool.get_converter(opset),
         "QLinearGlobalAveragePool": QLinearGlobalAveragePool.get_converter(opset),

tests/python/frontend/onnx/test_forward.py

@@ -6990,6 +6990,31 @@ def verify_qlinearsigmoid(a_shape):
     verify_qlinearsigmoid([])
 
 
+@tvm.testing.parametrize_targets
+def test_qlinearsoftmax(target, dev):
+    """test_qlinearsoftmax"""
+
+    def verify_qlinearsoftmax(a_shape):
+
+        _ = np.random.random(a_shape).astype("float32")
+
+        input_nodes = [helper.make_tensor_value_info("a", TensorProto.FLOAT, list(a_shape))]
+
+        node = helper.make_node("Softmax", ["a"], ["B"])
+        graph = helper.make_graph(
+            [node],
+            "qlinearsoftmax_test",
+            inputs=input_nodes,
+            outputs=[helper.make_tensor_value_info("B", TensorProto.FLOAT, list(a_shape))],
+        )
+        model = helper.make_model(graph, producer_name="qlinearsoftmax_test")
+        quantize_and_verify_with_ort(model, ["a"], [a_shape], target, dev)
+
+    verify_qlinearsoftmax([4, 2])
+    verify_qlinearsoftmax([5])
+    verify_qlinearsoftmax([3, 4, 5])
+
+
 @tvm.testing.parametrize_targets("llvm")
 def test_random_bernoulli(target, dev):
     """test_random_bernoulli"""