From 4faee2e44c2029d7f5fac1d927f8c843b22189ca Mon Sep 17 00:00:00 2001
From: "G. Ramalingam" <grama@microsoft.com>
Date: Fri, 14 Jul 2023 14:44:59 -0700
Subject: [PATCH] Fix issue in constant-propagation inside function subgraph
 (#16330)

### Description

The SequenceMap function-op has a graph-attribute. ORT's
constant-folding optimization may identify constant-expressions inside
the subgraph and promote them to constants, stored as initializers in
the main graph. When it does this, the optimization updates the subgraph
to remove the corresponding nodes.

When we expand a SequenceMap node by inlining its function-expansion, we
need to use this updated subgraph. However, the existing code uses the
original graph-attribute (GraphProto), instead of regenerating it from
the modified subgraph. This results in producing a graph with duplicate
definitions for the constant-folded variable, resulting in an error
during graph-resolve.

This PR fixes this issue (just a single line fix), and adds a test-case
to cover this scenario.

---------

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
Co-authored-by: Suryaprakash Shanmugam <suryaprakash.shanmugam@intel.com>
---
 onnxruntime/core/graph/graph.cc               |  2 +-
 .../openvino/ov_versions/capability.cc        | 10 +++++---
 .../openvino/ov_versions/data_ops.cc          |  1 +
 onnxruntime/test/framework/function_test.cc   | 24 +++++++++++++++++++
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc
index 80425cb7ad..3c45e910ae 100644
--- a/onnxruntime/core/graph/graph.cc
+++ b/onnxruntime/core/graph/graph.cc
@@ -585,7 +585,7 @@ bool Node::TryGetFunctionProto(ONNX_NAMESPACE::FunctionProto& onnx_function_prot
     // Check if this node has a schema defined function proto.
     if (op_->HasContextDependentFunction()) {
       NodeProto node_proto;
-      ToProto(node_proto);
+      ToProto(node_proto, true);
       std::vector<TypeProto> input_types;
       for (size_t i = 0, n = InputDefs().size(); i < n; i++) {
         auto p_node_arg = InputDefs().at(i);
diff --git a/onnxruntime/core/providers/openvino/ov_versions/capability.cc b/onnxruntime/core/providers/openvino/ov_versions/capability.cc
index ae87961dc8..251f475525 100644
--- a/onnxruntime/core/providers/openvino/ov_versions/capability.cc
+++ b/onnxruntime/core/providers/openvino/ov_versions/capability.cc
@@ -75,11 +75,15 @@ std::vector<std::unique_ptr<ComputeCapability>> GetCapability::Execute() {
     }
 
     const auto& nodes = graph_viewer_.GetNodesInTopologicalOrder();
+
+    // Handle cases where lone, reoccuring Ops in smaller models cannot be supported in OpenVINO
+    // If only a node of the same lone,unsupported type is present, then do not proceed with the subgraph
+    const auto& node = graph_viewer_.GetNode(nodes[0]);
+    if (data_ops_->IsOpSupportedOnlyInModel(node->OpType()))
+      return result;
+
     // Nodes that work well in models but not as a single node
     if (nodes.size() == 1) {
-      const auto& node = graph_viewer_.GetNode(nodes[0]);
-      if (data_ops_->IsOpSupportedOnlyInModel(node->OpType()))
-        return result;
       // If reshape is not an intermediate node, shape needs to be an initializer
       if (data_ops_->SpecialConditionForClusterSizeOne(ng_required_initializers, node)) {
         return result;
diff --git a/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc b/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc
index 30ba5a3198..f30f251f75 100644
--- a/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc
+++ b/onnxruntime/core/providers/openvino/ov_versions/data_ops.cc
@@ -30,6 +30,7 @@ namespace openvino_ep {
 
 // Ops which are supported only in models(as intermediate nodes) and not in unit tests
 std::set<std::string> ops_supported_only_in_model = {
+    "Add",
     "Cast",
     "Concat",
     "ConstantOfShape",
diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc
index c68a5f4abb..e126979532 100644
--- a/onnxruntime/test/framework/function_test.cc
+++ b/onnxruntime/test/framework/function_test.cc
@@ -506,5 +506,29 @@ TEST(FunctionTest, UnusedFunctionInputs) {
   Check(code, "x", {1.0, 2.0, 3.0}, "y", {1.0, 4.0, 9.0});
 }
 
+// Test constant-folding inside a sub-graph is handled correctly
+// for functions that are inlined.
+TEST(FunctionTest, ConstantFoldingInSubGraph) {
+  const char* code = R"(
+    <ir_version: 8, opset_import: [ "" : 17 ]>
+    agraph (float[N] X) => (float[M] Y)  {
+        seq1 = SequenceConstruct(X, X, X)
+        seq2 = SequenceMap (seq1) <body =
+            add1 (float[K] Z) => (float[K] W) {
+                C1 = Constant <value = float {1.0}> ()
+                C2 = Constant <value = float {1.0}> ()
+                # C is a constant, which will be constant-folded into an initializer out of the sub-graph.
+                C = Add (C1, C2)
+                # After optimization, only following Add will be left in this sub-graph.
+                W = Add (Z, C)
+            }
+        >
+        Y = ConcatFromSequence <axis=0> (seq2)
+    }
+  )";
+
+  Check(code, "X", {1.0, 2.0, 3.0}, "Y", {3.0, 4.0, 5.0, 3.0, 4.0, 5.0, 3.0, 4.0, 5.0});
+}
+
 }  // namespace test
 }  // namespace onnxruntime