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

onnx_diagnostic/export/control_flow.py

@@ -1,42 +1,18 @@
 import contextlib
-import inspect
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
-import onnx
-import onnx.helper as oh
+from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
 import torch
-from torch._higher_order_ops.utils import materialize_as_graph
-from torch._higher_order_ops.utils import check_input_alias_and_mutation_return_outputs
-from .api import to_onnx
+from torch._higher_order_ops.utils import (
+    materialize_as_graph,
+    check_input_alias_and_mutation_return_outputs,
+    # _maybe_reenter_make_fx,
+)
 
 _TEST_EXPORT = False
-_REGISTERED_SCHEMA = {}  # type: ignore[var-annotated]
-_DISPATCHER = None
-
-
-def create_global_dispatcher():
-    global _DISPATCHER
-
-    if not _DISPATCHER:
-        from experimental_experiment.torch_interpreter import Dispatcher
-
-        class ControlFlowDispatcher(Dispatcher):
-            def __init__(self):
-                super().__init__({})
-
-            def register(self, aten_name: str, converter: Callable):
-                assert aten_name not in self.registered_functions, (
-                    f"Name {aten_name!r} is already registered in "
-                    f"{sorted(self.registered_functions)}"
-                )
-                self.registered_functions[aten_name] = converter
-
-        _DISPATCHER = ControlFlowDispatcher()
-    return _DISPATCHER
 
 
 @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
