027、Linalg Dialect:线性代数算子的结构化表示
从一次矩阵乘法调试说起
上周帮同事定位一个模型推理性能问题。他手写了一个自定义算子,用MLIR的Standard Dialect表达矩阵乘法,结果在A100上跑出来的吞吐量只有cuBLAS的1/3。我让他把IR dump出来一看——好家伙,整个matmul被拆成了几十个load、multiply、add、store的循环,LLVM后端根本认不出这是个矩阵乘法,自然没法触发tensor core的mma指令。
这就是典型的“算子碎片化”问题。你明明想表达一个高层次的线性代数操作,但Standard Dialect只认识标量和简单向量操作,所有语义都被降级成了底层指令序列。Linalg Dialect就是来解决这个问题的——它让你能保留“这是一个矩阵乘法”这个高层次语义,直到最后一步才决定是映射到tensor core还是SIMD指令。
Linalg的核心思想:结构化迭代
Linalg的全称是“Linear Algebra Dialect”,但它不是简单地把BLAS接口搬进MLIR。它的核心抽象是结构化迭代器——用一组迭代域(iteration domains)和访问模式(access patterns)来描述一个计算。
看一个最简单的例子,两个向量相加:
#map0 = affine_map<(i) -> (i)>
#map1 =



