Merge attrs and attr_protos in IRFunction | chore(irbuilder)

ghstack-source-id: 3582671
Pull Request resolved: #625

Author: justinchuby

onnxscript/converter.py

@@ -1326,6 +1326,7 @@ def translate_function_def(self, fn: ast.FunctionDef) -> irbuilder.IRFunction:
                 self.ir_builder.add_attr_parameter(
                     self.current_fn,
                     x.arg,
+                    ta.pytype_to_attrtype(typeinfo),
                     default_value,
                 )
                 self.bind(x.arg, values.AttrRef(x.arg, typeinfo, self.source_of(x)))

onnxscript/irbuilder.py

@@ -108,10 +108,16 @@ def _opt_var_to_str(x):
 
 
 class IRAttributeValue:
-    """An attribute value (representing an actual parameter)."""
+    """An attribute value (representing an actual parameter).
 
-    def __init__(self, attrproto) -> None:
+    Attributes:
+        attr_proto: The attribute proto
+        has_default: Whether the attribute has a default value.
+    """
+
+    def __init__(self, attrproto, has_default: bool) -> None:
         self.attr_proto = attrproto
+        self.has_default = has_default
 
     def __str__(self):
         if self.attr_proto.HasField("ref_attr_name"):
@@ -191,9 +197,7 @@ def __init__(self, name: str, domain: str = "") -> None:
         self.outputs: list[IRVar] = []
         self.stmts: list[IRStmt] = []
         # attribute parameters
-        self.attrs: list[str] = []
-        # attribute parameters with default value
-        self.attr_protos: list[IRAttributeValue] = []
+        self.attrs: list[IRAttributeValue] = []
         self.called_functions: dict[str, onnx.FunctionProto] = {}
         self.docstring: str = ""
         # a dictionary of nested function-definitions
@@ -207,11 +211,10 @@ def assigned_names(self) -> Sequence[str]:
 
     def __str__(self):
         attrs = _format(self.attrs, "<", ", ", ">") if self.attrs else ""
-        attr_protos = _format(self.attr_protos, "<", ", ", ">") if self.attr_protos else ""
         inputs = _format([x.typed_str() for x in self.inputs], "(", ", ", ")")
         outputs = _format([x.typed_str() for x in self.outputs], "(", ", ", ")")
         stmts = _format(self.stmts, "\n{\n   ", "\n   ", "\n}\n")
-        return f"{self.name} {attrs}{attr_protos}{inputs} => {outputs}{stmts}"
+        return f"{self.name} {attrs}{inputs} => {outputs}{stmts}"
 
     def append_docstring(self, docstring):
         self.docstring += docstring
@@ -225,11 +228,8 @@ def append_input(self, name: IRVar) -> None:
     def append_output(self, name: IRVar) -> None:
         self.outputs.append(name)
 
-    def add_attr_parameter(self, attr: str | IRAttributeValue) -> None:
-        if isinstance(attr, IRAttributeValue):
-            self.attr_protos.append(attr)
-        else:
-            self.attrs.append(attr)
+    def add_attr_parameter(self, attr: IRAttributeValue) -> None:
+        self.attrs.append(attr)
 
     def debug_print(self):
         if logger.isEnabledFor(logging.DEBUG):
@@ -398,19 +398,19 @@ def to_function_proto(self) -> onnx.FunctionProto:
             onnx.helper.make_opsetid(domain, version) for domain, version in opsets.items()
         ]
 
-        # attribute_proto is introduced in version onnx==1.13.0.
+        # attribute_proto is introduced in version onnx==1.14.0.
         # If this attribute is available, onnxscript uses it to
         # default values for attributes. The function has then two
         # lists, one list for attributes without default values,
         # another one for attributes with default values.
         # If this *attribute_proto* is not available,
-        # all attributes with a default value are moved to the first
+        # all attributes are moved to the first
         # list, default values are removed.
         # TODO: remove this when onnx with attribute_proto is released.
         if hasattr(onnx.FunctionProto, "attribute_proto"):
-            atts = self.attrs
+            attribute_names = [attr.name for attr in self.attrs if not attr.has_default]
         else:
-            atts = self.attrs + [a.attr_proto.name for a in self.attr_protos]
+            attribute_names = [attr.name for attr in self.attrs]
 
         f = helper.make_function(
             self.domain,
@@ -419,11 +419,13 @@ def to_function_proto(self) -> onnx.FunctionProto:
             outputs=[y.name for y in self.outputs],
             nodes=nodes,
             opset_imports=opset_imports,  # TODO
-            attributes=atts,
+            attributes=attribute_names,
             doc_string=self.docstring,
         )
         if hasattr(onnx.FunctionProto, "attribute_proto"):
-            f.attribute_proto.extend([a.attr_proto for a in self.attr_protos])
+            f.attribute_proto.extend(
+                [attr.attr_proto for attr in self.attrs if attr.has_default]
+            )
         return f
 
 
