Auto generate OpSchema for functions | feat(op_schema) (#626)

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* #692
* #674
* __->__ #626
* #684

This change adds the capability to auto generate `OpSchema`.

### Changes

- Implement the `opschema` property in `OnnxFunction`
- Test on all torch_lib functions

### Next PR

Support trace_only functions

## Example

```python
from onnxscript.function_libs.torch_aten.ops import core, nn


print("core.aten_abs.opschema: ", core.aten_abs.opschema)

print("nn.aten_cross_entropy_loss.opschema: ", nn.aten_cross_entropy_loss.opschema)
```

Results

```
core.aten_abs.opschema:  OpSchema(
    name='aten_abs',
    domain='onnxscript.atenlib',
    since_version=1,
    doc='abs(Tensor self) -> Tensor',
    type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TReal', allowed_type_strs=['tensor(float)', 'tensor(int8)', 'tensor(int16)', 'tensor(int32)', 'tensor(int64)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
    inputs=[OpSchema.FormalParameter(name='self', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    outputs=[OpSchema.FormalParameter(name='return_val', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    attributes={}
)
nn.aten_cross_entropy_loss.opschema:  OpSchema(
    name='aten_cross_entropy_loss',
    domain='onnxscript.atenlib',
    since_version=1,
    doc='cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor',
    type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TFloatOrBFloat16', allowed_type_strs=['tensor(float)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description=''), OpSchema.TypeConstraintParam(type_param_str='T1', allowed_type_strs=['tensor(float)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
    inputs=[OpSchema.FormalParameter(name='self', type_str='TFloatOrBFloat16', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>), OpSchema.FormalParameter(name='weight', type_str='T1', description='', param_option=<FormalParameterOption.Optional: 1>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    outputs=[OpSchema.FormalParameter(name='result_10', type_str='TFloatOrBFloat16', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    attributes={'ignore_index': OpSchema.Attribute(name='ignore_index', type=<AttrType.INT: 2>, description='', default_value=name: "ignore_index"
i: -100
type: INT
, required=False), 'label_smoothing': OpSchema.Attribute(name='label_smoothing', type=<AttrType.FLOAT: 1>, description='', default_value=name: "label_smoothing"
f: 0.0
type: FLOAT
, required=False), 'reduction': OpSchema.Attribute(name='reduction', type=<AttrType.INT: 2>, description='', default_value=name: "reduction"
i: 1
type: INT
, required=False), 'target': OpSchema.Attribute(name='target', type=<AttrType.INTS: 7>, description='', default_value=, required=True)}
)
```

Fixes #476

Author: justinchuby

onnxscript/autocast.py

@@ -86,7 +86,7 @@ def cast(x, typeinfo) -> tensor.Tensor:
     return cast_inputs(get_type_info, cast, op_schema, *args)
 
 
-def static_cast_inputs(converter, op_schema: OpSchema, *args):
+def static_cast_inputs(converter, op_schema: Optional[OpSchema], *args):
     """Used for autocast during script-translation."""
     if op_schema is None:
         return args

onnxscript/tests/function_libs/torch_lib/ops_test.py

@@ -106,6 +106,19 @@ def test_script_function_passes_checker(self, _, func_with_wrangler):
         function_proto = func.to_function_proto()
         onnx.checker.check_function(function_proto)  # type: ignore[attr-defined]
 
+    @parameterized.parameterized.expand(
+        list(ops_test_data.OPINFO_FUNCTION_MAPPING_SCRIPTED.items())
+    )
+    @unittest.skipIf(
+        version_utils.onnx_older_than("1.15"),
+        "OpSchema is not writable before ONNX 1.15",
+    )
+    def test_script_function_has_op_schema(self, _, func_with_wrangler):
+        func, _ = _split_function_and_wrangler(func_with_wrangler)
+        schema = func.opschema
+        self.assertIsNotNone(schema)
+        self.assertEqual(schema.name, func.name)
+
 
 def run_test_output_match(
     test_suite: unittest.TestCase,

onnxscript/values.py

@@ -8,12 +8,14 @@
 import logging
 import types
 from enum import IntFlag
-from typing import Any, Optional, Sequence, _GenericAlias  # type: ignore[attr-defined]
+from typing import _GenericAlias  # type: ignore[attr-defined]
+from typing import Any, Optional, Sequence
 
 import onnx
 import onnx.defs
 
-from onnxscript import irbuilder, sourceinfo
+from onnxscript import irbuilder, sourceinfo, type_annotation
+from onnxscript._internal import version_utils
 
 _ATTRIBUTE_TYPE_TO_PYTHON_TYPE = {
     onnx.defs.OpSchema.AttrType.FLOAT: float,
@@ -34,6 +36,7 @@
 
 # A special value to indicate that the default value is not specified
 _EmptyDefault = object()
+_ONNX_OP_SCHEMA_WRITABLE = not version_utils.onnx_older_than("1.14")
 
 
 class Opset:
@@ -173,7 +176,7 @@ def __init__(
     ) -> None:
         self.opset = opset
         self.opname = opname
-        self.opschema = opschema
+        self._opschema = opschema
         self._param_schemas: Optional[tuple[ParamSchema, ...]] = None
 
     def __call__(self, *args, **kwargs):
@@ -190,9 +193,13 @@ def __call__(self, *args, **kwargs):
     def is_single_op(self) -> bool:
         return isinstance(self.opname, str)
 
+    @property
+    def opschema(self) -> Optional[onnx.defs.OpSchema]:
+        return self._opschema
+
     def get_schema(self) -> Optional[onnx.defs.OpSchema]:
         """Returns the ONNX OpSchema for this op."""
-        if self.opschema:
+        if self.opschema is not None:
             return self.opschema
         return self.opset[self.opname]
 
@@ -249,6 +256,100 @@ class OnnxClosure:
     function: Any
 
 
+@dataclasses.dataclass
+class TypeConstraint:
+    """Represents a type constraint for an ONNX op.
+
+    Attributes:
+        name: The name of the type constraint.
+        allowed_types: The allowed types for the type constraint.
+    """
+
+    name: str
+    allowed_types: list[str]
+    description: str = ""
+
+    def as_tuple(self) -> tuple[str, list[str], str]:
+        """Returns the type constraint as a tuple."""
+        return (self.name, self.allowed_types, self.description)
+
+
+def op_schema_from_function_ir(
+    function_ir: irbuilder.IRFunction, opset: Opset
+) -> onnx.defs.OpSchema:
+    """Construct an ONNX OpSchema from an IRFunction."""
+
+    # Find all distinct types in the inputs and outputs
+    distinct_types = {arg.typeinfo for arg in function_ir.inputs}.union(
+        {arg.typeinfo for arg in function_ir.outputs}
+    )
+    # Create a mapping from type to a unique name
+    type_to_constraint = {}
+    for i, type_ in enumerate(distinct_types):
+        name = f"T{i}"
+        type_to_constraint[type_] = TypeConstraint(
+            name=type_annotation.get_type_constraint_name(type_) or name,
+            allowed_types=type_annotation.pytype_to_type_strings(type_),
+        )
+
+    formal_inputs = [
+        onnx.defs.OpSchema.FormalParameter(
+            arg.name,
+            type_to_constraint[arg.typeinfo].name,
+            param_option=(
+                onnx.defs.OpSchema.FormalParameterOption.Optional
+                if type_annotation.is_optional(arg.typeinfo)
+                else onnx.defs.OpSchema.FormalParameterOption.Single
+            ),
+            # TODO(justinchu): Check this is_homogeneous thing
+            is_homogeneous=True,
+        )
+        for arg in function_ir.inputs
+    ]
+    formal_outputs = [
+        onnx.defs.OpSchema.FormalParameter(
+            arg.name,
+            type_to_constraint[arg.typeinfo].name,
+            param_option=(
+                onnx.defs.OpSchema.FormalParameterOption.Optional
+                if type_annotation.is_optional(arg.typeinfo)
+                else onnx.defs.OpSchema.FormalParameterOption.Single
+            ),
+            # TODO(justinchu): Check this is_homogeneous thing
+            is_homogeneous=True,
+        )
+        for arg in function_ir.outputs
+    ]
+
+    return onnx.defs.OpSchema(
+        function_ir.name,
+        opset.domain,
+        since_version=opset.version,
+        doc=function_ir.docstring,
+        inputs=formal_inputs,
+        outputs=formal_outputs,
+        type_constraints=[constraint.as_tuple() for constraint in type_to_constraint.values()],
+        attributes=[
+            *[
+                onnx.defs.OpSchema.Attribute(
+                    attr.name,
+                    type=onnx.defs.OpSchema.AttrType(attr.type),
+                )
+                for attr in function_ir.attrs
+                if not attr.has_default
+            ],
+            *[
+                onnx.defs.OpSchema.Attribute(
+                    attr.name,
+                    default_value=attr.attr_proto,
+                )
+                for attr in function_ir.attrs
+                if attr.has_default
+            ],
+        ],
+    )
+
+
 class OnnxFunction(Op):
     """Represents an ONNX op for which a function-body has been defined in onnxscript.
 
@@ -276,12 +377,26 @@ def __init__(
         self.source = source
         self.kwargs = kwargs
         self._param_schemas: Optional[tuple[ParamSchema, ...]] = None
+        self._opschema: Optional[onnx.defs.OpSchema] = None
 
     @property
     def name(self):
         """Returns the function name."""
         return self.opname
 
+    @property
+    def opschema(self) -> Optional[onnx.defs.OpSchema]:
+        """Construct an OpSchema from function_ir."""
+        if self._opschema is not None:
+            return self._opschema
+
+        if not _ONNX_OP_SCHEMA_WRITABLE:
+            return None
+
+        self._opschema = op_schema_from_function_ir(self.function_ir, self.opset)
+
+        return self._opschema
+
     def __getitem__(self, instance):
         """Returns a lambda to evaluate function using given evaluator instance.
 
@@ -311,6 +426,9 @@ def param_schemas(self) -> tuple[ParamSchema, ...]:
         if self._param_schemas is not None:
             return self._param_schemas
 
+        # NOTE: We generate the parameter schemas from the function_ir instead
+        # of relying on the auto generated OpSchema because we need to preserve the keyword
+        # argument order from the Python function definition, which is lost in OpSchema.
         function_ir = self.function_ir
         # The first len(func_ir.inputs) arguments are onnx inputs
         inputs = function_ir.inputs