[pass] Update LiftSubgraphInitializersToMainGraphPass to disallow variable shadowing (#2348)

Variable shadowing (reusing value names) is disallowed in ONNX across
the main graph and subgraphs according to the spec
(https://github.com/onnx/onnx/pull/6955/files). This change updates to
the logic to check and raise on such cases. A subsequent PR will
implement #1432 to allow users to fix names explicitly.

Author: justinchuby

onnxscript/ir/passes/common/constant_manipulation.py

@@ -137,11 +137,37 @@ class LiftSubgraphInitializersToMainGraphPass(ir.passes.InPlacePass):
     This pass lifts the initializers of a subgraph to the main graph.
     It is used to ensure that the initializers are available in the main graph
     for further processing or optimization.
+
+    Initializers that are also graph inputs will not be lifted.
+
+    Preconditions:
+        - All initializers in the model must have unique names across the main graph and subgraphs.
     """
 
+    def requires(self, model: ir.Model) -> None:
+        """Ensure all initializer names are unique."""
+        registered_initializer_names: set[str] = set()
+        duplicated_initializers: list[ir.Value] = []
+        for graph in model.graphs():
+            for initializer in graph.initializers.values():
+                if initializer.name is None:
+                    raise ir.passes.PreconditionError(
+                        f"Initializer name is None. Please ensure all initializers have unique names: {initializer!r}"
+                    )
+                if initializer.name in registered_initializer_names:
+                    duplicated_initializers.append(initializer)
+                else:
+                    registered_initializer_names.add(initializer.name)
+        if duplicated_initializers:
+            raise ir.passes.PreconditionError(
+                "Found duplicated initializers in the model. "
+                "Initializer name must be unique across the main graph and subgraphs. "
+                "Please ensure all initializers have unique names. Duplicated: "
+                f"{duplicated_initializers!r}"
+            )
+
     def call(self, model: ir.Model) -> ir.passes.PassResult:
         count = 0
-        registered_initializer_names: dict[str, int] = {}
         for graph in model.graphs():
             if graph is model.graph:
                 continue
@@ -156,15 +182,6 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
                     continue
                 # Remove the initializer from the subgraph
                 graph.initializers.pop(name)
-                # To avoid name conflicts, we need to rename the initializer
-                # to a unique name in the main graph
-                if name in registered_initializer_names:
-                    name_count = registered_initializer_names[name]
-                    initializer.name = f"{name}_{name_count}"
-                    registered_initializer_names[name] = name_count + 1
-                else:
-                    assert initializer.name is not None
-                    registered_initializer_names[initializer.name] = 1
                 model.graph.register_initializer(initializer)
                 count += 1
                 logger.debug(

onnxscript/ir/passes/common/constant_manipulation_test.py

@@ -248,12 +248,12 @@ def test_not_lifting_constants_to_initializers_when_it_is_output(self):
 class TestLiftSubgraphInitializersToMainGraphPass(unittest.TestCase):
     @parameterized.parameterized.expand(
         [
-            ("then_initializer", "else_initializer"),
-            ("initializer", "initializer"),
+            ("unique_init_names", "then_initializer", "else_initializer"),
+            ("duplicated_init_names", "initializer", "initializer"),
         ]
     )
     def test_pass_with_lifting_constants_to_initializers_within_subgraph(
-        self, then_initializer_name, else_initializer_name
+        self, _: str, then_initializer_name: str, else_initializer_name: str
     ):
         input_value = ir.Value(
             name="input", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
@@ -311,6 +311,13 @@ def test_pass_with_lifting_constants_to_initializers_within_subgraph(
             graph=main_graph,
             ir_version=10,
         )
+        if then_initializer_name == else_initializer_name:
+            with self.assertRaisesRegex(
+                ir.passes.PreconditionError,
+                "Initializer name must be unique across the main graph and subgraphs",
+            ):
+                constant_manipulation.LiftSubgraphInitializersToMainGraphPass()(model)
+            return
         result = constant_manipulation.LiftSubgraphInitializersToMainGraphPass()(model)
         self.assertTrue(result.modified)
 
@@ -325,12 +332,12 @@ def test_pass_with_lifting_constants_to_initializers_within_subgraph(
 
     @parameterized.parameterized.expand(
         [
-            ("then_initializer", "else_initializer"),
-            ("initializer", "initializer"),
+            ("unique_init_names", "then_initializer", "else_initializer"),
+            ("duplicated_init_names", "initializer", "initializer"),
         ]
     )
     def test_pass_does_not_lift_initialized_inputs_in_subgraph(
-        self, then_initializer_name, else_initializer_name
+        self, _: str, then_initializer_name: str, else_initializer_name: str
     ):
         input_value = ir.Value(
             name="input", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
@@ -390,6 +397,13 @@ def test_pass_does_not_lift_initialized_inputs_in_subgraph(
             graph=main_graph,
             ir_version=10,
         )
+        if then_initializer_name == else_initializer_name:
+            with self.assertRaisesRegex(
+                ir.passes.PreconditionError,
+                "Initializer name must be unique across the main graph and subgraphs",
+            ):
+                constant_manipulation.LiftSubgraphInitializersToMainGraphPass()(model)
+            return
         result = constant_manipulation.LiftSubgraphInitializersToMainGraphPass()(model)
         self.assertTrue(result.modified)