Partial Conversion

Partial Conversion

在 MLIR 的异构多级降级(Lowering)管线中,将整个模块(Module)从一个高层方言无脑、一刀切地全部翻译到底层方言往往是不切实际的。

例如,当我们试图把自研的 NPU 高层算子图逐步向硬件靠拢时,我们可能只想把矩阵乘法和卷积降级到硬件厂商的私有方言或LLVM裸指针形态,而图中的某些复杂控制流(如scf.if)或标量加减法(如arith.addi),我们希望保留它们,让 MLIR 官方的通用优化 Pass 继续去处理。

为了在一条流水线内实现“精准动手术、局部降级、新老方言共存”的控制,MLIR 引入了极其硬核且灵活的编译机制——Partial Conversion(局部转换 / 部分转换)

本文将带你深度剖析 Partial Conversion 的工作机理、如何编写一个局部降级通道,以及它在异构 AI 基础设施开发中的关键价值。


1. 什么是 Partial Conversion?三大转换模式对比

MLIR 官方的方言转换框架(Dialect Conversion Framework)提供了三种截然不同的转译战略。理解它们的边界是掌握 Partial Conversion 的第一步:

转换模式核心行为匹配失败后果适用场景
Full Conversion


(完全转换) | 图中所有属于源方言的算子必须百分之百被消灭、翻译干净。 | 只要有一个算子没被完全消除,编译直接报错崩溃。 | 编译管线的最后一公里(如最终生成纯 LLVM IR 或 汇编前)。 |
|Partial Conversion


(局部转换) | 允许未知、无法匹配或故意保留的算子继续合法地留在图里。 | 匹配失败的算子会被优雅跳过,新老方言在图里和平共存。 |中层异构降级,逐步剥离自研算子,保留官方通用控制流。 |
|Analysis Conversion


(分析转换) | 只执行模式匹配和类型检查,不真正动手修改 IR 图结构。 | 不会改变图,只收集元数据。 | 编译期的合规性静态分析。 |

Partial Conversion 的核心灵魂:
它是一种渐进式、外科手术式的优化手段。它允许编译器专家只去啃那些自己关心的硬骨头算子,而把通用的、复杂的控制流和标量计算“战略性留白”,留在随后的通用通道里解决。


2. 文本 IR 视角:Partial Conversion 后的“杂交”生态

为了让你产生最直观的工程体感,我们来看看当一个包含复杂混合算子的函数,经历了一次针对MemRef访存的Partial Conversion Pass后,IR 内部发生的长相剧变。

Conversion 前:统一、高层的中间方言

此时全图算子语义都处于相对高层的抽象形态:

