onnx-diagnostic 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Patches, Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.8.3"
+__version__ = "0.8.5"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -198,15 +198,19 @@ def get_parser_print() -> ArgumentParser:
     )
     parser.add_argument(
         "fmt",
-        choices=["pretty", "raw", "text", "printer"],
+        choices=["dot", "pretty", "printer", "raw", "shape", "text"],
         default="pretty",
         help=textwrap.dedent(
             """
             Prints out a model on the standard output.
-            raw     - just prints the model with print(...)
-            printer - onnx.printer.to_text(...)
+
+            dot     - converts the graph into dot
             pretty  - an improved rendering
+            printer - onnx.printer.to_text(...)
+            raw     - just prints the model with print(...)
+            shape   - prints every node node with input and output shapes
             text    - uses GraphRendering
+
             """.strip(
                 "\n"
             )
@@ -232,6 +236,14 @@ def _cmd_print(argv: List[Any]):
         from .helpers.graph_helper import GraphRendering
 
         print(GraphRendering(onx).text_rendering())
+    elif args.fmt == "shape":
+        from experimental_experiment.xbuilder import GraphBuilder
+
+        print(GraphBuilder(onx).pretty_text())
+    elif args.fmt == "dot":
+        from .helpers.dot_helper import to_dot
+
+        print(to_dot(onx))
     else:
         raise ValueError(f"Unexpected value fmt={args.fmt!r}")
 
@@ -517,12 +529,12 @@ def get_parser_validate(name: str = "validate") -> ArgumentParser:
         nargs="*",
         help=textwrap.dedent(
             """
-        Applies patches before exporting, it can be a boolean
-        to enable to disable the patches or be more finetuned
-        (default is True). It is possible to disable patch for torch
-        by adding:
-            --patch "patch_sympy=False" --patch "patch_torch=False"
-        """.strip(
+            Applies patches before exporting, it can be a boolean
+            to enable to disable the patches or be more finetuned
+            (default is True). It is possible to disable patch for torch
+            by adding:
+                --patch "patch_sympy=False" --patch "patch_torch=False"
+            """.strip(
                 "\n"
             )
         ),
@@ -1286,7 +1298,13 @@ def get_parser_sbs() -> ArgumentParser:
         "--first",
         action=BooleanOptionalAction,
         default=False,
-        help="First runs the whole model.",
+        help="First runs the whole model (default is False).",
+    )
+    parser.add_argument(
+        "--sbs",
+        action=BooleanOptionalAction,
+        default=True,
+        help="Runs the side-by-side (default is True).",
     )
     parser.add_argument(
         "-2",
@@ -1342,6 +1360,20 @@ def get_parser_sbs() -> ArgumentParser:
         default="replay",
         help="If the replay is triggered, this defines the folder where everything is dumped.",
     )
+    parser.add_argument(
+        "-p",
+        "--replay-prefix-model",
+        action=BooleanOptionalAction,
+        default=False,
+        help=textwrap.dedent(
+            """
+            There are two ways to recompute an intermediate output, the first one is to "
+            produce the minimal model between torch and onnx.
+            The second one is to dump onnx models from the inputs
+            to the considered intermediate results. This enables the second one.
+            """
+        ),
+    )
 
     return parser
 
@@ -1417,6 +1449,10 @@ def _cmd_sbs(argv: List[Any]):
         print("-- done")
         del sess
 
+    if not args.sbs:
+        print("-- done")
+        return
+
     print(f"-- load onnx {args.onnx!r}")
     begin = time.perf_counter()
     onx = onnx.load(args.onnx)
@@ -1431,6 +1467,7 @@ def _cmd_sbs(argv: List[Any]):
                 set(args.replay_op_types.split(",")) if args.replay_op_types else None
             ),
             dump_folder=args.replay_folder,
+            dump_prefix_model=args.replay_prefix_model,
         )
 
     print("-- starts side-by-side")

onnx_diagnostic/export/api.py

@@ -3,6 +3,52 @@ import torch
 from .onnx_plug import EagerDirectReplacementWithOnnx
 
 
+def get_main_dispatcher(
+    use_control_flow_dispatcher: bool = False,
+    onnx_plugs: Optional[List[EagerDirectReplacementWithOnnx]] = None,
+) -> Any:  # Dispatcher
+    """Creates a custom dispatcher for the custom exporter."""
+    from experimental_experiment.torch_interpreter import Dispatcher
+
+    if use_control_flow_dispatcher:
+        from .control_flow_onnx import create_global_dispatcher
+
+        control_flow_dispatcher = create_global_dispatcher()
+    else:
+        control_flow_dispatcher = None
+
+    class MainDispatcher(Dispatcher):
+        def __init__(self, previous_dispatcher=None):
+            super().__init__({})
+            self.previous_dispatcher = previous_dispatcher
+
+        @property
+        def supported(self):
+            if self.previous_dispatcher:
+                return set(self.registered_functions) | self.previous_dispatcher.supported
+            return set(self.registered_functions)
+
+        def find_function(self, name: Any):
+            if self.previous_dispatcher:
+                find = self.previous_dispatcher.find_function(name)
+                if find:
+                    return find
+            return Dispatcher.find_function(self, name)
+
+        def find_method(self, name: Any):
+            if self.previous_dispatcher:
+                find = self.previous_dispatcher.find_method(name)
+                if find:
+                    return find
+            return Dispatcher.find_method(self, name)
+
+    main_dispatcher = MainDispatcher(control_flow_dispatcher)
+    if onnx_plugs:
+        for plug in onnx_plugs:
+            main_dispatcher.registered_functions[plug.target_name] = plug.custom_converter()
+    return main_dispatcher
+
+
 def to_onnx(
     mod: Union["torch.nn.Module", "torch.fx.GraphModule"],  # noqa: F821
     args: Optional[Sequence["torch.Tensor"]] = None,  # noqa: F821
@@ -18,6 +64,7 @@ def to_onnx(
     exporter_kwargs: Optional[Dict[str, Any]] = None,
     save_ep: Optional[str] = None,
     optimize: bool = True,
+    optimizer_for_ort: bool = True,
     use_control_flow_dispatcher: bool = False,
     onnx_plugs: Optional[List[EagerDirectReplacementWithOnnx]] = None,
     inline: bool = True,
@@ -42,6 +89,7 @@ def to_onnx(
     :param exporter_kwargs: additional parameters sent to the exporter
     :param save_ep: saves the exported program
     :param optimize: optimizes the model
+    :param optimizer_for_ort: optimizes the model for onnxruntime
     :param use_control_flow_dispatcher: use the dispatcher created to supported
         custom loops (see :func:`onnx_diagnostic.export.control_flow_onnx.loop_for_onnx`)
     :param onnx_plugs: the code was modified to replace some parts with onnx translation
@@ -80,53 +128,15 @@ def to_onnx(
         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_onnx import create_global_dispatcher
-
-                control_flow_dispatcher = create_global_dispatcher()
-            else:
-                control_flow_dispatcher = None
-
-            class MainDispatcher(Dispatcher):
-                def __init__(self, previous_dispatcher=None):
-                    super().__init__({})
-                    self.previous_dispatcher = previous_dispatcher
-
-                @property
-                def supported(self):
-                    if self.previous_dispatcher:
-                        return (
-                            set(self.registered_functions) | self.previous_dispatcher.supported
-                        )
-                    return set(self.registered_functions)
-
-                def find_function(self, name: Any):
-                    if self.previous_dispatcher:
-                        find = self.previous_dispatcher.find_function(name)
-                        if find:
-                            return find
-                    return Dispatcher.find_function(self, name)
-
-                def find_method(self, name: Any):
-                    if self.previous_dispatcher:
-                        find = self.previous_dispatcher.find_method(name)
-                        if find:
-                            return find
-                    return Dispatcher.find_method(self, name)
-
-            main_dispatcher = MainDispatcher(control_flow_dispatcher)
-            if onnx_plugs:
-                for plug in onnx_plugs:
-                    main_dispatcher.registered_functions[plug.target_name] = (
-                        plug.custom_converter()
-                    )
-        else:
-            main_dispatcher = None
+        if options is None and optimize:
+            options = OptimizationOptions(
+                patterns="default+onnxruntime" if optimizer_for_ort else "default"
+            )
+        main_dispatcher = (
+            get_main_dispatcher(use_control_flow_dispatcher, onnx_plugs)
+            if onnx_plugs or use_control_flow_dispatcher
+            else None
+        )
 
         return _to_onnx(
             mod,
@@ -149,11 +159,15 @@ def to_onnx(
 
     if exporter in ("dynamo", "onnx-dynamo"):
         import os
+        from ..helpers import flatten_object
         import onnxscript.rewriter.ort_fusions as ort_fusions
 
         assert (
             not output_dynamic_shapes
         ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
+        assert (
+            optimize
+        ), f"torch.onnx.export always optimizes the model but optimize={optimize}"
         custom_translation_table = {}
         if onnx_plugs:
             for plug in onnx_plugs:
@@ -173,21 +187,34 @@ def to_onnx(
             custom_translation_table=custom_translation_table,
             **(exporter_kwargs or {}),
         )
-        if not inline and optimize:
+        if not inline and optimize and optimizer_for_ort:
             ort_fusions.optimize_for_ort(epo.model)
 
         if onnx_plugs:
             import onnx_ir as ir
             import onnx_ir.passes.common as common_passes
 
-            irfunctions = [ir.from_proto(plug.function_proto) for plug in onnx_plugs]
+            opset = (
+                18
+                if target_opset is None
+                else (target_opset if isinstance(target_opset, int) else target_opset[""])
+            )
+
+            irfunctions = [
+                ir.from_proto(
+                    plug.get_function_proto(
+                        opset, *flatten_object((args, kwargs), drop_keys=True)
+                    )
+                )
+                for plug in onnx_plugs
+            ]
             for func in irfunctions:
                 epo.model.functions[func.identifier()] = func
             if inline:
                 common_passes.InlinePass()(epo.model)
                 common_passes.RemoveUnusedOpsetsPass()(epo.model)
 
-        if inline and optimize:
+        if inline and optimize and optimizer_for_ort:
             ort_fusions.optimize_for_ort(epo.model)
         if filename:
             epo.save(filename, external_data=True)
@@ -212,6 +239,10 @@ def to_onnx(
             f"Only a specified set of inputs is supported for exporter={exporter!r}, "
             f"but it is {list(kwargs)}"  # type: ignore[arg-type]
         )
+        assert optimizer_for_ort and optimize, (
+            f"ModelBuilder only produces model optimized for onnxruntime but "
+            f"optimizer_for_ort={optimizer_for_ort} and optimize={optimize}"
+        )
         flat_inputs = flatten_object(kwargs, drop_keys=True)
         first = flat_inputs[0]
         first_float = [

onnx_diagnostic/export/control_flow_research.py

@@ -92,10 +92,11 @@ def simple_loop_for(
 
     from torch._higher_order_ops.utils import setup_compilation_env
 
-    with setup_compilation_env() as backend:
-        return torch.compile(_loop_for_op_wrapper, backend=backend, fullgraph=True)(
-            n_iter, body_fn, operands
-        )
+    with setup_compilation_env() as _backend:
+        return _loop_for_op_wrapper(n_iter, body_fn, *operands)
+        # return torch.compile(_loop_for_op_wrapper, backend=backend, fullgraph=True)(
+        #    n_iter, body_fn, operands
+        # )
 
 
 def trace_loop_for(proxy_mode, func_overload, n_iter, body_fn, operands):
@@ -127,9 +128,13 @@ def loop_for_op_dense(n_iter, body_fn, operands):
     ), f"Dense implementation operands must be a list of tensors and ints {operands}"
     mode = _get_current_dispatch_mode()
     assert mode is None, "Mode should never be enabled for CPU/CUDA key"
-    return _loop_for_onnx_fn(body_fn, n_iter, None, *operands)
+    return _loop_for_onnx_fn(body_fn, n_iter, None, operands)
 
 
 @simple_loop_for_op.py_impl(ProxyTorchDispatchMode)
 def inner(mode, n_iter, body_fn, operands):
     return trace_loop_for(mode, simple_loop_for_op, n_iter, body_fn, operands)
+
+
+simple_loop_for_op.fallthrough(torch._C.DispatchKey.AutogradCPU)
+simple_loop_for_op.fallthrough(torch._C.DispatchKey.AutogradCUDA)