【MLIR】Linalg中ElementwiseOpFusion的优化模式技术分析（总）
【MLIR】Linalg中ElementwiseOpFusion的优化模式技术分析（总）
./mlir-opt -h | grep linalg
通过此命令可以查看MLIR中关于linalg的所有Pass，本篇主要分析：linalg-fuse-elementwise-ops（基于llvm 21.1.8版本）。
1. 介绍
1.1 代码介绍
linalg-fuse-elementwise-ops是 Linalg 中关于 Elementwise 类算子融合的优化 Pass。从源代码中全文检索此关键字，在mlir/include/mlir/Dialect/Linalg/Passes.td:73中找到了LinalgElementwiseOpFusionPass定义。定义非常简单，只声明了依赖的三种方言，如下：
继续检索LinalgElementwiseOpFusionPass关键字，在mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp:2284中找到了具体实现，代码如下：
该类继承自LinalgElementwiseOpFusionPassBase类，跳转进去后发现是使用 mlir-tblgen工具生成的代码，如下：
该类最重要的作用是实现了Pass机制中的虚函数runOnOperation，也就是该优化的核心功能，见上图中红框部分populate关键词开头的几个函数。
1.2 Pass核心流程图
graph LR
A[LinalgElementwiseOpFusionPass<br/>Linalg元素级操作融合Pass] --> B[runOnOperation<br/>执行优化变换]
B --> B1[populateElementwiseOpsFusionPatterns<br/>注册元素级融合模式]
B --> B2[populateFoldReshapeOpsByExpansionPatterns<br/>注册扩展融合模式]
B --> B3[populateFoldReshapeOpsByCollapsingPatterns<br/>注册折叠融合模式]
B --> B4[Canonicalization & Constant Folding<br/>规范化和常量折叠]
subgraph S1["元素级融合 - 合并连续的Generic操作"]
C1[FuseElementwiseOps<br/>匹配并融合元素级操作]
C2[FoldFillWithGenericOp<br/>将Fill操作内联到Generic]
C3[FoldScalarOrSplatConstant<br/>折叠标量/Splat常量]
C4[RemoveOutsDependency<br/>移除未使用的输出依赖]
D1{areElementwiseOpsFusable<br/>检查融合前置条件}
D2[fuseElementwiseOps<br/>执行融合变换]
E1[getPreservedProducerResults<br/>确定需要保留的Producer结果]
E2[generateFusedElementwiseOpRegion<br/>生成融合后的操作体]
E3[getIndexingMapOfProducerOperands<br/>计算融合坐标系中的索引映射]
F1[isOpOperandCanBeDropped<br/>判断操作数能否被删除]
end
subgraph S2["维度扩展融合 - 通过扩展迭代空间消除Reshape"]
G1[FoldReshapeWithGenericOpByExpansion<br/>折叠Producer的Reshape]
G2[FoldWithProducerReshapeOpByExpansion<br/>折叠Consumer的Reshape]
G3[FoldPadWithProducerReshapeOpByExpansion<br/>处理Pad+Reshape组合]
H1{isFusableWithReshapeByDimExpansion<br/>检查是否可通过扩展融合}
H2[fuseWithReshapeByExpansion<br/>执行扩展融合]
I1[ExpansionInfo::compute<br/>计算维度扩展映射]
I2[createExpandedOp<br/>创建扩展后的Linalg操作]
J1[createExpandedGenericOp<br/>创建扩展的GenericOp]
J2[createExpandedTransposeOp<br/>创建扩展的TransposeOp]
J3[updateExpandedGenericOpRegion<br/>修正扩展后的index操作]
end
subgraph S3["维度折叠融合 - 通过折叠迭代空间消除Reshape"]
K1[FoldWithProducerReshapeOpByCollapsing<br/>折叠Consumer的ExpandShape]
K2[FoldReshapeWithGenericOpByCollapsing<br/>折叠Producer的CollapseShape]
K3[FoldPadWithProducerReshapeOpByCollapsing<br/>处理Pad+ExpandShape组合]
L1[getCollapsableIterationSpaceDims<br/>计算可折叠的迭代维度]
L2[collapseOpIterationDims<br/>执行维度折叠]
M1[CollapsingInfo::initialize<br/>初始化折叠映射]
M2[createCollapsedOp<br/>创建折叠后的Linalg操作]
M3[generateCollapsedIndexingRegion<br/>通过除法和取模恢复原始索引]
N1[cloneToCollapsedOp<br/>克隆并调整操作]
N2[collapseOperandsAndResults<br/>折叠操作数和结果张量]
P1[isDimSequencePreserved<br/>检查维度序列在映射中是否保持]
end
B1 --> C1
B1 --> C2
B1 --> C3
B1 --> C4
C1 --> D1
C1 --> D2
D2 --> E1
D2 --> E2
D2 --> E3
E1 --> F1
E2 --> E3
B2 --> G1
B2 --> G2
B2 --> G3
G1 --> H1
G1 --> H2
G2 --> H1
G2 --> H2
G3 --> H2
H2 --> I1
H2 --> I2
I2 --> J1
I2 --> J2
J1 --> J3
B3 --> K1
B3 --> K2
B3 --> K3
K1 --> L1
K1 --> L2
K2 --> L1
K2 --> L2
K3 --> L2
L2 --> M1
L2 --> M2
L2 --> M3
M2 --> N1
N1 --> N2
L1 --> P1
style A fill:#e1f5fe,stroke:#01579b,stroke-width:3px
style B fill:#f3e5f5,stroke:#4a148c,stroke-width:2px
style B1 fill:#c8e6c9,stroke:#2e7d32,stroke-width:2px
style B2 fill:#ffe0b2,stroke:#e65100,stroke-width:2px
style B3 fill:#f8bbd0,stroke:#880e4f,stroke-width:2px
style C1 fill:#ffcdd2
style G1 fill:#fff9c4
style K1 fill:#e1bee7
style D2 fill:#a5d6a7,stroke:#2e7d32,stroke-width:2px
style H2 fill:#ffcc80,stroke:#e65100,stroke-width:2px
style L2 fill:#ce93d8,stroke:#880e4f,stroke-width:2px
浙公网安备 33010602011771号