Create the OpLike protocol and refactor Op | feat(values)

- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 853c076
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

onnxscript/function_libs/torch_lib/tracing.py

@@ -5,14 +5,18 @@
 import inspect
 import textwrap
 import types
-from typing import Optional
+import typing
+from typing import Optional, Tuple
 
 import onnx
 
 import onnxscript
 from onnxscript import converter as ons_converter
 from onnxscript._internal import version_utils
 
+if typing.TYPE_CHECKING:
+    from onnxscript import irbuilder
+
 _ONNX_OP_SCHEMA_WRITABLE = not version_utils.onnx_older_than("1.14")
 
 
@@ -33,7 +37,7 @@ def _get_src_and_ast(f: types.FunctionType) -> tuple[str, ast.FunctionDef]:
     return src, f_ast
 
 
-class TraceOnlyFunction:
+class TraceOnlyFunction(onnxscript.values.OpLike):
     """TraceOnlyFunction.
 
     Attributes:
@@ -44,9 +48,11 @@ class TraceOnlyFunction:
     def __init__(self, opset: onnxscript.values.Opset, func: types.FunctionType):
         self._opset = opset
         self._func = func
-        self._opschema: Optional[onnx.defs.OpSchema] = None
         # Set the signature of the class to function's
         self.__signature__ = inspect.signature(func)
+        # Cached computed fields
+        self._opschema: Optional[onnx.defs.OpSchema] = None
+        self._param_schemas: Optional[Tuple[onnxscript.values.ParamSchema, ...]] = None
 
     def __call__(self, *args, **kwargs):
         return self._func(*args, **kwargs)
@@ -72,13 +78,32 @@ def opset(self) -> onnxscript.values.Opset:
     @property
     def opschema(self) -> Optional[onnx.defs.OpSchema]:
         """Return the opschema."""
-
         if self._opschema is not None:
             return self._opschema
-
         if not _ONNX_OP_SCHEMA_WRITABLE:
             return None
 
+        # FIXME(justinchuby): outputs are empty. Need to fix.
+        self._opschema = onnxscript.values.op_schema_from_function_ir(
+            self._function_ir(), self._opset
+        )
+
+        return self._opschema
+
+    def param_schemas(self) -> tuple[onnxscript.values.ParamSchema, ...]:
+        """Generate param_schemas for the TraceOnlyFunction."""
+        if self._param_schemas is None:
+            self._param_schemas = onnxscript.values.param_schemas_from_function_ir(
+                self._function_ir()
+            )
+
+        return self._param_schemas
+
+    def _function_ir(self) -> irbuilder.IRFunction:
+        """Return the IRFunction of the function.
+
+        This IRFunction contains only the function signature.
+        """
         src, func_ast = _get_src_and_ast(self._func)
         module = inspect.getmodule(self._func)
         closure = inspect.getclosurevars(self._func)
@@ -90,9 +115,4 @@ def opschema(self) -> Optional[onnx.defs.OpSchema]:
             source=src,
         )
 
-        function_ir = converter.translate_function_signature(func_ast)
-
-        # FIXME(justinchuby): outputs are empty. Need to fix.
-        self._opschema = onnxscript.values.op_schema_from_function_ir(function_ir, self._opset)
-
-        return self._opschema
+        return converter.translate_function_signature(func_ast)

onnxscript/irbuilder.py

@@ -202,7 +202,7 @@ def __str__(self):
 
         args = _format(self.args, "(", ", ", ")", _opt_var_to_str)
         domain = self.callee.opset.domain
-        opname = self.callee.opname
+        opname = self.callee.name
         callee = f"{domain}.{opname}" if (domain != "") else opname
         return f"{lhs} = {callee} {attrs}{args}"
 
@@ -212,7 +212,7 @@ def debug_print(self):
 
     def to_node_proto(self, node_name: str) -> onnx.NodeProto:
         n = helper.make_node(
-            self.callee.opname,
+            self.callee.name,
             [_opt_var_to_str(x) for x in self.args],
             [str(x) for x in self.result],
             domain=self.callee.opset.domain,