Update opset imports in version_converter

Author: shubhambhokare1

onnxscript/version_converter/_version_converter.py

@@ -220,6 +220,38 @@ class _VersionConverter:
     def __init__(self, target_version: int):
         self.target_version = target_version
 
+    def _maybe_set_opset_version(
+        self, model_or_function: ir.Model | ir.Function, domain: str, version: int | None
+    ):
+        """Set the opset version for the domain."""
+        current_version = model_or_function.opset_imports.get(domain)
+        if version is None or current_version is None:
+            return
+        if domain == "":
+            model_or_function.opset_imports[domain] = max(version, current_version)
+            return
+        elif domain == "ai.onnx":
+            model_or_function.opset_imports[domain] = max(version, current_version)
+            return
+        else:
+            return
+
+    def _update_opset_imports(self, model: ir.Model) -> None:
+        """Collect all opsets used and add opset imports to the model and functions."""
+        for node in ir.traversal.RecursiveGraphIterator(model.graph):
+            domain = node.domain
+            self._maybe_set_opset_version(model, domain, node.version)
+
+        for function in model.functions.values():
+            for node in ir.traversal.RecursiveGraphIterator(function):
+                domain = node.domain
+                self._maybe_set_opset_version(function, domain, node.version)
+            for domain, version in function.opset_imports.items():
+                # Add all opsets used in the function to the model, because ONNX Runtime
+                # does not handle adding the opset imports to the model after inlining during inference.
+                # This should happen after all opsets are collected for the function from its nodes.
+                self._maybe_set_opset_version(model, domain, version)
+
     def _upgrade_version(self, node: ir.Node, opset_version: int, up_conversion: bool) -> None:
         if up_conversion is True:
             node.version = opset_version + 1
@@ -320,6 +352,8 @@ def visit_model(self, model: ir.Model) -> None:
                 return None
         self.model_version = model_version
         self.visit_graph(model.graph)
+        # Finally, update the opset imports for the model
+        self._update_opset_imports(model)
         return None
 
 

onnxscript/version_converter/_version_converter_test.py

@@ -115,6 +115,7 @@ def test_version_convert_compatible(self):
         )
         model = ir.serde.deserialize_model(model_proto)
         target_version = 19
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(0).op_type, "Constant")
@@ -123,6 +124,7 @@ def test_version_convert_compatible(self):
         self.assertEqual(model.graph.node(1).version, 19)
         self.assertEqual(model.graph.node(4).op_type, "MatMul")
         self.assertEqual(model.graph.node(4).version, 19)
+        self.assertEqual(model.opset_imports[""], 19)
 
 
 class VersionConverter19to20Test(unittest.TestCase):
@@ -142,6 +144,7 @@ def test_version_convert_compatible(self):
         )
         model = ir.serde.deserialize_model(model_proto)
         target_version = 20
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(0).op_type, "Constant")
@@ -153,6 +156,7 @@ def test_version_convert_compatible(self):
         self.assertEqual(model.graph.node(3).op_type, "DFT")
         self.assertEqual(model.graph.node(3).version, 20)
         self.assertEqual(len(model.graph.node(3).inputs), 2)
+        self.assertEqual(model.opset_imports[""], 20)
 
     def test_version_convert_gridsample_linear(self):
         model_proto = onnx.parser.parse_model(
@@ -175,6 +179,7 @@ def test_version_convert_gridsample_linear(self):
         self.assertEqual(model.graph.node(4).attributes["mode"].value, "bilinear")
 
         target_version = 20
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(0).op_type, "Constant")
@@ -184,6 +189,7 @@ def test_version_convert_gridsample_linear(self):
         self.assertEqual(model.graph.node(4).op_type, "GridSample")
         self.assertEqual(model.graph.node(4).version, 20)
         self.assertEqual(model.graph.node(4).attributes["mode"].value, "linear")
+        self.assertEqual(model.opset_imports[""], 20)
 
     def test_version_convert_gridsample_cubic(self):
         model_proto = onnx.parser.parse_model(
@@ -206,6 +212,7 @@ def test_version_convert_gridsample_cubic(self):
         self.assertEqual(model.graph.node(4).attributes["mode"].value, "bicubic")
 
         target_version = 20
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(0).op_type, "Constant")
@@ -215,6 +222,7 @@ def test_version_convert_gridsample_cubic(self):
         self.assertEqual(model.graph.node(4).op_type, "GridSample")
         self.assertEqual(model.graph.node(4).version, 20)
         self.assertEqual(model.graph.node(4).attributes["mode"].value, "cubic")
+        self.assertEqual(model.opset_imports[""], 20)
 
     def test_version_convert_inline(self):
         model_proto = onnx.parser.parse_model(
@@ -238,6 +246,7 @@ def test_version_convert_inline(self):
         )
         model = ir.serde.deserialize_model(model_proto)
         target_version = 20
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(0).op_type, "Constant")
@@ -250,6 +259,7 @@ def test_version_convert_inline(self):
         self.assertEqual(model.graph.node(6).op_type, "DFT")
         self.assertEqual(model.graph.node(6).version, 20)
         self.assertEqual(len(model.graph.node(6).inputs), 2)
+        self.assertEqual(model.opset_imports[""], 20)
 
 
 class VersionConverter20to21Test(unittest.TestCase):
@@ -267,6 +277,7 @@ def test_version_groupnorm(self):
         )
         model = ir.serde.deserialize_model(model_proto)
         target_version = 21
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(3).op_type, "Reshape")
@@ -283,6 +294,7 @@ def test_version_groupnorm(self):
         self.assertEqual(model.graph.node(8).version, 21)
         self.assertEqual(model.graph.node(9).op_type, "GroupNormalization")
         self.assertEqual(model.graph.node(9).version, 21)
+        self.assertEqual(model.opset_imports[""], 21)
 
     def test_version_groupnorm_no_bias(self):
         model_proto = onnx.parser.parse_model(
@@ -298,10 +310,12 @@ def test_version_groupnorm_no_bias(self):
         )
         model = ir.serde.deserialize_model(model_proto)
         target_version = 21
+        self.assertEqual(model.opset_imports[""], 18)
         version_converter.convert_version(model, target_version=target_version)
 
         self.assertEqual(model.graph.node(0).op_type, "GroupNormalization")
         self.assertEqual(model.graph.node(0).version, 20)
+        self.assertEqual(model.opset_imports[""], 21)
 
 
 class VersionConverter23to24Test(unittest.TestCase):