Add local functions to TorchOnnxEvaluator (#122)

* Add local functions

* mypy

* mypy

* mypy

* mypy

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.6.1
 +++++
 
+* :pr:`122`: add local functions to TorchOnnxEvaluator
 * :pr:`120`: enables TorchOnnxEvaluator in command line ``python -m onnx_diagnostic validate ...``
 * :pr:`115`, :pr:`116`, :pr:`117`, :pr:`118`, :pr:`119`:
   first steps for TorchOnnxEvaluator

_doc/api/reference/torch_ops/index.rst

@@ -22,6 +22,12 @@ OpRunValue
 .. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunValue
     :members:
 
+OpRunFunction
++++++++++++++
+
+.. autoclass:: onnx_diagnostic.reference.torch_ops.OpRunFunction
+    :members:
+
 Other functions
 +++++++++++++++
 

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -119,7 +119,6 @@ def test_op_binary_cmp(self):
             ir_version=9,
             opset_imports=[oh.make_opsetid("", 18)],
         )
-        onnx.checker.check_model(model)
         self._finalize_test(
             model,
             torch.abs(torch.rand(3, 4, 5, dtype=torch.float32)),
@@ -853,6 +852,51 @@ def test_op_trilu_k(self):
             torch.tensor([2], dtype=torch.int64),
         )
 
