onnx_plug

Author: xadupre

_doc/api/export/index.rst

@@ -8,6 +8,7 @@ onnx_diagnostic.export
     api
     control_flow
     dynamic_shapes
+    onnx_plug
     shape_helper
     validate
 

_doc/api/export/onnx_plug.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.export.onnx_plug
+================================
+
+.. automodule:: onnx_diagnostic.export.onnx_plug
+    :members:
+    :no-undoc-members:

_unittests/ut_export/test_onnx_plug.py

@@ -0,0 +1,103 @@
+import unittest
+import onnx.helper as oh
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.export.onnx_plug import EagerDirectReplacementWithOnnx
+from onnx_diagnostic.export.api import to_onnx
+
+
+class TestOnnxPlus(ExtTestCase):
+    def test_onnx_plug_verify(self):
+        def _test_customadd(x, y):
+            return x + y
+
+        def _test_customadd_shape(x, y):
+            return torch.empty(torch.broadcast_shapes(x.shape, y.shape), dtype=x.dtype)
+
+        def make_function_proto():
+            return oh.make_function(
+                "onnx_plug",
+                "_test_customadd",
+                ["x", "y"],
+                ["z"],
+                [oh.make_node("Add", ["x", "y"], ["z"])],
+                opset_imports=[oh.make_opsetid("", 22)],
+            )
+
+        rep = EagerDirectReplacementWithOnnx(
+            _test_customadd, _test_customadd_shape, make_function_proto(), 2, 1
+        )
+
+        x = torch.randn((3, 4), dtype=torch.float32)
+        y = torch.randn((3, 1), dtype=torch.float32)
+        self.assertEqualArray(_test_customadd(x, y), x + y)
+        res = rep.verify(x, y)
+        self.assertEqualAny(res.eager_outputs, (x + y,))
+        self.assertEqual(len(res.diffs), 1)
+        self.assertEqual(res.diffs[0]["abs"], 0)
+
+    def test_onnx_plug_export(self):
+        def _test_customsub(x, y):
+            return x - y
+
+        def _test_customsub_shape(x, y):
+            return torch.empty(torch.broadcast_shapes(x.shape, y.shape), dtype=x.dtype)
+
+        def make_function_proto():
+            return oh.make_function(
+                "onnx_plug",
+                "_test_customsub",
+                ["x", "y"],
+                ["z"],
+                [oh.make_node("Sub", ["x", "y"], ["z"])],
+                opset_imports=[oh.make_opsetid("", 22)],
+            )
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                y = x.sum(axis=1, keepdim=True)
+                d = torch.ops.onnx_plug._test_customsub(x, y)
+                return torch.abs(d)
+
+        replacements = [
+            EagerDirectReplacementWithOnnx(
+                _test_customsub, _test_customsub_shape, make_function_proto(), 2, 1
+            )
+        ]
+
+        x = torch.randn((3, 4), dtype=torch.float32)
+        model = Model()
+        expected = model(x)
+        ds = ({0: "d1", 1: "d2"},)
+        ep = torch.export.export(model, (x,), dynamic_shapes=self.use_dyn_not_str(ds))
+        self.assertIn("torch.ops.onnx_plug._test_customsub.default", str(ep))
+        got = ep.module()(x)
+        self.assertEqualArray(expected, got)
+
+        with self.subTest(exporter="custom"):
+            onx = to_onnx(
+                model,
+                (x,),
+                dynamic_shapes=ds,
+                exporter="custom",
+                onnx_plugs=replacements,
+                target_opset=22,
+            )
+            self.assert_onnx_disc("test_onnx_plug_export_custom", onx.model_proto, model, (x,))
+
+        with self.subTest(exporter="onnx-dynamo"):
+            onx = to_onnx(
+                model,
+                (x,),
+                dynamic_shapes=ds,
+                exporter="onnx-dynamo",
+                onnx_plugs=replacements,
+                target_opset=22,
+            )
+            self.assert_onnx_disc(
+                "test_onnx_plug_export_onnx_dynamo", onx.model_proto, model, (x,)
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnx_diagnostic/export/api.py

@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 import torch
+from .onnx_plug import EagerDirectReplacementWithOnnx
 
 
 def to_onnx(
@@ -18,6 +19,7 @@ def to_onnx(
     save_ep: Optional[str] = None,
     optimize: bool = True,
     use_control_flow_dispatcher: bool = False,
+    onnx_plugs: Optional[List[EagerDirectReplacementWithOnnx]] = None,
 ) -> Any:
     """
     Common API for exporters. By default, the models are optimized to use the
@@ -41,6 +43,7 @@ def to_onnx(
     :param optimize: optimizes the model
     :param use_control_flow_dispatcher: use the dispatcher created to supported
         custom loops (see :func:`onnx_diagnostic.export.control_flow.loop_for`)
+    :param onnx_plugs: the code was modified to replace some parts with onnx translation
     :return: the output of the selected exporter, usually a structure including
         an onnx model
 
@@ -55,6 +58,10 @@ def to_onnx(
             exporter=exporter,
             filename=filename,
         )
+
+    Some examples using control flows are available in
+    :func:`onnx_diagnostic.export.control_flow.loop_for` or
+    :class:`onnx_diagnostic.export.onnx_plug.EagerDirectReplacementWithOnnx`.
     """
     if exporter == "custom":
         from experimental_experiment.torch_interpreter import (
@@ -63,16 +70,38 @@ def to_onnx(
         )
         from experimental_experiment.xbuilder import OptimizationOptions
 
-        if use_control_flow_dispatcher:
-            from .control_flow import create_global_dispatcher
-
-            dispatcher = create_global_dispatcher()
-
         options = None
         if exporter_kwargs is not None:
             options = exporter_kwargs.pop("options", None)
         if options is None:
             options = OptimizationOptions(patterns="default+onnxruntime")
+        if onnx_plugs or use_control_flow_dispatcher:
+            from experimental_experiment.torch_interpreter import Dispatcher
+
+            if use_control_flow_dispatcher:
+                from .control_flow import create_global_dispatcher
+
+                control_flow_dispatcher = create_global_dispatcher()
+            else:
+                control_flow_dispatcher = None
+
+            class MainDispatcher(Dispatcher):
+                def __init__(self):
+                    super().__init__({})
+
+            main_dispatcher = MainDispatcher()
+            if control_flow_dispatcher:
+                main_dispatcher.registered_functions.update(
+                    control_flow_dispatcher.registered_functions
+                )
+            if onnx_plugs:
+                for plug in onnx_plugs:
+                    main_dispatcher.registered_functions[plug.target_name] = (
+                        plug.custom_converter()
+                    )
+
+        else:
+            main_dispatcher = None
 
         return _to_onnx(
             mod,
@@ -89,7 +118,7 @@ def to_onnx(
             export_options=ExportOptions(save_ep=save_ep),
             options=options,
             **(exporter_kwargs or {}),
-            dispatcher=dispatcher if use_control_flow_dispatcher else None,
+            dispatcher=main_dispatcher,
         )
 
     if exporter in ("dynamo", "onnx-dynamo"):
@@ -99,6 +128,10 @@ def to_onnx(
         assert (
             not output_dynamic_shapes
         ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
+        custom_translation_table = {}
+        if onnx_plugs:
+            for plug in onnx_plugs:
+                custom_translation_table[plug.torch_op] = plug.onnx_dynamo_converter()
         epo = torch.onnx.export(
             mod,
             args=args or tuple(),
@@ -111,6 +144,7 @@ def to_onnx(
             verbose=verbose,
             dump_exported_program=bool(save_ep),
             artifacts_dir=os.path.dirname(filename) if filename else ".",
+            custom_translation_table=custom_translation_table,
             **(exporter_kwargs or {}),
         )
         if optimize:

onnx_diagnostic/export/control_flow.py

@@ -36,7 +36,7 @@ def register(self, aten_name: str, converter: Callable):
 
 @contextlib.contextmanager
 def enable_code_export_control_flow():
-    """Enables the code means to be exported."""
+    """Enables the code meant to be exported."""
     global _TEST_EXPORT
     old = _TEST_EXPORT
     _TEST_EXPORT = True