Torch-MLIR 项目
Torch-MLIR 项目旨在为 PyTorch 生态系统向 MLIR 生态系统提供一流的编译器支持。
该项目参与了 LLVM 孵化器流程:因此,它不属于任何官方 LLVM 版本。虽然孵化状态并不一定反映代码的完整性或稳定性,但它确实表明该项目尚未被认可为 LLVM 的组件。
PyTorch PyTorch 是一个开源的机器学习框架,能够实现从研究和原型设计到生产级部署的无缝过渡。
MLIR MLIR 项目提供了一种构建可扩展和可重用编译器架构的新方法,这些方法可以解决软件碎片化问题,降低开发特定领域编译器的成本,改进异构硬件的编译,并促进现有编译器之间的兼容性。
Torch-MLIR 多个供应商已将 MLIR 作为其系统中的中间层,从而使他们能够将 PyTorch、JAX 和 TensorFlow 等框架映射到 MLIR,并随后降低到其目标硬件。我们已经观察到半打从 PyTorch 到 MLIR 的自定义下层处理,这使得硬件供应商可以专注于其独特价值,而不必再为 MLIR 实现另一个 PyTorch 前端。最终目标是类似于当前硬件供应商添加 LLVM 目标支持,而不是每个供应商都实现 Clang 或 C++ 前端。
从 PyTorch 到 Torch MLIR 方言的所有路径
我们有几条路径可以降到 Torch MLIR 方言。
- TorchScript 这是降到 Torch MLIR 方言中测试最多的路径。
- LazyTensorCore 阅读更多详细信息 这里。
- 我们还具有基本的 TorchDynamo/PyTorch 2.0 支持,请参阅我们的 长期路线图 和 关于 PyTorch 2.0 的思考 了解更多详细信息。
项目沟通
- LLVM Discord 上的
#torch-mlir
频道 - 这是最活跃的沟通渠道 - Github 问题 这里
- LLVM Discourse 上的
torch-mlir
部分
会议
社区会议 / 开发者时间:
- 每月的第一个和第三个星期一上午9点 PST
- 每月的第二个和第四个星期一下午5点 PST
办公时间:
- 每周四上午8:30 PST
会议链接请见 这里。
安装 torch-mlir 快照
在撰写本文时,我们发布了 torch-mlir 的预构建快照,用于 Python 3.11 和 Python 3.10。
如果您有支持的 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
pip install torch-mlir -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
演示
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)]
TorchScript ResNet18
将 PyTorch ResNet18 模型转换为 MLIR 并在 CPU 后端运行的独立脚本:
# 如果您尚未检出代码,请获取最新示例
wget https://raw.githubusercontent.com/llvm/torch-mlir/main/projects/pt1/examples/torchscript_resnet18.py
# 作为独立脚本运行 ResNet18。
python projects/pt1/examples/torchscript_resnet18.py
从 https://upload.wikimedia.org/wikipedia/commons/2/26/YellowLabradorLooking_new.jpg 加载图片
正在下载: "https://download.pytorch.org/models/resnet18-f37072fd.pth" 到 /home/mlir/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth
100.0%
PyTorch 预测
[('拉布拉多猎犬', 70.66319274902344), ('金毛犬', 4.956596374511719), ('切萨皮克湾猎犬', 4.195662975311279)]
torch-mlir 预测
[('拉布拉多猎犬', 70.66320037841797), ('金毛犬', 4.956601619720459), ('切萨皮克湾猎犬', 4.195651531219482)]
Lazy Tensor Core
查看示例 这里。
仓库布局
该项目遵循典型的基于 MLIR 项目的惯例:
include/torch-mlir
和lib
结构用于 C++ MLIR 编译器方言/传递。test
目录用于存放测试代码。tools
目录用于torch-mlir-opt
等工具。python
顶级目录用于存放 Python 代码
开发者
如果您想从源码开发和构建 torch-mlir,请查看 开发说明。