Skip to content

Commit 9bbf182

Browse files
HonryHonry
committed
[ORT] Extend reshape fusion for dynamic shape dimensions
Enable reshape fusion for models with symbolic (dynamic) dimensions, covering patterns found in MobileNetV2, vision encoders, and SLMs. Changes: - Match_Shape: allow dimension matching by symbolic dim_param, not just static values. Needed for SLMs (e.g. Qwen) where position_ids takes Shape from input_ids — same dim_param names, no static values. - Match_Shape: allow matching when only the gathered dimension agrees between Shape source and reshape root. Needed for vision encoders (e.g. moondream2) where Shape is taken from pre-projection tensor but Reshape operates on post-projection output. - Match_Shape: add GlobalAveragePool passthrough — accept Shape from pre-pool input when gathering batch dim. Needed for MobileNetV2. - ReshapeHelper: handle zero-dim shapes (allow_zero=true) by computing unknown dim from non-zero dimension product only. Needed for MobileNetV2 reshape fusion output. - Add corresponding unit tests.
1 parent bf7e27a commit 9bbf182

3 files changed

Lines changed: 273 additions & 21 deletions

File tree

‎onnxruntime/core/optimizer/reshape_fusion.cc‎

Lines changed: 83 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -132,17 +132,78 @@ static bool Match_Linear_Subgraph_1(Graph& graph, const Node& concat, const Node
132132
return true;
133133
}
134134

135-
static bool Match_Shape(Graph& graph, const Node& concat, const Node& shape, const NodeArg& root_input, const logging::Logger& logger) {
135+
static bool Match_Shape(Graph& graph, const Node& concat, const Node& shape, const NodeArg& root_input,
136+
int64_t expected_gather_index, const logging::Logger& logger) {
136137
const NodeArg& shape_input = *(shape.InputDefs()[0]);
137138
if (shape_input.Name() == root_input.Name()) {
138139
return true;
139140
}
140141

142+
// Allow a narrow passthrough pattern: input -> GlobalAveragePool -> Reshape, where shape is taken from
143+
// the pre-pool input and only batch dim (index 0) is gathered.
144+
const Node* root_input_producer = graph.GetProducerNode(root_input.Name());
145+
if (root_input_producer != nullptr && root_input_producer->OpType() == "GlobalAveragePool" &&
146+
root_input_producer->InputDefs().size() > 0 && root_input_producer->InputDefs()[0] != nullptr &&
147+
root_input_producer->InputDefs()[0]->Name() == shape_input.Name() && expected_gather_index == 0) {
148+
return true;
149+
}
150+
141151
const ONNX_NAMESPACE::TensorShapeProto* shape_input_shape = shape_input.Shape();
142152
const ONNX_NAMESPACE::TensorShapeProto* root_input_shape = root_input.Shape();
143153

144-
if (shape_input_shape != nullptr && root_input_shape != nullptr)
145-
return optimizer_utils::CompareShape(*shape_input_shape, *root_input_shape);
154+
if (shape_input_shape != nullptr && root_input_shape != nullptr) {
155+
// First try the existing static-value comparison.
156+
if (optimizer_utils::CompareShape(*shape_input_shape, *root_input_shape)) {
157+
return true;
158+
}
159+
// Also allow matching when both shapes have the same rank and each dimension
160+
// is either equal by static value or by matching symbolic dim_param.
161+
// This covers cases like position_ids [batch_size, sequence_length] taking
162+
// shape from input_ids [batch_size, sequence_length] — same symbolic dims
163+
// but no static values.
164+
if (shape_input_shape->dim_size() == root_input_shape->dim_size() &&
165+
shape_input_shape->dim_size() > 0) {
166+
bool all_dims_match = true;
167+
for (int i = 0; i < shape_input_shape->dim_size(); ++i) {
168+
const auto& dim_a = shape_input_shape->dim(i);
169+
const auto& dim_b = root_input_shape->dim(i);
170+
if (utils::HasDimValue(dim_a) && utils::HasDimValue(dim_b)) {
171+
if (dim_a.dim_value() != dim_b.dim_value()) {
172+
all_dims_match = false;
173+
break;
174+
}
175+
} else if (utils::HasDimParam(dim_a) && utils::HasDimParam(dim_b)) {
176+
if (dim_a.dim_param() != dim_b.dim_param()) {
177+
all_dims_match = false;
178+
break;
179+
}
180+
} else {
181+
// One has value, the other has param, or neither has anything — no match.
182+
all_dims_match = false;
183+
break;
184+
}
185+
}
186+
if (all_dims_match) {
187+
return true;
188+
}
189+
}
190+
}
191+
192+
// Allow match when only the gathered dimension needs to agree.
193+
// E.g., Shape extracts batch dim from a pre-projection tensor; that dim is
194+
// identical to the batch dim of the post-projection tensor being reshaped.
195+
if (shape_input_shape != nullptr && root_input_shape != nullptr &&
196+
expected_gather_index < shape_input_shape->dim_size() &&
197+
expected_gather_index < root_input_shape->dim_size()) {
198+
const auto& dim_a = shape_input_shape->dim(static_cast<int>(expected_gather_index));
199+
const auto& dim_b = root_input_shape->dim(static_cast<int>(expected_gather_index));
200+
if ((utils::HasDimValue(dim_a) && utils::HasDimValue(dim_b) &&
201+
dim_a.dim_value() == dim_b.dim_value()) ||
202+
(utils::HasDimParam(dim_a) && utils::HasDimParam(dim_b) &&
203+
dim_a.dim_param() == dim_b.dim_param())) {
204+
return true;
205+
}
206+
}
146207

147208
const Node* p_node_before_shape = graph_utils::GetInputNode(shape, 0);
148209
if (p_node_before_shape == nullptr) {
@@ -188,11 +249,12 @@ bool ReshapeFusion::Match_One_Element_Output_Subgraph_1(Graph& graph, const Node
188249
return true;
189250
}
190251

191-
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), int64_t(shape_value.size()), false)) {
252+
const int64_t expected_gather_index = static_cast<int64_t>(shape_value.size());
253+
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather.InputDefs()[1]), expected_gather_index, false)) {
192254
return false;
193255
}
194256

