Auto OpSchema for trace_only functions | feat(torchlib)

ghstack-source-id: 3b97c3f
Pull Request resolved: #674

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

Author: justinchuby

onnxscript/_internal/ast_utils.py

@@ -0,0 +1,23 @@
+"""Utilities for working with Python ASTs."""
+
+import ast
+import inspect
+import textwrap
+import types
+
+
+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

onnxscript/function_libs/torch_lib/registration.py

@@ -67,7 +67,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 | onnxscript.values.TracedOnnxFunction]:
     """Register a torch op.
 
     Args:
@@ -81,12 +81,16 @@ 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 | onnxscript.values.TracedOnnxFunction:
+        # Compile the function
+        custom_opset = onnxscript.values.Opset(domain="onnxscript.atenlib", version=1)
+
+        processed_func: onnxscript.OnnxFunction | onnxscript.values.TracedOnnxFunction
         if trace_only:
-            processed_func = func
+            processed_func = onnxscript.values.TracedOnnxFunction(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/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,

onnxscript/main.py

@@ -8,36 +8,14 @@
 import ast
 import inspect
 import sys
-import textwrap
 import types
 from typing import Any, Callable, Optional, Sequence, cast
 
 import onnx.helper
 
 import onnxscript
 from onnxscript import converter, irbuilder, values
-
-
-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
-
-
-def get_ast(f: types.FunctionType) -> ast.FunctionDef:
-    _, f_ast = get_src_and_ast(f)
-    return f_ast
+from onnxscript._internal import ast_utils
 
 
 def script_check(
@@ -104,7 +82,7 @@ def transform(f: types.FunctionType) -> onnxscript.OnnxFunction:
         if not inspect.isfunction(f):
             raise TypeError("The ONNXScript decorator should be applied to functions only.")
 
-        src, f_ast = get_src_and_ast(f)  # pylint: disable=redefined-outer-name
+        src, f_ast = ast_utils.get_src_and_ast(f)
         # The script should be compiled using the globals/locals at the definition site.
         # This allows the script to reference names defined outside the script,
         # which is used for a few different purposes.

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,