Move to_model_proto

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>

Author: gramalingam

onnxscript/converter.py

@@ -1475,7 +1475,6 @@ def translate_function_def(self, stmt: ast.FunctionDef) -> irbuilder.IRFunction:
             self._current_fn = self.ir_builder.new_function(stmt.name, domain, True)
             self._analyzer = analysis.AstAnalyzer(stmt, self._message, self.globals)
             fn_ir = self._translate_function_def_common(stmt)
-            fn_ir.debug_print()
             self.this_module.add_function_def(fn_ir)
             self._analyzer = None
             return fn_ir

onnxscript/converter_test.py

@@ -235,7 +235,7 @@ def test_unary_op(self):
     def test_subfunction_check_model(self):
         from tests.models import subfunction
 
-        model = subfunction.MyElu.function_ir.to_model_proto(producer_name="p2o")
+        model = subfunction.MyElu.to_model_proto(producer_name="p2o")
         model = onnx.shape_inference.infer_shapes(model)
         onnx.checker.check_model(model)
 

onnxscript/irbuilder.py

@@ -9,30 +9,23 @@
 
 import onnx
 import onnx_ir as ir
-from onnx import helper
-from onnx.defs import onnx_opset_version
 
-import onnxscript
 import onnxscript.type_annotation
 from onnxscript import values
-from onnxscript.onnx_types import ONNXType
 from onnxscript.sourceinfo import SourceInfo
 
 logger = logging.getLogger("onnxscript")
 
 
-def _format(seq: Sequence[Any], prefix: str, sep: str, suffix: str, formatter=str):
-    """Formats a sequence of objects into a string."""
-    return prefix + sep.join([formatter(x) for x in seq]) + suffix
-
-
 def select_ir_version(version: int, domain: str = "") -> int:
     """Selects a suitable ONNX ir_version for a given opset version."""
     if domain == "":
         domain = "ai.onnx"
-    if (domain, version) not in helper.OP_SET_ID_VERSION_MAP:
-        return max(v for k, v in helper.OP_SET_ID_VERSION_MAP.items() if k[0] == "ai.onnx")
-    return helper.OP_SET_ID_VERSION_MAP[domain, version]
+    if (domain, version) not in onnx.helper.OP_SET_ID_VERSION_MAP:
+        return max(
+            v for k, v in onnx.helper.OP_SET_ID_VERSION_MAP.items() if k[0] == "ai.onnx"
+        )
+    return onnx.helper.OP_SET_ID_VERSION_MAP[domain, version]
 
 
 TypeAnnotationValue = onnxscript.type_annotation.TypeAnnotationValue
@@ -76,9 +69,7 @@ def assigned_names(self) -> Sequence[str]:
 
     @property
     def inputs(self) -> Sequence[ir.Value]:
-        return (
-            self.ir_function.inputs
-        )  # [var for var in self.ordered_inputs_and_attrs if isinstance(var, IRVar)]
+        return self.ir_function.inputs
 
     @property
     def attrs(self) -> Sequence[ir.Attr]:
@@ -116,106 +107,9 @@ def add_attr_parameter(self, attr: ir.Attr) -> None:
         self.ordered_inputs_and_attrs.append(attr)
         self.ir_function.attributes.add(attr)
 
-    def debug_print(self):
-        if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(str(self.ir_function))
-
     def add_nested_function(self, fun: IRFunction) -> None:
         self.nested_functions[fun.name] = fun
 