195-
if (!Match_Shape(graph, concat, shape, root_input, logger)) {
257+
if (!Match_Shape(graph, concat, shape, root_input, expected_gather_index, logger)) {
196258
return false;
197259
}
198260
return true;
@@ -340,19 +402,23 @@ index corresponding to the index of the argument, or a custom subgraph in which
340402
have only one output edge. Note the resulting shape value should contain no more than one
341403
value of -1.
342404
405+
The Shape node may feed from Sub-graph Root itself, or from a different tensor upstream
406+
as long as the gathered dimension matches (same static value or same symbolic dim_param)
407+
between the Shape source and Sub-graph Root.
408+
343409
Before fusion:
344-
[Sub-graph Root]
345-
| / \
346-
| Shape Shape
347-
| | |
348-
| Gather(indices=0) a[] Gather(indices=2) b[] or subgraph
349-
| \ / / /
350-
| Unsqueeze / Unsqueeze /
351-
| \ / ___________/ /
352-
| \ / / _______________________/
353-
| \ / / /
354-
\ Concat
355-
\ /
410+
[Sub-graph Root] [Ancestor of Root or Root itself]
411+
| / \
412+
| Shape Shape
413+
| | |
414+
| Gather(indices=0) a[] Gather(indices=2) b[] or subgraph
415+
| \ / / /
416+
| Unsqueeze / Unsqueeze /
417+
| \ / ___________/ /
418+
| \ / / _______________________/
419+
| \ / / /
420+
\ Concat
421+
\ /
356422
Reshape
357423
358424
After fusion:

‎onnxruntime/core/providers/cpu/tensor/reshape_helper.h‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ class ReshapeHelper {
2121
auto nDims = requested_shape.size();
2222
ptrdiff_t unknown_dim = -1;
2323
int64_t size = 1;
24+
int64_t size_for_inference = 1;
25+
bool has_zero_dim = false;
2426
for (size_t i = 0; i < nDims; ++i) {
2527
ORT_ENFORCE(requested_shape[i] >= -1, "A dimension cannot be less than -1, got ", requested_shape[i]);
2628
if (requested_shape[i] == -1) {
@@ -39,17 +41,43 @@ class ReshapeHelper {
3941
", requested shape:", TensorShape(requested_shape));
4042
}
4143
size *= dim;
44+
if (dim == 0) {
45+
has_zero_dim = true;
46+
} else {
47+
size_for_inference *= dim;
48+
}
4249
}
4350
}
4451

4552
const auto requested_shape_size = size;
4653

4754
if (unknown_dim != -1) {
4855
// calculate unknown dimension
49-
ORT_ENFORCE(requested_shape_size != 0 && (input_shape_size % requested_shape_size) == 0,
50-
"The input tensor cannot be reshaped to the requested shape. Input shape:", input_shape,
51-
", requested shape:", TensorShape(requested_shape));
52-
requested_shape[unknown_dim] = input_shape_size / requested_shape_size;
56+
if (has_zero_dim) {
57+
ORT_ENFORCE(input_shape_size == 0,
58+
"The input tensor cannot be reshaped to the requested shape. Input shape:", input_shape,
59+
", requested shape:", TensorShape(requested_shape));
60+
ORT_ENFORCE(size_for_inference != 0,
61+
"The input tensor cannot be reshaped to the requested shape. Input shape:", input_shape,
62+
", requested shape:", TensorShape(requested_shape));
63+
64+
int64_t input_shape_non_zero_size = 1;
65+
for (size_t i = 0; i < input_shape.NumDimensions(); ++i) {
66+
if (input_shape[i] != 0) {
67+
input_shape_non_zero_size *= input_shape[i];
68+
}
69+
}
70+
71+
ORT_ENFORCE((input_shape_non_zero_size % size_for_inference) == 0,
72+
"The input tensor cannot be reshaped to the requested shape. Input shape:", input_shape,
73+
", requested shape:", TensorShape(requested_shape));
74+
requested_shape[unknown_dim] = input_shape_non_zero_size / size_for_inference;
75+
} else {
76+
ORT_ENFORCE(size != 0 && (input_shape_size % size) == 0,
77+
"The input tensor cannot be reshaped to the requested shape. Input shape:", input_shape,
78+
", requested shape:", TensorShape(requested_shape));
79+
requested_shape[unknown_dim] = input_shape_size / size;
80+
}
5381
} else {
5482
// check if the output shape is valid.
5583
ORT_ENFORCE(input_shape_size == requested_shape_size,

‎onnxruntime/test/optimizer/graph_transform_test.cc‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8602,6 +8602,164 @@ TEST_F(GraphTransformationTests, ReshapeFusionOpsetTest) {
86028602
}
86038603
#endif
86048604

8605+
// Test reshape fusion when Shape feeds from a tensor with the same symbolic dim_param names
8606+
// as the reshape root (e.g., position_ids reshaped using Shape from input_ids in Qwen).
8607+
TEST_F(GraphTransformationTests, ReshapeFusionDimParamMatch) {
8608+
auto build_test_case = [&](ModelTestBuilder& builder) {
8609+
// input_ids: [batch_size, sequence_length] — symbolic dims
8610+
auto* input_ids = builder.MakeSymbolicInput<float>({std::string("batch_size"), std::string("sequence_length")});
8611+
// position_ids: [batch_size, sequence_length] — same symbolic dims
8612+
auto* position_ids = builder.MakeSymbolicInput<float>({std::string("batch_size"), std::string("sequence_length")});
8613+
8614+
auto* shape_out = builder.MakeIntermediate();
8615+
auto* gather_out = builder.MakeIntermediate();
8616+
auto* unsqueeze_out = builder.MakeIntermediate();
8617+
auto* concat_out = builder.MakeIntermediate();
8618+
auto* out = builder.MakeOutput();
8619+
8620+
auto* scalar_int_0 = builder.MakeInitializer<int64_t>({}, {0});
8621+
auto* unsqueeze_axes = builder.MakeInitializer<int64_t>({1}, {0});
8622+
auto* const_neg1 = builder.MakeInitializer<int64_t>({1}, {-1});
8623+
8624+
// Shape(input_ids) -> Gather(0) -> Unsqueeze -> Concat([unsqueeze, -1]) -> Reshape(position_ids)
8625+
builder.AddNode("Shape", {input_ids}, {shape_out});
8626+
builder.AddNode("Gather", {shape_out, scalar_int_0}, {gather_out});
8627+
builder.AddNode("Unsqueeze", {gather_out, unsqueeze_axes}, {unsqueeze_out});
8628+
builder.AddNode("Concat", {unsqueeze_out, const_neg1}, {concat_out}).AddAttribute("axis", static_cast<int64_t>(0));
8629+
builder.AddNode("Reshape", {position_ids, concat_out}, {out});
8630+
};
8631+
8632+
auto pre_graph_checker = [&](Graph& graph) {
8633+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Shape"] == 1);
8634+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gather"] == 1);
8635+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Unsqueeze"] == 1);
8636+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Concat"] == 1);
8637+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Reshape"] == 1);
8638+
return Status::OK();
8639+
};
8640+
8641+
auto post_graph_checker = [&](Graph& graph) {
8642+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Shape"] == 0);
8643+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gather"] == 0);
8644+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Unsqueeze"] == 0);
8645+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Concat"] == 0);
8646+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Reshape"] == 1);
8647+
return Status::OK();
8648+
};
8649+
8650+
std::unique_ptr<GraphTransformer> transformer = std::make_unique<ReshapeFusion>();
8651+
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer),
8652+
TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker));
8653+
}
8654+
8655+
// Test reshape fusion when Shape feeds from a different tensor than the reshape root,
8656+
// but the gathered dimension matches by symbolic dim_param (moondream2 vision encoder pattern:
8657+
// Shape from pre-projection tensor, Reshape on post-projection output).
8658+
TEST_F(GraphTransformationTests, ReshapeFusionCrossTensorDimMatch) {
8659+
auto build_test_case = [&](ModelTestBuilder& builder) {
8660+
// pre_proj: [batch_size, seq_len, 768] — pre-projection (layernorm output)
8661+
auto* pre_proj = builder.MakeSymbolicInput<float>(
8662+
{std::string("batch_size"), std::string("seq_len"), int64_t(768)});
8663+
// post_proj: [batch_size, seq_len, 2304] — post-projection (qkv linear output, different last dim)
8664+
auto* post_proj = builder.MakeSymbolicInput<float>(
8665+
{std::string("batch_size"), std::string("seq_len"), int64_t(2304)});
8666+
8667+
auto* shape_out = builder.MakeIntermediate();
8668+
auto* gather_out = builder.MakeIntermediate();
8669+
auto* unsqueeze_out = builder.MakeIntermediate();
8670+
auto* concat_out = builder.MakeIntermediate();
8671+
auto* out = builder.MakeOutput();
8672+
8673+
auto* scalar_int_0 = builder.MakeInitializer<int64_t>({}, {0});
8674+
auto* unsqueeze_axes = builder.MakeInitializer<int64_t>({1}, {0});
8675+
auto* const_neg1 = builder.MakeInitializer<int64_t>({1}, {-1});
8676+
auto* const_3 = builder.MakeInitializer<int64_t>({1}, {3});
8677+
auto* const_256 = builder.MakeInitializer<int64_t>({1}, {256});
8678+
8679+
// Shape(pre_proj) -> Gather(0) -> Unsqueeze -> Concat([unsqueeze, -1, 3, 256]) -> Reshape(post_proj)
8680+
builder.AddNode("Shape", {pre_proj}, {shape_out});
8681+
builder.AddNode("Gather", {shape_out, scalar_int_0}, {gather_out});
8682+
builder.AddNode("Unsqueeze", {gather_out, unsqueeze_axes}, {unsqueeze_out});
8683+
builder.AddNode("Concat", {unsqueeze_out, const_neg1, const_3, const_256}, {concat_out})
8684+
.AddAttribute("axis", static_cast<int64_t>(0));
8685+
builder.AddNode("Reshape", {post_proj, concat_out}, {out});
8686+
};
8687+
8688+
auto pre_graph_checker = [&](Graph& graph) {
8689+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Shape"] == 1);
8690+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gather"] == 1);
8691+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Unsqueeze"] == 1);
8692+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Concat"] == 1);
8693+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Reshape"] == 1);
8694+
return Status::OK();
8695+
};
8696+
8697+
auto post_graph_checker = [&](Graph& graph) {
8698+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Shape"] == 0);
8699+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gather"] == 0);
8700+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Unsqueeze"] == 0);
8701+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Concat"] == 0);
8702+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Reshape"] == 1);
8703+
return Status::OK();
8704+
};
8705+
8706+
std::unique_ptr<GraphTransformer> transformer = std::make_unique<ReshapeFusion>();
8707+
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer),
8708+
TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker));
8709+
}
8710+
8711+
// Test reshape fusion with GlobalAveragePool passthrough pattern (MobileNetV2):
8712+
// Shape is taken from the pre-pool input, Reshape operates on the pool output,
8713+
// and only batch dim (gather index 0) is extracted.
8714+
TEST_F(GraphTransformationTests, ReshapeFusionGlobalAveragePoolPassthrough) {
8715+
auto build_test_case = [&](ModelTestBuilder& builder) {
8716+
auto* input = builder.MakeInput<float>({{1, 1280, 7, 7}});
8717+
auto* pool_out = builder.MakeIntermediate();
8718+
auto* shape_out = builder.MakeIntermediate();
8719+
auto* gather_out = builder.MakeIntermediate();
8720+
auto* unsqueeze_out = builder.MakeIntermediate();
8721+
auto* concat_out = builder.MakeIntermediate();
8722+
auto* out = builder.MakeOutput();
8723+
8724+
auto* scalar_int_0 = builder.MakeInitializer<int64_t>({}, {0});
8725+
auto* unsqueeze_axes = builder.MakeInitializer<int64_t>({1}, {0});
8726+
auto* const_neg1 = builder.MakeInitializer<int64_t>({1}, {-1});
8727+
8728+
// input -> GlobalAveragePool -> pool_out
8729+
builder.AddNode("GlobalAveragePool", {input}, {pool_out});
8730+
// Shape(input) -> Gather(0) -> Unsqueeze -> Concat([unsqueeze, -1]) -> Reshape(pool_out)
8731+
builder.AddNode("Shape", {input}, {shape_out});
8732+
builder.AddNode("Gather", {shape_out, scalar_int_0}, {gather_out});
8733+
builder.AddNode("Unsqueeze", {gather_out, unsqueeze_axes}, {unsqueeze_out});
8734+
builder.AddNode("Concat", {unsqueeze_out, const_neg1}, {concat_out}).AddAttribute("axis", static_cast<int64_t>(0));
8735+
builder.AddNode("Reshape", {pool_out, concat_out}, {out});
8736+
};
8737+
8738+
auto pre_graph_checker = [&](Graph& graph) {
8739+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["GlobalAveragePool"] == 1);
8740+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Shape"] == 1);
8741+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gather"] == 1);
8742+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Unsqueeze"] == 1);
8743+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Concat"] == 1);
8744+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Reshape"] == 1);
8745+
return Status::OK();
8746+
};
8747+
8748+
auto post_graph_checker = [&](Graph& graph) {
8749+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["GlobalAveragePool"] == 1);
8750+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Shape"] == 0);
8751+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Gather"] == 0);
8752+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Unsqueeze"] == 0);
8753+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Concat"] == 0);
8754+
TEST_RETURN_IF_NOT(CountOpsInGraph(graph)["Reshape"] == 1);
8755+
return Status::OK();
8756+
};
8757+
8758+
std::unique_ptr<GraphTransformer> transformer = std::make_unique<ReshapeFusion>();
8759+
ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 14, *logger_, std::move(transformer),
8760+
TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker));
8761+
}
8762+
86058763
TEST_F(GraphTransformationTests, DynamicQuantizeMatMulTest) {
86068764
constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/dynamic_quantize_matmul.onnx";
86078765
std::shared_ptr<Model> p_model;

0 commit comments

Comments
 (0)