Add ONNX-MLIR dialect support

README.md

@@ -206,6 +206,50 @@ running
 python -c "from tvm.runtime import Module"
 ```
 
+### Step 2-tris: Install ONNX-MLIR Compiler
+Since `ONNX-MLIR` compiler need to be installed from the source code, it build
+the specific commit of [llvm-project](https://github.com/llvm/llvm-project) from
+source, further builds [ONNX-MLIR](https://github.com/onnx/onnx-mlir) from source.
+Its installation can take several minutes (or even hours) for being performed.
+For this reason we decided to not include it in the default auto-installer. However,
+if you want to  squeeze out the maximum of the performance from your model on your machine,
+we highly recommend installing ONNX-MLIR as well. With `nebullvm` it is super-easy! Just run
+
+```
+python -c "from nebullvm.installers.installers import install_onnx_mlir; install_onnx_mlir()"
+```
+and wait for the compiler to be installed! You can check that everything worked by running
+```
+python -c "from nebullvm.inference_learners import onnx_mlir"
+```
+The default installation location of `ONNX-MLIR` is in home directory,
+if you want to install it in different directory, installation directory can be
+changed by specifying the custom installation directory in `install_onnx_mlir` function
+```
+python -c "from nebullvm.installers.installers import install_onnx_mlir; 
+install_onnx_mlir(absolute_path_of_installation_directory)"
+```
+and wait for the compiler to be installed!
+In the subsequent run, specify ONNX-MLIR installation directory by exporting
+the environment variable `MLIR_INSTALLATION_ROOT="absolute_path_of_installation_directory"` by running
+
+```
+export MLIR_INSTALLATION_ROOT="absolute_path_of_installation_directory"
+```
+
+from your command line or adding
+
+```
+import os
+os.environ["MLIR_INSTALLATION_ROOT"] = "absolute_path_of_installation_directory"
+```
+in your python code before importing `nebullvm` for the first time.
+You can check that everything worked by running
+```
+export MLIR_INSTALLATION_ROOT="absolute_path_of_installation_directory"
+python -c "from nebullvm.inference_learners import onnx_mlir"
+```
+
 </details>
 
 <details> 

nebullvm/base.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Tuple, List, Dict, Union
+from typing import Dict, List, Tuple, Union
 
 
 class DataType(str, Enum):
@@ -106,6 +106,7 @@ class ModelCompiler(Enum):
     OPENVINO = "openvino"
     APACHE_TVM = "tvm"
     ONNX_RUNTIME = "onnxruntime"
+    ONNX_MLIR = "onnx_mlir"
 
 
 class QuantizationType(Enum):

nebullvm/config.py

@@ -37,3 +37,7 @@
     "description_file": "description.xml",
     "weights": "weights.bin",
 }
+
+ONNX_MLIR_FILENAMES = {
+    "model_name": "mlir_model.so",
+}

nebullvm/inference_learners/onnx_mlir.py

@@ -0,0 +1,220 @@
+import copy
+import os
+import shutil
+import sys
+import warnings
+from abc import ABC
+from pathlib import Path
+from typing import Dict, Generator, List, Tuple, Type, Union
+
+import cpuinfo
+import numpy as np
+import tensorflow as tf
+import torch
+from nebullvm.base import DeepLearningFramework, ModelParams
+from nebullvm.config import ONNX_MLIR_FILENAMES
+from nebullvm.inference_learners.base import (
+    BaseInferenceLearner,
+    LearnerMetadata,
+    PytorchBaseInferenceLearner,
+    TensorflowBaseInferenceLearner,
+)
+
+try:
+    # Set the ONNX_MLIR_HOME as the environment variable and append in the path,
+    # directory path where the MLIR is built
+
+    # retrieve the ONNX-MLIR installation directory from environment variable
+    # if exists otherwise set to home directory
+    MLIR_INSTALLATION_ROOT = os.environ.get(
+        "MLIR_INSTALLATION_ROOT", Path.home()
+    )
+
+    os.environ["ONNX_MLIR_HOME"] = os.path.join(
+        MLIR_INSTALLATION_ROOT,
+        "onnx-mlir",
+        "build",
+        "Debug",
+    )
+
+    sys.path.append(
+        os.path.join(
+            os.environ.get("ONNX_MLIR_HOME", ""),
+            "lib",
+        )
+    )
+    import PyRuntime
+except ImportError:
+    # Disable the ONNX-MLIR auto-installer for the time being as it takes long to
+    #  install, can be installed by explicitly running the install_onnx_mlir function
+
+    # TODO: Remove the False flag for allowing onnx-mlir to be installed by
+    #  the Auto-Installer.
+    if False and not NO_COMPILER_INSTALLATION:
+        warnings.warn(
+            "No valid onnx-mlir installation found. Trying to install it..."
+        )
+        from nebullvm.installers.installers import install_onnx_mlir
+
+        install_onnx_mlir(
+            working_dir=MLIR_INSTALLATION_ROOT,
+        )
+        import PyRuntime
+    else:
+        warnings.warn(
+            "Not found any valid onnx-mlir installation. "
+            "ONNX-MLIR will not be available in the following steps. "
+            "ONNX-MLIR should be explicitly installed using install_onnx_mlir."
+        )
+
+
+class ONNXMlirInferenceLearner(BaseInferenceLearner, ABC):
+    """Model converted from ONNX to Shared Object file using ONNX-MLIR dialect
+    and run with ONNX-MLIR's PyRuntime
+    created at onnx-mlir/build/Debug/lib/PyRuntime.cpython-<target>.so.
+
+    Attributes:
+        onnx_mlir_model_path (str or Path): Path to the shared object mlir model.
+        network_parameters (ModelParams): The model parameters as batch
+                size, input and output sizes.
+    """
+
+    def __init__(
+        self,
+        onnx_mlir_model_path: Union[str, Path],
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.onnx_mlir_model_path = onnx_mlir_model_path
+        self._session = PyRuntime.ExecutionSession(
+            os.path.abspath(str(self.onnx_mlir_model_path)),
+        )
+
+    def save(self, path: Union[str, Path], **kwargs):
+        """Save the model.
+
+        Args:
+            path (Path or str): Path to the directory where the model will
+                be stored.
+            kwargs (Dict): Dictionary of key-value pairs that will be saved in
+                the model metadata file.
+        """
+        metadata = LearnerMetadata.from_model(self, **kwargs)
+        metadata.save(path)
+
+        shutil.copy(
+            self.onnx_mlir_model_path,
+            os.path.join(str(path), ONNX_MLIR_FILENAMES["model_name"]),
+        )
+
+    @classmethod
+    def load(cls, path: Union[Path, str], **kwargs):
+        """Load the model.
+
+        Args:
+            path (Path or str): Path to the directory where the model is
+                stored.
+            kwargs (Dict): Dictionary of additional arguments for consistency
+                with other Learners.
+
+        Returns:
+            ONNXInferenceLearner: The optimized model.
+        """
+        if len(kwargs) > 0:
+            warnings.warn(
+                f"No extra keywords expected for the load method. "
+                f"Got {kwargs}."
+            )
+        onnx_mlir_model_path = os.path.join(
+            str(path), ONNX_MLIR_FILENAMES["model_name"]
+        )
+        metadata = LearnerMetadata.read(path)
+
+        return cls(
+            network_parameters=ModelParams(**metadata.network_parameters),
+            onnx_mlir_model_path=onnx_mlir_model_path,
+        )
+
+    def _predict_arrays(self, input_arrays: Generator[np.ndarray, None, None]):
+        outputs = self._session.run(list(input_arrays))
+        return outputs
+
+
+class PytorchONNXMlirInferenceLearner(
+    ONNXMlirInferenceLearner, PytorchBaseInferenceLearner
+):
+    """Model run with ONNX-MLIR's PyRuntime using a Pytorch interface.
+
+    Attributes:
+        onnx_mlir_model_path (str or Path): Path to the shared object mlir model.
+        network_parameters (ModelParams): The model parameters as batch
+                size, input and output sizes.
+    """
+
+    def predict(self, *input_tensors: torch.Tensor) -> Tuple[torch.Tensor]:
+        """Predict on the input tensors.
+
+        Note that the input tensors must be on the same batch. If a sequence
+        of tensors is given when the model is expecting a single input tensor
+        (with batch size >= 1) an error is raised.
+
+        Args:
+            input_tensors (Tuple[Tensor]): Input tensors belonging to the same
+                batch. The tensors are expected having dimensions
+                (batch_size, dim1, dim2, ...).
+
+        Returns:
+            Tuple[Tensor]: Output tensors. Note that the output tensors does
+                not correspond to the prediction on the input tensors with a
+                1 to 1 mapping. In fact the output tensors are produced as the
+                multiple-output of the model given a (multi-) tensor input.
+        """
+        input_arrays = (
+            input_tensor.cpu().detach().numpy()
+            for input_tensor in input_tensors
+        )
+        outputs = self._predict_arrays(input_arrays)
+        return tuple(torch.from_numpy(output) for output in outputs)
+
+
+class TensorflowONNXMlirInferenceLearner(
+    ONNXMlirInferenceLearner, TensorflowBaseInferenceLearner
+):
+    """Model run with ONNX-MLIR's PyRuntime using a tensorflow interface.
+
+    Attributes:
+        onnx_mlir_model_path (str or Path): Path to the shared object mlir model.
+        network_parameters (ModelParams): The model parameters as batch
+                size, input and output sizes.
+    """
+
+    def predict(self, *input_tensors: tf.Tensor) -> Tuple[tf.Tensor]:
+        """Predict on the input tensors.
+
+        Note that the input tensors must be on the same batch. If a sequence
+        of tensors is given when the model is expecting a single input tensor
+        (with batch size >= 1) an error is raised.
+
+        Args:
+            input_tensors (Tuple[Tensor]): Input tensors belonging to the same
+                batch. The tensors are expected having dimensions
+                (batch_size, dim1, dim2, ...).
+
+        Returns:
+            Tuple[Tensor]: Output tensors. Note that the output tensors does
+                not correspond to the prediction on the input tensors with a
+                1 to 1 mapping. In fact the output tensors are produced as the
+                multiple-output of the model given a (multi-) tensor input.
+        """
+        input_arrays = (input_tensor.numpy() for input_tensor in input_tensors)
+        outputs = self._predict_arrays(input_arrays)
+        # noinspection PyTypeChecker
+        return tuple(tf.convert_to_tensor(output) for output in outputs)
+
+
+ONNX_MLIR_INFERENCE_LEARNERS: Dict[
+    DeepLearningFramework, Type[ONNXMlirInferenceLearner]
+] = {
+    DeepLearningFramework.PYTORCH: PytorchONNXMlirInferenceLearner,
+    DeepLearningFramework.TENSORFLOW: TensorflowONNXMlirInferenceLearner,
+}

nebullvm/installers/install_onnx_mlir.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# Installation steps to build the ONNX-MLIR from source
+
+# Set non interactive mode for apt-get
+export DEBIAN_FRONTEND=noninteractive
+
+# Build ONNX-MLIR
+
+if [ ! -d "onnx-mlir" ]
+then
+
+  git clone --recursive https://github.com/onnx/onnx-mlir.git onnx-mlir
+fi
+
+
+if [ -z "$NPROC" ]
+then
+  NPROC=4
+fi
+
+
+# Export environment variables pointing to LLVM-Projects.
+export MLIR_DIR=$(pwd)/llvm-project/build/lib/cmake/mlir
+
+# Get the python interpreter path
+export PYTHON_LOCATION=$(which python3)
+
+mkdir onnx-mlir/build && cd onnx-mlir/build
+
+if [[ -z "$PYTHON_LOCATION" ]]; then
+  cmake -G Ninja \
+        -DCMAKE_CXX_COMPILER=/usr/bin/c++ \
+        -DMLIR_DIR=${MLIR_DIR} \
+        ..
+else
+  echo "Using python path " $PYTHON_LOCATION
+  echo "Using MLIR_DIR " $MLIR_DIR
+
+  cmake -G Ninja \
+      -DCMAKE_CXX_COMPILER=/usr/bin/c++ \
+      -DPython3_ROOT_DIR=${PYTHON_LOCATION} \
+      -DPython3_EXECUTABLE=${PYTHON_LOCATION} \
+      -DMLIR_DIR=${MLIR_DIR} \
+      ..
+
+fi
+
+cmake --build . --parallel $NPROC
+
+# Run lit tests:
+export LIT_OPTS=-v
+cmake --build . --parallel $NPROC --target check-onnx-lit