【架构分析】Torch-MLIR 技术详解

目录

概述

架构

举个栗子

IR转换详解

OpConversion代码分析

核心时序

Linalg Generic Op说明

Code Gen VS LibCall 


概述

Torch-MLIR 作为一个编译器支持将pytorch生态转换到MLIR生态,本文基于Torch-MLIR 开源项目 对它的主要工作流程,以及涉及到的MLIR主要概念进行了分析说明

架构

【架构分析】Torch-MLIR 技术详解_第1张图片

上面的架构图说明从Pytorch生态有多条pipeline可以转换到MLIR的世界,本文主要介绍

TorchScript -> MLIR Converter -> Torch Dialect -> Linalg Tensor 这条pipeline

举个栗子

torch-mlir 工作流程

动画展示了一个Pytorch模型 -> TorchScript IR 经过Torch-MLIR Importer后进入到MLIR世界的详细过程

IR转换详解

【架构分析】Torch-MLIR 技术详解_第2张图片【架构分析】Torch-MLIR 技术详解_第3张图片【架构分析】Torch-MLIR 技术详解_第4张图片【架构分析】Torch-MLIR 技术详解_第5张图片【架构分析】Torch-MLIR 技术详解_第6张图片

上图说明了Pytorch模型代码 ->TorchScript IR -> TorchDialect IR 的转换对应关系

OpConversion代码分析

【架构分析】Torch-MLIR 技术详解_第7张图片【架构分析】Torch-MLIR 技术详解_第8张图片【架构分析】Torch-MLIR 技术详解_第9张图片【架构分析】Torch-MLIR 技术详解_第10张图片【架构分析】Torch-MLIR 技术详解_第11张图片【架构分析】Torch-MLIR 技术详解_第12张图片【架构分析】Torch-MLIR 技术详解_第13张图片【架构分析】Torch-MLIR 技术详解_第14张图片

上图以OpConversionPattern代码为例,说明了TorchDialect IR 中的torch.aten.mm Op 转换到LinalgDialect Op的过程

核心时序

 Torch-MLIR从Python应用、初始化、MLIR Converion Framework注册与回调、Dialect & Op Convert Pass 的核心时序图如上

Linalg Generic Op说明

理解Linalg Dialect的Op才能对Torch Dialect Op做对应的转换,而Linalg Generic Op又是最重要的一个Op,用一个详细的例子对它说明如下:

【架构分析】Torch-MLIR 技术详解_第15张图片【架构分析】Torch-MLIR 技术详解_第16张图片【架构分析】Torch-MLIR 技术详解_第17张图片【架构分析】Torch-MLIR 技术详解_第18张图片【架构分析】Torch-MLIR 技术详解_第19张图片【架构分析】Torch-MLIR 技术详解_第20张图片【架构分析】Torch-MLIR 技术详解_第21张图片【架构分析】Torch-MLIR 技术详解_第22张图片

Code Gen VS LibCall 

【架构分析】Torch-MLIR 技术详解_第23张图片

Code Gen:用一堆IR生成另外一堆表达同样计算语义的IR

LibCall: 用一个API调用时序一堆IR的计算,比如用cublasSgemm 计算矩阵乘法

你可能感兴趣的:(深度学习,人工智能,编译器)