欢迎光临
我们一直在努力

027、Linalg Dialect:线性代数算子的结构化表示

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 =

赞(0)
未经允许不得转载:171主机测评 » 027、Linalg Dialect:线性代数算子的结构化表示
分享到: 更多 (0)

评论 抢沙发

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址