diff --git a/onnxruntime/core/optimizer/reshape_fusion.cc b/onnxruntime/core/optimizer/reshape_fusion.cc index 3e2389495a..0721fc1cef 100644 --- a/onnxruntime/core/optimizer/reshape_fusion.cc +++ b/onnxruntime/core/optimizer/reshape_fusion.cc @@ -39,9 +39,99 @@ Status ReshapeFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, c return Status::OK(); } + +/** +Provide check for Reshape Fusion for DistilBert. The following are subgraphs that +match the pattern for DistilBert + +DistilBert reshape pattern: + [Root] + / \ _ _ _ _ _ _ _ + Shape \ + | MatMul(w * w) + Gather(indices=0) | + \ | + Unsqueeze Add(w) + \ / + Concat _ _ _ _ / + \ / + Reshape + + - -> Shape + | +Check the subgraph that matches [root] -> MatMul(w * w) -> Add(w) -> Reshape. +*/ +static bool Match_Linear_Subgraph_1(Graph& graph, const Node& concat, const Node& root, const logging::Logger& logger) { + if (!optimizer_utils::CheckOutputEdges(graph, concat, 1)) { + return false; + } + auto reshape_itr = concat.OutputNodesBegin(); + if ((*reshape_itr).OpType().compare("Reshape") != 0) { + return false; + } + const Node& reshape = *reshape_itr; + + std::vector linear_path{ + {0, 0, "Add", {7}, kOnnxDomain}, + {0, 0, "MatMul", {1, 9}, kOnnxDomain}}; + std::vector edges; + if (!graph_utils::FindPath(reshape, true, linear_path, edges, logger)) { + return false; + } + + const Node& linear_path_add = edges[0]->GetNode(); + const Node& linear_path_matmul = edges[1]->GetNode(); + + const Node* p_node_before_matmul = graph_utils::GetInputNode(linear_path_matmul, 0); + if (p_node_before_matmul != nullptr && p_node_before_matmul->Index() != root.Index()) { + return false; + } + + if (linear_path_add.InputDefs().size() < 2) { + return false; + } + const NodeArg& linear_path_add_b = *(linear_path_add.InputDefs()[1]); + if (!graph_utils::IsInitializer(graph, linear_path_add_b.Name(), true)) { + return false; + } + + if (!optimizer_utils::IsShapeKnownOnAllDims(linear_path_add_b, 1)) { + return false; + } + int64_t hidden_size = linear_path_add_b.Shape()->dim(0).dim_value(); + + if (!optimizer_utils::ValidateShape(*(linear_path_matmul.InputDefs()[1]), {hidden_size, hidden_size})) { + return false; + } + + return true; +} + +static bool Match_Shape(Graph& graph, const Node& concat, const Node& shape, const NodeArg& root_input, const logging::Logger& logger) { + const NodeArg& shape_input = *(shape.InputDefs()[0]); + if (shape_input.Name() == root_input.Name()) { + return true; + } + + const ONNX_NAMESPACE::TensorShapeProto* shape_input_shape = shape_input.Shape(); + const ONNX_NAMESPACE::TensorShapeProto* root_input_shape = root_input.Shape(); + + if (shape_input_shape != nullptr && root_input_shape != nullptr) + return optimizer_utils::CompareShape(*shape_input_shape, *root_input_shape); + + const Node* p_node_before_shape = graph_utils::GetInputNode(shape, 0); + if (p_node_before_shape == nullptr) { + return false; + } + if (Match_Linear_Subgraph_1(graph, concat, *p_node_before_shape, logger)) { + return true; + } + return false; +} + /** * Find the subgraph that matches [root] -> Shape -> Gather -> Unsqueeze. - * If checkOneElementOnly is set to true, this function only checks if the matched subgraph produces a + * If checkOneElementOnly is set to true, this function only checks if the matched subgraph produces a * one element output(skip the Gather input indices check). */ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const NodeArg& root_input, const Node& concat, @@ -57,11 +147,6 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node const Node& gather = edges[1]->GetNode(); const Node& shape = edges[2]->GetNode(); - const NodeArg& shape_input = *(shape.InputDefs()[0]); - if (shape_input.Name() != root_input.Name()) { - return false; - } - std::vector axes; if (!(graph_utils::GetRepeatedNodeAttributeValues(unsqueeze, "axes", axes) && axes.size() == 1 && axes[0] == 0)) { return false; @@ -74,6 +159,10 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(shape_value.size()), false)) { return false; } + + if (!Match_Shape(graph, concat, shape, root_input, logger)) { + return false; + } return true; } @@ -81,7 +170,7 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node } /** - * Find the subgraph that matches [root] -> Shape -> Slice -> Squeeze. Check the inputs of slice + * Find the subgraph that matches [root] -> Shape -> Slice -> Squeeze. Check the inputs of slice * to make sure the graph produces output with exactly one element. */ bool ReshapeFusion::Match_One_Element_Output_Subgraph_2(Graph& graph, const NodeArg& root_input, const Node& cur_node, @@ -117,7 +206,7 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_2(Graph& graph, const Node } /** - * Check if the i-th input of the current node contains exactly one element by checking + * Check if the i-th input of the current node contains exactly one element by checking * its inferred shape. */ bool ReshapeFusion::Is_One_Element_Input(const Node& cur_node, int index) { @@ -140,8 +229,8 @@ bool ReshapeFusion::Is_One_Element_Input(const Node& cur_node, int index) { } /** - * Search all known patterns of one element subgraphs, which include - - * 1. A concat input with inferred shape that can only contain one element. + * Search all known patterns of one element subgraphs, which include - + * 1. A concat input with inferred shape that can only contain one element. * 2. [root] -> Shape -> Gather(any 1d indice) -> Unsqueeze -> [Concat] * 3. [root] -> Shape -> Slice (slice to one element) -> Squeeze -> (Div/Mul) -> Unsqueeze -> [Concat] * | @@ -191,7 +280,7 @@ bool ReshapeFusion::Is_One_Element_Output_Subgraph(Graph& graph, const NodeArg& auto input_count = binary_node.InputArgCount().front(); for (int i = 0; i < input_count; ++i) { - // For each input, look for "one-element subgraph -> concat" or "shape -> slice -> squeeze" path for + // For each input, look for "one-element subgraph -> concat" or "shape -> slice -> squeeze" path for // a potential match. if (!ReshapeFusion::Is_One_Element_Input(binary_node, i) && !ReshapeFusion::Match_One_Element_Output_Subgraph_2(graph, root_input, binary_node, i, logger)) { @@ -207,7 +296,7 @@ bool ReshapeFusion::Is_One_Element_Output_Subgraph(Graph& graph, const NodeArg& Apply Reshape Fusion. The following are subgraphs before and after fusion: (a[] and b[] are int64[] constant initializers; Concat may have any number of arguments, each of which is a constant initializer or a Shape->Gather->Unsqueeze chain with the -index corresponding to the index of the argument, or a custom subgraph in which nodes +index corresponding to the index of the argument, or a custom subgraph in which nodes have only one output edge. Note the resulting shape value should contain no more than one value of -1. diff --git a/onnxruntime/core/optimizer/utils.cc b/onnxruntime/core/optimizer/utils.cc index af2d7c9e95..44c6a955af 100644 --- a/onnxruntime/core/optimizer/utils.cc +++ b/onnxruntime/core/optimizer/utils.cc @@ -195,6 +195,24 @@ bool ValidateShape(const NodeArg& node_arg, const std::initializer_list return true; } +bool CompareShape(const ONNX_NAMESPACE::TensorShapeProto& node_arg_shape, const ONNX_NAMESPACE::TensorShapeProto& node_arg_other_shape) { + if (node_arg_shape.dim_size() != node_arg_other_shape.dim_size()) + return false; + + if (node_arg_shape.dim_size() < 1) + return false; + + for (int i = 0; i < node_arg_shape.dim_size(); ++i) { + const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim = node_arg_shape.dim(i); + const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim_other = node_arg_other_shape.dim(i); + if (!utils::HasDimValue(dim) || !utils::HasDimValue(dim_other)) + return false; + if (dim.dim_value() != dim_other.dim_value()) + return false; + } + return true; +} + bool IsShapeKnownOnAllDims(const NodeArg& node_arg, int expected_dim_size) { auto shape = node_arg.Shape(); if (shape == nullptr || shape->dim_size() != expected_dim_size) { diff --git a/onnxruntime/core/optimizer/utils.h b/onnxruntime/core/optimizer/utils.h index 17aa0b4b53..2af8abfe31 100644 --- a/onnxruntime/core/optimizer/utils.h +++ b/onnxruntime/core/optimizer/utils.h @@ -50,6 +50,11 @@ bool AppendTensorFromInitializer(const Graph& graph, const NodeArg& input_arg, s */ bool ValidateShape(const NodeArg& node_arg, const std::initializer_list& expected_dim_values); +/** Compare Shape of node input or output. +@remarks exactly compare two TensorShapeProtos. Return true if they are same +*/ +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 @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). diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 6d7cb51abc..3e0694d3a7 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1558,6 +1558,42 @@ TEST_F(GraphTransformationTests, ReshapeFusionConcatSubgraphWithMul) { } } +TEST_F(GraphTransformationTests, ReshapeFusionDistilBertTest) { + auto model_uri = MODEL_FOLDER "fusion/reshape_fusion_distillbert.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + 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_TRUE(op_to_count["Shape"] == 0); + ASSERT_TRUE(op_to_count["Gather"] == 0); + ASSERT_TRUE(op_to_count["Unsqueeze"] == 0); + ASSERT_TRUE(op_to_count["Concat"] == 0); + ASSERT_TRUE(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(), 4); + + const int64_t* val = initializer->data(); + EXPECT_EQ(val[0], 0); + EXPECT_EQ(val[1], -1); + EXPECT_EQ(val[2], 2); + EXPECT_EQ(val[3], 4); + } + } +} + TEST_F(GraphTransformationTests, ExpandElimination) { auto model_uri = MODEL_FOLDER "expand_elimination.onnx"; std::shared_ptr model; diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_distillbert.onnx b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_distillbert.onnx new file mode 100644 index 0000000000..543970d34a Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_distillbert.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py index b6895820cf..21d81f8a82 100644 --- a/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py +++ b/onnxruntime/test/testdata/transform/fusion/reshape_fusion_gen.py @@ -195,7 +195,7 @@ graph = helper.make_graph( helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), - + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), helper.make_node("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"), @@ -227,9 +227,9 @@ graph = helper.make_graph( helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), - + helper.make_node("Shape", ["unsqueeze0_out"], ["dummy_out"], "dummy"), - + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), helper.make_node("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"), @@ -248,7 +248,7 @@ graph = helper.make_graph( helper.make_tensor('indices0', TensorProto.INT64, [], [0]), helper.make_tensor('indices1', TensorProto.INT64, [], [1]), helper.make_tensor('slice_starts', TensorProto.INT64, [1], [2]), - helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3]) + helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3]) ] ) @@ -262,7 +262,7 @@ graph = helper.make_graph( helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), - + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), helper.make_node("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"), helper.make_node("Pad", ["slice_out", "pads"], ["pad0_out"], "pad0", mode = "constant"), @@ -280,9 +280,9 @@ graph = helper.make_graph( [ # initializers helper.make_tensor('indices0', TensorProto.INT64, [], [0]), helper.make_tensor('indices1', TensorProto.INT64, [], [1]), - helper.make_tensor('pads', TensorProto.INT64, [2], [1, 0]), + helper.make_tensor('pads', TensorProto.INT64, [2], [1, 0]), helper.make_tensor('slice_starts', TensorProto.INT64, [1], [2]), - helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3]) + helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3]) ] ) @@ -296,7 +296,7 @@ graph = helper.make_graph( helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), - + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), helper.make_node("Slice", ["shape2_out", "slice_starts", "slice_ends"], ["slice_out"], "slice1"), helper.make_node("Squeeze", ["slice_out"], ["squeeze0_out"], "squeeze0", axes=[0]), @@ -316,9 +316,9 @@ graph = helper.make_graph( [ # initializers helper.make_tensor('indices0', TensorProto.INT64, [], [0]), helper.make_tensor('indices1', TensorProto.INT64, [], [1]), - helper.make_tensor('div_init', TensorProto.INT64, [], [1]), + helper.make_tensor('div_init', TensorProto.INT64, [], [1]), helper.make_tensor('slice_starts', TensorProto.INT64, [1], [2]), - helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3]) + helper.make_tensor('slice_ends', TensorProto.INT64, [1], [3]) ] ) @@ -332,7 +332,7 @@ graph = helper.make_graph( helper.make_node("Gather", ["shape1_out", "indices1"], ["gather1_out"], "gather1", axis=0), helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), - + helper.make_node("Shape", ["SubgraphRoot"], ["shape2_out"], "shape2"), helper.make_node("Slice", ["shape2_out", "slice_starts_0", "slice_ends_0"], ["slice0_out"], "slice0"), helper.make_node("Squeeze", ["slice0_out"], ["squeeze0_out"], "squeeze0", axes=[0]), @@ -358,8 +358,62 @@ graph = helper.make_graph( helper.make_tensor('slice_starts_0', TensorProto.INT64, [1], [2]), helper.make_tensor('slice_ends_0', TensorProto.INT64, [1], [3]), helper.make_tensor('slice_starts_1', TensorProto.INT64, [1], [1]), - helper.make_tensor('slice_ends_1', TensorProto.INT64, [1], [2]) + helper.make_tensor('slice_ends_1', TensorProto.INT64, [1], [2]) ] ) -save_model(graph, 'reshape_fusion_concat_subgraph_mul.onnx') \ No newline at end of file +save_model(graph, 'reshape_fusion_concat_subgraph_mul.onnx') + +matmul_weights = [ + -0.04888916015625, 0.0143280029296875, 0.066650390625,-0.0343017578125, + -0.0010356903076171875, -0.00048232078552246094, 0.07470703125, -0.04736328125, + 0.01454925537109375, -0.0086669921875, -0.051971435546875, -0.0201568603515625, + 0.040435791015625, -0.019256591796875, 0.0205078125, 0.0111541748046875, + 0.0071868896484375, -0.0298309326171875, -0.0306549072265625, -0.0225372314453125, + -0.04193115234375, 0.07073974609375, -0.048065185546875, 0.0198822021484375, + -0.035552978515625, -0.022796630859375, 0.03839111328125, 0.007099151611328125, + -0.0080108642578125, -0.0017957687377929688, 0.0266265869140625,-0.028289794921875, + 0.0032901763916015625, 0.0208740234375, -0.01529693603515625, -0.046600341796875, + -0.034637451171875, 0.011322021484375, -0.026458740234375, 0.04656982421875, + -0.0091705322265625, 0.017913818359375, -0.019256591796875, -0.001216888427734375, + -0.08245849609375, -0.023162841796875, -0.04132080078125, -0.03363037109375, + 0.0029315948486328125, 0.03173828125, -0.004024505615234375, 0.04534912109375, + -0.0036163330078125, -0.03912353515625, -0.00800323486328125, 0.058197021484375, + 0.05572509765625, 0.01165771484375, 0.06756591796875, 0.05816650390625, + -0.0654296875, -0.0241851806640625, 0.0205535888671875, -0.031707763671875 +] + +add_weight = [-0.23681640625, -0.16552734375, 0.2191162109375, -0.1756591796875, + -0.03460693359375, -0.05316162109375, -0.336181640625, -0.253662109375] + +graph = helper.make_graph( + [ # nodes + helper.make_node("Add", ["Input", "Bias"], ["add0_out"], "add0"), + helper.make_node("Shape", ["add0_out"], ["shape0_out"], "shape0"), + helper.make_node("Gather", ["shape0_out", "indices0"], ["gather0_out"], "gather0", axis=0), + helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), + helper.make_node("Concat", ["unsqueeze0_out", "dim_-1", "dim_2", "dim_4"], ["concat_out"], "concat", axis=0), + helper.make_node("MatMul", ["add0_out", "matmul_weight"], ["matmul_out"], "matmul"), + helper.make_node("Add", ["matmul_out", "add_weight"], ["add1_out"], "add1"), + helper.make_node("Reshape", ["add1_out", "concat_out"], ["Result"], "reshape"), + ], + "Reshape_Fusion", #name + [ # inputs + helper.make_tensor_value_info('Input', TensorProto.FLOAT, [1, 8]), + ], + [ # outputs + helper.make_tensor_value_info('Result', TensorProto.FLOAT, [1, -1, 2, 4]), + ], + [ # initializers + helper.make_tensor('Bias', TensorProto.FLOAT, [8], add_weight), + helper.make_tensor('dim_-1', TensorProto.INT64, [1], [-1]), + helper.make_tensor('dim_2', TensorProto.INT64, [1], [2]), + helper.make_tensor('dim_4', TensorProto.INT64, [1], [4]), + helper.make_tensor('indices0', TensorProto.INT64, [], [0]), + helper.make_tensor('matmul_weight', TensorProto.FLOAT, [8, 8], matmul_weights), + helper.make_tensor('add_weight', TensorProto.FLOAT, [8], add_weight), + ] +) + +save_model(graph, 'reshape_fusion_distillbert.onnx') +