1. 从ONNX到自研芯片为什么还要再造一个编译器第一次接触TPU-MLIR这个项目时我的反应和大多数人一样ONNX已经有了ONNX Runtime也能跑TVM、XLA这些框架也都在为什么还要费劲用MLIR再搭一套编译器这个问题如果不搞清楚后面看代码基本就是看天书。先说结论通用推理框架解决的是能跑而自研TPU需要的是跑得好。这两者之间的差距恰恰就是TPU-MLIR存在的理由。ONNX本质上是一个模型交换格式它定义了一套算子集合和计算图表示。ONNX Runtime作为通用推理引擎它的优化策略是面向CPU、CUDA这类通用硬件的。但自研TPU的硬件架构往往非常特殊——可能有自己的脉动阵列、自己的片上存储层次、自己的数据搬运指令、自己的量化格式。通用框架里的优化pass根本不知道这些硬件特性它只能做通用层面的图优化比如常量折叠、算子融合再往下就无能为力了。我举个具体的例子。假设你的TPU有一个专门的矩阵乘加单元支持INT8输入、INT32累加并且要求输入数据必须按照特定的tile格式排布在片上缓存里。ONNX Runtime跑到这个算子时它只能调用你提供的一个kernel至于数据怎么搬、什么时候搬、搬多少它一概不管。而TPU-MLIR要做的是从计算图层面就开始规划哪些算子可以合并成一个大的矩阵运算、数据在什么阶段做layout转换、量化参数怎么传播、片上内存怎么分配才能让数据搬运次数最少。这些决策必须在编译期完成运行时只负责执行。这就是MLIR的价值所在。MLIR不是又一个IR它是一套可扩展的编译器基础设施。你可以定义自己的dialect方言描述你硬件特有的操作你可以写自己的pass在合适的抽象层次上做优化你可以利用MLIR已有的基础设施比如affine dialect做循环优化、linalg dialect做张量运算的逐步lowering。TPU-MLIR正是基于这套基础设施构建了一条从ONNX到自研TPU的完整编译流水线。提示如果你之前只接触过ONNX Runtime这类推理框架建议先花点时间理解MLIR的dialect概念。这是理解TPU-MLIR架构的前提否则后面看到一堆.td文件和pass会完全摸不着头脑。从工程角度看TPU-MLIR解决的是一类非常具体的问题模型部署的最后一公里。训练框架产出的模型是浮点的、算子粒度是粗的、内存布局是框架决定的而芯片能执行的是定点的、算子粒度是细的、内存布局是硬件规定的。这中间的鸿沟靠手工写kernel填不平靠通用框架也填不平必须有一个专门的编译器来做系统性的转换和优化。2. TPU-MLIR的编译流水线从ONNX一路降到芯片指令理解了为什么之后接下来看怎么做。TPU-MLIR的编译流程不是一步到位的它是一条多级lowering的流水线每一级都在不同的抽象层次上做该做的事。我把这条流水线拆成几个关键阶段来讲。2.1 前端导入ONNX模型怎么变成MLIR的IRTPU-MLIR的前端入口是ONNX。你给它一个.onnx文件它首先做的是把ONNX的计算图翻译成MLIR的top dialect。这个top dialect是TPU-MLIR自己定义的高层抽象算子粒度基本和ONNX对齐但表达方式已经是MLIR的SSA形式了。这一步看起来简单实际上有不少细节。ONNX的算子版本很多同一个算子在不同opset版本里语义可能不一样。比如Resize算子opset 10和opset 11的行为就有差异。TPU-MLIR在导入时需要根据模型的opset版本做相应的语义适配。另外ONNX允许一些比较野的写法比如动态shape、控制流算子这些在导入阶段就要做规范化处理。导入完成后你会得到一个top dialect的MLIR文件。这个文件可以用mlir-opt工具查看结构比ONNX的protobuf可读性好很多。我建议在这个阶段就仔细检查一遍确认算子映射是否正确、shape信息是否完整。很多后续问题其实在前端导入时就已经埋下了。2.2 高层优化在top dialect上做图级变换拿到top dialect的IR之后TPU-MLIR会跑一系列图优化pass。这些pass的目标是简化计算图、减少算子数量、为后续的量化做准备。常见的优化包括常量折叠把编译期就能算出来的子图直接算掉减少运行时计算量。算子融合比如ConvBNReLU这种经典组合融合成一个算子。融合之后不仅减少kernel启动开销更重要的是让后续量化能在一个更大的范围内做。死代码消除去掉那些输出不被使用的算子。shape推断补全所有tensor的shape信息为后续的内存分配做准备。这些优化在ONNX Runtime里也有但TPU-MLIR做这些优化的目的不太一样。它不是为了通用性能提升而是为了让计算图更适合后续的量化和小算子合并。比如算子融合的粒度TPU-MLIR会倾向于融合得更大一些因为自研TPU通常有比较大的片上缓存大算子可以减少数据搬运。2.3 量化从浮点到定点的关键一跃量化是TPU-MLIR里最核心也最容易出问题的环节。自研TPU通常只支持定点运算所以浮点模型必须转成定点。TPU-MLIR支持多种量化模式包括F32、F16、BF16以及INT8的对称和非对称量化。量化的基本思路是对每个tensor统计其数值范围然后计算scale和zero_point把浮点值映射到定点值。听起来简单但实际操作中有几个坑第一个坑是量化粒度。Per-tensor量化是整个tensor共用一个scalePer-channel量化是每个通道一个scale。Per-channel精度更高但硬件支持程度取决于你的TPU。如果TPU的矩阵乘单元不支持per-channel的scale那你在编译器里做了per-channel量化到了硬件上还是得退化成per-tensor反而多了一层转换开销。第二个坑是量化参数传播。Conv的输出scale和输入的scale、权重的scale是有数学关系的。如果每一层都独立统计scale误差会逐层累积。TPU-MLIR的做法是在校准阶段用一批代表性数据跑一遍统计每层的激活值分布然后用这些统计信息来定scale。校准数据的选取很关键如果校准集和实际推理数据分布差异大量化精度会明显下降。第三个坑是混合精度。有些层对精度敏感比如第一层和最后一层强行INT8量化会导致精度暴跌。TPU-MLIR允许你指定某些层保持F16或BF16其余层用INT8。这个混合精度的配置需要根据实际模型来调没有万能公式。注意量化不是一键INT8就完事了。我见过太多案例模型量化后精度掉十几个点最后发现是校准集选得不对或者某个关键层不该量化。建议在量化后一定用验证集跑一遍对比浮点模型的输出逐层排查精度损失。2.4 Lowering到TPU dialect硬件相关的优化量化完成后IR会从top dialect lowering到TPU dialect。这一步是硬件相关的也是TPU-MLIR和通用编译器最大的区别所在。在TPU dialect层面算子已经被映射到硬件实际支持的指令上了。比如一个Conv算子在top dialect里就是一个Conv但在TPU dialect里它可能被拆成数据搬运指令、矩阵乘指令、累加指令、激活指令。具体怎么拆取决于你的TPU架构。这个阶段还会做内存分配。自研TPU的片上内存通常很有限比如只有几MB的SRAM。编译器需要决定每个tensor放在哪里、什么时候加载、什么时候释放。TPU-MLIR会做liveness分析找出每个tensor的生命周期然后做内存复用。如果内存不够还需要把部分数据spill到DDR这又会引入额外的搬运开销。2.5 代码生成从TPU dialect到可执行文件最后一步是把TPU dialect的IR翻译成芯片能执行的指令流。这一步的输出通常是一个二进制文件或者一组指令序列加载到TPU上就能跑。代码生成的质量直接影响性能。同样的计算图不同的指令调度策略可能导致几倍的性能差异。TPU-MLIR在这个阶段会做指令重排、流水线调度、双缓冲等优化。比如数据搬运和计算可以重叠执行编译器需要识别出哪些搬运可以和哪些计算并行然后插入合适的同步指令。3. 动手跑通第一个模型环境搭建与实操步骤光看架构不够得实际跑一遍才能有体感。这一章我带你从零开始把一个ONNX模型编译到TPU上。整个过程我尽量给出可复现的步骤但需要说明的是TPU-MLIR的具体命令和配置会随版本变化以下基于我实际操作时的版本。3.1 环境准备依赖安装与编译TPU-MLIR的构建依赖LLVM/MLIR。如果你直接从源码编译第一步是clone LLVM项目切换到TPU-MLIR要求的commit然后编译。这个过程比较耗时建议机器配置至少16核、32GB内存否则编译LLVM能等到天荒地老。# 克隆LLVMTPU-MLIR通常要求特定commit git clone https://github.com/llvm/llvm-project.git cd llvm-project git checkout tpu-mlir要求的commit # 编译MLIR mkdir build cd build cmake -G Ninja ../llvm \ -DCMAKE_BUILD_TYPERelease \ -DLLVM_ENABLE_PROJECTSmlir \ -DLLVM_TARGETS_TO_BUILDhost \ -DLLVM_ENABLE_ASSERTIONSON ninja编译完LLVM之后再clone TPU-MLIR配置好LLVM的路径编译TPU-MLIR本身。如果你不想从源码编译也可以找预编译的docker镜像很多团队会提供。但预编译镜像的版本可能和你需要的对不上长期来看还是自己编译更可控。环境搭好之后验证一下工具是否可用tpu-mlir-opt --version如果输出了版本信息说明基本环境没问题。3.2 模型准备从PyTorch导出ONNXTPU-MLIR的输入是ONNX所以如果你有PyTorch模型需要先导出。导出时有几个注意事项import torch import torch.onnx model YourModel() model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, model.onnx, opset_version11, # 建议用11或13 input_names[input], output_names[output], dynamic_axesNone # TPU通常不支持动态shape固定住 )opset版本建议用11或13太新的版本TPU-MLIR可能还没适配。dynamic_axes最好设成None因为大多数自研TPU不支持动态shape固定shape能让编译器做更激进的优化。导出之后用ONNX的工具检查一下模型import onnx model onnx.load(model.onnx) onnx.checker.check_model(model) print(onnx.helper.printable_graph(model.graph))这一步能帮你发现一些低级错误比如算子不支持、shape不匹配等。3.3 编译流程一步步把ONNX变成TPU可执行文件TPU-MLIR的编译通常分几步走。以下是一个典型的流程# 第一步ONNX转MLIR tpu-mlir-opt --import-onnx model.onnx -o model.mlir # 第二步跑优化pass tpu-mlir-opt --optimize model.mlir -o model_opt.mlir # 第三步量化需要校准数据 tpu-mlir-opt --quantize --calibration-data calib/ model_opt.mlir -o model_quant.mlir # 第四步lowering到TPU dialect tpu-mlir-opt --lower-to-tpu model_quant.mlir -o model_tpu.mlir # 第五步代码生成 tpu-mlir-translate --gen-binary model_tpu.mlir -o model.bin每一步的输出都可以用文本编辑器打开看这是MLIR的一大优势——IR是可读的。我强烈建议在每一步之后都检查一下IR确认变换是否符合预期。特别是量化那一步看看scale和zero_point是否合理有没有出现全零或者异常大的scale。3.4 验证与调试怎么确认编译结果是对的编译出二进制文件只是第一步还得确认它跑出来的结果是对的。TPU-MLIR通常提供一个模拟器或者参考实现可以在CPU上模拟TPU的执行结果。验证的基本思路是用同一组输入分别跑浮点模型和编译后的模型对比输出差异。如果差异在可接受范围内比如余弦相似度大于0.99说明编译流程基本正确。如果差异很大就需要逐层排查。排查的方法是在MLIR的每一级IR上插入中间输出对比每一层的输出。TPU-MLIR支持在IR里标记某些tensor为输出这样编译后的模型会把这些中间结果也输出出来。通过逐层对比你能定位到是哪一层引入了误差。常见的误差来源包括量化参数不合理、layout转换出错、算子融合后语义变了、内存复用导致数据被覆盖。这些问题在IR层面通常都能看出来。4. 踩坑实录量化精度、内存分配与算子适配的典型问题前面讲的都是应该怎么做但实际操作中一定会遇到各种意外。这一章我分享几个我踩过的坑以及排查思路。4.1 量化后精度暴跌从校准集到逐层排查有一次我编译一个分类模型浮点模型准确率78%量化后掉到62%。第一反应是量化太激进了但调了量化配置也没用。后来逐层对比输出发现是某一层的Conv输出scale异常大导致后续量化时大量数值被截断。根因是那一层的输入数据分布有个别极端值outlier。校准集里恰好包含了这些极端值导致统计出来的最大值远大于正常值scale被拉得很大正常数值的量化精度就下降了。解决办法有两个一是换校准集去掉那些包含极端值的样本二是用KL散度校准代替min-max校准KL散度校准会自动截断一部分极端值让scale更合理。TPU-MLIR支持多种校准方法可以在配置里指定。提示校准集不需要很大但一定要有代表性。我通常从验证集里随机抽100-500张确保覆盖所有类别。如果模型有多个输入分支每个分支都要有对应的校准数据。4.2 内存不够用片上缓存的分配策略自研TPU的片上内存通常很小编译大模型时经常遇到内存不够的问题。TPU-MLIR在内存分配阶段会报错告诉你哪个tensor分配失败。遇到这种情况首先看是不是有tensor的生命周期重叠了。如果两个tensor在时间上不重叠它们可以复用同一块内存。TPU-MLIR的liveness分析通常能处理这种情况但如果IR里的依赖关系不准确分析结果就会偏保守。另一个思路是调整算子融合策略。融合得越大中间tensor越少但每个tensor占的内存越大。有时候把一个大融合拆成两个小融合反而能降低峰值内存。这个需要根据具体模型来试。如果实在放不下就只能把部分数据放到DDR。但这会引入搬运开销性能会下降。TPU-MLIR允许你指定哪些tensor可以放DDR通常把权重放DDR、激活值放SRAM是比较合理的策略。4.3 算子不支持自定义dialect的扩展方法ONNX的算子有几百个但你的TPU可能只支持其中一部分。遇到不支持的算子时TPU-MLIR会报错。这时候有几种处理方式第一种是算子替换。比如ONNX里的HardSwish如果你的TPU不支持可以用几个基础算子组合出来。TPU-MLIR的pattern rewrite机制可以自动做这种替换你只需要写一条rewrite rule。第二种是自定义算子。如果这个算子在模型里很关键组合替换性能太差那就需要在TPU dialect里定义一个新算子然后在代码生成阶段实现它。这需要你同时改编译器和硬件驱动工作量比较大。第三种是回退到CPU。如果这个算子只在模型开头或结尾出现计算量不大可以把它放到CPU上跑其余部分在TPU上跑。TPU-MLIR支持这种异构执行但需要你手动指定分割点。实际操作中我优先用第一种实在不行才考虑第二种。第三种虽然简单但异构执行的同步开销有时候比算子本身的计算量还大。5. 从编译器视角看模型部署一些反直觉的经验做了一段时间TPU-MLIR之后我积累了一些和常规认知不太一样的经验分享出来供参考。5.1 算子融合不是越多越好大多数框架的优化指南都会告诉你算子融合能减少kernel启动开销、提升性能。这在GPU上是成立的但在自研TPU上不一定。原因在于自研TPU的片上缓存通常比GPU的寄存器文件大但比GPU的L2小。融合后的算子需要把中间结果全部保留在片上如果融合后的工作集超过了片上缓存容量就会导致频繁的spill和reload性能反而下降。我的经验是融合的粒度应该和片上缓存容量匹配。具体来说融合后所有中间tensor的总大小不应该超过片上缓存的70%留30%给数据搬运的缓冲。这个比例不是绝对的需要根据你的TPU架构来调。5.2 量化不是精度越低越好INT8量化是主流但并不是所有模型都适合INT8。有些模型对数值精度非常敏感比如涉及softmax、layernorm的层INT8量化后误差会明显放大。TPU-MLIR支持混合精度你可以让大部分层用INT8少数敏感层用F16。虽然F16的计算吞吐比INT8低但整体精度能保住。实际部署时精度不达标比性能不达标更致命所以宁可牺牲一点性能也要保证精度。判断哪些层该用高精度我的方法是先全部用INT8跑一遍逐层对比输出找出误差最大的几层把它们改成F16再跑一遍。通常迭代两三次就能找到合适的混合精度配置。5.3 编译时间也是成本TPU-MLIR编译一个大模型可能需要几十分钟甚至几个小时。如果每次调参都要重新编译开发效率会非常低。我的做法是把编译流程拆成两段前端导入和优化只做一次把优化后的IR存下来量化和lowering可以反复做因为这两步才是需要调参的。这样每次调参只需要跑后半段时间能省一半以上。另外TPU-MLIR的pass支持并行执行如果你的机器核多可以在cmake里开启并行编译选项。LLVM的编译本身就支持多线程TPU-MLIR的pass也可以配置成多线程跑。5.4 版本管理比想象中重要TPU-MLIR依赖LLVM/MLIR而LLVM的API变动非常频繁。今天能编译的代码下周更新了LLVM可能就编译不过了。所以一定要把LLVM的commit号固定住不要用main分支。我的做法是在项目里维护一个llvm_commit.txt记录当前使用的LLVM commit。每次构建时从这个文件读取commit号checkout对应的版本。这样能保证构建的可复现性也方便团队协作。TPU-MLIR本身的版本也要管理。不同版本的TPU-MLIR支持的ONNX算子集可能不一样编译出来的结果也可能有差异。建议在模型部署的整个生命周期里锁定一个TPU-MLIR版本不要随意升级。6. 这条编译流水线还能怎么用TPU-MLIR虽然是为自研TPU设计的但它的架构思路其实可以复用到很多场景。比如你做的是NPU、DSP或者FPGA加速器同样面临从ONNX到自定义硬件的编译问题。TPU-MLIR的分层设计——前端导入、高层优化、量化、硬件相关lowering、代码生成——这套流程是通用的。你只需要替换掉硬件相关的dialect和pass其余部分可以复用。再比如如果你只是想在CPU上做模型优化TPU-MLIR的量化pass和算子融合pass也可以单独拿出来用。MLIR的pass是模块化的你可以只跑你需要的那些pass不需要走完整的流水线。从更宏观的角度看MLIR正在成为编译器领域的事实标准。掌握MLIR的dialect定义、pass编写、pattern rewrite这些技能不仅仅是为了用TPU-MLIR更是为了在未来的编译器开发中有更多的可能性。自研芯片越来越多对编译器的需求只会增不会减这个方向的技术积累是值得的。我在实际项目中的体会是编译器开发最难的不是写代码而是理解硬件和模型之间的语义鸿沟。TPU-MLIR提供了一套工具和框架但真正把模型跑好还是需要对硬件架构和模型结构都有深入的理解。工具能帮你做转换但做什么样的转换、怎么转换还是得靠人来决策。