+    def test_local_function(self):
+        new_domain = "custom"
+
+        linear_regression = oh.make_function(
+            new_domain,
+            "LinearRegression",
+            ["x", "a", "b"],
+            ["y"],
+            [
+                oh.make_node("MatMul", ["x", "a"], ["xa"]),
+                oh.make_node("Add", ["xa", "b"], ["y"]),
+            ],
+            [oh.make_opsetid("", 18)],
+            [],
+        )
+
+        graph = oh.make_graph(
+            [
+                oh.make_node("LinearRegression", ["X", "A", "B"], ["Y1"], domain=new_domain),
+                oh.make_node("Abs", ["Y1"], ["Y"]),
+            ],
+            "example",
+            [
+                oh.make_tensor_value_info("X", TFLOAT, ["a", "b"]),
+                oh.make_tensor_value_info("A", TFLOAT, ["a", "b"]),
+                oh.make_tensor_value_info("B", TFLOAT, ["a", "b"]),
+            ],
+            [oh.make_tensor_value_info("Y", TFLOAT, ["a", "b"])],
+        )
+
+        model = oh.make_model(
+            graph,
+            opset_imports=[oh.make_opsetid("", 18), oh.make_opsetid(new_domain, 1)],
+            functions=[linear_regression],
+            ir_version=10,
+        )
+        self.assertNotEmpty(model.functions)
+        self._finalize_test(
+            model,
+            torch.rand((3, 3), dtype=torch.float32),
+            torch.rand((3, 3), dtype=torch.float32),
+            torch.rand((3, 3), dtype=torch.float32),
+            atol=1e-6,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_evaluator.py

@@ -43,6 +43,7 @@ class TorchOnnxEvaluator:
     :param proto: a proto
     :param providers: where to run the model
     :param opsets: needed if proto is a graph
+    :param functions: known local functions
 
     The class holds the following attributes:
 
@@ -66,10 +67,12 @@ def __init__(
         proto: Union[onnx.FunctionProto, onnx.GraphProto, onnx.ModelProto],
         providers: Tuple[str, ...] = ("CPUExecutionProvider",),
         opsets: Optional[Dict[str, int]] = None,
+        local_functions: Optional[Dict[Tuple[str, str], "TorchOnnxEvaluator"]] = None,
     ):
         self.providers = providers
         self.constants: Dict[str, torch.Tensor] = {}
         self.kernels: List[Optional[torch_ops.OpRun]] = []
+        self.functions = local_functions.copy() if local_functions else {}
         self.CPU = torch.tensor([0]).to("cpu").device
         if "CUDAExecutionProvider" in providers:
             self.CUDA = torch.tensor([0]).to("cuda").device
@@ -83,6 +86,10 @@ def __init__(
             assert opsets is None, "proto is a model, opsets must be None in that case"
             assert not proto.graph.sparse_initializer, "sparse_initializer not support yet"
             self.opsets = {d.domain: d.version for d in proto.opset_import}
+            for f in proto.functions:
+                self.functions[f.domain, f.name] = TorchOnnxEvaluator(
+                    f, providers=providers, local_functions=self.functions
+                )
             self._build_initializers(proto.graph.initializer)
             self._build_initializers(proto.graph.node)
             self._build_kernels(proto.graph.node)
@@ -138,24 +145,33 @@ def _build_kernels(self, nodes: Sequence[onnx.NodeProto]):
         kernels = get_kernels()
         self.kernels.clear()
         for node in nodes:
+            if (node.domain, node.op_type) in self.functions:
+                kernel = torch_ops.OpRunFunction(
+                    self.functions[node.domain, node.op_type], node, self.opsets[node.domain]
+                )
+                self.kernels.append(kernel)
+                continue
+
             if node.op_type == "Constant" and node.domain == "":
                 # Treated as a constant.
                 self.kernels.append(None)
                 continue
+
             opset = self.opsets[node.domain]
             key = node.domain, node.op_type, opset
             while key not in kernels and opset > 0:
                 opset -= 1
                 key = node.domain, node.op_type, opset
-            assert (
-                key in kernels
-            ), f"Missing kernel for node type {node.op_type!r} from domain {node.domain!r}"
+            assert key in kernels, (
+                f"Missing kernel for node type {node.op_type!r} from domain {node.domain!r}, "
+                f"local functions={sorted(self.functions)}"
+            )
             cls = kernels[key]
             if cls.device_dependent():
-                kernel = cls(node, opset, self.default_device)  # type: ignore[call-arg]
+                kernel2: torch_ops.OpRun = cls(node, opset, self.default_device)  # type: ignore[call-arg]
             else:
-                kernel = cls(node, opset)
-            self.kernels.append(kernel)
+                kernel2 = cls(node, opset)  # type: ignore[assignment]
+            self.kernels.append(kernel2)
 
     def run(
         self,
@@ -165,7 +181,7 @@ def run(
         """
         Runs the ONNX model.
 
-        :param outputs: outputs required:
+        :param outputs: outputs required
         :param feeds: inputs
         :return: output tensors.
         """
@@ -218,7 +234,10 @@ def run(
             for name in self.last_used[it]:
                 self.runtime_info[name].clean_value()
 
-        res = [self.runtime_info[o].value.tensor for o in outputs]  # type: ignore[assignment, union-attr]
+        assert all(
+            self.runtime_info[o].value is not None for o in outputs
+        ), "Not implemented yet when one output is None."
+        fres = [self.runtime_info[o].value.tensor for o in outputs]  # type: ignore[union-attr]
 
         # clean previous execution
         for k in feeds:
@@ -227,5 +246,71 @@ def run(
             self.runtime_info[o].clean_value()
 
         if use_numpy:
-            return [None if a is None else a.detach().cpu().numpy() for a in res]  # type: ignore[union-attr]
-        return res  # type: ignore[return-value]
+            return [None if a is None else a.detach().cpu().numpy() for a in fres]
+        return fres
+
+    def run_with_values(
+        self, *args: Optional[torch_ops.OpRunValue]
+    ) -> Union[torch_ops.OpRunValue, Tuple[torch_ops.OpRunValue, ...]]:
+        """
+        Runs the ONNX model.
+
+        :param args: inputs
+        :return: output OpRunValue
+        """
+        assert all(
+            isinstance(a, torch_ops.OpRunValue) for a in args
+        ), f"Unexpected type in args: {[type(a) for a in args]}"
+        outputs = self.output_names
+
+        # sets constants
+        for k, v in self.constants.items():
+            r = self.runtime_info[k]
+            if not r.has_value:
+                r.set_value(
+                    torch_ops.OpRunValue(
+                        v.to(self.CUDA) if r.is_shape and self.on_cuda else v, True
+                    )
+                )
+
+        # inputs
+        for k, v in zip(self.input_names, args):
+            r = self.runtime_info[k]
+            r.set_value(torch_ops.OpRunValue(None if v is None else v.tensor))
+
+        # node execution
+        for it, kernel in enumerate(self.kernels):
+            if kernel is not None:
+                # kernel execution
+                inputs = [(self.runtime_info[i].value if i else None) for i in kernel.input]
+                res = kernel.run(*inputs)
+                if isinstance(res, tuple):
+                    # outputs
+                    assert all(isinstance(o, torch_ops.OpRunValue) for o in res), (
+                        f"Unexpected output type {[type(o) for o in res]} "
+                        f"for kernel {type(kernel)}."
+                    )
+                    for name, t in zip(kernel.output, res):
+                        self.runtime_info[name].set_value(t)
+                else:
+                    assert isinstance(
+                        res, torch_ops.OpRunValue
+                    ), f"Unexpected output type {type(res)} for kernel {type(kernel)}."
+                    self.runtime_info[kernel.output[0]].set_value(res)
+
+            # free intermediate results
+            for name in self.last_used[it]:
+                self.runtime_info[name].clean_value()
+
+        assert all(
+            self.runtime_info[o].value is not None for o in outputs
+        ), "Not implemented yet when one output is None."
+        res2 = [torch_ops.OpRunValue(self.runtime_info[o].value.tensor) for o in outputs]  # type: ignore[assignment, union-attr]
+
+        # clean previous execution
+        for k in self.input_names:
+            self.runtime_info[k].clean_value()
+        for o in self.output_names:
+            self.runtime_info[o].clean_value()
+
+        return res2[0] if len(res2) == 1 else tuple(res2)  # type: ignore[index, return-value, arg-type]

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -1,4 +1,4 @@
-from ._op_run import OpRun, OpRunValue
+from ._op_run import OpRun, OpRunFunction, OpRunValue
 from .access_ops import Gather_1, Slice_13
 from .binary_ops import (
     And_1,
@@ -29,6 +29,7 @@
     Unsqueeze_13,
 )
 from .unary_ops import (
+    Abs_1,
     Cos_1,
     Exp_1,
     Log_1,

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -90,7 +90,9 @@ def __str__(self) -> str:
             )
         return f"{self.op_type}({', '.join(self.input)}) -> {', '.join(self.output)}"
 
-    def run(self, *args) -> Union[OpRunValue, Tuple[OpRunValue, ...]]:
+    def run(
+        self, *args: Optional[OpRunValue]
+    ) -> Union[OpRunValue, Tuple[Optional[OpRunValue], ...]]:
         "Kernel implementation."
         raise NotImplementedError(
             f"Method run is not implemented for kernel {self.__class__.__name__!r}"
@@ -157,3 +159,24 @@ def get_attribute_tensor(self, node: onnx.NodeProto, name: str) -> Optional[torc
         if att is None:
             return None
         return to_tensor(att.t)
+
+
+class OpRunFunction(OpRun):
+    """
+    Defines a kernel based on a local functions.
+    """
+
+    def __init__(
+        self,
+        runtime: "onnx_diagnostic.reference.TorchOnnxEvaluator",  # noqa: F821
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+    ):
+        super().__init__(node, version)
+        self.runtime = runtime
+        self.input_names = runtime.input_names
+
+    def run(
+        self, *args: Optional[OpRunValue]
+    ) -> Union[OpRunValue, Tuple[Optional[OpRunValue], ...]]:
+        return self.runtime.run_with_values(*args)

onnx_diagnostic/reference/torch_ops/unary_ops.py

@@ -2,6 +2,13 @@
 from . import OpRun, OpRunValue
 
 
+class Abs_1(OpRun):
+    """Abs"""
+
+    def run(self, x: OpRunValue) -> OpRunValue:
+        return OpRunValue(torch.abs(x.tensor))
+
+
 class Cos_1(OpRun):
     """Cos"""
 

onnx_diagnostic/torch_onnx/runtime_info.py

@@ -167,28 +167,49 @@ def set_is_shape(node: onnx.NodeProto, values: Dict[str, RuntimeValue]) -> List[
 
 
 def first_used_last_used(
-    proto: onnx.ModelProto, constant_as_initializer: bool = False
+    proto: Union[onnx.FunctionProto, onnx.GraphProto, onnx.ModelProto],
+    constant_as_initializer: bool = False,
 ) -> Dict[str, RuntimeValue]:
     """
     Builds first used, last used information for every result
     in the model.
 
-    :param proto: model
+    :param proto: model, graph or function
     :param constant_as_initializer: outputs of node Constant is tagged as INITIALIZER
     :return: dictionary of RuntimeValue
     """
     values = {}
-    for init in proto.graph.initializer:
+    if isinstance(proto, onnx.ModelProto):
+        initializer = proto.graph.initializer
+        sparse_initializer = proto.graph.sparse_initializer
+        _input = proto.graph.input
+        output = proto.graph.output
+        _node = proto.graph.node
+    elif isinstance(proto, onnx.GraphProto):
+        initializer = proto.initializer
+        sparse_initializer = proto.sparse_initializer
+        _input = proto.input
+        output = proto.output
+        _node = proto.node
+    else:
+        initializer = []
+        sparse_initializer = []
+        _input = proto.input
+        output = proto.output
+        _node = proto.node
+
+    for init in initializer:
         values[init.name] = RuntimeValue(
             init.name, kind=RuntimeValueKind.INITIALIZER, created=-1
         )
-    for init in proto.graph.sparse_initializer:
+    for init in sparse_initializer:
         values[init.name] = RuntimeValue(
             init.name, created=-1, kind=RuntimeValueKind.INITIALIZER
         )
-    for inp in proto.graph.input:
-        values[inp.name] = RuntimeValue(inp.name, created=-1, kind=RuntimeValueKind.INPUT)
-    for it, node in enumerate(proto.graph.node):
+    for inp in _input:
+        n = inp if isinstance(inp, str) else inp.name
+        values[n] = RuntimeValue(n, created=-1, kind=RuntimeValueKind.INPUT)
+    for it, node in enumerate(_node):
         for i in node.input:
             if values[i].first_used is None:
                 values[i].first_used = it
@@ -212,7 +233,8 @@ def first_used_last_used(
             )
         set_is_shape(node, values)
 
-    for out in proto.graph.output:
-        values[out.name].kind = RuntimeValueKind.OUTPUT
-        values[out.name].last_used = len(proto.graph.node)
+    for out in output:
+        n = out if isinstance(out, str) else out.name
+        values[n].kind = RuntimeValueKind.OUTPUT
+        values[n].last_used = len(_node)
     return values

pyproject.toml

@@ -44,6 +44,14 @@ disable_error_code = ["assignment", "arg-type", "name-defined", "union-attr"]
 module = ["onnx_diagnostic.helpers.ort_session"]
 disable_error_code = ["union-attr"]
 
+[[tool.mypy.overrides]]
+module = ["onnx_diagnostic.reference.torch_ops.*"]
+disable_error_code = ["override"]
+
+[[tool.mypy.overrides]]
+module = ["onnx_diagnostic.reference.torch_ops._op_run"]
+disable_error_code = ["name-defined"]
+
 [[tool.mypy.overrides]]
 module = ["onnx_diagnostic.torch_export_patches.*"]
 disable_error_code = ["arg-type", "assignment", "attr-defined", "index", "misc", "name-defined", "operator", "return-value", "union-attr"]