-    def to_model_proto(
-        self,
-        functions=None,
-        io_types: Optional[ONNXType] = None,
-        input_types: Optional[Sequence[ONNXType]] = None,
-        output_types: Optional[Sequence[ONNXType]] = None,
-        value_infos: dict[str, ONNXType] | None = None,
-        opset_version: int | None = None,
-        **kwargs,
-    ) -> onnx.ModelProto:
-        """Converts this instance into a `onnx.ModelProto`.
-
-        Args:
-            functions: A list of functions to include in the model.
-                By default, all functions called at least once are included.
-            io_types: When specified, all the inputs/outputs of the model
-                are set to be of this type.
-            input_types: When specified, all the inputs of the model
-                are set to be of the corresponding type in this list.
-            output_types: When specified, all the outputs of the model
-                are set to be of the corresponding type in this list.
-            value_infos: A dictionary mapping intermediate variable names to ONNX types.
-                Used to set value_info for intermediate variables.
-            opset_version: The standard opset version to use for the model if it
-                cannot be inferred. Otherwise defaults to the current opset version.
-            kwargs: Additional parameters given to function :func:`onnx.helper.make_model`.
-
-        Returns:
-            An instance of :class:`onnx.ModelProto`.
-        """
-        value_infos = (
-            [
-                onnx.helper.make_value_info(name, type.to_type_proto())
-                for name, type in value_infos.items()
-            ]
-            if value_infos
-            else None
-        )
-        sub_functions = self.get_called_functions()
-        graph = self.to_graph_proto(use_default_type=False)
-        if value_infos:
-            graph.value_info.extend(value_infos)
-        if io_types is not None:
-            for input in graph.input:
-                if not input.HasField("type"):
-                    input.type.CopyFrom(io_types.to_type_proto())
-            for output in graph.output:
-                if not output.HasField("type"):
-                    output.type.CopyFrom(io_types.to_type_proto())
-        if input_types is not None:
-            for input, type in zip(graph.input, input_types):
-                input.type.CopyFrom(type.to_type_proto())
-        if output_types is not None:
-            for output, type in zip(graph.output, output_types):
-                output.type.CopyFrom(type.to_type_proto())
-        if functions is None:
-            functions = sub_functions.values()
-        else:
-
-            def to_proto(f):
-                if isinstance(f, onnx.FunctionProto):
-                    return f
-                if isinstance(f, onnxscript.OnnxFunction):
-                    return f.to_function_proto()
-                raise TypeError("Expected a value of type FunctionProto of OnnxFunction")
-
-            functions = [to_proto(f) for f in functions]
-
-        opsets = self.ir_function.opset_imports.copy()
-
-        for proto in functions:
-            if proto.domain not in opsets:
-                opsets[proto.domain] = 1
-            # TODO(rama): Handle conflicts with appropriate error/warning message.
-            for opset in proto.opset_import:
-                if opset.domain not in opsets:
-                    opsets[opset.domain] = opset.version
-
-        if "" not in opsets:
-            # No operator is using the standard opset.
-            # Use the specified version if provided or the default value.
-            opsets[""] = opset_version if opset_version is not None else onnx_opset_version()
-
-        if "ir_version" not in kwargs:
-            kwargs["ir_version"] = select_ir_version(opsets[""])
-        opset_imports = [
-            onnx.helper.make_opsetid(domain, version) for domain, version in opsets.items()
-        ]
-
-        return helper.make_model(
-            graph, opset_imports=opset_imports, functions=functions, **kwargs
-        )
-
     def get_called_functions(self) -> dict[str, onnx.FunctionProto]:
         called_functions: dict[str, values.OnnxFunction] = {}
 
@@ -297,7 +191,7 @@ def add_stmt(
         attrs: Sequence[ir.Attr],
     ) -> Sequence[ir.Value]:
         output_values = [ir.Value(name=o) for o in results]
