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 0000000000..c62056a4fe Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_input_is_graph_input.onnx differ 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 0000000000..ae17f6a087 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_1.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_2.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_2.onnx new file mode 100644 index 0000000000..1832a8c1ba Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_multiple_values_in_initializer_tensor_2.onnx differ 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 0000000000..9235b18a5b Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_overridable_initializer.onnx differ