Use the new IR for building graphs in torchlib (#1354)

Use the new IR for building graphs in torchlib. I created a flag
`TORCHLIB_EXPERIMENTAL_USE_IR` to control the feature. When enabled the
classes in `graph_building` will be swapped to use the new IR. The
exporter should see the same interfaces and not feel anything. After the
transition period we can decide where this should live.

Set `TORCHLIB_EXPERIMENTAL_USE_IR=1` to enable this feature.

Fix #997

Author: justinchuby

.github/workflows/main.yaml

@@ -31,6 +31,7 @@ jobs:
           - py311-onnx-weekly
           - py311-ort-nightly
           - py311-experimental-torchlib-tracing
+          - py311-experimental-torchlib-onnx-ir
           - py310
           - py39
           - py38
@@ -62,6 +63,9 @@ jobs:
           - name: py311-experimental-torchlib-tracing
             python-version: "3.11"
             nox-tag: test-experimental-torchlib-tracing
+          - name: py311-experimental-torchlib-onnx-ir
+            python-version: "3.11"
+            nox-tag: test-experimental-torchlib-onnx-ir
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4

noxfile.py

@@ -126,3 +126,24 @@ def test_experimental_torchlib_tracing(session):
         *session.posargs,
         env={"TORCHLIB_EXPERIMENTAL_PREFER_TRACING": "1"},
     )
+
+
+@nox.session(tags=["test-experimental-torchlib-onnx-ir"])
+def test_experimental_torchlib_onnx_ir(session):
+    """Test TorchLib using the ONNX IR to build graphs."""
+    session.install(
+        *COMMON_TEST_DEPENDENCIES,
+        PYTORCH,
+        TORCHVISON,
+        ONNX,
+        *ONNX_RUNTIME_NIGHTLY_DEPENDENCIES,
+    )
+    session.install("-r", "requirements/ci/requirements-ort-nightly.txt")
+    session.install(".", "--no-deps")
+    session.run("pip", "list")
+    session.run(
+        "pytest",
+        "tests/function_libs/torch_lib/ops_test.py",
+        *session.posargs,
+        env={"TORCHLIB_EXPERIMENTAL_USE_IR": "1"},
+    )

onnxscript/_internal/runtime_typing.py

@@ -5,12 +5,15 @@
 installed.
 """
 
+import typing
 import warnings
 
 __all__ = [
     "checked",
 ]
 
+T = typing.TypeVar("T", bound=typing.Callable[..., typing.Any])
+
 try:
     from beartype import beartype as checked
     from beartype import roar as _roar
@@ -25,12 +28,12 @@
     )
 except ImportError:
 
-    def checked(func):  # type: ignore[no-redef]
+    def checked(func: T) -> T:  # type: ignore[no-redef]
         return func
 
 except Exception as e:  # pylint: disable=broad-exception-caught
     # Warn errors that are not import errors (unexpected).
     warnings.warn(f"{e}", stacklevel=2)
 
-    def checked(func):  # type: ignore[no-redef]
+    def checked(func: T) -> T:  # type: ignore[no-redef]
         return func

onnxscript/function_libs/torch_lib/_flags.py

@@ -43,3 +43,7 @@ def _load_boolean_flag(
     "TORCHLIB_EXPERIMENTAL_PREFER_TRACING",
     this_will="trace all traceable functions to fold if branches and collapse constant expressions",
 )
+EXPERIMENTAL_USE_IR: bool = _load_boolean_flag(
+    "TORCHLIB_EXPERIMENTAL_USE_IR",
+    this_will="use the ONNX IR instead of the PyTorch Graph for graph building",
+)

onnxscript/function_libs/torch_lib/graph_building/__init__.py

@@ -0,0 +1,54 @@
+"""APIs for building an ONNX graph from a PyTorch model.
+
+This module exposes only three classes that will be used to build an ONNX graph
+by the ONNX exporter in PyTorch:
+
+- :class:`TorchScriptTensor`: Represents a symbolic value in the ONNX graph.
+- :class:`TorchScriptGraph`: Stores the graph being built.
+- :class:`TorchScriptTracingEvaluator`: An evaluator that will record all operators
+    applied on the ``TorchScriptTensor``. It has a reference to the ``TorchScriptGraph``
+    being built, will write to it, and will handle eager evaluations of Torch Lib
+    functions when desired.
+
+The usage is in https://github.com/pytorch/pytorch/blob/136f8378e1b5a8cb7127977b8d068fbf9c3e1247/torch/onnx/_internal/fx/fx_onnx_interpreter.py#L698-L702,
+and it is very simple::
+
+    with onnxscript.evaluator.default_as(onnxscript_tracer):  # onnxscript_tracer is a TorchScriptTracingEvaluator
+        output: Union[
+            onnxscript_graph_building.TorchScriptTensor,
+            Tuple[onnxscript_graph_building.TorchScriptTensor, ...],
+        ] = symbolic_fn(*onnx_args, **onnx_kwargs)
+
+Here, we set the default evaluator to be ``onnxscript_tracer`` so
+that ONNX Script will dispatch all operators calls to the evaluator. The ``symbolic_fn``
+can be a pure Python function (e.g. trace-only) or an ONNX Script function. Either way,
+they are recorded by ``onnxscript_tracer`` and onto the graph.
+
+The outputs, as ``TorchScriptTensor``, are then handed by to the exporter. On line
+https://github.com/pytorch/pytorch/blob/136f8378e1b5a8cb7127977b8d068fbf9c3e1247/torch/onnx/_internal/fx/fx_onnx_interpreter.py#L707
+the exporter fills in type and shape information from PyTorch by calling the setters
+on ``TorchScriptTensor.dtype`` and ``TorchScriptTensor.shape``.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "TorchScriptTensor",
+    "TorchScriptGraph",
+    "TorchScriptTracingEvaluator",
+]
+
+from onnxscript.function_libs.torch_lib import _flags
+
+if _flags.EXPERIMENTAL_USE_IR:
+    from ._graph_building_ir import (
+        TorchScriptGraph,
+        TorchScriptTensor,
+        TorchScriptTracingEvaluator,
+    )
+else:
+    from ._graph_building_torch import (  # type: ignore[assignment]
+        TorchScriptGraph,
+        TorchScriptTensor,
+        TorchScriptTracingEvaluator,
+    )