@@ -463,25 +465,35 @@ def add_input(
         v = IRVar(varname, type, info)
         fn.append_input(v)
 
-    def add_attr_parameter(self, fn: IRFunction, varname: str, default_value) -> None:
+    def add_attr_parameter(
+        self,
+        fn: IRFunction,
+        varname: str,
+        attribute_type: onnx.AttributeProto.AttributeType,
+        default_value,
+    ) -> None:
         if default_value is not None:
-            a = IRAttributeValue(helper.make_attribute(varname, default_value))
-            fn.add_attr_parameter(a)
+            fn.add_attr_parameter(
+                IRAttributeValue(
+                    helper.make_attribute(varname, default_value), has_default=True
+                )
+            )
         else:
-            fn.add_attr_parameter(varname)
+            proto = onnx.AttributeProto()
+            proto.name = varname
+            proto.type = attribute_type
+            fn.add_attr_parameter(IRAttributeValue(proto, has_default=False))
 
     def add_output(self, fn: IRFunction, varname: str, type, info) -> None:
         v = IRVar(varname, type, info)
         fn.append_output(v)
 
     def make_attr(self, attrname: str, attrval: Any) -> IRAttributeValue:
-        return IRAttributeValue(helper.make_attribute(attrname, attrval))
+        return IRAttributeValue(helper.make_attribute(attrname, attrval), has_default=True)
 
     def make_attr_ref(self, attrname: str, refname: str, pytype: type) -> IRAttributeValue:
         a = onnx.AttributeProto()
         a.name = attrname
         a.ref_attr_name = refname
-        type_ = ta.pytype_to_attrtype(pytype)
-        assert type_ is not None
-        a.type = type_
-        return IRAttributeValue(a)
+        a.type = ta.pytype_to_attrtype(pytype)
+        return IRAttributeValue(a, has_default=False)

onnxscript/type_annotation.py

@@ -59,7 +59,7 @@ def _is_primitive_attr_type(typeinfo: TypeAnnotationValue) -> bool:
 
 def pytype_to_attrtype(
     pytype: TypeAnnotationValue,
-) -> typing.Optional[onnx.AttributeProto.AttributeType]:
+) -> onnx.AttributeProto.AttributeType:
     pytype = _remove_annotation(pytype)
     if pytype in _PYTYPE_TO_ATTRTYPE_MAP:
         return _PYTYPE_TO_ATTRTYPE_MAP[pytype]
@@ -74,7 +74,7 @@ def pytype_to_attrtype(
         elt_type = get_args(pytype)[0]
         if elt_type in _LISTTYPE_TO_ATTRTYPE_MAP:
             return _LISTTYPE_TO_ATTRTYPE_MAP[elt_type]
-    return None
+    return onnx.AttributeProto.UNDEFINED
 
 
 def _is_tensor_type(typeinfo: TypeAnnotationValue) -> bool:

onnxscript/values.py

@@ -311,11 +311,10 @@ def param_schemas(self) -> tuple[ParamSchema, ...]:
         # The first len(func_ir.inputs) arguments are onnx inputs
         inputs = function_ir.inputs
         # The rest is onnx attributes
-        attributes = function_ir.attrs
         # Construct a dictionary of attributes with their names specified in the function
         # definition
         attr_name_to_protos = collections.OrderedDict(
-            (attr.name, attr) for attr in function_ir.attr_protos
+            (attr.name, attr) for attr in function_ir.attrs
         )
 
         # args with default value are attributes
@@ -325,26 +324,30 @@ def param_schemas(self) -> tuple[ParamSchema, ...]:
                 required = False
             else:
                 required = True
-            param_schema = ParamSchema(
-                name=arg.name, type=arg.typeinfo, is_input=True, required=required
+            schemas.append(
+                ParamSchema(name=arg.name, type=arg.typeinfo, is_input=True, required=required)
             )
-            schemas.append(param_schema)
-
-        for attr_name in attributes:
-            # Attributes without default values
-            # FIXME(justinchuby): Where can we find the type?
-            param_schema = ParamSchema(name=attr_name, type=None, is_input=False)
-            schemas.append(param_schema)
 
         for name, attr_value in attr_name_to_protos.items():
-            param_schema = ParamSchema(
-                name=name,
-                type=_ATTRIBUTE_TYPE_TO_PYTHON_TYPE[attr_value.type],
-                default=_get_attribute_value(attr_value.attr_proto),
-                is_input=False,
-                # All function attributes are required
-            )
-            schemas.append(param_schema)
+            if not attr_value.has_default:
+                schemas.append(
+                    ParamSchema(
+                        name=name,
+                        type=_ATTRIBUTE_TYPE_TO_PYTHON_TYPE[attr_value.type],
+                        is_input=False,
+                        required=True,
+                    )
+                )
+            else:
+                schemas.append(
+                    ParamSchema(
+                        name=name,
+                        type=_ATTRIBUTE_TYPE_TO_PYTHON_TYPE[attr_value.type],
+                        default=_get_attribute_value(attr_value.attr_proto),
+                        is_input=False,
+                        required=True,
+                    )
+                )
 
         self._param_schemas = tuple(schemas)
         return self._param_schemas  # type: ignore[return-value]