Address PadFusion shortcomings

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<int64_t>& pads_values, const uint32_t pads_size) {
+void UpdatePaddingAttribute(Node& child_node, const std::vector<int64_t>& 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<int64_t>& pads_v
   }
 
   auto child_pads = child_node.GetMutableAttributes()["pads"].mutable_ints();
-  uint32_t child_pads_size = static_cast<uint32_t>(child_pads->size());
+  const size_t child_pads_size = static_cast<size_t>(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<int>(child_index),
+                    child_pads->Get(static_cast<int>(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<int>(mirrored_child_index),
+                    child_pads->Get(static_cast<int>(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<uint32_t>(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<size_t>(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);

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<int64_t>& pads_data) {
+    auto build_test_case = [&](ModelTestBuilder& builder) {
+      auto* data_arg = builder.MakeInput<float>(std::nullopt);
+      auto* weight_arg = builder.MakeInitializer<float>({1, 1, 1, 1}, {1.0f});
+      auto* pads_arg = builder.MakeInitializer<int64_t>(
+          {static_cast<int64_t>(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<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformerL1");
+    ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<PadFusion>()));
+    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<float>(std::nullopt);
+    auto* weight_arg = builder.MakeInitializer<float>({1, 1, 1, 1}, {1.0f});
+    auto* pads_arg = builder.MakeInitializer<int64_t>({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<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformerL1");
+  ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<PadFusion>()));
+  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<float>(std::nullopt);
+    auto* weight_arg = builder.MakeInitializer<float>({1, 1, 1, 1}, {1.0f});
+    auto* pads_arg = builder.MakeInitializer<int64_t>({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<int64_t>{0, 0, 0, 0, 0, 0});
+  };
+
+  auto pre_graph_checker = [](Graph&) { return Status::OK(); };
+  auto post_graph_checker = [](Graph& graph) {
+    std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformerL1");
+  ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<PadFusion>()));
+  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<int64_t>& pads_data) {
+    auto build_test_case = [&](ModelTestBuilder& builder) {
+      auto* data_arg = builder.MakeInput<float>(std::nullopt);
+      auto* weight_arg = builder.MakeInitializer<float>({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<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformerL1");
+    ASSERT_STATUS_OK(rule_transformer->Register(std::make_unique<PadFusion>()));
+    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";