From 90294b9c433f8d2319bfa38e0a34caa707725b15 Mon Sep 17 00:00:00 2001 From: satyajandhyala Date: Sun, 28 Mar 2021 09:24:12 -0700 Subject: [PATCH] =?UTF-8?q?Fix=20Transpose=20and=20MatMul=20fusion=20code?= =?UTF-8?q?=20to=20check=20the=20input=20datatypes=20as=20=E2=80=A6=20(#71?= =?UTF-8?q?47)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix Transpose and MatMul fusion code to check the input datatypes as FusedMatMul only supports floating point datatypes. * Added testcases to make sure that the int32/int64 datatypes prevent Transport-MatMul fusion. --- .../core/optimizer/matmul_transpose_fusion.cc | 17 ++++++++++ .../test/optimizer/graph_transform_test.cc | 4 ++- ...tmul_4d_fusion_invalid_datatype_int32.onnx | Bin 0 -> 287 bytes ...tmul_4d_fusion_invalid_datatype_int64.onnx | Bin 0 -> 287 bytes .../transform/fusion/transpose_matmul_gen.py | 31 ++++++++++++++++++ 5 files changed, 51 insertions(+), 1 deletion(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int32.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/transpose_matmul_4d_fusion_invalid_datatype_int64.onnx 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 0000000000000000000000000000000000000000..76386889014baba434990fb10f0945bf2938bd15 GIT binary patch literal 287 zcmd;J7ZS+N%d03V%`3^wP1P+)EiSQo&&XxX#h#g0P+AgiAS6^$l$cjskYAjd5)Tzq z;tYX_X$f+%6r>jAIxsjeIxsshEnsBl(&fUU2WFt55Li&==#8$|db6dB@Fwx1+@^H8efn>eoic0 a3<5k(f?T}G`MG+znaM@@#rbI^0*n9%EJkVo literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..63728e2ff920a6b2711a77ca1d63ed460f18498c GIT binary patch literal 287 zcmd;J7ZS+N%d03V%`3^wP1P+)EiSQo&&XxX#h#g0P+AgiAS6^$l$cjskYAjd5)Tzq z;tYX_X$f+%6r>jAIxsjeIxsshEnsBl(&fUU2WFt55Li&==#8$|db6dB@Fwx1+@^H8efn>eoic0 a3<5k(f?T}G`MG+znaM@@#rbI^0*n9&BSvii literal 0 HcmV?d00001 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)