AI编译器的中间表示设计：从MLIR的Dialect生态看编译器IR的可扩展性
AI编译器的中间表示设计从MLIR的Dialect生态看编译器IR的可扩展性一、当自定义算子需要同时支持5种硬件后端传统编译器的IR扩展困境设计一个支持多硬件CUDA、ROCm、Metal、Vulkan、CPU的AI推理编译器时最大的困难不是生成各硬件的代码而是在IR层面表达不同硬件的算子语义。LLVM IR虽然通用但不理解矩阵乘法、卷积、Attention等AI特定的操作——这些操作在LLVM中被展开为基础运算丧失了优化机会。更麻烦的是每个硬件有独特的优化Pass。CUDA的Tensor Core优化与Metal的SIMD group优化完全不同。在单一IR上叠加所有Pass会导致不可维护的复杂度。MLIR通过Dialect机制解决这个问题——每种抽象层次有独立的Dialect它们可以在同一IR中共存、互操作。二、MLIR的Dialect分层架构graph TB subgraph 高层Dialect A[Torch Dialect] -- B[Linalg Dialect] A -- C[TOSA Dialect] C -- B end subgraph 中层Dialect B -- D[Affine Dialect] B -- E[SCF Dialect] B -- F[GPU Dialect] end subgraph 低层Dialect D -- G[LLVM Dialect] E -- G F -- G G -- H[LLVM IR] end subgraph 自定义Dialect I[Custom Ops Dialect] J[Quantization Dialect] end I -- B J -- F subgraph 硬件特定优化 F -- K[CUDA Passes] F -- L[Metal Passes] F -- M[Vulkan Passes] endDialect的核心理念IR中的每个Operation属于一个Dialect。Dialect定义了自己的类型系统、操作集和优化Pass。不同Dialect的操作可以在同一函数中共存——编译器逐步将高层操作Lowering到低层。三、自定义Dialect设计与Lowering流程// 注意MLIR主要使用C API这里以Rust概念性的DSL展示核心设计 /// MLIR Dialect的核心概念概念性Rust表示 /// 实际MLIR使用TableGen定义Dialect /// 自定义AI Compiler Dialect定义矩阵操作 mod ai_compiler_dialect { /// MatMul操作矩阵乘法 /// 属性transpose_a, transpose_b控制转置 struct MatMulOp { /// 输入矩阵A [M, K] lhs: TensorValue, /// 输入矩阵B [K, N] rhs: TensorValue, /// 是否转置A transpose_a: bool, /// 是否转置B transpose_b: bool, /// 结果矩阵 [M, N] result: TensorValue, } /// FusedMatMulAdd操作融合的 Linear Bias /// 融合为一个操作避免中间结果的显存写回 struct FusedMatMulAddOp { input: TensorValue, // [M, K] weight: TensorValue, // [K, N] bias: TensorValue, // [N] result: TensorValue, // [M, N] } /// Attention操作QKV计算 /// 保留高层语义——Lowering时根据硬件选择最优实现 struct AttentionOp { query: TensorValue, // [B, H, S, D] key: TensorValue, // [B, H, T, D] value: TensorValue, // [B, H, T, D] mask: OptionTensorValue, scale: f32, /// 注意力输出 output: TensorValue, // [B, H, S, D] } } /// 优化Pass在AI Compiler Dialect层进行 mod optimization_passes { /// Pass: 融合 MatMul Add → FusedMatMulAdd /// 减少kernel launch次数和中间结果带宽 fn fuse_linear_bias_pass(func: mut IrFunction) - usize { let mut fused_count 0; for op in func.operations_mut() { // 模式匹配寻找MatMul → Add的连续操作 if let Some(matmul) op.as_matmul() { // 检查matmul的唯一使用者是否为Add if let Some(add_op) matmul.result().single_user() { if let Some(add) add_op.as_add() { // 检查add的另一个操作数是否为常量bias if add.rhs().is_constant() { // 执行融合删除MatMul和Add插入FusedMatMulAdd let fused FusedMatMulAddOp::new( matmul.lhs(), matmul.rhs(), add.rhs(), ); func.replace_op_pair(matmul, add, fused); fused_count 1; } } } } } fused_count } /// Pass: Attention → MatMul分解CPU后端 /// CPU没有专门的Attention硬件需分解为基础矩阵运算 fn decompose_attention_for_cpu(func: mut IrFunction) - usize { let mut decomposed 0; for op in func.operations_mut() { if let Some(attention) op.as_attention() { // Q × K^T → scores [B, H, S, T] let scores MatMulOp::new( attention.query(), attention.key(), false, // 不转置Q true, // 转置K ); // scores × scale let scaled MulOp::new(scores.result(), attention.scale()); // softmax(scores) let probs SoftmaxOp::new(scaled.result()); // probs × V let output MatMulOp::new( probs.result(), attention.value(), false, false, ); // 替换Attention为分解的操作序列 func.replace_single_with_sequence( attention, vec![scores, scaled, probs, output], ); decomposed 1; } } decomposed } /// Pass: 针对GPU的Attention优化 /// GPU后端使用FlashAttention算法减少显存访问 fn lower_attention_for_gpu(func: mut IrFunction) - usize { let mut lowered 0; for op in func.operations_mut() { if let Some(attention) op.as_attention() { // FlashAttention: 分块计算避免O(N²)显存写入 // 在GPU Dialect层面实现分块策略 let block_size 128; // 根据共享内存大小确定 let flash_attn FlashAttentionOp::new( attention.query(), attention.key(), attention.value(), block_size, attention.scale(), ); func.replace_op(attention, flash_attn); lowered 1; } } lowered } } /// Lowering流程从高层次到LLVM IR mod lowering_pipeline { /// 完整的Lowering Pipeline fn build_lowering_pipeline(target: TargetHardware) - LoweringPipeline { let mut pipeline LoweringPipeline::new(); // 阶段1AI Compiler Dialect内的优化 // 在高层语义上进行算子融合和优化 pipeline.add_pass(fuse_linear_bias_pass); pipeline.add_pass(fuse_attention_patterns_pass); // 阶段2Lowering到通用Dialect // 根据目标硬件选择不同的Lowering策略 match target { TargetHardware::CPU { pipeline.add_pass(decompose_attention_for_cpu); pipeline.add_pass(lower_to_affine); } TargetHardware::GPU(GpuType::NVIDIA) { pipeline.add_pass(lower_attention_for_gpu); pipeline.add_pass(lower_to_gpu); pipeline.add_pass(nvidia_tensor_core_opt); } TargetHardware::GPU(GpuType::Apple) { pipeline.add_pass(lower_attention_for_gpu); pipeline.add_pass(lower_to_gpu); pipeline.add_pass(metal_simd_group_opt); } TargetHardware::GPU(GpuType::AMD) { pipeline.add_pass(lower_attention_for_gpu); pipeline.add_pass(lower_to_gpu); pipeline.add_pass(rocm_matrix_core_opt); } } // 阶段3Lowering到LLVM Dialect pipeline.add_pass(lower_to_llvm); // 阶段4LLVM优化Pass pipeline.add_pass(llvm_optimize); pipeline } } /// Dialect互操作的核心类型转换 mod type_conversion { /// Tensor类型在不同Dialect间转换 /// 关键不同Dialect的Tensor可能有不同的内存布局 fn convert_tensor_layout( tensor: TensorValue, source_dialect: str, target_dialect: str, ) - ResultTensorValue, ConversionError { match (source_dialect, target_dialect) { // AI Compiler → Linalg抽象Tensor到Named Op (ai_compiler, linalg) { // AI Compiler的Tensor可能是NCHW或NHWC // Linalg的Named Op需要特定布局 // 这里插入layout转换操作 let layout_converted insert_layout_cast(tensor, Layout::RowMajor); Ok(layout_converted) } // Linalg → GPU需要分配GPU内存 (linalg, gpu) { // 创建GPU分配操作 let gpu_alloc GpuAllocOp::new(tensor.element_type(), tensor.shape()); // 插入host到device的拷贝 let copy GpuMemcpyOp::new(tensor, gpu_alloc.result()); Ok(copy.result()) } _ Err(ConversionError::NoConversionPath), } } } /// 自定义Dialect的表定义概念性 /// 实际使用MLIR的TableGen (.td文件) mod tablegen_example { // 以下为MLIR TableGen的概念对应 // def AICompiler_Dialect : Dialect { // let name aicompiler; // let summary AI Compiler custom operations; // let description [{ // High-level operations for AI model compilation. // Includes fused operators and hardware-agnostic primitives. // }]; // } // def AIC_MatMulOp : AIC_Opmatmul { // let summary Matrix multiplication; // let arguments (ins // AnyTensor:$lhs, // AnyTensor:$rhs, // OptionalAttrBoolAttr:$transpose_a, // OptionalAttrBoolAttr:$transpose_b // ); // let results (outs AnyTensor:$result); // // let verifier [{ return verifyMatMulOp(*this); }]; // // // Canonicalization: 通过模式重写进行标准化 // let hasCanonicalizer 1; // } // 模式重写规则MatMul(a, b, transpose_atrue) → Transpose(a)的MatMul // def TransposeMatMulPattern : Pat // (AIC_MatMulOp $a, $b, true, false), // (AIC_MatMulOp (AIC_TransposeOp $a), $b, false, false) // ; } /// MLIR的核心概念总结 fn mlir_key_concepts() { // 1. Operation是最小单元——每个op有唯一的OperationNamedialect.op_name // 2. Dialect是相关Operation的集合——定义类型系统、操作和Pass // 3. Region/Block提供控制流——不同于纯DAG支持循环和分支 // 4. Attribute是编译期常量——维度、步长、数据类型 // 5. Type系统可扩展——每个Dialect定义自己的类型 // MLIR的IR示例文本格式 let _mlir_example r# func.func matmul_add(%a: tensor128x256xf32, %b: tensor256x512xf32, %c: tensor512xf32) - tensor128x512xf32 { %0 aicompiler.matmul %a, %b {transpose_a false, transpose_b false} : tensor128x256xf32, tensor256x512xf32 - tensor128x512xf32 %1 aicompiler.fused_matmul_add %a, %b, %c : tensor128x256xf32, tensor256x512xf32, tensor512xf32 - tensor128x512xf32 return %1 : tensor128x512xf32 } #; } // 辅助类型定义 enum TargetHardware { CPU, GPU(GpuType), } enum GpuType { NVIDIA, Apple, AMD, } enum Layout { RowMajor, ColMajor, NCHW, NHWC, } struct TensorValue; struct IrFunction; struct LoweringPipeline; impl IrFunction { fn operations_mut(self) - [IrOp] { [] } } struct IrOp; impl IrOp { fn as_matmul(self) - OptionMatMulOpRef { None } fn as_add(self) - OptionAddOpRef { None } fn as_attention(self) - OptionAttentionOpRef { None } } struct MatMulOpRef { lhs: (), rhs: (), result: (), } impl MatMulOpRef { fn result(self) - TensorValue { TensorValue } fn lhs(self) - TensorValue { TensorValue } fn rhs(self) - TensorValue { TensorValue } } struct AddOpRef { rhs: (), } impl AddOpRef { fn rhs(self) - TensorValue { TensorValue } } struct AttentionOpRef { query: (), key: (), value: (), scale: (), } impl AttentionOpRef { fn query(self) - TensorValue { TensorValue } fn key(self) - TensorValue { TensorValue } fn value(self) - TensorValue { TensorValue } fn scale(self) - f32 { 1.0 } } impl TensorValue { fn single_user(self) - OptionIrOp { None } } #[derive(Debug, thiserror::Error)] enum ConversionError { #[error(No conversion path between dialects)] NoConversionPath, } // Stub functions fn insert_layout_cast(_tensor: TensorValue, _layout: Layout) - TensorValue { TensorValue } struct LoweringPipeline; impl LoweringPipeline { fn new() - Self { Self } fn add_passF: Fn(mut IrFunction) - usize(mut self, _: F) {} } fn fuse_attention_patterns_pass(_: mut IrFunction) - usize { 0 } fn lower_to_affine(_: mut IrFunction) - usize { 0 } fn lower_to_gpu(_: mut IrFunction) - usize { 0 } fn nvidia_tensor_core_opt(_: mut IrFunction) - usize { 0 } fn metal_simd_group_opt(_: mut IrFunction) - usize { 0 } fn rocm_matrix_core_opt(_: mut IrFunction) - usize { 0 } fn lower_to_llvm(_: mut IrFunction) - usize { 0 } fn llvm_optimize(_: mut IrFunction) - usize { 0 } struct FusedMatMulAddOp; impl FusedMatMulAddOp { fn new(_: TensorValue, _: TensorValue, _: TensorValue) - Self { Self } } struct MatMulOp; impl MatMulOp { fn new(_: TensorValue, _: TensorValue, _: bool, _: bool) - Self { Self } fn result(self) - TensorValue { TensorValue } } struct MulOp; impl MulOp { fn new(_: TensorValue, _: f32) - Self { Self } fn result(self) - TensorValue { TensorValue } } struct SoftmaxOp; impl SoftmaxOp { fn new(_: TensorValue) - Self { Self } fn result(self) - TensorValue { TensorValue } } struct FlashAttentionOp; impl FlashAttentionOp { fn new(_: TensorValue, _: TensorValue, _: TensorValue, _: usize, _: f32) - Self { Self } } struct GpuAllocOp; impl GpuAllocOp { fn new(_: (), _: [usize]) - Self { Self } fn result(self) - TensorValue { TensorValue } } struct GpuMemcpyOp; impl GpuMemcpyOp { fn new(_: TensorValue, _: TensorValue) - Self { Self } fn result(self) - TensorValue { TensorValue } }核心设计Dialect定义操作语义Attention保持高层语义直到Lowering阶段分阶段LoweringAI Compiler Dialect → Linalg/Affine → GPU → LLVM硬件特定的Lowering路径同一操作在不同硬件上有不同实现Pass可组合按目标平台组合不同的Pass序列Dialect间类型转换的正确性证明与Shape推断。MLIR的Dialect互操作中最棘手的不是Pass流程而是形状Shape推断的跨Dialect传播。当aicompiler.attention被lower到linalg.matmul linalg.softmax时输出Tensor的shape需要从输入推断——但这要求Attention的Lowering逻辑不仅生成操作序列还要显式传播shape约束。MLIR通过ShapeDialect和shape.func来显式建模维度关系%c shape.mul %m, %k表示c的shape是由m和k的shape乘积决定的。这种显式的shape计算在高层Dialect中引入了额外的IR节点但在Lowering到LLVM时被常量折叠消除因为所有shape在编译期为常量。另一个深入问题是Dialect间的语义等价验证——当tosa.matmul和linalg.matmul在IR中同时存在时编译器必须知道它们在数学上等价才能进行跨Dialect的算子融合。MLIR通过Interface机制实现TosaOp和LinalgOp可以实现相同的MatMulInterfacePass通过Interface查询而非具体Dialect匹配操作实现面向接口编程式的编译优化。这种设计的关键优势是添加新Dialect时只需实现已有Interface现有的跨Dialect优化Pass无需修改即可工作。但Interface粒度选择的代价是编译器的实现复杂度——过细的Interface如HasSingleResult虽然复用性高但信息量低过粗的Interface如HasMatMulSemantics信息量高但可复用性差。这是编译器IR设计中抽象层次通胀的经典权衡。四、MLIR的工程实用性评估MLIR的优势Dialect扩展性不需要修改核心IR添加新Dialect即可渐进式Lowering逐步降低抽象层次每层有独立的优化空间硬件Backend隔离添加新硬件支持只需添加新的Lowering路径与LLVM共享生态最终阶段利用LLVM的优化和代码生成MLIR的挑战学习曲线陡峭Dialect、Pass、Pattern Rewriting等概念需要时间理解编译时间多层Lowering和Pass增加了编译时间TableGen依赖Dialect定义使用TableGen DSL——额外工具链Rust生态不完整目前主要C APIRust绑定(Melior)在早期阶段适用场景多硬件后端的编译器通过Dialect隔离硬件差异需要自定义算子的深度学习框架DSL到硬件的编译链五、总结MLIR通过Dialect机制实现IR的可扩展性——每层抽象独立定义可组合的Pass完成渐进Lowering高层Dialect保留领域语义——Attention/Convolution的操作语义在Lowering前可优化硬件特定的Lowering路径解决多后端支持问题——每个硬件有独立的优化Pass序列算子融合在IR层面而非源码层实现——利用use-def链进行模式匹配和替换MLIR的设计核心是分离关注点——Dialect定义语义Pass实现转换Type系统保证正确性
)
)
