Auto OpSchema for trace_only functions | feat(torchlib)

ghstack-source-id: b25bb51
Pull Request resolved: #674

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

Author: justinchuby

onnxscript/function_libs/torch_lib/registration.py

@@ -6,6 +6,7 @@
 from typing import Any, Callable, Optional
 
 import onnxscript
+from onnxscript.function_libs.torch_lib import tracing
 
 
 class OverloadedFunction:
@@ -67,7 +68,7 @@ def torch_op(
     registry: Optional[Registry] = None,
     trace_only: bool = False,
     private: bool = False,
-) -> Callable[[Callable[..., Any]], onnxscript.OnnxFunction | Callable[..., Any]]:
+) -> Callable[[FunctionType], onnxscript.OnnxFunction | tracing.TraceOnlyFunction]:
     """Register a torch op.
 
     Args:
@@ -81,12 +82,14 @@ def torch_op(
     if registry is None:
         registry = default_registry
 
-    def wrapper(func: Callable[..., Any]) -> onnxscript.OnnxFunction | Callable[..., Any]:
+    def wrapper(func: FunctionType) -> onnxscript.OnnxFunction | tracing.TraceOnlyFunction:
+        # Compile the function
+        custom_opset = onnxscript.values.Opset(domain="onnxscript.atenlib", version=1)
+
+        processed_func: onnxscript.OnnxFunction | tracing.TraceOnlyFunction
         if trace_only:
-            processed_func = func
+            processed_func = tracing.TraceOnlyFunction(custom_opset, func)
         else:
-            # Compile the function
-            custom_opset = onnxscript.values.Opset(domain="onnxscript.atenlib", version=1)
             assert isinstance(func, FunctionType)
             processed_func = onnxscript.script(opset=custom_opset)(func)
 

onnxscript/function_libs/torch_lib/tracing.py

@@ -0,0 +1,98 @@
+"""Module for trace_only functions."""
+from __future__ import annotations
+
+import ast
+import inspect
+import textwrap
+import types
+from typing import Optional
+
+import onnx
+
+import onnxscript
+from onnxscript import converter as ons_converter
+from onnxscript._internal import version_utils
+
+_ONNX_OP_SCHEMA_WRITABLE = not version_utils.onnx_older_than("1.14")
+
+
+def _get_src_and_ast(f: types.FunctionType) -> tuple[str, ast.FunctionDef]:
+    try:
+        src = inspect.getsource(f)
+    except OSError as e:
+        raise RuntimeError(
+            f"Decorator script does not work on dynamically "
+            f"compiled function {f.__name__}."
+        ) from e
+    src = textwrap.dedent(src)
+    top_level_ast = ast.parse(src)
+    assert isinstance(top_level_ast, ast.Module)
+    assert len(top_level_ast.body) == 1
+    f_ast = top_level_ast.body[0]
+    assert isinstance(f_ast, ast.FunctionDef)
+    return src, f_ast
+
+
+class TraceOnlyFunction:
+    """TraceOnlyFunction.
+
+    Attributes:
+        name: Name of the op. E.g. "aten::add".
+        func: Function.
+    """
+
+    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)
+
+    def __call__(self, *args, **kwargs):
+        return self._func(*args, **kwargs)
+
+    def __repr__(self):
+        return f"TraceOnlyFunction({self!r})"
+
+    @property
+    def name(self) -> str:
+        """Return the name of the op."""
+        return self._func.__name__
+
+    @property
+    def source(self) -> str:
+        """Return the source of the op."""
+        return inspect.getsource(self._func)
+
+    @property
+    def opset(self) -> onnxscript.values.Opset:
+        """Return the opset."""
+        return self._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
+
+        src, func_ast = _get_src_and_ast(self._func)
+        module = inspect.getmodule(self._func)
+        closure = inspect.getclosurevars(self._func)
+        global_names = module.__dict__.copy()
+        global_names.update(closure.nonlocals)
+        converter = ons_converter.Converter(
+            opset=self._opset,
+            global_names=global_names,
+            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

onnxscript/tests/function_libs/torch_lib/ops_test.py

@@ -119,6 +119,19 @@ def test_script_function_has_op_schema(self, _, func_with_wrangler):
         self.assertIsNotNone(schema)
         self.assertEqual(schema.name, func.name)
 
+    @parameterized.parameterized.expand(
+        list(ops_test_data.OPINFO_FUNCTION_MAPPING_TRACE_ONLY.items())
+    )
+    @unittest.skipIf(
+        version_utils.onnx_older_than("1.15"),
+        "OpSchema is not writable before ONNX 1.15",
+    )
+    def test_trace_only_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

@@ -431,7 +431,7 @@ class OnnxFunction(Op):
 
     def __init__(
         self,
-        opset: Opset,
+        opset: Optional[Opset],
         pyfun: types.FunctionType,
         irfun: irbuilder.IRFunction,
         source: str,
@@ -455,6 +455,8 @@ def __init__(
         self.kwargs = kwargs
         self._param_schemas: Optional[tuple[ParamSchema, ...]] = None
         self._opschema: Optional[onnx.defs.OpSchema] = None
+        # Set the signature of the class to function's
+        self.__signature__ = inspect.signature(pyfun)
 
     @property
     def opschema(self) -> Optional[onnx.defs.OpSchema]: