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

一、引言

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 pile

下图是pytorch 官方介绍 PyTorch 2.0 版本提供的pytorch 2.x 之后的编译过程，即torch pile。 图的原文 可以发现torch pile主要实现两个过程：前端和后端 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 pile 和 torch_mlir.fx.export_and_import。

第一条路径是旧项目 pt1 代码的一部分 （torch_mlir.torchscript pile），允许用户测试编译器的 输出到不同的 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 piler_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)