在 MLIR 中,TableGen (
.td
文件) 可用于定义简化(重写)模式(patterns),这使得编写和管理复杂的模式匹配规则更加简洁和易于维护。TableGen 提供了一种声明性语法来定义模式,用于匹配 MLIR 操作并将其重写为新的操作。1. 定义模式 (Pattern) 的基础结构典型的 .td 文件结构2. 定义重写模式示例:简单加法结合律重写3. 自定义操作的重写4. 运用模式组合多个操作5. 应用模式生成 .inc 文件:在 C++ 中使用生成的模式:6. 常用模板和规则7. 自定义模式的参数化8. 优势
下面是如何使用
.td
文件定义 MLIR 重写模式的详细步骤和示例:1. 定义模式 (Pattern
) 的基础结构
MLIR 的
.td
文件通过 TableGen 的 Pattern
机制,定义从一种操作到另一种操作的转换规则。这些模式可以在 pass 中应用,用于优化或转换 IR。典型的 .td
文件结构
// 包含一些基本的模式定义和宏 include "mlir/IR/PatternBase.td" // 引入我们想要匹配的操作的 ODS 定义 include "MyDialect/MyDialectOps.td"
2. 定义重写模式
在
.td
文件中,使用 Pattern
声明模式,它指定了要匹配的操作和如何将其转换为新操作。示例:简单加法结合律重写
假设我们有一个
AddOp
,我们希望通过模式匹配将 (a + 0)
转换为 a
。def SimplifyAddZero : Pat< (AddOp $op, $a, 0), // 匹配 AddOp,其中一个操作数是0 (replaceWithValue $a)>; // 将其替换为另一个操作数
AddOp
是一个我们要匹配的操作。
$a
是一个通配符,用于表示加法的一个操作数。
0
是我们想要匹配的常量值。
replaceWithValue $a
是重写部分,表示匹配后用$a
进行替换。
3. 自定义操作的重写
你可以为自定义操作定义匹配模式。例如,假设你有一个
MulOp
,你希望简化乘以 1 的情况:def SimplifyMulOne : Pat< (MulOp $op, $a, 1), // 匹配乘法操作 MulOp,其中一个操作数是1 (replaceWithValue $a)>; // 将其简化为另一个操作数
4. 运用模式组合多个操作
你也可以通过组合多个操作,来进行更复杂的重写。
def CombineAddMul : Pat< (AddOp (MulOp $a, $b), (MulOp $c, $d)), // 匹配 Add(Mul(a, b), Mul(c, d)) (MulOp (AddOp $a, $c), (AddOp $b, $d)) // 转换为 Mul(Add(a, c), Add(b, d)) >;
5. 应用模式
定义好
.td
文件后,MLIR 将生成 C++ 代码,可以将其用于 RewritePatternSet
中来应用这些模式。在 C++ 中,可以将这些模式应用到 PatternRewriteDriver
中:生成 .inc
文件:
- 在 CMake 构建系统中,使用
tablegen
来生成模式的 C++ 实现:
mlir_tablegen(MyPatterns.td -gen-rewriters -outdir ${CMAKE_CURRENT_BINARY_DIR}) include_directories(${CMAKE_CURREN_BINARY_DIR})
- 生成的
.inc
文件包含 TableGen 定义的重写模式,可以在 C++ 中包含它们:
#include "MyDialect/MyPatterns.inc"
在 C++ 中使用生成的模式:
void populateWithGeneratedPatterns(RewritePatternSet &patterns); void runOnFunction() override { RewritePatternSet patterns(&getContext()); populateWithGeneratedPatterns(patterns); // 添加生成的模式 applyPatternsAndFoldGreedily(getFunction(), std::move(patterns)); // 应用模式 }
6. 常用模板和规则
- Pat:用于定义重写模式。
- Dag:用于描述操作的树结构。
- replaceWithValue:用于指定替换时应返回的值。
- replaceWithOp:用于用新的操作替换匹配的操作。
- MatchDagOp:用于匹配操作并提取值。
7. 自定义模式的参数化
有时,你可能需要创建更具通用性的模式。可以通过参数化的方式定义模式:
class SimplifyAddWithConstant<ValueType> : Pat< (AddOp $a, ValueType:$const), (replaceWithValue $a) >;
你可以通过不同类型的常量实例化这个模式。
8. 优势
- 声明性:TableGen 提供了声明性方式来定义复杂的模式匹配和转换规则,比直接编写 C++ 重写模式更简洁。
- 可维护性:集中定义模式,便于管理和扩展新的模式。
- 生成代码:TableGen 自动生成 C++ 代码,减少手动编写模式匹配逻辑的错误。