Torch-MLIR 项目

Torch-MLIR 项目旨在为 MLIR 生态系统提供来自 PyTorch 生态的第一流编译器支持。

本项目正在参与 LLVM 孵化器流程：因此，它不属于任何官方的 LLVM 发布版本。虽然孵化状态并不一定反映代码的完整性和稳定性，但它确实表明该项目尚未被认可为 LLVM 的一部分。

PyTorch PyTorch 是一个开源机器学习框架，能够无缝地从研究和原型设计过渡到生产环境中的部署。

MLIR MLIR 项目提供了一种构建可扩展且可重用编译器架构的新方法，解决了软件碎片化问题，降低了开发领域特定编译器的成本，提升了针对异构硬件的编译效率，并促进了现有编译器之间的兼容性。

Torch-MLIR 多家厂商已将 MLIR 作为其系统中的中间层，从而能够将 PyTorch、JAX 和 TensorFlow 等框架映射到 MLIR，再进一步下放到目标硬件上。我们已经观察到大约六个从 PyTorch 到 MLIR 的自定义下放实现，这使得硬件厂商可以更加专注于自身的核心价值，而无需为 MLIR 再次实现一个新的 PyTorch 前端。最终目标是类似于当前硬件厂商为 LLVM 添加目标支持的方式，而不是每个厂商都去实现 Clang 或 C++ 前端。

从 PyTorch 到 Torch MLIR Dialect 的所有路径

我们将 PyTorch 下降到 Torch MLIR Dialect 的路径主要有两条：

• 以 ONNX 作为入口点。
• 以 Fx 作为入口点。

项目沟通

• LLVM Discord 上的 #torch-mlir 频道——这是最活跃的沟通渠道。
• GitHub 问题 这里。
• LLVM Discourse 中的 torch-mlir 版块。

安装 torch-mlir 快照

截至撰写本文时，我们发布了适用于 Python 3.11 和 Python 3.10 的 预构建 torch-mlir 快照。

如果您使用的是受支持的 Python 版本，可以通过以下命令初始化一个虚拟环境：

python3.11 -m venv mlir_venv
source mlir_venv/bin/activate

或者，如果您希望使用 conda 在多个 Python 版本之间切换，可以创建一个包含 Python 3.11 的 conda 环境：

conda create -n torch-mlir python=3.11
conda activate torch-mlir
python -m pip install --upgrade pip

然后，我们可以安装对应版本的 torch 和 torchvision 夜间构建版的 torch-mlir：

pip install --pre torch-mlir torchvision \
  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels

使用 torch-mlir

Torch-MLIR 主要是一个集成到编译器中，用于连接 PyTorch 和 ONNX 的项目。如果您正在考虑进行新的集成，可以参考现有的下游项目：

• IREE
• Blade

尽管该项目的大部分功能都是通过测试路径来验证的，但终端用户也可以在不进行进一步集成的情况下直接使用一些 API：

FxImporter ResNet18

# 如果您尚未克隆代码，请获取最新示例
wget https://raw.githubusercontent.com/llvm/torch-mlir/main/projects/pt1/examples/fximporter_resnet18.py

# 以独立脚本形式运行 ResNet18。
python projects/pt1/examples/fximporter_resnet18.py

# 输出
从 https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg 加载图像
...
PyTorch 预测
[('拉布拉多寻回犬', 70.65674591064453), ('金毛寻回犬', 4.988346099853516), ('萨路基犬, 瞪羚猎犬', 4.477451324462891)]
torch-mlir 预测
[('拉布拉多寻回犬', 70.6567153930664), ('金毛寻回犬', 4.988325119018555), ('萨路基犬, 瞪羚猎犬', 4.477458477020264)]

仓库布局

该项目遵循典型的基于 MLIR 的项目的惯例：

• include/torch-mlir 和 lib 目录用于存放 C++ MLIR 编译器方言和优化 pass。
• test 目录用于存放测试代码。
• tools 目录用于存放 torch-mlir-opt 等工具。
• python 顶级目录用于存放 Python 代码。

开发者

如果您希望从源码开发并构建 torch-mlir，请参阅 开发说明。

torch-mlir 快速上手指南

torch-mlir 旨在为 PyTorch 生态系统提供一流的编译器支持，将其无缝对接到 MLIR 生态系统中。它允许将 PyTorch 模型转换为 MLIR 方言，以便下游编译器（如 IREE、BladeDISC）进一步优化并部署到各种硬件后端。

环境准备

• 操作系统：Linux (推荐), macOS, Windows (支持程度视具体版本而定)
• Python 版本：推荐使用 Python 3.10 或 Python 3.11
• 前置依赖：
  • pip (包管理工具)
  • venv 或 conda (用于创建隔离的虚拟环境)
  • 网络连接 (需访问 GitHub 和 PyTorch 官方源)

注意：目前项目处于 LLVM 孵化阶段，建议使用预构建的快照版本进行体验，无需从源码编译。

安装步骤

以下提供两种常见的环境配置方式，任选其一即可。

方式一：使用 venv (推荐)

1. 创建并激活虚拟环境：

   python3.11 -m venv mlir_venv
   source mlir_venv/bin/activate
   

2. 安装 torch-mlir 及对应的 PyTorch 夜间构建版：

   pip install --pre torch-mlir torchvision \
     --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
     -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
   

方式二：使用 Conda

1. 创建并激活 Conda 环境：

   conda create -n torch-mlir python=3.11
   conda activate torch-mlir
   python -m pip install --upgrade pip
   

2. 执行相同的安装命令：

   pip install --pre torch-mlir torchvision \
     --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
     -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
   

国内加速提示：若下载速度较慢，可尝试将 https://download.pytorch.org/whl/nightly/cpu 替换为国内镜像源（如清华源或阿里源中对应的 PyTorch 夜间版地址），但需确保索引 URL 与 -f 指定的 wheel 资产地址兼容性。若不确定，建议先使用官方源测试连通性。

基本使用

torch-mlir 最直接的用法是通过 FxImporter 将 PyTorch 模型导入为 MLIR 格式。以下以 ResNet18 为例演示完整流程。

1. 获取示例脚本：

   wget https://raw.githubusercontent.com/llvm/torch-mlir/main/projects/pt1/examples/fximporter_resnet18.py
   

2. 运行示例： 确保在当前目录下执行，该脚本会自动下载图片、加载模型、通过 PyTorch 推理，并将其转换为 torch-mlir 后进行推理对比。

   python fximporter_resnet18.py
   

   (注：如果未克隆仓库，文件位于当前目录；若已克隆仓库，路径应为 python projects/pt1/examples/fximporter_resnet18.py)

3. 预期输出： 脚本将输出 PyTorch 原生预测结果与 torch-mlir 转换后的预测结果，两者应高度一致。

   load image from https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg
   ...
   PyTorch prediction
   [('Labrador retriever', 70.65674591064453), ('golden retriever', 4.988346099853516), ('Saluki, gazelle hound', 4.477451324462891)]
   torch-mlir prediction
   [('Labrador retriever', 70.6567153930664), ('golden retriever', 4.988325119018555), ('Saluki, gazelle hound', 4.477458477020264)]
   

此示例展示了从 PyTorch FX 图到 Torch MLIR Dialect 的基本转换流程。对于更复杂的集成（如对接 IREE 或 BladeDISC），请参考相应下游项目的文档。