Model Preparation

Overview

The Torq compiler can compile models expressed in MLIR, which is a generic framework for intermediate representations. To compile a model with the Torq compiler, it must first be converted to MLIR.

This describes how to convert models from TFLite and Torch to MLIR.

Convert a TFLite model to TOSA and MLIR

IREE provides a command line tool that can be used to convert a TFLite model to a binary MLIR file expressed in the TOSA dialect.

  • If not yet done, activate the Python environment as explained in Getting Started (skip this step if using the Docker container).

  • Navigate to the root directory of the Release Package, or run the Docker container. For the Docker container, the release package is located at:

    $ cd /opt/release

  • Convert the model to TOSA using the following command:

    Model Source: This model - MobileNetV2 is generated from tf.keras.applications using tf_model_generator.py. The dataset for int8 quantization is done using random data.

    $ iree-import-tflite tests/hf/Synaptics_MobileNetV2/MobileNetV2_int8.tflite -o mobilenetv2.tosa

Note

The tests/hf/ directory is only included in the Release Package and is not available in the compiler GitHub repository.

  • More details on the IREE TFLite tools can be found in the official IREE website.

    Tip

    The generated .tosa file is a binary MLIR file. To convert it to a text representation, run:

    $ iree-opt mobilenetv2.tosa -o mobilenetv2.mlir

Convert Torch Model to MLIR

  • Torch models can be converted to MLIR using the torch-mlir toolchain.

  • This process involves exporting a PyTorch model and converting it to MLIR in various dialects such as TORCH, TOSA, LINALG_ON_TENSORS, or STABLEHLO.

  • The resulting MLIR file can then be used as input for the Torq compiler, depending on the supported dialects and features.

  • Export torch-mlir Python Packages

    $ export PYTHONPATH=`pwd`/build/tools/torch-mlir/python_packages/torch_mlir:`pwd`/test/python/fx_importer

  • Create Torch Test Model and Output to different MLIR Dialect

    import torch
from torch_mlir import torchscript

class SimpleModule(torch.nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.ops.aten.abs(x)


if __name__ == '__main__':
    test_input = torch.ones(2, 16)
    graph = SimpleModule()
    graph.eval()
    module = torchscript.compile(graph,
                                test_input,
                                torchscript.OutputType.TORCH,
                                use_tracing=False,
                                verbose=False)
    print(module.operation.get_asm())

    with open("./aten-torch.mlir", "w") as fp:
        fp.write(module.operation.get_asm())

  • output type could be

    • TORCH

    • LINALG_ON_TENSORS

    • TOSA

    • STABLEHLO

Convert ONNX Model to MLIR

IREE provides an ONNX importer that converts ONNX models into a text-based MLIR representation. The importer is available in the Torq compiler Python environment.

  • If not using the Docker container, activate the Python environment as explained in Getting Started.

  • Run the IREE importer to convert an ONNX model to MLIR:

    $ python -m iree.compiler.tools.import_onnx path/to/model.onnx -o path/to/model.mlir --data-prop

    Tip

    Use the -h flag to view all available importer options.

    $ python -m iree.compiler.tools.import_onnx -h

The output MLIR file can be compiled with Torq by specifying onnx-torq as the input dialect.

References