torch-mlir

torch-mlir 项目介绍

torch-mlir 项目旨在为从 PyTorch 生态系统到 MLIR 生态系统提供一流的编译器支持。

项目背景

PyTorch 是一个开源的机器学习框架，帮助从研究到产品级部署实现无缝过渡。MLIR 项目则提供了一种新颖的方法来构建可扩展和可重用的编译器架构，以解决软件碎片化问题，降低开发领域专用编译器的成本，改善对异构硬件的编译，并促进现有编译器之间的兼容性。

torch-mlir 项目正在参与 LLVM 孵化器计划，因此尚未成为任何官方 LLVM 版本的一部分。尽管孵化状态并不一定反映代码的完整性或稳定性，但这仍表明该项目尚未被官方认可为 LLVM 的组成部分。

项目的目标

torch-mlir 项目力图通过将 PyTorch 等框架映射到 MLIR，并进一步降低到目标硬件，实现类似于当前硬件供应商添加 LLVM 目标支持的效果。这使得硬件供应商可以更加专注于自己的独特优势，而无需再为 MLIR 创建新的 PyTorch 前端。

通信及社区

torch-mlir 项目主要通过以下渠道进行沟通：

LLVM Discord 社区的 #torch-mlir 频道，这是最活跃的交流渠道。
Github 的问题区。
LLVM Discourse 中的 torch-mlir 版块。

安装 torch-mlir 快照版本

目前可为 Python 3.11 和 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 夜版：

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

# 预测输出示例
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)]