当前位置：首页 > news >正文

介绍 torch-mlir 从 pytorch 生态到 mlir 生态

news 来源：原创 2025/6/9 6:46:08

一、引言

The Torch-MLIR project provides core infrastructure for bridging the PyTorch ecosystem and the MLIR ecosystem. For example, Torch-MLIR enables PyTorch models to be lowered to a few different MLIR dialects. Torch-MLIR does not attempt to provide a production end-to-end flow for PyTorch programs by itself, but is a useful component for constructing one.

这是torch-mlir官方的介绍。总结下来就是 torch-mlir 提供了一个方案，用来连接 pytorch 生态和 mlir 生态。

二、torch.compile

下图是pytorch 官方介绍 PyTorch 2.0 版本提供的pytorch 2.x 之后的编译过程，即torch.compile。
图的原文
在这里插入图片描述
可以发现torch.compile主要实现两个过程：前端和后端
1、前端获取计算图输出中间表示，供后端生成二进制可执行文件。主要通过 TorchDynamo+AOT Autograd 输出中间表示 FX Graph。
其中TorchDynamo处理后，得到的是可以带控制逻辑的高层计算图，再经过AOT Autograd 处理就得到了一种通用的aten ir 中间表示，是不含控制逻辑底层计算图，aten ir 可以等价于 pytorch 中定义的prim ops + aten ops。
2、后端翻译生成二进制文件。pytorch 提供的一个工具是 TorchInductor，这里跳过，不是本文讨论的重点。

而 torch-mlir 主要实现的部分是紧接着 aten ir，可以将 aten ir 继续下降为 mlir，即接入 MLIR生态。

三、Torch-MLIR

torch-mlir 的实现也可以分为前端和后端。前后端的中间交互是 torch-mlir 自己定义的一个 dialect ：Torch dialect。而 backend contract是Torch dialect的子集，既可以向上兼容 aten ir，又可以向下对接mlir 生态中的目标 dialect。

torch-mlir 的前端与pytorch 本身的接口有关，torch-mlir 封装 pytorch 的 API 为 backend contract。即也是将python program 下降为 aten ir。
torch-mlir 的后端是将 backend contract 下降为 mlir 生态中的目标 dialect。如 Linalg、TOSA、MHLO等。
下图为torch-mlir的官方提供的架构图

目前torch-mlir 项目有两个主要的 API ：torch_mlir.torchscript.compile 和 torch_mlir.fx.export_and_import。
第一条路径是旧项目 pt1 代码的一部分（torch_mlir.torchscript.compile），允许用户测试编译器的输出到不同的 MLIR 方言
第二条路径 (torch_mlir.fx.export_and_import)允许用户导入任意 Python 可调用对象（nn.Module、函数或方法）的合并 torch.export.ExportedProgram 实例，并输出到 torch dialect mlir 模块。该路径与 PyTorch 的路线图一致。

from torch_mlir import fx
from torch_mlir.compiler_utils import run_pipeline_with_repro_report

class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = torch.nn.Linear(10, 10)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        return x

# front
aten_ir = fx.export_and_import(Basic(), torch.randn(3, 4))
print(aten_ir)

# backend to linalg
run_pipeline_with_repro_report(
   aten_ir,
   (
       "builtin.module("
       "func.func(torch-decompose-complex-ops),"
       "torch-backend-to-linalg-on-tensors-backend-pipeline)"
   ),
   "Lowering TorchFX IR -> Linalg IR",
   enable_ir_printing=False,
)

# linalg dialect
print(aten_ir)