MLIR/TVM/XLA深度学习编译器深度对比与实战一、引言AI 芯片百花齐放NVIDIA GPU、Apple M系列、Google TPU、华为昇腾、高通 Hexagon…每种芯片都有独特的指令集和内存模型。写一次到处优化成为奢望。深度学习编译器正是破解这一困局的关键——将高层计算图自动编译为底层高效代码。本文将深度对比三大编译器Google XLA、Apache TVM、LLVM MLIR并通过实战案例展示从模型定义到自动调优的完整流程。二、编译器分层架构┌─────────────────────────────────┐ │ 前端 (Frontend) │ PyTorch/TensorFlow/ONNX ├─────────────────────────────────┤ │ 高层IR (HLO/Relay) │ 算子融合、图优化 ├─────────────────────────────────┤ │ 中层IR (Linalg/StableHLO) │ 循环变换、内存布局 ├─────────────────────────────────┤ │ 底层IR (LLVM IR/SPIR-V/PTX) │ 向量化、指令选择 ├─────────────────────────────────┤ │ 后端 (Backend) │ GPU/CPU/TPU/NPU └─────────────────────────────────┘三、XLA (Accelerated Linear Algebra)XLA 是 Google 开发的 JIT 编译器深度集成于 TensorFlow/JAX。importtorchimporttorch_xlaimporttorch_xla.core.xla_modelasxm# 方式1: PyTorch/XLA (TPU/GPU)devicexm.xla_device()modelMyModel().to(device)# JIT编译torch.jit.scriptdefcompiled_forward(x):returnmodel(x)# 训练循环fordataindataloader:datadata.to(device)outputcompiled_forward(data)losscriterion(output,target)loss.backward()xm.optimizer_step(optimizer)# 方式2: JAX (原生XLA支持)importjaximportjax.numpyasjnpjax.jit# 自动编译为XLAdeftrain_step(params,batch):defloss_fn(params):logitsmodel.apply(params,batch[x])return-jnp.mean(jax.nn.log_softmax(logits)*batch[y])gradjax.grad(loss_fn)(params)returngrad# 查看编译后的HLOprint(jax.xla_computation(train_step)(params,batch).as_hlo_text())XLA核心优化# XLA的算子融合示例# 原始代码:# y matmul(W, x)# y y b# y relu(y)## XLA将其融合为单个kernel: FusedMatMulBiasRelu# 查看JAX的HLO IRimportjax computationjax.xla_computation(my_function)(x)print(computation.as_hlo_text())# 输出:# HloModule ...# %fused_computation {# %param_0 f32[1024,512] parameter(0)# %param_1 f32[512,256] parameter(1)# %dot f32[1024,256] dot(%param_0, %param_1)# %broadcast f32[1024,256] broadcast(%bias)# %add f32[1024,256] add(%dot, %broadcast)# ROOT %relu f32[1024,256] maximum(%add, 0)# }四、TVM (Tensor Virtual Machine)TVM 是 Apache 开源的端到端深度学习编译器。4.1 从模型到部署importtvmfromtvmimportrelay,auto_schedulerimporttvm.contrib.graph_executorasruntimeimportonnx# 1. 导入模型支持ONNX/PyTorch/TF/Kerasonnx_modelonnx.load(resnet18.onnx)mod,paramsrelay.frontend.from_onnx(onnx_model)# 2. 图级别优化算子融合、常量折叠withtvm.transform.PassContext(opt_level3):modrelay.transform.InferType()(mod)modrelay.transform.FuseOps(fuse_opt_level3)(mod)# 算子融合modrelay.transform.FoldConstant()(mod)# 常量折叠modrelay.transform.AlterOpLayout()(mod)# 布局优化# 3. 自动调优 (AutoTVM / AutoScheduler)targettvm.target.Target(cuda -archsm_80)# A100tasks,task_weightsauto_scheduler.extract_tasks(mod[main],params,target)tunerauto_scheduler.TaskScheduler(tasks,task_weights)tune_optionauto_scheduler.TuningOptions(num_measure_trials200,runnerauto_scheduler.LocalRunner(repeat10,enable_cpu_cache_flushTrue),measure_callbacks[auto_scheduler.RecordToFile(resnet18.json)],)tuner.tune(tune_option)# 4. 应用最佳调优配置withauto_scheduler.ApplyHistoryBest(resnet18.json):withtvm.transform.PassContext(opt_level3,config{relay.backend.use_auto_scheduler:True}):librelay.build(mod,targettarget,paramsparams)# 5. 部署运行devtvm.cuda(0)moduleruntime.GraphModule(lib[default](dev))module.set_input(input,input_data)module.run()outputmodule.get_output(0)4.2 手动调度示例importtvmfromtvmimportte# 矩阵乘法的手动调度M,N,K1024,1024,1024# 定义计算Ate.placeholder((M,K),nameA)Bte.placeholder((K,N),nameB)kte.reduce_axis((0,K),namek)Cte.compute((M,N),lambdai,j:te.sum(A[i,k]*B[k,j],axisk))# 创建调度ste.create_schedule(C.op)# 分块Tilingblock_x,block_y32,32xo,yo,xi,yis[C].tile(C.op.axis[0],C.op.axis[1],block_x,block_y)# 向量化s[C].vectorize(yi)# 缓存共享内存AAs.cache_read(A,shared,[C])BBs.cache_read(B,shared,[C])# 绑定到GPUs[AA].compute_at(s[C],xo)s[BB].compute_at(s[C],xo)# 编译functvm.build(s,[A,B,C],targetcuda)print(func.imported_modules[0].get_source())4.3 性能对比后端(TVM编译)ResNet50MobileNetV2BERTPyTorch Eager45ms12ms85msTVM AutoScheduler22ms5.5ms42msTVM TensorRT15ms4.2ms30ms加速比3x2.8x2.8x五、MLIR多层中间表示MLIR 是 LLVM 项目的子项目提供可组合的编译器基础设施。// MLIR方言示例从高层到低层 // 1. StableHLO方言XLA兼容 func.func main(%arg0: tensor1x3x224x224xf32) - tensor1x1000xf32 { %0 stablehlo.convolution(%arg0, %filter) dim_numbers [b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f], window {stride [2, 2], pad [[1, 1], [1, 1]]} : (tensor1x3x224x224xf32, tensor64x3x7x7xf32) - tensor1x64x112x112xf32 %1 stablehlo.batch_norm_inference %0, %scale, %offset, %mean, %variance : tensor1x64x112x112xf32 return %1 : tensor1x64x112x112xf32 } // 2. Linalg方言线性代数操作 func.func matmul(%A: memref1024x512xf32, %B: memref512x256xf32, %C: memref1024x256xf32) { linalg.matmul ins(%A, %B : memref1024x512xf32, memref512x256xf32) outs(%C : memref1024x256xf32) return } // 3. SCF方言结构化控制流 scf.for %i %c0 to %N step %c1 { %val memref.load %A[%i] : memref1024xf32 %squared arith.mulf %val, %val : f32 memref.store %squared, %B[%i] : memref1024xf32 }MLIR Python实战frommlir.irimport*frommlir.dialectsimportfunc,arith,scf,memref,linalgdefbuild_matmul():用MLIR Python API构建矩阵乘法withContext()asctx,Location.unknown():moduleModule.create()withInsertionPoint(module.body):M,N,K1024,256,512# 函数定义ftypeFunctionType.get([MemRefType.get([M,K],F32Type.get()),MemRefType.get([K,N],F32Type.get()),MemRefType.get([M,N],F32Type.get())],[])func_opfunc.FuncOp(matmul,ftype)entry_blockfunc_op.add_entry_block()withInsertionPoint(entry_block):a,b,centry_block.arguments# i循环zeroarith.ConstantOp.create_index(0)onearith.ConstantOp.create_index(1)i_loopscf.ForOp(zero,arith.ConstantOp.create_index(M),one)withInsertionPoint(i_loop.body):j_loopscf.ForOp(zero,arith.ConstantOp.create_index(N),one)withInsertionPoint(j_loop.body):# 初始化累加器accmemref.AllocaOp(MemRefType.get([1],F32Type.get()),[],[])k_loopscf.ForOp(zero,arith.ConstantOp.create_index(K),one)withInsertionPoint(k_loop.body):# C[i,j] A[i,k] * B[k,j]a_valmemref.LoadOp(a,[i_loop.induction_variable,k_loop.induction_variable])b_valmemref.LoadOp(b,[k_loop.induction_variable,j_loop.induction_variable])prodarith.MulFOp(a_val,b_val)oldmemref.LoadOp(acc,[zero])new_valarith.AddFOp(old,prod)memref.StoreOp(new_val,acc,[zero])scf.YieldOp([])final_valmemref.LoadOp(acc,[zero])memref.StoreOp(final_val,c,[i_loop.induction_variable,j_loop.induction_variable])scf.YieldOp([])scf.YieldOp([])func.ReturnOp([])print(module)returnmodule build_matmul()六、三大编译器对比特性XLATVMMLIR所属GoogleApacheLLVM目标用户JAX/TF开发者芯片/框架厂商编译器开发者输入格式HLORelay/ONNX自定义方言自动调优❌✅ AutoTVM/Ansor基础passes硬件后端TPU/GPU/CPU全平台全平台学习曲线低JAX透明中高DIY生产成熟度⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐典型用户Google内部华为/阿里/字节Apple/Google七、选择建议场景推荐JAX/TF用户TPU部署XLA自定义芯片/极致性能TVM构建新编译器/DSLMLIRNVIDIA GPU通用优化TVM TensorRT移动端部署TVM (ARM Mali/Adreno)浏览器推理TVM (WebGPU/WebAssembly)八、总结深度学习编译器的核心价值XLA— 零配置加速JAX/TF用户首选TVM— 自动调优 全平台覆盖极致性能MLIR— 构建下一代编译器的基础设施三者关系XLA专注TPU生态TVM覆盖全硬件MLIR提供编译器构建框架。实际项目中TVM是通用性最好的选择。
MLIR/TVM/XLA深度学习编译器深度对比与实战一、引言AI 芯片百花齐放NVIDIA GPU、Apple M系列、Google TPU、华为昇腾、高通 Hexagon…每种芯片都有独特的指令集和内存模型。写一次到处优化成为奢望。深度学习编译器正是破解这一困局的关键——将高层计算图自动编译为底层高效代码。本文将深度对比三大编译器Google XLA、Apache TVM、LLVM MLIR并通过实战案例展示从模型定义到自动调优的完整流程。二、编译器分层架构┌─────────────────────────────────┐ │ 前端 (Frontend) │ PyTorch/TensorFlow/ONNX ├─────────────────────────────────┤ │ 高层IR (HLO/Relay) │ 算子融合、图优化 ├─────────────────────────────────┤ │ 中层IR (Linalg/StableHLO) │ 循环变换、内存布局 ├─────────────────────────────────┤ │ 底层IR (LLVM IR/SPIR-V/PTX) │ 向量化、指令选择 ├─────────────────────────────────┤ │ 后端 (Backend) │ GPU/CPU/TPU/NPU └─────────────────────────────────┘三、XLA (Accelerated Linear Algebra)XLA 是 Google 开发的 JIT 编译器深度集成于 TensorFlow/JAX。importtorchimporttorch_xlaimporttorch_xla.core.xla_modelasxm# 方式1: PyTorch/XLA (TPU/GPU)devicexm.xla_device()modelMyModel().to(device)# JIT编译torch.jit.scriptdefcompiled_forward(x):returnmodel(x)# 训练循环fordataindataloader:datadata.to(device)outputcompiled_forward(data)losscriterion(output,target)loss.backward()xm.optimizer_step(optimizer)# 方式2: JAX (原生XLA支持)importjaximportjax.numpyasjnpjax.jit# 自动编译为XLAdeftrain_step(params,batch):defloss_fn(params):logitsmodel.apply(params,batch[x])return-jnp.mean(jax.nn.log_softmax(logits)*batch[y])gradjax.grad(loss_fn)(params)returngrad# 查看编译后的HLOprint(jax.xla_computation(train_step)(params,batch).as_hlo_text())XLA核心优化# XLA的算子融合示例# 原始代码:# y matmul(W, x)# y y b# y relu(y)## XLA将其融合为单个kernel: FusedMatMulBiasRelu# 查看JAX的HLO IRimportjax computationjax.xla_computation(my_function)(x)print(computation.as_hlo_text())# 输出:# HloModule ...# %fused_computation {# %param_0 f32[1024,512] parameter(0)# %param_1 f32[512,256] parameter(1)# %dot f32[1024,256] dot(%param_0, %param_1)# %broadcast f32[1024,256] broadcast(%bias)# %add f32[1024,256] add(%dot, %broadcast)# ROOT %relu f32[1024,256] maximum(%add, 0)# }四、TVM (Tensor Virtual Machine)TVM 是 Apache 开源的端到端深度学习编译器。4.1 从模型到部署importtvmfromtvmimportrelay,auto_schedulerimporttvm.contrib.graph_executorasruntimeimportonnx# 1. 导入模型支持ONNX/PyTorch/TF/Kerasonnx_modelonnx.load(resnet18.onnx)mod,paramsrelay.frontend.from_onnx(onnx_model)# 2. 图级别优化算子融合、常量折叠withtvm.transform.PassContext(opt_level3):modrelay.transform.InferType()(mod)modrelay.transform.FuseOps(fuse_opt_level3)(mod)# 算子融合modrelay.transform.FoldConstant()(mod)# 常量折叠modrelay.transform.AlterOpLayout()(mod)# 布局优化# 3. 自动调优 (AutoTVM / AutoScheduler)targettvm.target.Target(cuda -archsm_80)# A100tasks,task_weightsauto_scheduler.extract_tasks(mod[main],params,target)tunerauto_scheduler.TaskScheduler(tasks,task_weights)tune_optionauto_scheduler.TuningOptions(num_measure_trials200,runnerauto_scheduler.LocalRunner(repeat10,enable_cpu_cache_flushTrue),measure_callbacks[auto_scheduler.RecordToFile(resnet18.json)],)tuner.tune(tune_option)# 4. 应用最佳调优配置withauto_scheduler.ApplyHistoryBest(resnet18.json):withtvm.transform.PassContext(opt_level3,config{relay.backend.use_auto_scheduler:True}):librelay.build(mod,targettarget,paramsparams)# 5. 部署运行devtvm.cuda(0)moduleruntime.GraphModule(lib[default](dev))module.set_input(input,input_data)module.run()outputmodule.get_output(0)4.2 手动调度示例importtvmfromtvmimportte# 矩阵乘法的手动调度M,N,K1024,1024,1024# 定义计算Ate.placeholder((M,K),nameA)Bte.placeholder((K,N),nameB)kte.reduce_axis((0,K),namek)Cte.compute((M,N),lambdai,j:te.sum(A[i,k]*B[k,j],axisk))# 创建调度ste.create_schedule(C.op)# 分块Tilingblock_x,block_y32,32xo,yo,xi,yis[C].tile(C.op.axis[0],C.op.axis[1],block_x,block_y)# 向量化s[C].vectorize(yi)# 缓存共享内存AAs.cache_read(A,shared,[C])BBs.cache_read(B,shared,[C])# 绑定到GPUs[AA].compute_at(s[C],xo)s[BB].compute_at(s[C],xo)# 编译functvm.build(s,[A,B,C],targetcuda)print(func.imported_modules[0].get_source())4.3 性能对比后端(TVM编译)ResNet50MobileNetV2BERTPyTorch Eager45ms12ms85msTVM AutoScheduler22ms5.5ms42msTVM TensorRT15ms4.2ms30ms加速比3x2.8x2.8x五、MLIR多层中间表示MLIR 是 LLVM 项目的子项目提供可组合的编译器基础设施。// MLIR方言示例从高层到低层 // 1. StableHLO方言XLA兼容 func.func main(%arg0: tensor1x3x224x224xf32) - tensor1x1000xf32 { %0 stablehlo.convolution(%arg0, %filter) dim_numbers [b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f], window {stride [2, 2], pad [[1, 1], [1, 1]]} : (tensor1x3x224x224xf32, tensor64x3x7x7xf32) - tensor1x64x112x112xf32 %1 stablehlo.batch_norm_inference %0, %scale, %offset, %mean, %variance : tensor1x64x112x112xf32 return %1 : tensor1x64x112x112xf32 } // 2. Linalg方言线性代数操作 func.func matmul(%A: memref1024x512xf32, %B: memref512x256xf32, %C: memref1024x256xf32) { linalg.matmul ins(%A, %B : memref1024x512xf32, memref512x256xf32) outs(%C : memref1024x256xf32) return } // 3. SCF方言结构化控制流 scf.for %i %c0 to %N step %c1 { %val memref.load %A[%i] : memref1024xf32 %squared arith.mulf %val, %val : f32 memref.store %squared, %B[%i] : memref1024xf32 }MLIR Python实战frommlir.irimport*frommlir.dialectsimportfunc,arith,scf,memref,linalgdefbuild_matmul():用MLIR Python API构建矩阵乘法withContext()asctx,Location.unknown():moduleModule.create()withInsertionPoint(module.body):M,N,K1024,256,512# 函数定义ftypeFunctionType.get([MemRefType.get([M,K],F32Type.get()),MemRefType.get([K,N],F32Type.get()),MemRefType.get([M,N],F32Type.get())],[])func_opfunc.FuncOp(matmul,ftype)entry_blockfunc_op.add_entry_block()withInsertionPoint(entry_block):a,b,centry_block.arguments# i循环zeroarith.ConstantOp.create_index(0)onearith.ConstantOp.create_index(1)i_loopscf.ForOp(zero,arith.ConstantOp.create_index(M),one)withInsertionPoint(i_loop.body):j_loopscf.ForOp(zero,arith.ConstantOp.create_index(N),one)withInsertionPoint(j_loop.body):# 初始化累加器accmemref.AllocaOp(MemRefType.get([1],F32Type.get()),[],[])k_loopscf.ForOp(zero,arith.ConstantOp.create_index(K),one)withInsertionPoint(k_loop.body):# C[i,j] A[i,k] * B[k,j]a_valmemref.LoadOp(a,[i_loop.induction_variable,k_loop.induction_variable])b_valmemref.LoadOp(b,[k_loop.induction_variable,j_loop.induction_variable])prodarith.MulFOp(a_val,b_val)oldmemref.LoadOp(acc,[zero])new_valarith.AddFOp(old,prod)memref.StoreOp(new_val,acc,[zero])scf.YieldOp([])final_valmemref.LoadOp(acc,[zero])memref.StoreOp(final_val,c,[i_loop.induction_variable,j_loop.induction_variable])scf.YieldOp([])scf.YieldOp([])func.ReturnOp([])print(module)returnmodule build_matmul()六、三大编译器对比特性XLATVMMLIR所属GoogleApacheLLVM目标用户JAX/TF开发者芯片/框架厂商编译器开发者输入格式HLORelay/ONNX自定义方言自动调优❌✅ AutoTVM/Ansor基础passes硬件后端TPU/GPU/CPU全平台全平台学习曲线低JAX透明中高DIY生产成熟度⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐典型用户Google内部华为/阿里/字节Apple/Google七、选择建议场景推荐JAX/TF用户TPU部署XLA自定义芯片/极致性能TVM构建新编译器/DSLMLIRNVIDIA GPU通用优化TVM TensorRT移动端部署TVM (ARM Mali/Adreno)浏览器推理TVM (WebGPU/WebAssembly)八、总结深度学习编译器的核心价值XLA— 零配置加速JAX/TF用户首选TVM— 自动调优 全平台覆盖极致性能MLIR— 构建下一代编译器的基础设施三者关系XLA专注TPU生态TVM覆盖全硬件MLIR提供编译器构建框架。实际项目中TVM是通用性最好的选择。
双梁正问题 【 torch 实战】（Python代码实现）)
