071、使用MLIR实现一个简单的矩阵乘法（MatMul）算子从一次诡异的性能回退说起上个月调试一个AI推理引擎，发现同样的MatMul计算，在LLVM后端生成的代码比手写汇编慢了40%。排查到最后，问题出在MLIR的linalg.generic op上——我天真地以为把矩阵乘法映射成通用循环就能自动获得最佳优化。结果LLVM把循环展开成了灾难性的寄存器溢出。这个教训让我意识到：MLIR的MatMul实现，远不是写几行循环那么简单。先搭骨架：定义我们的MatMul Op在MLIR中实现算子，第一步是定义Operation。我习惯从最简单的开始——一个接受两个输入矩阵、输出一个结果矩阵的MatMul op。// 定义在自定义Dialect中，这里用Toy Dialect做演示 // 注意：实际生产环境建议用linalg或tosa，但为了理解原理，我们手搓一个 def Toy_MatMulOp : Toy_Op"matmul" { let arguments = (ins RankedTensorOf[F64]:$lhs, // 左矩阵，形状[M, K] RankedTensorOf[F64]:$rhs // 右矩阵，形状[K, N] ); let results = (outs RankedTensorOf[F64]:$result // 结果矩阵，形状[M 071、使用MLIR实现一个简单的矩阵乘法（MatMul）算子
071、使用MLIR实现一个简单的矩阵乘法（MatMul）算子从一次诡异的性能回退说起上个月调试一个AI推理引擎，发现同样的MatMul计算，在LLVM后端生成的代码比手写汇编慢了40%。排查到最后，问题出在MLIR的linalg.generic op上——我天真地以为把矩阵乘法映射成通用循环就能自动获得最佳优化。结果LLVM把循环展开成了灾难性的寄存器溢出。这个教训让我意识到：MLIR的MatMul实现，远不是写几行循环那么简单。先搭骨架：定义我们的MatMul Op在MLIR中实现算子，第一步是定义Operation。我习惯从最简单的开始——一个接受两个输入矩阵、输出一个结果矩阵的MatMul op。// 定义在自定义Dialect中，这里用Toy Dialect做演示 // 注意：实际生产环境建议用linalg或tosa，但为了理解原理，我们手搓一个 def Toy_MatMulOp : Toy_Op"matmul" { let arguments = (ins RankedTensorOf[F64]:$lhs, // 左矩阵，形状[M, K] RankedTensorOf[F64]:$rhs // 右矩阵，形状[K, N] ); let results = (outs RankedTensorOf[F64]:$result // 结果矩阵，形状[M RELATED READING
延伸阅读
更多一线实战笔记与深度复盘，助您持续精进
