diff --git a/onnxruntime/core/optimizer/matmul_transpose_fusion.cc b/onnxruntime/core/optimizer/matmul_transpose_fusion.cc index 6fac01a033..703babffe9 100644 --- a/onnxruntime/core/optimizer/matmul_transpose_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_transpose_fusion.cc @@ -168,6 +168,15 @@ static Node* ReorderCastAndTranspose(Graph& graph, Node* cast, return &new_transpose; } +// Check whether the element_type is an allowed FusedMatMul data type or not. +static bool IsAllowedFusedMatMulDataType(ONNX_NAMESPACE::TensorProto_DataType element_type) +{ + return element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT || + element_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 || + element_type == ONNX_NAMESPACE::TensorProto_DataType_DOUBLE || + element_type == ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16; +} + /********************************************************************************************* Case I: The followin is a scenario where Transpose output feeds MatMul. The Transpose input can be either on the left or right. @@ -277,9 +286,17 @@ Status MatmulTransposeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_ } NodeArg* left_input = node.MutableInputDefs()[0]; + auto left_type = left_input->TypeAsProto()->tensor_type().elem_type(); + if (!IsAllowedFusedMatMulDataType(static_cast(left_type))) { + continue; + } auto left = GetTransposeNodeFromOutput(graph, *left_input); NodeArg* right_input = node.MutableInputDefs()[1]; + auto right_type = right_input->TypeAsProto()->tensor_type().elem_type(); + if (!IsAllowedFusedMatMulDataType(static_cast(right_type))) { + continue; + } auto right = GetTransposeNodeFromOutput(graph, *right_input); if (!left) { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 4d5453a047..2a0872e48b 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -918,10 +918,12 @@ TEST_F(GraphTransformationTests, TransposeMatmulFusionOnThreeTranspose) { ASSERT_TRUE(static_cast(node.GetAttributes().at("transB").i())); } -TEST_F(GraphTransformationTests, TransposeMatmulNoFusionOnInvalidPerm) { +TEST_F(GraphTransformationTests, TransposeMatmulNoFusionOnInvalidInput) { const std::vector model_uris = { MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_perm.onnx", MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_default_perm.onnx", + MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_datatype_int32.onnx", + MODEL_FOLDER "fusion/transpose_matmul_4d_fusion_invalid_datatype_int64.onnx", }; for (const auto& model_uri : model_uris) { std::shared_ptr p_model; diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int32.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int32.onnx new file mode 100644 index 0000000000..7638688901 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int32.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int64.onnx b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int64.onnx new file mode 100644 index 0000000000..63728e2ff9 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int64.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py index dda0efb108..1f2bc05551 100644 --- a/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py +++ b/onnxruntime/test/testdata/transform/fusion/transpose_matmul_gen.py @@ -221,3 +221,34 @@ def gen_transpose_fusion_with_cast(model_path): gen_transpose_fusion_with_cast( "transpose_cast_matmul_4d_fusion") + +def gen_transpose_fusion_invalid_datatype(model_path, datatype): + nodes = [ + helper.make_node( + "Transpose", + ["input_0"], + ["transposed_input_0"], + perm = [0, 1, 3, 2]), + helper.make_node( + "MatMul", + ["transposed_input_0", "input_1"], + ["output"]) + ] + + inputs = [ + helper.make_tensor_value_info( + "input_0", datatype, [2, 3, 'K', 'M']), + helper.make_tensor_value_info( + "input_1", datatype, [2, 3, 'K', 'N']) + ] + + outputs = [ + helper.make_tensor_value_info( + "output", datatype, [2, 3, 'M', 'N']) + ] + + save(model_path, nodes, inputs, outputs, []) + + +gen_transpose_fusion_invalid_datatype("transpose_matmul_4d_fusion_invalid_datatype_int32.onnx", TensorProto.INT32) +gen_transpose_fusion_invalid_datatype("transpose_matmul_4d_fusion_invalid_datatype_int64.onnx", TensorProto.INT64)