diff --git a/onnxruntime/core/optimizer/pad_fusion.cc b/onnxruntime/core/optimizer/pad_fusion.cc index b364b770a425c..179a08e0c6f9e 100644 --- a/onnxruntime/core/optimizer/pad_fusion.cc +++ b/onnxruntime/core/optimizer/pad_fusion.cc @@ -50,7 +50,7 @@ bool VerifyNotCastChild(const Node& child_node) { return true; } -void UpdatePaddingAttribute(Node& child_node, const std::vector& pads_values, const uint32_t pads_size) { +void UpdatePaddingAttribute(Node& child_node, const std::vector& pads_values, const size_t pads_size) { auto reset_pads = true; if (child_node.GetAttributes().find("pads") != child_node.GetAttributes().end()) { /* pads can be empty, overwrite pads attribute in this case */ @@ -62,13 +62,15 @@ void UpdatePaddingAttribute(Node& child_node, const std::vector& pads_v } auto child_pads = child_node.GetMutableAttributes()["pads"].mutable_ints(); - uint32_t child_pads_size = static_cast(child_pads->size()); + const size_t child_pads_size = static_cast(child_pads->size()); - for (uint32_t pads_index = 2, child_index = 0; pads_index < pads_size / 2; pads_index++, child_index++) { - child_pads->Set(child_index, child_pads->Get(child_index) + pads_values[pads_index]); - uint32_t mirrored_child_index = child_index + (child_pads_size / 2); - uint32_t mirrored_pad_index = pads_index + (pads_size / 2); - child_pads->Set(mirrored_child_index, child_pads->Get(mirrored_child_index) + pads_values[mirrored_pad_index]); + for (size_t pads_index = 2, child_index = 0; pads_index < pads_size / 2; pads_index++, child_index++) { + child_pads->Set(static_cast(child_index), + child_pads->Get(static_cast(child_index)) + pads_values[pads_index]); + const size_t mirrored_child_index = child_index + (child_pads_size / 2); + const size_t mirrored_pad_index = pads_index + (pads_size / 2); + child_pads->Set(static_cast(mirrored_child_index), + child_pads->Get(static_cast(mirrored_child_index)) + pads_values[mirrored_pad_index]); } if (child_node.OpType() == "AveragePool") { @@ -158,7 +160,14 @@ Status PadFusion::Apply(Graph& graph, Node& pad_node, RewriteRuleEffect& rule_ef pads_values.assign(pad_node.GetAttributes().at("pads").ints().begin(), pad_node.GetAttributes().at("pads").ints().end()); } - uint32_t pads_size = static_cast(pads_values.size()); + const size_t pads_size = pads_values.size(); + // Per ONNX Pad spec, pads has 2*rank elements. This fusion only applies when the leading two + // dimensions (N, C) have zero padding, so we need at least rank 2 (pads_size >= 4) and an even + // number of entries. + if (pads_size < 4 || (pads_size % 2) != 0) { + return Status::OK(); + } + // check if padding is applied only on feature dims if (pads_values[0] != 0 || pads_values[1] != 0 || pads_values[pads_size / 2] != 0 || pads_values[pads_size / 2 + 1] != 0) { @@ -173,7 +182,21 @@ Status PadFusion::Apply(Graph& graph, Node& pad_node, RewriteRuleEffect& rule_ef Node& child_node = *graph.GetNode(pad_node.OutputNodesBegin()->Index()); // We don't need to cast the pad_constant_value because this fusion requires that constant_pad_value // to be zero. See PadFusion::SatisfyCondition for details. - Node& target_padding_node = (child_node.OpType() == "Cast") ? *graph.GetNode(child_node.OutputNodesBegin()->Index()) : child_node; + Node& target_padding_node = (child_node.OpType() == "Cast") + ? *graph.GetNode(child_node.OutputNodesBegin()->Index()) + : child_node; + + // If the target node already has an explicit pads attribute, its length must match the expected + // pads length for the target op (2 * spatial_rank), where spatial_rank = pads_size / 2 - 2. + // Otherwise the fused padding values would be written into mismatched positions. + const auto& target_attrs = target_padding_node.GetAttributes(); + auto target_pads_iter = target_attrs.find("pads"); + if (target_pads_iter != target_attrs.end() && !target_pads_iter->second.ints().empty()) { + if (static_cast(target_pads_iter->second.ints().size()) != pads_size - 4) { + return Status::OK(); + } + } + UpdatePaddingAttribute(target_padding_node, pads_values, pads_size); graph_utils::RemoveNodeOutputEdges(graph, pad_node); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 924ceaa19a47d..3f66d4d0a6696 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -2103,6 +2103,141 @@ TEST_F(GraphTransformationTests, FusePadWithAvgPoolWithPadNoInclude) { ASSERT_EQ(op_to_count["AveragePool"], 1); } +// Verify PadFusion does not fuse and does not crash when the Pad node's `pads` initializer carries +// fewer than four elements (which violates the 2*rank requirement of the ONNX Pad spec). The data +// input uses an unspecified rank so ONNX shape inference cannot pre-validate the pads length. +TEST_F(GraphTransformationTests, PadFusionRejectsShortPadsInitializer) { + auto run_case = [&](const std::vector& pads_data) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* data_arg = builder.MakeInput(std::nullopt); + auto* weight_arg = builder.MakeInitializer({1, 1, 1, 1}, {1.0f}); + auto* pads_arg = builder.MakeInitializer( + {static_cast(pads_data.size())}, pads_data); + auto* pad_out = builder.MakeIntermediate(); + auto* conv_out = builder.MakeOutput(); + + builder.AddNode("Pad", {data_arg, pads_arg}, {pad_out}); + builder.AddNode("Conv", {pad_out, weight_arg}, {conv_out}); + }; + + auto pre_graph_checker = [](Graph&) { return Status::OK(); }; + auto post_graph_checker = [](Graph& graph) { + // Fusion must not have run; both nodes remain. + std::map op_to_count = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(op_to_count["Pad"] == 1, "Pad should not be fused away"); + ORT_RETURN_IF_NOT(op_to_count["Conv"] == 1, "Conv should remain"); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformerL1"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 18, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); + }; + + // Cover empty, 1, 2, and 3 element pads vectors as well as an odd-length vector. + run_case({}); + run_case({0}); + run_case({0, 0}); + run_case({0, 0, 0}); + run_case({0, 0, 0, 0, 0}); +} + +// Verify PadFusion bails out when the Pad node's pads contain a negative value. +TEST_F(GraphTransformationTests, PadFusionRejectsNegativePads) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* data_arg = builder.MakeInput(std::nullopt); + auto* weight_arg = builder.MakeInitializer({1, 1, 1, 1}, {1.0f}); + auto* pads_arg = builder.MakeInitializer({8}, {0, 0, 1, 1, 0, 0, -1, 1}); + auto* pad_out = builder.MakeIntermediate(); + auto* conv_out = builder.MakeOutput(); + + builder.AddNode("Pad", {data_arg, pads_arg}, {pad_out}); + builder.AddNode("Conv", {pad_out, weight_arg}, {conv_out}); + }; + + auto pre_graph_checker = [](Graph&) { return Status::OK(); }; + auto post_graph_checker = [](Graph& graph) { + std::map op_to_count = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(op_to_count["Pad"] == 1, "Pad with negative pads should not be fused"); + ORT_RETURN_IF_NOT(op_to_count["Conv"] == 1, "Conv should remain"); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformerL1"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 18, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); +} + +// Verify PadFusion bails out when the child Conv already has an explicit `pads` attribute whose +// length does not match the spatial rank implied by the Pad node. +TEST_F(GraphTransformationTests, PadFusionRejectsMismatchedChildPadsRank) { + auto build_test_case = [&](ModelTestBuilder& builder) { + // Pad pads length 8 implies spatial rank 2 (expected child pads length 4), + // but Conv carries a `pads` attribute of length 6, simulating a rank mismatch + // that could arise from a malformed model. Data uses unspecified rank so shape + // inference does not reject the model before the transformer runs. + auto* data_arg = builder.MakeInput(std::nullopt); + auto* weight_arg = builder.MakeInitializer({1, 1, 1, 1}, {1.0f}); + auto* pads_arg = builder.MakeInitializer({8}, {0, 0, 1, 1, 0, 0, 1, 1}); + auto* pad_out = builder.MakeIntermediate(); + auto* conv_out = builder.MakeOutput(); + + builder.AddNode("Pad", {data_arg, pads_arg}, {pad_out}); + auto& conv = builder.AddNode("Conv", {pad_out, weight_arg}, {conv_out}); + conv.AddAttribute("pads", std::vector{0, 0, 0, 0, 0, 0}); + }; + + auto pre_graph_checker = [](Graph&) { return Status::OK(); }; + auto post_graph_checker = [](Graph& graph) { + std::map op_to_count = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(op_to_count["Pad"] == 1, "Pad should not be fused into a rank-mismatched Conv"); + ORT_RETURN_IF_NOT(op_to_count["Conv"] == 1, "Conv should remain"); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformerL1"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 18, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); +} + +// Verify PadFusion does not crash for the opset-10 attribute form of Pad when the `pads` attribute +// has fewer than four elements. +TEST_F(GraphTransformationTests, PadFusionRejectsShortPadsAttributeOpset10) { + auto run_case = [&](const std::vector& pads_data) { + auto build_test_case = [&](ModelTestBuilder& builder) { + auto* data_arg = builder.MakeInput(std::nullopt); + auto* weight_arg = builder.MakeInitializer({1, 1, 1, 1}, {1.0f}); + auto* pad_out = builder.MakeIntermediate(); + auto* conv_out = builder.MakeOutput(); + + auto& pad_node = builder.AddNode("Pad", {data_arg}, {pad_out}); + pad_node.AddAttribute("pads", pads_data); + pad_node.AddAttribute("mode", "constant"); + builder.AddNode("Conv", {pad_out, weight_arg}, {conv_out}); + }; + + auto pre_graph_checker = [](Graph&) { return Status::OK(); }; + auto post_graph_checker = [](Graph& graph) { + std::map op_to_count = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(op_to_count["Pad"] == 1, "Pad should not be fused away"); + ORT_RETURN_IF_NOT(op_to_count["Conv"] == 1, "Conv should remain"); + return Status::OK(); + }; + + auto rule_transformer = std::make_unique("RuleTransformerL1"); + ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique())); + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 10, *logger_, std::move(rule_transformer), + TransformerLevel::Level1, 1, pre_graph_checker, post_graph_checker)); + }; + + run_case({}); + run_case({0, 0}); + run_case({0, 0, 0}); +} + TEST_F(GraphTransformationTests, FuseMatmulBNWithInBetweenNodes) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "fusion/fuse-matmul-bn-with-reshape.onnx";