onnx-diagnostic 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -0,0 +1,326 @@
+from typing import Any, List, Optional, Union, Tuple
+import onnx
+import torch
+from ...api import TensorLike
+from ...helpers import string_type
+from ...helpers.torch_helper import to_tensor
+
+
+class OpRunValue(TensorLike):
+    """Defines a value for the runtime, a tensor or a sequence."""
+
+    __slots__ = ("cached", "is_constant", "sequence", "tensor")
+
+    @classmethod
+    def is_sequence(cls) -> bool:
+        "Tells if it is sequence."
+        raise NotImplementedError("is_sequence must be overwritten.")
+
+
+class OpRunTensor(OpRunValue):
+    """
+    Wrapper around a tensor.
+
+    :param tensor: torch.Tensor
+    :param is_constant: is it a constant
+    :param may_cpu: change the device the tensor is if
+        more appropriate
+    """
+
+    def __init__(self, tensor, is_constant: bool = False, may_cpu: bool = False):
+        assert isinstance(tensor, torch.Tensor), (
+            f"Unexpected type {type(tensor)}, "
+            f"__name__={getattr(tensor, '__name__', 'no name')}"
+        )
+        assert tensor is None or tensor.numel() != 1 or tensor.item() != -666666
+        self.tensor = (
+            tensor.cpu()
+            if may_cpu
+            and len(tensor.shape) == 1
+            and tensor.numel() < 8
+            and tensor.dtype == torch.int64
+            and tensor.get_device() >= 0
+            else tensor
+        )
+        self.is_constant = is_constant
+        self.cached: Optional[Tuple[int, ...]] = None
+
+    @classmethod
+    def is_sequence(cls) -> bool:
+        "Tells if it is sequence."
+        return False
+
+    def to(self, to: Any) -> "OpRunTensor":
+        "Changes the device."
+        return OpRunTensor(self.tensor.to(to))
+
+    def string_type(self) -> str:
+        "Returns information about the value as a string."
+        s = string_type(self.tensor, with_shape=True, with_min_max=True, with_device=True)
+        if self.is_constant:
+            return f"CST({s})"
+        return s
+
+    def __repr__(self) -> str:
+        "usual"
+        if self.is_constant:
+            return (
+                f"{self.__class__.__name__}"
+                f"({string_type(self.tensor, with_shape=True)}, is_constant=True)"
+            )
+        return f"{self.__class__.__name__}({string_type(self.tensor, with_shape=True)})"
+
+    @property
+    def tensor_or_sequence(self) -> Union[torch.Tensor, List[torch.Tensor]]:
+        "Returns either a tensor or a sequence."
+        return self.tensor
+
+    @property
+    def shape(self):
+        "shape (torch shape)"
+        return self.tensor.shape
+
+    @property
+    def dtype(self):
+        "dtype (torch dtype)"
+        return self.tensor.dtype
+
+    def _tensor_as_tuple_int(self) -> Tuple[int, ...]:
+        return tuple(map(int, self.tensor))
+
+    def numel(self) -> int:
+        "Returns the number of elements."
+        return 0 if self.tensor is None else self.tensor.numel()
+
+    def get_device(self) -> int:
+        "Returns the device id."
+        return -1 if self.tensor is None else self.tensor.get_device()
+
+    @property
+    def device(self):
+        "Returns the device."
+        return -1 if self.tensor is None else self.tensor.device
+
+    @property
+    def as_tuple_int(self) -> Tuple[int, ...]:
+        "value as int"
+        if self.is_constant:
+            if self.cached is None:
+                self.cached = self._tensor_as_tuple_int()
+            return self.cached
+        return self._tensor_as_tuple_int()
+
+    def copy(self) -> "OpRunTensor":
+        "Shallow copy."
+        return self.__class__(self.tensor)
+
+
+class OpRunSequence(OpRunValue):
+    """Defines a sequence."""
+
+    def __init__(
+        self, sequence: Optional[List[torch.Tensor]] = None, dtype: torch.dtype = torch.float32
+    ):
+        self.tensor = torch.tensor(-666666, dtype=dtype)
+        self.is_shape = False
+        self.sequence = sequence or []
+        self.cached: Optional[Tuple[int, ...]] = None
+        assert all(
+            isinstance(s, torch.Tensor) for s in self.sequence
+        ), f"Unexpected type in sequence {[type(s) for s in self.sequence]}"
+
+    @property
+    def dtype(self):
+        "dtype (torch dtype)"
+        return self.tensor.dtype
+
+    @property
+    def tensor_or_sequence(self) -> Union[torch.Tensor, List[torch.Tensor]]:
+        "Returns either a tensor or a sequence."
+        return self.sequence
+
+    @classmethod
+    def is_sequence(cls) -> bool:
+        "Tells if it is sequence."
+        return True
+
+    def insert_at(
+        self, tensor: torch.Tensor, position: Optional[OpRunTensor] = None
+    ) -> "OpRunSequence":
+        "Inserts a value at a given position."
+        assert isinstance(tensor, OpRunTensor), f"Unexpected type {type(tensor)} for tensor"
+        new_seq = OpRunSequence()
+        seq = self.sequence.copy()
+        new_seq.sequence = seq
+        if position is None:
+            seq.append(tensor.tensor)
+        else:
+            seq.insert(int(position.tensor.item()), tensor.tensor)
+        return new_seq
+
+    def copy(self) -> "OpRunSequence":
+        "Shallow copy."
+        return self.__class__(self.sequence, dtype=self.dtype)
+
+    def string_type(self) -> str:
+        "Returns a string which can be printed."
+        return string_type(self.sequence, with_shape=True)
+
+
+class OpRun:
+    """
+    Main class. Every kernel should inherit from it.
+    It does not copy the proto.
+    """
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        """
+        Returns True if the kernel needs a device to be efficiently initialized.
+        """
+        return False
+
+    @classmethod
+    def has_subgraphs(cls) -> bool:
+        """Returns True if the kernel has subgraphs."""
+        return False
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        assert isinstance(
+            node, onnx.NodeProto
+        ), f"node must be a NodeProto but node is {type(node)}"
+        self.op_type = node.op_type
+        self.domain = node.domain
+        self.input = node.input
+        self.output = node.output
+        if version is None:
+            name = self.__class__.__name__.split("_")
+            assert (
+                len(name) == 2
+            ), f"Cannot guess version from name={self.__class__.__name__!r}"
+            version = int(name[1])
+        self.version = version
+        self.name = node.name
+
+    def __str__(self) -> str:
+        "usual"
+        if self.domain:
+            return (
+                f"{self.op_type}[{self.domain}]({', '.join(self.input)}) "
+                f"-> {', '.join(self.output)}"
+            )
+        return f"{self.op_type}({', '.join(self.input)}) -> {', '.join(self.output)}"
+
+    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}"
+        )
+
+    def _find_attribute(self, node: onnx.NodeProto, name: str):
+        for att in node.attribute:
+            if att.name == name:
+                return att
+        return None
+
+    def get_attribute_float(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[float] = None
+    ) -> Optional[float]:
+        """
+        Returns an attribute as an int.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else float(att.f)
+
+    def get_attribute_int(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[int] = None
+    ) -> Optional[int]:
+        """
+        Returns an attribute as an int.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else int(att.i)
+
+    def get_attribute_ints(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[Tuple[int, ...]] = None
+    ) -> Optional[Tuple[int, ...]]:
+        """
+        Returns an attribute as a tuple of ints.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else tuple(map(int, att.ints))
+
+    def get_attribute_string(
+        self, node: onnx.NodeProto, name: str, default_value: Optional[str] = None
+    ) -> Optional[str]:
+        """
+        Returns an attribute as a tuple of ints.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        return default_value if att is None else att.s.decode("utf-8")
+
+    def get_attribute_tensor(self, node: onnx.NodeProto, name: str) -> Optional[torch.Tensor]:
+        """
+        Returns an attribute as a torch tensor.
+
+        :param node: NodeProto
+        :param name: name
+        :param default_value: default_value
+        :return: value
+        """
+        att = self._find_attribute(node, name)
+        if att is None:
+            return None
+        return to_tensor(att.t)
+
+    def same_device(self, *tensors: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        """Puts all tensors on the same device."""
+        devices = [t.get_device() for t in tensors]
+        if len(set(devices)) == 1:
+            return tuple(tensors)
+        index = devices.index(max(devices))
+        device = tensors[index].device
+        return tuple(t.to(device) for t in tensors)
+
+
+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/access_ops.py

@@ -0,0 +1,84 @@
+from typing import Optional
+import onnx
+import torch
+from . import OpRun, OpRunTensor
+
+
+class Gather_1(OpRun):
+    "Gather"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        axis = self.get_attribute_int(node, "axis", 0)
+        assert isinstance(axis, int), f"Unexpected value for attribute axis={axis!r}"
+        self.axis = axis
+
+    def run(self, x, indices):
+        if indices.tensor.numel() == 0:
+            return torch.empty((0,), dtype=x.tensor.dtype, device=x.tensor.device)
+        ind = [slice(0, s) for s in x.shape]
+        ind[self.axis] = indices.tensor
+        return OpRunTensor(x.tensor[tuple(ind)])
+
+
+class ScatterND_16(OpRun):
+    "ScatterND"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.reduction = self.get_attribute_string(node, "reduction", "none")
+
+    def run(
+        self, data: OpRunTensor, indices: OpRunTensor, updates: OpRunTensor
+    ) -> OpRunTensor:
+        # This implementation is not efficient.
+        grids = torch.meshgrid(*[torch.arange(s) for s in indices.shape[:-1]], indexing="ij")
+        stacked = torch.stack(grids, dim=-1)
+        index = stacked.reshape(-1, len(indices.shape) - 1)
+        output = data.tensor.clone()
+        for i in index:
+            if self.reduction == "add":
+                output[indices.tensor[i]] += updates.tensor[i]
+            elif self.reduction == "mul":
+                output[indices.tensor[i]] *= updates.tensor[i]
+            elif self.reduction == "max":
+                output[indices.tensor[i]] = torch.maximum(
+                    output[indices.tensor[i]], updates.tensor[i]
+                )
+            elif self.reduction == "min":
+                output[indices.tensor[i]] = torch.minimum(
+                    output[indices.tensor[i]], updates.tensor[i]
+                )
+            else:
+                output[indices.tensor[i]] = updates.tensor[i]
+        return OpRunTensor(output)
+
+
+class Slice_13(OpRun):
+    "Slice"
+
+    def run(
+        self,
+        data: OpRunTensor,
+        starts: OpRunTensor,
+        ends: OpRunTensor,
+        axes: Optional[OpRunTensor] = None,
+        steps: Optional[OpRunTensor] = None,
+    ) -> OpRunTensor:
+        if axes is None:
+            if steps is None:
+                slices = [slice(s, e) for s, e in zip(starts.tensor, ends.tensor)]
+            else:
+                slices = [
+                    slice(s, e, d) for s, e, d in zip(starts.tensor, ends.tensor, steps.tensor)
+                ]
+        else:
+            if steps is None:
+                slices = [slice(0, a) for a in data.shape]
+                for s, e, a in zip(starts.tensor, ends.tensor, axes.tensor):
+                    slices[a] = slice(s, e)
+            else:
+                slices = [slice(0, a) for a in data.shape]
+                for s, e, a, d in zip(starts.tensor, ends.tensor, axes.tensor, steps.tensor):
+                    slices[a] = slice(s, e, d)
+        return OpRunTensor(data.tensor[tuple(slices)])

onnx_diagnostic/reference/torch_ops/binary_ops.py

@@ -0,0 +1,108 @@
+import torch
+from . import OpRun, OpRunTensor
+
+
+class OpRunBinary(OpRun):
+    "Binary Op"
+
+    def run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        if x.get_device() != y.get_device():
+            if x.get_device() >= 0:
+                y = y.to(x.device)
+            else:
+                x = x.to(y.device)
+        return self._run(x, y)
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        raise NotImplementedError(f"Operator {self.__class__.__name__!r} is not complete.")
+
+
+class And_1(OpRunBinary):
+    """And"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor & y.tensor)
+
+
+class Add_1(OpRunBinary):
+    """Add"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor + y.tensor)
+
+
+class Div_1(OpRunBinary):
+    """Div"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor / y.tensor)
+
+
+class Equal_1(OpRunBinary):
+    """Equal"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor == y.tensor)
+
+
+class Greater_1(OpRunBinary):
+    """Greater"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor > y.tensor)
+
+
+class GreaterOrEqual_1(OpRunBinary):
+    """GreaterOrEqual"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor >= y.tensor)
+
+
+class Less_1(OpRunBinary):
+    """Less"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor < y.tensor)
+
+
+class LessOrEqual_1(OpRunBinary):
+    """LessOrEqual"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor <= y.tensor)
+
+
+class MatMul_1(OpRunBinary):
+    """MatMul"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor @ y.tensor)
+
+
+class Mul_1(OpRunBinary):
+    """Mul"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor * y.tensor)
+
+
+class Or_1(OpRunBinary):
+    """Or"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor | y.tensor)
+
+
+class Pow_12(OpRunBinary):
+    """Pow"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.pow(x.tensor, y.tensor))
+
+
+class Sub_1(OpRunBinary):
+    """Sub"""
+
+    def _run(self, x: OpRunTensor, y: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor - y.tensor)

onnx_diagnostic/reference/torch_ops/controlflow_ops.py

@@ -0,0 +1,118 @@
+from typing import Any, Dict, Optional
+import onnx
+import torch
+from . import OpRun, OpRunTensor
+
+
+class OpRunControlFlow(OpRun):
+    """Common ancestor for control flows."""
+
+    @classmethod
+    def has_subgraphs(cls) -> bool:
+        """Returns True if the kernel has subgraphs."""
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        parent: Optional["onnx_diagnostic.reference.TorchOnnxEvaluator"] = None,  # noqa: F821
+    ):
+        super().__init__(node, version)
+        assert (
+            parent is not None
+        ), f"parent must be specified for operator {self.__class__.__name__!r}"
+        for att in node.attribute:
+            if att.type == onnx.AttributeProto.GRAPH:
+                rt = parent.__class__(
+                    att.g,
+                    providers=parent.providers,
+                    opsets=parent.opsets,
+                    local_functions=parent.functions,
+                    verbose=parent.verbose,
+                )
+                setattr(self, att.name, rt)
+
+
+class If_1(OpRunControlFlow):
+    "If"
+
+    def run(self, cond, context: Optional[Dict[str, Any]] = None):
+        rt = self.then_branch if cond.tensor.item() else self.else_branch  # type: ignore[attr-defined]
+        return rt.run_with_values(context=context)
+
+
+class Loop_16(OpRunControlFlow):
+    "Loop"
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        parent: Optional["onnx_diagnostic.reference.TorchOnnxEvaluator"] = None,  # noqa: F821
+    ):
+        super().__init__(node, version, parent)
+        self.output_index = {n: i for i, n in enumerate(self.body.output_names)}
+        self.N = len(self.body.input_names) - 2
+        self.K = len(self.body.output_names) - self.N - 1
+
+    def run(self, M, cond, *args, context: Optional[Dict[str, Any]] = None):
+        if args:
+            v_initial = args[0]
+            args = args[1:]
+        else:
+            v_initial = None
+        assert M is None or hasattr(
+            M, "dtype"
+        ), f"M must be empty or an array but its type is {type(M)}."
+        body = self.body
+        loop_inputs = body.input_names
+        inputs = dict.fromkeys(loop_inputs)
+        if v_initial is not None:
+            inputs[loop_inputs[2]] = v_initial
+        cond_name = body.output_names[0]
+        if args:
+            begin = len(loop_inputs) - len(args)
+            all_inputs = loop_inputs[begin:]
+            for name, val in zip(all_inputs, args):
+                inputs[name] = val
+        if context is not None:
+            for a in context:
+                inputs[a] = context[a]
+
+        k_carried_away = [[] for i in range(self.K)]  # type: ignore
+        it = 0
+        while (cond is None or cond.tensor is None or cond.tensor.item()) and (
+            M is None or M.tensor is None or it < M.tensor.item()
+        ):
+            if len(body.input_names) > 0 and body.input_names[0] is not None:
+                inputs[body.input_names[0]] = OpRunTensor(
+                    torch.tensor(it, dtype=None if M is None else M.dtype)
+                )
+            if len(body.input_names) > 1 and body.input_names[1] is not None:
+                inputs[body.input_names[1]] = cond
+            outputs = list(
+                self.body.run_with_values(
+                    *[inputs[k] for k in self.body.input_names], context=context
+                )
+            )
+            if self.K > 0:
+                for k in range(self.K):
+                    k_carried_away[k].append(outputs[-self.K + k])
+            index_cond = self.output_index[cond_name]
+            cond = outputs[index_cond]
+            assert (
+                cond is not None
+            ), f"Condition {cond_name!r} returned by the subgraph cannot be None."
+            for i, o in zip(body.input_names[2:], body.output_names[1:]):
+                inputs[i] = outputs[self.output_index[o]]
+            it += 1
+
+        if it == 0:
+            outputs = [inputs[i] for i in body.input_names[2:]]
+        else:
+            outputs = outputs[1 : 1 + self.N]
+        outputs.extend([OpRunTensor(torch.cat(x, axis=0)) for x in k_carried_away])
+        while len(outputs) < len(self.body.output_names):
+            outputs.append(OpRunTensor(torch.empty(())))
+        return tuple(outputs)

onnx_diagnostic/reference/torch_ops/generator_ops.py

@@ -0,0 +1,35 @@
+from typing import Optional
+import onnx
+import torch
+from . import OpRun, OpRunTensor
+
+
+class Range_11(OpRun):
+    """Range"""
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        """
+        Returns True if the kernel needs a device to be efficiently initialized.
+        """
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__(node, version)
+        self.device = device
+
+    def run(self, starts: OpRunTensor, limit: OpRunTensor, delta: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(
+            torch.arange(
+                starts.tensor,
+                limit.tensor,
+                delta.tensor,
+                dtype=starts.dtype,
+                device=self.device,
+            )
+        )