Fix function inliner bug re. outer-scope names (microsoft#14734)

### Description Fix the function inliner logic for renaming variables. Typically, a FunctionProto does not contain references to outer-scope names. However, special cases, such as the function-expansion of SequenceMap, can generate such FunctionProtos. Extend the renaming logic to ensure that references to outer-scope names are not renamed. ### Motivation and Context Fixes onnx/onnx#4892 Signed-off-by: Ganesan Ramalingam <[email protected]>
chenfucn · Mar 2, 2023 · 2facc5e · 2facc5e
1 parent 603026f
commit 2facc5e
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/onnxruntime/core/graph/function_utils.cc b/onnxruntime/core/graph/function_utils.cc
@@ -348,7 +348,7 @@ class Inliner {
     name = new_name;
   }
 
-  void rename(std::string& name) {
+  void rename(std::string& name, bool is_new_def) {
     if (name.empty()) return;
     for (auto i = rename_scopes.size(); i > 0; --i) {
       const auto& map = rename_scopes[i - 1];
@@ -358,7 +358,10 @@ class Inliner {
         return;
       }
     }
-    make_unique(name);
+    if (is_new_def) {
+      make_unique(name);
+    }
+    // Otherwise, it is a reference to an outer-scope variable that should not be renamed.
   }
 
   template <bool isOutput>
@@ -396,10 +399,10 @@ class Inliner {
       n.set_name(prefix + n.name());
 
     for (auto& x : *n.mutable_input()) {
-      rename(x);
+      rename(x, false);
     }
     for (auto& y : *n.mutable_output()) {
-      rename(y);
+      rename(y, true);
     }
     auto& attributes = *n.mutable_attribute();
     for (auto attr_iter = attributes.begin(); attr_iter != attributes.end();) {

diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc
@@ -352,5 +352,25 @@ TEST(FunctionTest, Variadics) {
   ASSERT_STATUS_OK(session_object.Load(model_uri));
   ASSERT_STATUS_OK(session_object.Initialize());
 }
+
+// Test use of outer-scope names inside sub-graphs in functions that are inlined.
+TEST(FunctionTest, OuterScopeName) {
+  const char* code = R"(
+        <ir_version: 8, opset_import: [ "" : 17 ]>
+        agraph (float[N] x) => (float[N] y)
+        {
+            xseq = SequenceConstruct (x)
+            zeros = Constant <value = float[3] {0.0, 0.0, 0.0}> ()
+            yseq = SequenceMap (xseq) <body =
+              zeropad (float[3] lx) => (float[6] ly) {
+                ly = Concat <axis = 0> (lx, zeros)
+              }>
+            zero = Constant <value = int64{0}> ()
+            y = SequenceAt (yseq, zero)
+        }
+        )";
+
+  Check(code, "x", {1.0, 2.0, 3.0}, "y", {1.0, 2.0, 3.0, 0.0, 0.0, 0.0});
+}
 }  // namespace test
 }  // namespace onnxruntime