拓冰网站定制 
191、MLIR与XLA（加速线性代数）的对比与融合从一次诡异的性能回退说起去年冬天，我在调试一个TPU上的BERT推理管线。模型在XLA编译下跑得挺好，延迟稳定在3.2ms。团队决定迁移到MLIR生态，用StableHLO做中间表示，再走Linalg lower到目标后端。结果一跑，延迟飙到4.7ms，整整多了47%。更诡异的是，同样的计算图，同样的硬件，只是中间表示层换了一下，性能就崩了。我盯着profile结果看了三天。发现XLA在layout优化上做了个很脏但很有效的trick：它会在HLO level就偷偷把某些reshape和transpose合并成物理内存重映射，而MLIR的StableHLO到Linalg的lowering路径上，这些优化被拆成了多个小步，每一步都引入了额外的拷贝。这就是典型的“中间表示层次越丰富，优化机会越容易被稀释”的案例。这个坑让我意识到，MLIR和XLA不是简单的替代关系，它们背后是两种完全不同的哲学。XLA的“黑盒”哲学XLA（Accelerated Linear Algebra）本质上是一个端到端的JIT编译器。你给它一个HLO（High Level Operations）图，它直接给你吐出针对特定硬件的机器码。HLO是XLA的中间表示，但XLA的优化器是高度耦合的——它知道整个pipeline的每个环节，从HLO fusion到buffer assignment到codegen，全部在一个pass pipeline里完成。这种设计的好处是极致优化。XLA的fusion p 191、MLIR与XLA（加速线性代数）的对比与融合
191、MLIR与XLA（加速线性代数）的对比与融合从一次诡异的性能回退说起去年冬天，我在调试一个TPU上的BERT推理管线。模型在XLA编译下跑得挺好，延迟稳定在3.2ms。团队决定迁移到MLIR生态，用StableHLO做中间表示，再走Linalg lower到目标后端。结果一跑，延迟飙到4.7ms，整整多了47%。更诡异的是，同样的计算图，同样的硬件，只是中间表示层换了一下，性能就崩了。我盯着profile结果看了三天。发现XLA在layout优化上做了个很脏但很有效的trick：它会在HLO level就偷偷把某些reshape和transpose合并成物理内存重映射，而MLIR的StableHLO到Linalg的lowering路径上，这些优化被拆成了多个小步，每一步都引入了额外的拷贝。这就是典型的“中间表示层次越丰富，优化机会越容易被稀释”的案例。这个坑让我意识到，MLIR和XLA不是简单的替代关系，它们背后是两种完全不同的哲学。XLA的“黑盒”哲学XLA（Accelerated Linear Algebra）本质上是一个端到端的JIT编译器。你给它一个HLO（High Level Operations）图，它直接给你吐出针对特定硬件的机器码。HLO是XLA的中间表示，但XLA的优化器是高度耦合的——它知道整个pipeline的每个环节，从HLO fusion到buffer assignment到codegen，全部在一个pass pipeline里完成。这种设计的好处是极致优化。XLA的fusion p