OpenXLA IREE 开发者实用技巧指南

OpenXLA IREE 开发者实用技巧指南

iree A retargetable MLIR-based machine learning compiler and runtime toolkit. 项目地址: https://gitcode.com/gh_mirrors/ir/iree

前言

OpenXLA IREE 是一个基于 MLIR 的端到端机器学习编译器和运行时工具链。本文将介绍一些实用的开发者技巧，帮助开发者更好地使用和调试 IREE 编译器。

编译器选项设置

IREE 编译器工具（如 iree-compile）通过命令行参数接收选项。使用 --help 可以查看所有可用选项：

$ iree-compile --help

在 Python 绑定中，可以通过 extra_args 参数传递选项：

import iree.compiler as ireec

compiled_module = ireec.tools.compile_str(
    input_mlir,
    target_backends=["llvm-cpu"],
    extra_args=["--mlir-timing"])

分析 .vmfb 文件

IREE 编译器生成的 .vmfb 文件（Virtual Machine FlatBuffer）是 IREE 运行时可加载执行的二进制格式。

文件结构分析

.vmfb 文件默认可以作为 zip 文件解压：

$ unzip -d simple_abs_cpu ./simple_abs_cpu.vmfb

解压后可以看到包含：

• 模块元数据文件（module.fb）
• 目标平台特定的二进制文件（如 ELF 共享对象）

使用专用工具分析

iree-dump-module 工具可以显示 .vmfb 文件的详细信息：

$ iree-dump-module simple_abs.vmfb

输出内容包括：

• 模块版本信息
• 依赖的类型和函数
• 导入/导出的函数列表

转储可执行文件

IREE 提供了一系列标志来保存编译过程中的中间文件，这对理解编译过程和调试非常有帮助。

常用转储标志

| 标志 | 功能 | |------|------| | --iree-hal-dump-executable-files-to | 保存所有相关文件 | | --iree-hal-dump-executable-sources-to | 保存 HAL 编译前的源文件 | | --iree-hal-dump-executable-intermediates-to | 保存中间文件（.o, .mlir 等） | | --iree-hal-dump-executable-binaries-to | 保存二进制文件（.so, .spv, .ptx） |

CPU 目标示例

$ mkdir -p /tmp/iree/simple_abs/
$ iree-compile simple_abs.mlir \
  --iree-hal-target-device=local \
  --iree-hal-local-target-device-backends=llvm-cpu \
  --iree-llvmcpu-link-embedded=false \
  --iree-hal-dump-executable-files-to=/tmp/iree/simple_abs \
  -o /tmp/iree/simple_abs/simple_abs_cpu.vmfb

这将生成包括 LLVM IR、目标代码等在内的完整编译流水线文件。

GPU 目标示例

对于 Vulkan 目标：

$ iree-compile simple_abs.mlir \
  --iree-hal-target-device=vulkan \
  --iree-hal-dump-executable-files-to=/tmp/iree/simple_abs \
  -o /tmp/iree/simple_abs/simple_abs_vulkan.vmfb

对于 CUDA 目标：

$ iree-compile simple_abs.mlir \
  --iree-hal-target-device=cuda \
  --iree-hal-dump-executable-files-to=/tmp/iree/simple_abs \
  -o /tmp/iree/simple_abs/simple_abs_cuda.vmfb

性能基准测试

IREE 提供了两种级别的基准测试方式：

模块级基准测试

测试完整 IREE 运行时环境下的性能：

$ iree-compile simple_abs.mlir \
  --iree-hal-target-device=local \
  --iree-hal-local-target-device-backends=llvm-cpu \
  --iree-hal-dump-executable-benchmarks-to=/tmp/iree/simple_abs/ \
  -o /dev/null

$ iree-benchmark-module \
  /tmp/iree/simple_abs/module_abs_dispatch_0_benchmark.vmfb

底层二进制基准测试

直接测试生成的二进制文件，绕过 VM 和 HAL 层：

$ iree-benchmark-executable \
  --device=local-sync \
  --executable_format=embedded-elf-x86_64 \
  --executable_file=/tmp/iree/simple_abs/module_abs_dispatch_0_embedded_elf_x86_64.so \
  --entry_point=0 \
  --binding=f32=-2.5 \
  --binding=f32=0 \
  --workgroup_count=1,1,1

分阶段编译

IREE 的编译过程分为多个阶段，开发者可以指定从某个阶段开始或结束编译：

graph LR
  A[Input] --> B[ABI]
  B --> C[Flow]
  C --> D[Stream]
  D --> E[HAL]
  E --> F[VM]

使用 --compile-to 和 --compile-from 参数可以控制编译阶段：

$ iree-compile --compile-to=flow input.mlir -o flow.mlir
$ iree-compile --compile-from=flow flow.mlir -o output.vmfb

结语

本文介绍了 OpenXLA IREE 开发中的一些实用技巧，包括编译器选项设置、文件分析、中间文件转储和基准测试等。掌握这些技巧可以帮助开发者更高效地使用 IREE 进行机器学习模型的编译和优化。

iree A retargetable MLIR-based machine learning compiler and runtime toolkit. 项目地址: https://gitcode.com/gh_mirrors/ir/iree