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]
end
Dialect的核心理念：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: Option<TensorValue>,
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();
// 阶段1：AI Compiler Dialect内的优化
// 在高层语义上进行算子融合和优化
pipeline.add_pass(fuse_linear_bias_pass);
pipeline.add_pass(fuse_attention_patterns_pass);
// 阶段2：Lowering到通用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);
}
}
// 阶段3：Lowering到LLVM Dialect
pipeline.add_pass(lower_to_llvm);
// 阶段4：LLVM优化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,
) -> Result<TensorValue, 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_Op<"matmul"> {
// let summary = "Matrix multiplication";
// let arguments = (ins
// AnyTensor:$lhs,
// AnyTensor:$rhs,
// OptionalAttr<BoolAttr>:$transpose_a,
// OptionalAttr<BoolAttr>:$transpose_b
// );
// let results = (outs AnyTensor:$result);
//
// let verifier = [{ return verifyMatMulOp(*this); }];
//
// // Canonicalization: 通过模式重写进行标准化
// let hasCanonicalizer = 1;
// }
// 模式重写规则：MatMul(a, b, transpose_a=true) → 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有唯一的OperationName（dialect.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: tensor<128x256xf32>,
%b: tensor<256x512xf32>,
%c: tensor<512xf32>) -> tensor<128x512xf32> {
%0 = aicompiler.matmul %a, %b {transpose_a = false, transpose_b = false}
: tensor<128x256xf32>, tensor<256x512xf32> -> tensor<128x512xf32>
%1 = aicompiler.fused_matmul_add %a, %b, %c
: tensor<128x256xf32>, tensor<256x512xf32>, tensor<512xf32>
-> tensor<128x512xf32>
return %1 : tensor<128x512xf32>
}
"#;
}
// 辅助类型定义
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) -> Option<&MatMulOpRef> { None }
fn as_add(&self) -> Option<&AddOpRef> { None }
fn as_attention(&self) -> Option<&AttentionOpRef> { 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) -> Option<&IrOp> { 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_pass<F: 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阶段
- 分阶段Lowering：AI 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可以实现相同的MatMulInterface，Pass通过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++ API，Rust绑定(Melior)在早期阶段
适用场景：
- 多硬件后端的编译器：通过Dialect隔离硬件差异
- 需要自定义算子的深度学习框架
- DSL到硬件的编译链
