From 3fab8ebfe9166fa03aece11eb98fb6a80b6d294d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 2 May 2020 00:20:00 -0700 Subject: [PATCH] (MaximKalininMS) Fix Reshape Fusion and Crash in Reshape (#3777) * Fix a crash in Reshape Reshape doesn't handle 0 input dimension properly, which leads to a division by zero * Fix reshape fusion https://github.com/microsoft/onnxruntime/pull/3554 introduced a bug: initializers can now come before Shape->Gather->Unsqueeze chains; if those initializers have more than 1 element, expected dimensions in the chains are now incorrect. Authored-by: Max Kalinin --- onnxruntime/core/optimizer/reshape_fusion.cc | 2 +- .../providers/cpu/tensor/reshape_helper.h | 2 +- .../test/optimizer/graph_transform_test.cc | 85 ++++++++++++++++ .../providers/cpu/tensor/tensor_op_test.cc | 26 +++++ .../transform/fusion/reshape_fusion_gen.py | 92 +++++++++++++++++- .../reshape_fusion_input_is_graph_input.onnx | Bin 0 -> 499 bytes ...ltiple_values_in_initializer_tensor_1.onnx | Bin 0 -> 430 bytes ...ltiple_values_in_initializer_tensor_2.onnx | Bin 0 -> 430 bytes ...eshape_fusion_overridable_initializer.onnx | Bin 0 -> 237 bytes 9 files changed, 204 insertions(+), 3 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/reshape_fusion_input_is_graph_input.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_2.onnx create mode 100644 onnxruntime/test/testdata/transform/fusion/reshape_fusion_overridable_initializer.onnx diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 399dcf744c..2183d4a00e 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -121,7 +121,7 @@ bool ReshapeFusion::Fuse_Subgraph1(Node& reshape, Graph& graph, const logging::L return false; } - if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(i), false)) { + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(shape_value.size()), false)) { return false; } diff --git a/onnxruntime/core/providers/cpu/tensor/reshape_helper.h b/onnxruntime/core/providers/cpu/tensor/reshape_helper.h index 65cc20c7e6..1baa5a523e 100644 --- a/onnxruntime/core/providers/cpu/tensor/reshape_helper.h +++ b/onnxruntime/core/providers/cpu/tensor/reshape_helper.h @@ -34,7 +34,7 @@ class ReshapeHelper { if (unknown_dim != -1) { // calculate unknown dimension - ORT_ENFORCE((input_shape.Size() % size) == 0, + ORT_ENFORCE(size != 0 && (input_shape.Size() % size) == 0, "The input tensor cannot be reshaped to the requested shape. Input shape:", input_shape, ", requested shape:", TensorShape(requested_shape)); requested_shape[unknown_dim] = input_shape.Size() / size; } else { diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 368947d6d6..e3820593df 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1153,6 +1153,91 @@ TEST_F(GraphTransformationTests, ReshapeFusionGraphInputsTest) { ASSERT_EQ(op_to_count["Concat"], 1); ASSERT_EQ(op_to_count["Reshape"], 1); } + +TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerDoesntApplyTest) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); + Graph& graph = p_model->MainGraph(); + std::map op_to_count_orig = CountOpsInGraph(graph); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + ASSERT_TRUE(ret.IsOK()); + + // The optimization does not apply. + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count_orig, op_to_count); +} + +TEST_F(GraphTransformationTests, ReshapeFusionMultipleValuesInInitializerAppliesTest) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_multiple_values_in_initializer_tensor_2.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count["Shape"], 0); + ASSERT_EQ(op_to_count["Gather"], 0); + ASSERT_EQ(op_to_count["Unsqueeze"], 0); + ASSERT_EQ(op_to_count["Concat"], 0); + ASSERT_EQ(op_to_count["Reshape"], 1); + for (const Node& node : graph.Nodes()) { + if (node.OpType() == "Reshape") { + const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, node.InputDefs()[1]->Name()); + ASSERT_TRUE(tensor_proto != nullptr); + + auto initializer = onnxruntime::make_unique(*tensor_proto, graph.ModelPath()); + EXPECT_EQ(tensor_proto->data_type(), ONNX_NAMESPACE::TensorProto_DataType_INT64); + EXPECT_EQ(initializer->size(), 3); + + const int64_t* val = initializer->data(); + EXPECT_EQ(val[0], 1); + EXPECT_EQ(val[1], 200); + EXPECT_EQ(val[2], 0); + } + } +} + +TEST_F(GraphTransformationTests, ReshapeFusionAnotherGraphInput) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_input_is_graph_input.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); + Graph& graph = p_model->MainGraph(); + std::map op_to_count_orig = CountOpsInGraph(graph); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + ASSERT_TRUE(ret.IsOK()); + + // The optimization does not apply. + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count_orig, op_to_count); +} + +TEST_F(GraphTransformationTests, ReshapeFusionOverridableInitializer) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_overridable_initializer.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, *logger_).IsOK()); + Graph& graph = p_model->MainGraph(); + std::map op_to_count_orig = CountOpsInGraph(graph); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level1); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_); + ASSERT_TRUE(ret.IsOK()); + + // The optimization does not apply. + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_EQ(op_to_count_orig, op_to_count); +} TEST_F(GraphTransformationTests, ExpandElimination) { auto model_uri = MODEL_FOLDER "expand_elimination.onnx"; diff --git a/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc b/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc index 57619960dd..d4619b60b8 100644 --- a/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc @@ -33,6 +33,32 @@ TEST(TensorOpTest, ReshapeWithEmptyDim) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // TensorRT doesn't support empty dimension } +TEST(TensorOpTest, ReshapeWithEmptyInput) { + OpTester test("Reshape"); + test.AddInput("data", {0, 10}, std::vector()); + test.AddInput("shape", {3}, {0, 10, 1}, false); + test.AddOutput("reshaped", {0, 10, 1}, std::vector()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // TensorRT doesn't support empty dimension +} + +TEST(TensorOpTest, ReshapeWithEmptyInputAndDynamicShape) { + { + OpTester test("Reshape"); + test.AddInput("data", {1, 0}, std::vector()); + test.AddInput("shape", {3}, {1, 0, -1}, false); + test.AddOutput("reshaped", {1, 0, 1}, {}); + test.Run(OpTester::ExpectResult::kExpectFailure, "The input tensor cannot be reshaped to the requested shape", {kTensorrtExecutionProvider}); // TensorRT doesn't support empty dimension + } + + { + OpTester test("Reshape"); + test.AddInput("data", {1, 0}, std::vector()); + test.AddInput("shape", {3}, {1, 1, -1}, false); + test.AddOutput("reshaped", {1, 1, 0}, {}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // TensorRT doesn't support empty dimension + } +} + TEST(TensorOpTest, ReshapeWithInitializer) { OpTester test("Reshape"); diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py index 156d9a2147..2b3f3558d0 100644 --- a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py +++ b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py @@ -70,7 +70,98 @@ graph = helper.make_graph( save_model(graph, 'reshape_fusion_internal_node_is_graph_output.onnx') +graph = helper.make_graph( + [ # nodes + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), + helper.make_node("Gather", ["shape2_out", "indices2"], ["gather2_out"], "gather2", axis=0), + helper.make_node("Unsqueeze", ["gather2_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + helper.make_node("Concat", ["a", "unsqueeze2_out"], ["concat_out"], "concat", axis=0), + helper.make_node("Reshape", ["SubgraphRoot", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('SubgraphRoot', TensorProto.FLOAT, [10, 20, 30]), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, ['unk_0', 'unk_1', 'unk_2']), + ], + [ # initializers + helper.make_tensor('a', TensorProto.INT64, [2], [1, 200]), + helper.make_tensor('indices2', TensorProto.INT64, [], [1]), + ] +) + +save_model(graph, 'reshape_fusion_multiple_values_in_initializer_tensor_1.onnx') + +graph = helper.make_graph( + [ # nodes + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), + helper.make_node("Gather", ["shape2_out", "indices2"], ["gather2_out"], "gather2", axis=0), + helper.make_node("Unsqueeze", ["gather2_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + + helper.make_node("Concat", ["a", "unsqueeze2_out"], ["concat_out"], "concat", axis=0), + helper.make_node("Reshape", ["SubgraphRoot", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('SubgraphRoot', TensorProto.FLOAT, [10, 20, 30]), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, ['unk_0', 'unk_1', 'unk_2']), + ], + [ # initializers + helper.make_tensor('a', TensorProto.INT64, [2], [1, 200]), + helper.make_tensor('indices2', TensorProto.INT64, [], [2]), + ] +) + +save_model(graph, 'reshape_fusion_multiple_values_in_initializer_tensor_2.onnx') + +graph = helper.make_graph( + [ # nodes + helper.make_node("Shape", ["AnotherInput"], ["shape2_out"], "shape2"), + helper.make_node("Gather", ["shape2_out", "indices2"], ["gather2_out"], "gather2", axis=0), + helper.make_node("Unsqueeze", ["gather2_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + + helper.make_node("Concat", ["a", "unsqueeze2_out"], ["concat_out"], "concat", axis=0), + helper.make_node("Reshape", ["SubgraphRoot", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('SubgraphRoot', TensorProto.FLOAT, [10, 20, 30]), + helper.make_tensor_value_info('AnotherInput', TensorProto.FLOAT, ['input_dim_0', 'input_dim_1', 'input_dim_2']), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, ['unk_0', 'unk_1', 'unk_2']), + ], + [ # initializers + helper.make_tensor('a', TensorProto.INT64, [2], [1, 200]), + helper.make_tensor('indices2', TensorProto.INT64, [], [2]), + ] +) + +save_model(graph, 'reshape_fusion_input_is_graph_input.onnx') + +graph = helper.make_graph( + [ # nodes + helper.make_node("Concat", ["a"], ["concat_out"], "concat", axis=0), + helper.make_node("Reshape", ["SubgraphRoot", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('SubgraphRoot', TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info('a', TensorProto.INT64, [3]), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, ['unk_0', 'unk_1', 'unk_2']), + ], + [ # initializers + helper.make_tensor('a', TensorProto.INT64, [3], [1, 1, 2*3*4]), + ] +) + +save_model(graph, 'reshape_fusion_overridable_initializer.onnx') graph = helper.make_graph( [ # nodes @@ -96,4 +187,3 @@ graph = helper.make_graph( save_model(graph, 'reshape_fusion_with_graph_inputs.onnx') - diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_input_is_graph_input.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_input_is_graph_input.onnx new file mode 100644 index 0000000000000000000000000000000000000000..c62056a4fe0e2cbd8ec77962388a4765a5d35789 GIT binary patch literal 499 zcmZ8e%SyvQ6m?RY#G6>d;6f@WW@kh&(N%F_Q&8NtxalTL)1evEnKqLF@oW4Xzs0X` z@~9?V44iZBJ@+-XIr{0rCAecQij3yZe60!$G9zmmB|;TLv&MdZR%$SYRuF8)mn@|+ z!p@46{@P{`d@y;k_k|2yqMeKfip!0n^s{LQ|J*-#*~aP;R6ghu7Tg0O zuw%rK3Z9app1Q`qIjOYn0C(Wd)MAyBbv6~EfR+)PQ!3S4F|=}8E$&-dC^yH01|*M4 zGQk7aHau(8L*FPyBskI-8S&pi=6<&c#^BiqyI|O$-?4sY^U2QLZz0q<*Kmj(R(?dX lWbesMmxf&$Ev}%M8DGG;wx5E9`-;D6V*|o1h+@NYegSCJjduV5 literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx new file mode 100644 index 0000000000000000000000000000000000000000..ae17f6a08759240249d8b0fef2ef1adc66cee0f8 GIT binary patch literal 430 zcmZ8dJx{|h5XC8t%_S<)i2)S}gg3$h5mVKn1PS(5EG&^5tR@0EC5|P;uVLZ8@Uw99 z0SUuN@80>nyJzOkzD95X{>rR3rPyZ6qNosbCblw43R6wp2J`W1$H6W1hQM<5oTpMp z=xs!m$+FSVw&CtW&4wQAi}zd~?|l--TJ#^|#3=n{r2J?X!e7r19=lm%A@#yqP?!S| z@K1=47Ah50b9IBcdslPA18#nyJzOkzD95X{>rR3rPyZ6qNosbCblw43R6wp2J`W1$H6W1hQM<5oTpMp z=xs!m$+FSVw&CtW&4wQAi}zd~?|l--TJ#^|#3=n{r2J?X!e7r19=lm%A@#yqP?!S| z@K1=47Ah50b9IBcdslPA18#kdS*6L>#- Cd3i1X literal 0 HcmV?d00001 diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_overridable_initializer.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_overridable_initializer.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9235b18a5b2897c21bb90f859a6f1b2ca8959265 GIT binary patch literal 237 zcmd;JvwFhFrNhOTD8!YVpO>6i5}#jMBE<$_E3r9)#k9D&SQ0BTizOHqFfwtOaPb6} zCZ!i87GwnF=a+CHX%u1$N-ZwUDUo6?N-fSvEJ#&i2MT}~LVQq0yjy8;W`3R)4+paV zyA?Adql6P