@@ -128,194 +104,31 @@ def make_custom_loop_for(
     :return: a name and the custom op definition, the name
         is used to cache the custom op
     """
-    global _DISPATCHER
     assert body_gm is not None, "body_gm cannot be None"
     assert body_mutated_inputs is not None, "body_mutated_inputs cannot be None"
     assert body_outputs is not None, "body_outputs cannot be None"
+
     srank = "_".join("x".join(map(str, s.shape)) for s in body_outputs)
     sred = "x".join(map(str, reduction_dim)) if reduction_dim else ""
-    name = f"loop_for_{body_fn.__name__}_{id(body_fn)}_{srank}_{sred}"
-    if name in _REGISTERED_SCHEMA:
-        return name, _REGISTERED_SCHEMA[name][0]
-    sig = inspect.signature(body_fn)
-    inputs = ", ".join([f"Tensor {p}" for p in sig.parameters])
-    schema = f"({inputs}) -> Tensor"
+    full_name = (
+        body_fn.__qualname__.replace("<locals>", "L")
+        .replace("<lambda>", "l")
+        .replace(".", "_")
+    )
+    name = f"loop_for_onnx_{full_name}_{srank}_{sred}"
+
+    schema = "(str body_fn, Tensor n_iter, Tensor[] body_inputs) -> Tensor"
     if len(body_outputs) > 1:
         schema += "[]"
-    custom_def = torch.library.CustomOpDef("onnx_higher_ops", name, schema, body_fn)
+    custom_def = torch.library.CustomOpDef("onnx_higher_ops", "loop_for", schema, body_fn)
     custom_def.register_kernel("cpu")(body_fn)
 
-    custom_def._abstract_fn = lambda *_args, _o=body_outputs: (
+    custom_def._abstract_fn = lambda _fn_id, *_args, _o=body_outputs: (
         tuple([torch.empty_like(s) for s in _o]) if len(_o) > 1 else torch.empty_like(_o[0])
     )
-
-    def _make_onx(
-        body_gm=body_gm, args=args, target_opset=None, verbose=0, exporter_kwargs=None
-    ):
-        return convert_into_onnx(
-            body_gm,
-            args,
-            exporter_kwargs=exporter_kwargs,
-            target_opset=target_opset,
-            verbose=verbose,
-        )
-
-    to_register = (
-        custom_def,
-        _make_onx,
-        (
-            lambda g, sts, outputs, *args, bc=_make_onx, rd=reduction_dim, name=name: (
-                convert_custom_loop_into_onnx(
-                    g,
-                    sts,
-                    outputs,
-                    *args,
-                    body_callable=bc,
-                    reduction_dim=rd,
-                    name=name,
-                )
-            )
-        ),
-    )
-    if _DISPATCHER is None:
-        create_global_dispatcher()
-    assert _DISPATCHER
-    _DISPATCHER.register(f"onnx_higher_ops::{name}", to_register[-1])
-    _REGISTERED_SCHEMA[name] = to_register
     return name, custom_def
 
 
-def convert_custom_loop_into_onnx(
-    g: Any,  # "GreaphBuilder"
-    sts: Dict[str, Any],
-    outputs: List[str],
-    *args: str,
-    body_callable: Callable[..., onnx.ModelProto],
-    reduction_dim: Optional[Sequence[int]] = None,
-    name: str = "loop_for",
-) -> Union[str, List[str]]:
-    """
-    Converts a custom op ``higher_ops::loop_for...`` into e sequence of node.
-
-    :param g: GreaphBuilder
-    :param sts: if not defined, torch does not know the output shapes
-    :param outputs: output names
-    :param args: input argument known at export time
-    :param body: GraphProto, the loop body
-    :param reduction_dim: the dimension to follow when aggregating the
-        list of tensors after the loop ran
-    :param name: to give the onnx nodes a name
-    :return: output names
-    """
-    assert body_callable is not None, "body_callable cannot be None"
-    # This should be part of a public API.
-    body = body_callable(
-        target_opset=g.main_opset,
-        verbose=g.verbose,
-        exporter_kwargs={"options": g.optimization_options},
-    )
-
-    graph = body.graph if isinstance(body, onnx.ModelProto) else body
-    assert isinstance(
-        graph, onnx.GraphProto
-    ), f"Unexpected type {type(body)} for body{g.get_debug_msg()}"
-    assert len(outputs) == 1, f"Only one outputs is expected but outputs={outputs!r}"
-    if len(graph.output) != 1:
-        outputs = [f"{outputs[0]}#{i}" for i in range(len(graph.output))]
-    input_names = [i.name for i in graph.input]
-    inputs = [
-        *graph.input[:1],
-        oh.make_tensor_value_info("cond_unused", onnx.TensorProto.BOOL, []),
-        *[
-            oh.make_tensor_sequence_value_info(
-                f"loop_in{i}", graph.output[i].type.tensor_type.elem_type, None
-            )
-            for i in range(len(graph.output))
-        ],
-        # hidden inputs are not added
-    ]
-    nodes = [
-        oh.make_node("Identity", ["cond_unused"], ["cond_out"]),
-        *[oh.make_node("Identity", [a], [r]) for a, r in zip(args[1:], input_names[1:])],
-        *graph.node,
-        *[
-            oh.make_node(
-                "SequenceInsert",
-                [f"loop_in{i}", graph.output[i].name],
-                [f"loop_out{i}"],
-            )
-            for i in range(len(graph.output))
-        ],
-    ]
-    graph_outputs = [
-        oh.make_tensor_value_info("cond_out", onnx.TensorProto.BOOL, []),
-        *[
-            oh.make_tensor_sequence_value_info(
-                f"loop_out{i}", graph.output[i].type.tensor_type.elem_type, None
-            )
-            for i in range(len(graph.output))
-        ],
-    ]
-    graph = oh.make_graph(
-        nodes, graph.name, inputs, graph_outputs, graph.initializer, graph.sparse_initializer
-    )
-
-    sequences = [g.op.SequenceEmpty() for _ in outputs]
-
-    outloop = [g.unique_name(f"loop_for{i}") for i in range(len(sequences))]
-
-    for i, s in enumerate(sequences):
-        g.set_sequence(s, graph.output[i].type.tensor_type.elem_type)
-    g.make_node("Loop", [args[0], "", *sequences], outloop, name=name, body=graph)
-    for i, o in enumerate(outloop):
-        g.set_sequence(o, graph.output[i].type.tensor_type.elem_type)
-    _res = [
-        g.op.ConcatFromSequence(
-            out,
-            outputs=[o],
-            name=name,
-            axis=0 if not reduction_dim or i >= len(reduction_dim) else reduction_dim[i],
-        )
-        for i, (out, o) in enumerate(zip(outloop, outputs))
-    ]
-    if not sts:
-        for i, o in enumerate(outputs):
-            g.set_type(o, graph.output[i].type.tensor_type.elem_type)
-            g.set_rank(o, len(graph.output[i].type.tensor_type.shape.dims))
-    return outputs if len(outputs) > 1 else outputs[0]
-
-
-def convert_into_onnx(
-    body_gm: torch.fx.GraphModule,
-    args: Sequence[torch.Tensor],
-    target_opset: Optional[int] = None,
-    verbose: int = 0,
-    exporter_kwargs: Optional[Dict[str, Any]] = None,
-) -> onnx.ModelProto:
-    """
-    Converts a torch.fx.GraphModule into ONNX.
-    It returns a ModelProto.
-
-    :param body_gm: a torch.fx.GraphModule
-    :param args: arguments known at export time
-    :param target_opset: targeted opset
-    :param verbose: verbosity level
-    :param exporter_kwargs: additional exporter arguments
-    :return: a ModelProto
-    """
-    # This does not work with onnx-dynamo.
-    # opset still needs to be defined
-    container = to_onnx(
-        body_gm,
-        args,
-        exporter="custom",
-        exporter_kwargs=exporter_kwargs,
-        target_opset=target_opset,
-        verbose=verbose,
-    )
-    return container.model_proto
-
-
 def loop_for(
     n_iter: Union[torch.SymInt, torch.Tensor],
     body_fn: Callable[..., Tuple[torch.Tensor]],
@@ -340,144 +153,6 @@ def loop_for(
         one of each output, each of them is concatenated into one
         tensor along one dimension, by default, it is the first
         dimension, but it can be defined otherwise
-
-    .. runpython::
-        :showcode:
-
-        import torch
-        import onnxruntime
-        from onnx_diagnostic.export.api import to_onnx
-        from onnx_diagnostic.export.control_flow import loop_for
-
-
-        class Model(torch.nn.Module):
-            def forward(self, n_iter, x):
-                def body(i, x):
-                    return x[: i.item() + 1].unsqueeze(1)
-
-                return loop_for(n_iter, body, (x,))
-
-
-        model = Model()
-        n_iter = torch.tensor(4, dtype=torch.int64)
-        x = torch.arange(10, dtype=torch.float32)
-        expected = model(n_iter, x)
-        print("expected:", expected)
-
-        onx = to_onnx(
-            model,
-            (n_iter, x),
-            dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC})),
-            exporter="custom",
-            use_control_flow_dispatcher=True,
-        ).model_proto
-
-        sess = onnxruntime.InferenceSession(
-            onx.SerializeToString(), providers=["CPUExecutionProvider"]
-        )
-        got = sess.run(None, dict(zip(["n_iter", "x"], [n_iter.numpy(), x.numpy()])))
-        print("got:", got)
-
-
-        # The loop is exported as a custom ops.
-        ep = torch.export.export(
-            model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
-        )
-        print(ep)
-
-    Another example with two outputs:
-
-    .. runpython::
-        :showcode:
-
-        import torch
-        import onnxruntime
-        from onnx_diagnostic.export.api import to_onnx
-        from onnx_diagnostic.export.control_flow import loop_for
-
-
-        class Model(torch.nn.Module):
-            def forward(self, n_iter, x):
-                def body(i, x):
-                    return x[: i.item() + 1].unsqueeze(1), x[: i.item() + 1].unsqueeze(1) + 1
-
-                two = loop_for(n_iter, body, (x,))
-                return two[0] + two[1]
-
-
-        model = Model()
-        n_iter = torch.tensor(4, dtype=torch.int64)
-        x = torch.arange(10, dtype=torch.float32)
-        expected = model(n_iter, x)
-        print("expected:", expected)
-
-        onx = to_onnx(
-            model,
-            (n_iter, x),
-            dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC})),
-            exporter="custom",
-            use_control_flow_dispatcher=True,
-        ).model_proto
-
-        sess = onnxruntime.InferenceSession(
-            onx.SerializeToString(), providers=["CPUExecutionProvider"]
-        )
-        got = sess.run(None, dict(zip(["n_iter", "x"], [n_iter.numpy(), x.numpy()])))
-        print("got:", got)
-
-
-        # The loop is exported as a custom ops.
-        ep = torch.export.export(
-            model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
-        )
-        print(ep)
-
-    A last example with ``reduction_dim``:
-
-    .. runpython::
-        :showcode:
-
-        import torch
-        import onnxruntime
-        from onnx_diagnostic.export.api import to_onnx
-        from onnx_diagnostic.export.control_flow import loop_for
-
-
-        class Model(torch.nn.Module):
-            def forward(self, n_iter, x):
-                def body(i, x):
-                    return x[: i.item() + 1].unsqueeze(1), x[: i.item() + 1].unsqueeze(0) + 1
-
-                two = loop_for(n_iter, body, (x,), reduction_dim=[0, 1])
-                return two[0] + two[1].T
-
-
-        model = Model()
-        n_iter = torch.tensor(4, dtype=torch.int64)
-        x = torch.arange(10, dtype=torch.float32)
-        expected = model(n_iter, x)
-        print("expected:", expected)
-
-        onx = to_onnx(
-            model,
-            (n_iter, x),
-            dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC})),
-            exporter="custom",
-            use_control_flow_dispatcher=True,
-        ).model_proto
-
-        sess = onnxruntime.InferenceSession(
-            onx.SerializeToString(), providers=["CPUExecutionProvider"]
-        )
-        got = sess.run(None, dict(zip(["n_iter", "x"], [n_iter.numpy(), x.numpy()])))
-        print("got:", got)
-
-
-        # The loop is exported as a custom ops.
-        ep = torch.export.export(
-            model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
-        )
-        print(ep)
     """
     assert args, "The function should have at least one arg."
     assert (
@@ -486,6 +161,12 @@ def loop_for(
         and len(n_iter.shape) == 0
     ), f"Only a tensor for one int64 is allowed for n_iter but it equal to {n_iter}."
     if is_exporting():
+        from torch.fx.experimental.proxy_tensor import _CURRENT_MAKE_FX_TRACER
+
+        # tracer = _CURRENT_MAKE_FX_TRACER.fx_tracer
+        root = _CURRENT_MAKE_FX_TRACER.fx_tracer.root
+        # graph = _CURRENT_MAKE_FX_TRACER.fx_tracer.graph
+
         body_gm: torch.fx.GraphModule = materialize_as_graph(
             body_fn, (torch.tensor(0, dtype=torch.int64), *args)
         )
@@ -505,7 +186,29 @@ def loop_for(
             body_mutated_inputs=body_mutated_inputs,
             body_outputs=body_outputs,
         )
-        fct = getattr(torch.ops.onnx_higher_ops, name)
-        return fct(n_iter, *args)
+        root.register_module(name, body_gm)
+        # body_graph = _maybe_reenter_make_fx(body_fn)(n_iter, *args)
+        return torch.ops.onnx_higher_ops.loop_for(name, n_iter, args)
 
     return _loop_for_fn(n_iter, body_fn, reduction_dim, args)
+
+
+"""
+        proxy_mode.tracer.root.register_module(cond_graph_name, cond_graph)
+        proxy_mode.tracer.root.register_module(body_graph_name, body_graph)
+
+        args = (cond_graph, body_graph, carried_inputs, additional_inputs)
+
+        proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args)
+
+        out_proxy = proxy_mode.tracer.create_proxy(
+            "call_function", op, proxy_args, {}, name=op._name
+        )
+
+        out = op(
+            cond_graph, body_graph, unspecialized_carried_inputs, additional_inputs
+        )
+        return track_tensor_tree(
+            out, out_proxy, constant=None, tracer=proxy_mode.tracer
+        )
+"""