Cleanup irbuilder

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/converter.py

@@ -120,6 +120,7 @@ def set_type_info(value: ir.Value, typeinfo: ta.TypeAnnotationValue) -> None:
         value.type = type_and_shape.type
         value.shape = type_and_shape.shape
     except AttributeError:
+        # TODO: This needs to be fixed.
         pass
     value.meta["typeinfo"] = typeinfo
 
@@ -166,7 +167,6 @@ def __init__(
         source: Optional[str] = None,
         default_opset: Optional[values.Opset] = None,
     ):
-        self.ir_builder = irbuilder.IRBuilder()
         self.source = source
         if global_names is not None:
             # We make a copy in case function eval modifies it.
@@ -1115,7 +1115,7 @@ def ret(exp, i, suffix):
                 t = None
             else:
                 t = self.returntype[i]
-            self._current_fn.append_output(
+            self._current_fn.outputs.append(
                 make_value(return_var.name, t, self._source_of(stmt))
             )
             return return_var
@@ -1324,7 +1324,7 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
             [],
         )
 
-        self._current_fn.append_output(
+        self._current_fn.outputs.append(
             make_value(
                 o_cond_out,
                 onnx_types.BOOL,
@@ -1342,7 +1342,7 @@ def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
                 ov = self._emit_copy(ov, pv)
             # TODO: retrieve variable type for the annotation if any.
             typeinfo = None
-            self._current_fn.append_output(
+            self._current_fn.outputs.append(
                 make_value(ov.name, typeinfo, self._source_of(loop_stmt))
             )
         body = self._exit_scope()
@@ -1389,7 +1389,7 @@ def _translate_block(
                     # To return an outer-scope variable, an ONNX Graph has to
                     # use an explicit copy via Identity.
                     output = self._emit_copy(output, pvar)
-                self._current_fn.append_output(
+                self._current_fn.outputs.append(
                     make_value(
                         output.name,
                         pv_val.typeinfo,
@@ -1413,7 +1413,7 @@ def _translate_block(
 
                 # TODO: retrieve the annotation if any.
                 typeinfo = None
-                self._current_fn.append_output(make_value(ovar.name, typeinfo, source))
+                self._current_fn.outputs.append(make_value(ovar.name, typeinfo, source))
         graph = self._exit_scope()
         return graph.graph
 

onnxscript/evaluator.py

@@ -447,7 +447,7 @@ def make_tensor_name() -> str:
     model = onnx.helper.make_model(  # noqa: TID251
         graph,
         opset_imports=[opset_id],
-        ir_version=irbuilder.select_ir_version(schema.since_version, domain=schema.domain),
+        ir_version=values.select_ir_version(schema.since_version, domain=schema.domain),
     )
     model = onnx.shape_inference.infer_shapes(model)
 

onnxscript/irbuilder.py

@@ -16,21 +16,8 @@
 
 logger = logging.getLogger("onnxscript")
 
-
-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]
-
-
 TypeAnnotationValue = onnxscript.type_annotation.TypeAnnotationValue
 
-
 class IRFunction(ir.Function):
     """Represents a function in the IR."""
 
@@ -43,11 +30,6 @@ def __init__(self, name: str, domain: str = "") -> None:
         self.nested_functions: dict[str, IRFunction] = {}
         self.outer_scope_variables: dict[Any, Any] = {}
 
-    @property
-    def docstring(self) -> str:
-        """Returns the docstring of this function."""
-        return self.doc_string or ""
-
     @property
     def assigned_names(self) -> Sequence[str]:
         """Returns the list of variables assigned to by this function."""
@@ -84,9 +66,6 @@ def append_parameter(self, parameter: ir.Value | ir.Attr) -> None:
                 raise TypeError(f"Expected ir.Value or ir.Attr, got {type(parameter)}")
             self.attributes.add(parameter)
 
-    def append_output(self, var: ir.Value) -> None:
-        self.outputs.append(var)
-
     def add_nested_function(self, fun: IRFunction) -> None:
         self.nested_functions[fun.name] = fun
 
@@ -117,31 +96,3 @@ def to_function_proto(self) -> onnx.FunctionProto:
         """Converts this instance into a `onnx.FunctionProto`."""
         return ir.to_proto(self)
 
-
-# IRBuilder: abstracts out details of the IR in the python-to-IR converter
-
-
-def set_type_info(value: ir.Value, typeinfo: TypeAnnotationValue) -> None:
-    """Sets the type information on an IR value."""
-    try:
-        type_and_shape = ir.from_proto(typeinfo.to_type_proto())
-        value.type = type_and_shape.type
-        value.shape = type_and_shape.shape
-    except AttributeError:
-        pass
-    value.meta["typeinfo"] = typeinfo
-
-
-def make_value(
-    varname: str, typeinfo: TypeAnnotationValue, sourceinfo: SourceInfo
-) -> ir.Value:
-    value = ir.Value(name=varname)
-    value.meta.setdefault("sourceinfo", sourceinfo)
-    if typeinfo is not None:
-        set_type_info(value, typeinfo)
-    return value
-
-
-class IRBuilder:
-    def __init__(self):
-        self.functions = {}

onnxscript/values.py

@@ -457,7 +457,7 @@ def _op_schema_from_function_ir(
         function_ir.name,
         opset.domain,
         since_version=opset.version,
-        doc=function_ir.docstring,
+        doc=function_ir.doc_string or "",
         inputs=formal_inputs,
         outputs=formal_outputs,
         type_constraints=[constraint.as_tuple() for constraint in type_to_constraint.values()],