解码 IREE：Tensorflow 和 TFlite 模型迈向 MLIR 前端 IR 的转换秘籍

IREE

本文计划着重讲一下IREE是如何实现tensorflow模型转化为MLIR方言的具体实现方法

前言

简单对比一下 MLIR 和 TVM

最近在研究模型推理这块，比较有名的就是TVM和MLIR这二个框架。

从定义来看，MLIR的定位更像是一个工程脚手架，构建一些列帮助开发者，构建完整编译器的工具。

TVM则更像是一个完整的编译器，并且TVM的发展历史要比MLIR更加深厚。

我是如何选择的：
本质上二者是有重合的地方，同时又有各自专注的方向。
在学习的时候，我其实更倾向于MLIR，因为ta前期在Tensorflow项目下，现在又在LLVM项目下挂载。 自带光环，并且是脚手架性质，有更大的自由度，但是从实现上来看，我无力一个人独自实现编译器，所以还是需要借助开源社区的一些优秀项目，来了解如何使用MLIR来实现一个端到端的AI编译器。

经过一番搜索和了解，我发现IREE项目是比较符合我的预期的，最近就准备研究一下IREE项目，其实我已经看了比较长的一段时间，最近打算把思路整理一下，汇总到博客上，与大家分享一下，如果有感兴趣的同学，希望可以一起讨论，共同进步

1. IREE的架构

从IREE的架构图上来看，首先第一步则是将原生模型，转化为MLIR的方言，这一部分功能称之为AI编译器的前端。就是我红色标注的那一部分。

因为一些原因，接触的tflite的模型比较多，所以今天就以tensorflow为基础，理一下IREE是如何将模型转为前端IR（MLIR）的。

2. IREE前端的实现

从IREE的文档里，我们可以看到。 ta们是通过命令行，来将tf模型转为MLIR前端IR的。

TFLITE_PATH=${WORKDIR}/model.tflite
IMPORT_PATH=${WORKDIR}/model.mlir
iree-import-tflite ${TFLITE_PATH} -o ${IMPORT_PATH}

具体的实现我们可以在IREE的代码里找到：integrations/tensorflow/python_projects/iree_tf
就是这个目录。在上游代码的 iree-tf 路径。

2.1 tflite的转MLIR的实现

我们来一起看一起TFlite模型转为mlir的代码：tflite转为MLIR方言

可以看到代码实现使用 python 来实现的，并且主体代码非常简单：

    parser.add_argument(
        "--output-format",
        dest="output_format",
        required=False,
        default="mlir-bytecode",
        help=argparse.SUPPRESS,
    )
    args = parser.parse_args()

    if args.output_format != "mlir-bytecode":
        logging.warning("output-format option is deprecated, emitting MLIR bytecode")

    tflite_to_tosa(
        flatbuffer=args.flatbuffer,
        bytecode=args.output_path,
        ordered_input_arrays=args.input_arrays,
        ordered_output_arrays=args.output_arrays,
    )
def tflite_to_tosa(
    flatbuffer,
    bytecode,
    use_external_constant=False,
    ordered_input_arrays=None,
    ordered_output_arrays=None,
):
    tflite_to_tosa_bytecode(
        flatbuffer,
        bytecode,
        use_external_constant,
        ordered_input_arrays,
        ordered_output_arrays,
    )

主要使用的 API 是 tflite_to_tosa_bytecode， 是直接从 tensorflow 代码库引进的API，也就是说，实际上这一步是 tensorflow 实现。 IREE 是借助 tensorflow 的工具，完成将 tflite 模型转为 mlir 模型的。

from tensorflow.mlir.experimental import tflite_to_tosa_bytecode

2.2 tensorflow模型

tensorflow模型转化为 MLIR的代码在 tensorflow模型转MLIR 。

我们依旧可以看到，整体代码比较简单，核心代码是

        convert_saved_model = pywrap_mlir.experimental_convert_saved_model_to_mlir
        convert_saved_model_v1 = pywrap_mlir.experimental_convert_saved_model_v1_to_mlir
        run_pass_pipeline = pywrap_mlir.experimental_run_pass_pipeline
        write_bytecode = pywrap_mlir.experimental_write_bytecode

    if import_type == "savedmodel_v2":
        result = convert_saved_model(
            saved_model_dir, exported_names=exported_names, show_debug_info=False
        )

其中核心的工作是 通过导入 tensorflow的API来完成的，和tflite的实现如出一辙。

from tensorflow.python import pywrap_mlir

        convert_saved_model = pywrap_mlir.experimental_convert_saved_model_to_mlir

2.3 总结

IREE的前端， 从模型转为MLIR方言，是通过 Tensorlfow 的官方支持完成的，这一部分工作 IREE 并没有自己做。

从tensorflow的 API文档中 我们也可以找到这一部分描述：将存储的模型转为MLIR方言。

但是同时我们需要注意到，这一部分功能的API目前是实验性质。

3. 实验样例

我这里已经编译好IREE的代码，

1. 下载模型
2. 模型编译为MLIR 的方言
3. MLRI方言作为IREE的输入，编译可执行文件 vmfb
4. 执行 vmfb 文件进行推理

WORKDIR="./"
TFLITE_URL="https://storage.googleapis.com/iree-model-artifacts/tflite-integration-tests/posenet_i8.tflite"
TFLITE_PATH=${WORKDIR}/model.tflite
IMPORT_PATH=${WORKDIR}/model.mlir
MODULE_PATH=${WORKDIR}/module.vmfb

# Fetch the sample model
wget ${TFLITE_URL} -O ${TFLITE_PATH}

# Import the sample model to an IREE compatible form
iree-import-tflite ${TFLITE_PATH} -o ${IMPORT_PATH}

# Compile for the CPU backend
iree-compile \
    --iree-hal-target-backends=llvm-cpu \
    ${IMPORT_PATH} \
    -o ${MODULE_PATH}

# excute
iree-run-module --device=local-task --module=module.vmfb --input="1x192x192x3xi8=0"

可以看到最后的执行结果。
ps: 1x192x192x3xi8=0" 代表的是这个尺寸的全0 tensor。