func.func @heterogeneous_graph(%arg0: memref<128xf32>, %cond: i1) { // 1. 纯标量控制流 cf.cond_br %cond, ^bb1, ^bb2 ^bb1: // 2. 复杂的、需要降级寻址的物理访存 %idx = arith.constant 0 : index %val = memref.load %arg0[%idx] : memref<128xf32> cf.br ^bb2 ^bb2: return }

Conversion 后:底层指针寻址与高层标量控制流的完美并存

在执行了Partial Conversion之后,原本高大上的memref.load被我们精准定向爆破,降级成了纯粹的LLVM裸指针和寻址公式;但与此同时,外层的cf.cond_br分支控制流因为不属于本次优化的准星范围,被完美且合法地保留了下来

func.func @heterogeneous_graph(%arg0: !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, %cond: i1) { // 高层的控制流完好无损地保留在图里! cf.cond_br %cond, ^bb1, ^bb2 ^bb1: // 核心:唯独这一块物理访存被外科手术式地翻译成了底层的裸指针结构 %base_ptr = llvm.extractvalue %arg0[1] : !llvm.struct<...> %zero = llvm.mlir.constant(0 : i64) %pixel_ptr = llvm.getelementptr %base_ptr[%zero] : (ptr, i64) -> ptr, f32 %val = llvm.load %pixel_ptr : ptr -> f32 cf.br ^bb2 ^bb2: return }

3. C++ 源码视角:手写一个 Partial Conversion Pass

在开发自研 AI 芯片后端、或者编写异构图切分通道时,工程师必须掌握如何利用 C++ 的ConversionTarget严格划定局部转换的势力范围。

下面展示了一个标准的工业级 Partial Conversion 通道的 C++ 写法范式:

#include"mlir/Transforms/DialectConversion.h"#include"mlir/Dialect/LLVMIR/LLVMDialect.h"#include"mlir/Dialect/MemRef/IR/MemRef.h"#include"mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"usingnamespacemlir;// 1. 继承官方的 ConversionPass 基类structPartialLoweringPass:publicPassWrapper<PartialLoweringPass,OperationPass<ModuleOp>>{voidrunOnOperation()override{ModuleOpmodule=getOperation();MLIRContext*context=&getContext();// 2. 实例化转换目标 (ConversionTarget) —— 划定法律边界ConversionTargettarget(*context);// 显式声明:整个 LLVM 方言在我们这个 Pass 执行完后是完全合法的target.addLegalDialect<LLVM::LLVMDialect>();// 显式声明:高层的控制流方言(ControlFlow)我们不想动它,也是完全合法的target.addLegalDialect<cf::ControlFlowDialect>();// 强力宣告:MemRef 方言中的所有算子,在此处通通被列为“非法(Illegal)”!// 编译器会像警犬一样在图里搜寻所有 MemRef 算子并强制要求其降级target.addIllegalDialect<memref::MemRefDialect>();// 3. 收拢并挂载转换模式 (Rewrite Patterns)RewritePatternSetpatterns(context);// 假设我们已经纯手写实现了一套 CustomMemRefToLLVMLowering 的 C++ 类populateMemRefToLLVMConversionPatterns(patterns);// 4. 惊艳的一击:启动局部转换引擎 (Partial Conversion)// 注意:如果是完全转换,这里会呼叫 applyFullConversionif(failed(applyPartialConversion(module,target,std::move(patterns)))){signalPassFailure();// 如果降级过程中匹配失败或内部报错,才宣布通道失败}}};

4. 大模型时代 Partial Conversion 的核心硬核价值

在动辄百亿、千亿参数的大模型(LLM)分布式Serving和推理管线中,Partial Conversion 是规避全图编译爆炸、实现“混合精度与异构硬件调度”的免死金牌:

① 混合泛型算子与专用硬件 Kernel 的无缝插拔

在大模型(如 Transformer 架构)中,矩阵乘法(MatMul)通常占了 90% 以上的计算量,各大厂商都会用手写的、极致压榨硬件的专用昂贵方言(如my_npu.custom_flash_attn)去承接。
但在矩阵乘法外面,还有一堆琐碎的、动态变化的自定义标量和控制流(如打印日志、动态的 Token 边界裁剪)。如果采用传统的全转换,你必须为这些琐碎算子也手写极其繁琐的硬件降级。
利用Partial Conversion,编译器可以只把核心的计算大块定向爆破到自研芯片后端,而把那些长尾的标量计算原封不动地踢给 CPU 或官方标准的LLVM运行时,极大地节省了底层的适配开发成本。

② 类型解耦与逐步降级(Type Converter)

在进行指针级降级时,高层的多维张量要变成低层的一维指针结构体,这中间涉及到输入参数和基本块(Block Arguments)的类型置换。
Partial Conversion 允许你挂载一个TypeConverter。当它在图的中途遇到跨方言连接边界(例如一个还没来得及降级的cf.br试图读取一个已经降级为LLVM ptr的变量时),Partial Conversion 引擎会在图内部隐式发射一个桥接算子(builtin.unrealized_conversion_cast)作为临时的“数据粘合带”。这样,哪怕整张图处于“半生不熟”的混交状态,MLIR 强大的验证器(Verifier)也不会报错崩溃,允许后续的 Pass 链条继续平滑推进。

总结

一句话概括:Partial Conversion 是 AI 编译器在面对纷繁复杂的异构生态时,能够实施精确制导、局部手术、新老方言完美共存的高级平衡艺术。

它打破了传统编译器非黑即白、一刀切的死板教条,用灵活的ConversionTarget在同一个计算图体内划出了清晰的合法与非法势力范围。深刻掌握 Partial Conversion 这一元编程控制手段,是每一个致力于大模型混合精度架构设计、异构多核算力芯片(CPU+GPU+NPU)协同调度的 AI 基础设施工程师(AI Infra)在复杂工业级图变换中挥洒自如的进阶利器。