Fuse MatMulIntegerToFloat only when scales are scalar (microsoft#6008)

MatMulIntegerToFloat fusion fuses per-row and per-column MatMulInteger, which is not supported by the MatMulIntegerToFloat kernel now. Limit the fusion to per-matrix only before we supporting the per-channel fully.

Author: yufenglee

onnxruntime/core/optimizer/matmul_integer_to_float.cc

@@ -34,7 +34,7 @@ static bool CheckBiasShape(const TensorShapeProto* bias_shape) {
 /**
 MatMulIntegerToFloatFusion will fuse subgraph like below into MatMulIntegerToFloat:
  
- A   A_Zero B B_Zero  A_Scale) B_Scale  Bias (Const, Optional)
+ A   A_Zero B B_Zero  A_Scale  B_Scale  Bias (Const, Optional)
   \    |    |    /        \      /             |
    \   |    |   /          \    /              |
     \  |    |  /            \  /               |  
@@ -84,6 +84,13 @@ Status MatMulIntegerToFloatFusion::ApplyImpl(Graph& graph, bool& modified, int g
       continue;
     }
 
+    // A_Scale is scalar and B_Scale is scalar or 1D tensor
+    auto mul_node_input_defs = p_mul_node_right->InputDefs();
+    if (!optimizer_utils::IsScalar(*mul_node_input_defs[0]) ||
+        !optimizer_utils::IsScalar(*mul_node_input_defs[1])) {
+      continue;
+    }
+
     Node& cast_node = *graph.GetNode(p_cast_node->Index());
     Node& matmulinteger_node = *graph.GetNode(p_matmulinteger_node->Index());
     Node& mul_node_right = *graph.GetNode(p_mul_node_right->Index());

onnxruntime/core/optimizer/utils.cc

@@ -24,7 +24,7 @@ bool IsFloatingPointDataType(const ONNX_NAMESPACE::TensorProto& tensor_proto) {
   return tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT || tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 || tensor_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE;
 }
 
-inline bool IsScalar(const NodeArg& input_arg) {
+bool IsScalar(const NodeArg& input_arg) {
   auto shape = input_arg.Shape();
   if (shape == nullptr) {
     // shape inferencing wasn't able to populate shape information for this NodeArg

onnxruntime/core/optimizer/utils.h

@@ -15,6 +15,9 @@ namespace optimizer_utils {
 // Check if TensorProto contains a floating point type.
 bool IsFloatingPointDataType(const ONNX_NAMESPACE::TensorProto& tensor_proto);
 
+// Check if input is a scalar
+bool IsScalar(const NodeArg& input_arg);
+
 /** Check whether a input is initializer with specified float value.
 @param expected_value is the expected value of the initializer.
 @param is_constant means whether the initializer is required to be constant.
@@ -60,7 +63,7 @@ bool ValidateShape(const NodeArg& node_arg, const std::initializer_list<int64_t>
 */
 bool CompareShape(const ONNX_NAMESPACE::TensorShapeProto& node_arg_shape, const ONNX_NAMESPACE::TensorShapeProto& node_arg_other_shape);
 
-/** Check check whether each dimension is known for shape of node_arg
+/** Check whether each dimension is known for shape of node_arg
 @returns false when shape is nullptr, or total dimension is not same as expected_dim_size length,
          or any dim is unknown (without dim value).
 */

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -3069,9 +3069,9 @@ TEST_F(GraphTransformationTests, MatMulIntegerToFloatTest) {
 
   std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
   EXPECT_EQ(op_to_count["DynamicQuantizeLinear"], 1);
-  EXPECT_EQ(op_to_count["MatMulInteger"], 0);
-  EXPECT_EQ(op_to_count["Cast"], 0);
-  EXPECT_EQ(op_to_count["Mul"], 0);
+  EXPECT_EQ(op_to_count["MatMulInteger"], 1);
+  EXPECT_EQ(op_to_count["Cast"], 1);
+  EXPECT_EQ(op_to_count["Mul"], 2);
   EXPECT_EQ(op_to_count["com.microsoft.MatMulIntegerToFloat"], 3);
   EXPECT_EQ(op_to_count["Add"], 1);
 }

onnxruntime/test/testdata/transform/fusion/matmul_integer_to_float.onnx

@@ -29,6 +29,7 @@ def GenerateModel(model_name):
     nodes.extend(MakeSubGraph("_1", True))
     nodes.extend(MakeSubGraph("_2", True))
     nodes.extend(MakeSubGraph("_3", False))
+    nodes.extend(MakeSubGraph("_4", False))
 
     initializers = []
     initializers.extend(MakeInitializer("_1"))
@@ -48,11 +49,15 @@ def GenerateModel(model_name):
             helper.make_tensor_value_info('b_quantized_2', TensorProto.UINT8, [2, 3]),
             helper.make_tensor_value_info('b_zp_2', TensorProto.UINT8, [1]),
             helper.make_tensor_value_info('b_scale_2', TensorProto.FLOAT, [1]),
+            helper.make_tensor_value_info('b_quantized_4', TensorProto.UINT8, [2, 3]),
+            helper.make_tensor_value_info('b_zp_4', TensorProto.UINT8, [3]),
+            helper.make_tensor_value_info('b_scale_4', TensorProto.FLOAT, [3]),
         ],
         [  # outputs
             helper.make_tensor_value_info('output_1', TensorProto.FLOAT, [3, 3]),
             helper.make_tensor_value_info('output_2', TensorProto.FLOAT, [3, 3]),
             helper.make_tensor_value_info('output_3', TensorProto.FLOAT, [3, 3]),
+            helper.make_tensor_value_info('output_4', TensorProto.FLOAT, [3, 3]),
         ],
         initializers)