-        attributes = attrs  # [ir.from_proto(a.attr_proto) for a in attrs]
+        attributes = attrs
         node = ir.Node(
             domain=callee.opset.domain,
             version=callee.opset.version,

onnxscript/values.py

@@ -30,11 +30,23 @@
 from onnxscript import irbuilder, sourceinfo, type_annotation
 from onnxscript._internal import ast_utils, deprecation
 from onnxscript.ir import _schemas
+from onnxscript.onnx_types import ONNXType
 
 _R = TypeVar("_R")
 _P = ParamSpec("_P")
 
 
+def select_ir_version(version: int, domain: str = "") -> int:
+    """Selects a suitable ONNX ir_version for a given opset version."""
+    if domain == "":
+        domain = "ai.onnx"
+    if (domain, version) not in onnx.helper.OP_SET_ID_VERSION_MAP:
+        return max(
+            v for k, v in onnx.helper.OP_SET_ID_VERSION_MAP.items() if k[0] == "ai.onnx"
+        )
+    return onnx.helper.OP_SET_ID_VERSION_MAP[domain, version]
+
+
 _ATTRIBUTE_TYPE_TO_PYTHON_TYPE = {
     onnx.defs.OpSchema.AttrType.FLOAT: float,
     onnx.defs.OpSchema.AttrType.INT: int,
@@ -609,7 +621,103 @@ def to_model_proto(self, **kwargs):
 
         # Merge kwargs specified in script-decorator with those specified in this call.
         merged_kw_args = {**self.kwargs, **kwargs}
-        return self.function_ir.to_model_proto(**merged_kw_args)
+        return self._to_model_proto(**merged_kw_args)
+
+    def _to_model_proto(
+        self,
+        functions=None,
+        io_types: Optional[ONNXType] = None,
+        input_types: Optional[Sequence[ONNXType]] = None,
+        output_types: Optional[Sequence[ONNXType]] = None,
+        value_infos: dict[str, ONNXType] | None = None,
+        opset_version: int | None = None,
+        **kwargs,
+    ) -> onnx.ModelProto:
+        """Converts this instance into a `onnx.ModelProto`.
+
+        Args:
+            functions: A list of functions to include in the model.
+                By default, all functions called at least once are included.
+            io_types: When specified, all the inputs/outputs of the model
+                are set to be of this type.
+            input_types: When specified, all the inputs of the model
+                are set to be of the corresponding type in this list.
+            output_types: When specified, all the outputs of the model
+                are set to be of the corresponding type in this list.
+            value_infos: A dictionary mapping intermediate variable names to ONNX types.
+                Used to set value_info for intermediate variables.
+            opset_version: The standard opset version to use for the model if it
+                cannot be inferred. Otherwise defaults to the current opset version.
+            kwargs: Additional parameters given to function :func:`onnx.helper.make_model`.
+
+        Returns:
+            An instance of :class:`onnx.ModelProto`.
+        """
+        value_infos = (
+            [
+                onnx.helper.make_value_info(name, type.to_type_proto())
+                for name, type in value_infos.items()
+            ]
+            if value_infos
+            else None
+        )
+
+        graph = self.function_ir.to_graph_proto(use_default_type=False)
+        if value_infos:
+            graph.value_info.extend(value_infos)
+        if io_types is not None:
+            for input in graph.input:
+                if not input.HasField("type"):
+                    input.type.CopyFrom(io_types.to_type_proto())
+            for output in graph.output:
+                if not output.HasField("type"):
+                    output.type.CopyFrom(io_types.to_type_proto())
+        if input_types is not None:
+            for input, type in zip(graph.input, input_types):
+                input.type.CopyFrom(type.to_type_proto())
+        if output_types is not None:
+            for output, type in zip(graph.output, output_types):
+                output.type.CopyFrom(type.to_type_proto())
+        if functions is None:
+            sub_functions = self.function_ir.get_called_functions()
+            functions = sub_functions.values()
+        else:
+
+            def to_proto(f):
+                if isinstance(f, onnx.FunctionProto):
+                    return f
+                if isinstance(f, OnnxFunction):
+                    return f.to_function_proto()
+                raise TypeError("Expected a value of type FunctionProto of OnnxFunction")
+
+            functions = [to_proto(f) for f in functions]
+
+        opsets = self.function_ir.ir_function.opset_imports.copy()
+
+        for proto in functions:
+            if proto.domain not in opsets:
+                opsets[proto.domain] = 1
+            # TODO(rama): Handle conflicts with appropriate error/warning message.
+            for opset in proto.opset_import:
+                if opset.domain not in opsets:
+                    opsets[opset.domain] = opset.version
+
+        if "" not in opsets:
+            # No operator is using the standard opset.
+            # Use the specified version if provided or the default value.
+            opsets[""] = (
+                opset_version if opset_version is not None else onnx.defs.onnx_opset_version()
+            )
+
+        if "ir_version" not in kwargs:
+            kwargs["ir_version"] = select_ir_version(opsets[""])
+        opset_imports = [
+            onnx.helper.make_opsetid(domain, version) for domain, version in opsets.items()
+        ]
+
+        return onnx.helper.make_model(
+            graph, opset_imports=opset_imports, functions=functions, **kwargs
+        )
 
 
 class TracedOnnxFunction(Op):

tests/common/onnx_script_test_case.py

@@ -144,7 +144,7 @@ def _create_model_from_param(
             # there is not way from the onnx test case's model and feed to get TypeProto
             # in order to build a model.
             # we have to resolve the TypeProto from script function.
-            local_function_model_proto = param.function.function_ir.to_model_proto(
+            local_function_model_proto = param.function.to_model_proto(
                 ir_version=ir_version
             )
             input_value_infos = []
@@ -202,7 +202,7 @@ def run_converter_test(
                 param, onnx_case_model, ir_version=ir_version
             )
         else:
-            model = param.function.function_ir.to_model_proto(
+            model = param.function.to_model_proto(
                 producer_name="call_clip", ir_version=ir_version
             )
         try: