TVM Relay 数据流模式语言（DataFlow Pattern）完全指南：从模式匹配到子图重写与划分

编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载TVM Relay 的数据流模式DataFlow Pattern简称 DFPattern是一套面向数据流图的模式匹配语言它让开发者可以用接近 Python 表达式的语法描述卷积 激活这类算子组合从而在 Relay IR 中定位、重写、划分子图。本文以 dataflow_pattern.rst 所对应的 Python API 为主体结合 python/tvm/relay/dataflow_pattern/init.py 的完整源码实现、C 侧的模式匹配器以及 test_dataflow_pattern.py 测试用例系统讲解模式语言的全部构件、组合语法、约束条件以及rewrite重写和partition子图划分两大核心工作流帮助你在自定义 Pass、算子融合和代码生成中熟练运用这套工具。一、什么是 Relay 数据流模式语言在 Relay 中做图优化如算子融合、算子替换时传统做法是用 ExprVisitor 手工遍历表达式树用大量 if/else 判断节点类型与拓扑结构代码冗长且容易出错。数据流模式语言提供了一种声明式的替代方案先用模式Pattern描述目标子图结构再让引擎在图上寻找所有匹配。模式与普通 Relay 表达式relay.Expr的区别在于模式是结构模板其中的某些位置可以留空通配符也可以附加数据类型、形状、属性等约束而表达式是具体的图。DFPattern 的基类定义在 include/tvm/relay/dataflow_pattern.h 中其_type_key为DFPatternNode所有模式节点如relay.dataflow_pattern.ExprPattern均继承自它。Python 侧的所有 API 都通过tvm._ffi._init_api(relay.dataflow_pattern, __name__)绑定到 C 实现见 python/tvm/relay/dataflow_pattern/_ffi.py而模式匹配的核心引擎位于 src/relay/ir/dataflow_matcher.cc其中DFPatternMatcher::Match负责把模式与表达式逐一比对DFPatternVisitor则在 src/relay/ir/dataflow_pattern_functor.cc 中实现对模式树的递归遍历。二、模式的基础构件语法糖函数tvm.relay.dataflow_pattern提供了一组以is_开头的语法糖函数用于快速构造各种基础模式。它们大多只是具体模式类的构造函数包装。语法糖函数底层模式类语义wildcard()WildcardPattern匹配任意表达式通配符is_var(name)VarPattern匹配 Relay 变量name为空时匹配任意 Varis_constant()ConstantPattern匹配 Relay 常量is_expr(expr)ExprPattern按结构等价匹配指定表达式is_op(op_name)ExprPattern匹配指定名字的算子如nn.conv2dis_tuple(fields)TuplePattern匹配元组fields为各字段模式is_tuple_get_item(tuple, indexNone)TupleGetItemPattern匹配取元组元素indexNone表示任意下标is_if(cond, true_branch, false_branch)IfPattern匹配If表达式三分支is_let(var, value, body)LetPattern匹配Let绑定的三个部分has_type(ttype, patternNone)TypePattern为模式附加类型约束has_dtype(dtype, patternNone)DataTypePattern为模式附加数据类型约束has_shape(shape, patternNone)ShapePattern为模式附加形状约束has_attr(attrs, patternNone)AttrPattern为模式附加属性约束dominates(parent, path, child)DominatorPattern匹配支配结构2.1 通配符与变量wildcard()是最常用的构件表示这里可以是任何表达式。is_var(name)除了能匹配特定名字的变量外还在partition流程中承担重要角色——它决定了被划分出的函数的参数如何生成。2.2 算子与表达式is_op通过tvm.relay.op.get(op_name)取得算子后包装成ExprPattern因此它匹配的是以该算子为头的 Call 节点from tvm.relay.dataflow_pattern import is_op, wildcard # 匹配任意输入的 relu 算子 relu_pat is_op(nn.relu)(wildcard())is_expr使用**结构等价structural equality**判断源码注释明确说明ExprPattern的匹配依赖表达式结构等价见 include/tvm/relay/dataflow_pattern.h。三、模式组合语法运算符重载基础模式通过DFPattern上的运算符重载进行组合这是模式语言最直观的部分Python 实现见 python/tvm/relay/dataflow_pattern/init.pyC 侧声明见 include/tvm/relay/dataflow_pattern.hpattern(*args)调用运算符构造CallPattern。args为子模式列表传入None表示匹配任意参数。注意__call__对args的处理CallPattern(self, args)中的args会被包装成数组用于匹配 Call 节点的参数列表。pattern | otherAltPattern两个模式二选一。pattern other/pattern - other/pattern * other/pattern / other分别等价于is_op(add)(self, other)、is_op(subtract)(self, other)、is_op(multiply)(self, other)、is_op(divide)(self, other)。例如测试用例 test_dataflow_pattern.py 中的模式add_pattern is_op(add)(wildcard(), wildcard())等价写法是wildcard() wildcard()。C 侧还额外提供了operator||对应 Python 的|与Optional方法。3.1 可选模式optionalDFPattern.optional(option_constructor)是AltPattern的快捷方式等价于self | option_constructor(self)。它用于描述可有可无的后续算子例如卷积后可能跟 bias 再跟 relupattern is_op(nn.conv2d)(wildcard(), wildcard()).optional( lambda x: is_op(nn.bias_add)(x, wildcard()) )测试 test_dataflow_pattern.py 展示了链式optional的写法pattern.optional(is_op(nn.relu)).optional(is_op(tanh))即卷积后面可以有 relurelu 后面还可以有 tanh两者都可缺省。3.2 支配模式dominatesDFPattern.dominates(parent, pathNone)构造DominatorPattern用于匹配单生产者、单最终消费者的链式结构parent是产出数据的节点child即self是该数据链上所有节点的最终使用者path是二者之间模糊路径通常匹配逐元素算子。若path为None默认使用wildcard()。测试 test_dataflow_pattern.py 中P is_op(nn.conv2d)(wildcard(), wildcard()) # parent I is_op(nn.relu)(wildcard()) # intermediate (path) pattern I.dominates(P) # 卷积支配 relu四、约束模式类型、数据格式、形状与属性模式可以叠加约束让匹配更加精确。所有这些约束方法都返回新的模式对象TypePattern、DataTypePattern、ShapePattern、AttrPattern因此可以链式调用。4.1 has_type类型约束DFPattern.has_type(ttype)要求被匹配的表达式具有指定的 Relay 类型tvm.ir.type.Type。底层构造TypePattern(pattern, ttype)。4.2 has_dtype数据类型约束DFPattern.has_dtype(dtype)约束张量数据类型dtype为字符串如float32。例如只匹配 float32 的卷积conv_fp32 is_op(nn.conv2d)(wildcard(), wildcard()).has_dtype(float32)4.3 has_shape形状约束DFPattern.has_shape(shape)约束张量形状shape为List[tvm.ir.PrimExpr]。例如conv_3x3 is_op(nn.conv2d)(wildcard(), wildcard()).has_shape([1, 3, 224, 224])4.4 has_attr算子属性约束DFPattern.has_attr(attrs)接受Dict[str, Object]内部通过make_node(DictAttrs, **attrs)构造属性字典然后包装成AttrPattern。注意源码注释特别说明目前只支持 Op 属性匹配不支持 Call 属性见 python/tvm/relay/dataflow_pattern/init.py。测试用例 test_dataflow_pattern.py 展示了多种属性约束# 匹配 NCHW 布局的卷积 is_conv2d is_op(nn.conv2d)(wildcard(), wildcard()).has_attr({data_layout: NCHW}) # 匹配 3x3 卷积 is_conv2d is_op(nn.conv2d)(wildcard(), wildcard()).has_attr({kernel_size: [3, 3]})属性值支持标量字符串、数字与数组如[3, 3]。在 C 匹配端MatchRetValue会递归展开数组并对运行时类型如runtime::Int与编译期 IR 类型做自动转换保证((0,0),(0,0))这类嵌套结构也能正确比较见 src/relay/ir/dataflow_matcher.cc。五、核心操作一match 模式匹配match(pattern, expr)返回布尔值判断表达式是否命中模式。DFPattern.match(expr)是它的实例方法形式最终都调用 FFI 函数ffi.match见 python/tvm/relay/dataflow_pattern/init.py。C 侧 src/relay/ir/dataflow_matcher.cc 的Match实现要点每次匹配前清空memo_记忆化缓存与matched_nodes_记录VisitDFPattern采用带回溯的记忆化匹配命中则把 (pattern, expr) 记入 memo未命中则通过ClearMap回滚到匹配前的水位线避免分支匹配失败污染后续结果匹配成功后matched_nodes_中保存了 pattern 与 expr 的对应关系供后续重写/划分使用。from tvm.relay.dataflow_pattern import is_op, wildcard, match pat is_op(add)(wildcard(), wildcard()) x, y relay.var(x), relay.var(y) assert match(pat, x y) # True assert not match(pat, x * y) # False六、核心操作二rewrite 模式重写6.1 DFPatternCallback 回调模式重写通过继承DFPatternCallback实现python/tvm/relay/dataflow_pattern/init.pyrequire_type回调前是否需要先执行类型推断InferTyperewrite_once为True时只执行一次回调pattern子类必须提供要匹配的模式callback(self, pre, post, node_map)命中时被调用三个参数分别是——pre为原始图中命中的表达式post为输入已被改写后的表达式node_map为tvm.ir.container.Map[DFPattern, List[Expr]]保存模式节点到命中的表达式列表的映射。回调返回值将替换匹配到的子图。测试 test_dataflow_pattern.py 中的经典示例——把加法改写成减法class TestRewrite(DFPatternCallback): def __init__(self): super(TestRewrite, self).__init__() self.pattern add_pattern def callback(self, pre, post, node_map): return post.args[0] - post.args[1] out rewrite(TestRewrite(), x y) assert sub_pattern.match(out)6.2 rewrite 函数rewrite(callbacks, expr, modNone)接受单个回调或回调列表对表达式执行重写python/tvm/relay/dataflow_pattern/init.py。实现细节mod参数可选用于关联 IRModule若未提供则内部创建空IRModule()回调在传入 C 前会被包装为_DFPatternCallback即 C 侧DFPatternCallback对象并断言pattern非空返回重写后的表达式。test_rewrite_functest_dataflow_pattern.py还展示了重写发生在函数调用内部的场景——对func(x, w) y执行重写后函数体中的加法同样被改写。七、核心操作三partition 子图划分partition(pattern, expr, attrsNone, checkNone)是数据流模式最强大的能力把表达式中所有命中模式的子图提取为独立的 Relay 函数并用对该函数的调用替换原子图python/tvm/relay/dataflow_pattern/init.py。DFPattern.partition(expr, attrs, check)是其实例方法形式。参数说明attrsOptional[Dict[str, Object]]添加到被划分函数上的属性字典checkCallable[[Expr], bool]对命中的表达式做更复杂检查的回调返回True才继续划分默认恒为True。7.1 实战划分 BatchNorm测试 test_partition_batchnorm 展示了完整的 BatchNorm 划分流程BN gamma * (x - mean) / relay.op.sqrt(var eps) beta # 定义匹配 BN 计算链的模式 class BatchnormCallback(DFPatternCallback): def __init__(self): super(BatchnormCallback, self).__init__() self.pattern ( is_op(add)( is_op(divide)( is_op(multiply)(wildcard(), is_op(subtract)(wildcard(), wildcard())), is_op(sqrt)(is_op(add)(wildcard(), is_constant())), ), wildcard(), ) ) partitioned BatchnormCallback().pattern.partition(BN)划分结果是一个带PartitionedFromPattern属性的函数调用属性值形如subtract_multiply_add_sqrt_divide_add_记录模式的操作序列这正是下游代码生成器如外部代码生成后端识别子图来源的依据。test_partition_double_batchnormtest_dataflow_pattern.py验证了嵌套 BatchNorm 会被划分成两个嵌套的函数调用test_partition_check第 1545 行附近则验证check回调可以阻止不符合额外条件的匹配被划分。7.2 划分与融合的关系partition是 Relay 中算子融合如FuseOps和外部代码生成BYOC的基础设施先定义哪些算子组合可以融合/下放再通过模式划分把它们聚合成函数最后交给后端。测试 test_partition_overused 与test_partition_fuzzy_tuple第 1489 行等用例覆盖了变量被多处复用、元组/函数参数等复杂拓扑下的划分行为。八、模式树遍历与 C 实现原理理解了 Python API 之后再看底层实现能更好地把握模式语言的能力边界。8.1 模式节点体系所有模式类在 Python 侧通过register_df_node注册到 C 对象系统python/tvm/relay/dataflow_pattern/init.py其 C 类型 key 均为relay.dataflow_pattern.*。完整的模式节点家族定义在 include/tvm/relay/dataflow_pattern.h 中与 Python 一一对应ExprPattern匹配字面表达式结构等价比较VarPattern匹配变量可选名字ConstantPattern匹配常量CallPattern匹配 Call 节点算子 参数列表args可为空表示任意参数FunctionPattern匹配函数参数列表 函数体IfPattern/LetPattern匹配控制流结构TuplePattern/TupleGetItemPattern匹配元组及其取元素操作AltPattern二选一WildcardPattern通配TypePattern/DataTypePattern/ShapePattern/AttrPattern四类约束DominatorPattern支配结构8.2 匹配器与访问器DFPatternMatchersrc/relay/ir/dataflow_matcher.cc负责模式与表达式的比对。核心是带 memo 的回溯匹配并维护node_map供重写/划分使用。AltPattern的匹配即left || right短路求值第 71-73 行。DFPatternVisitorsrc/relay/ir/dataflow_pattern_functor.cc对模式树本身做递归遍历用visited_集合去重防止环。例如CallPattern会依次访问其算子模式与每个参数模式第 46-53 行。匹配结果中模式节点 → 命中的表达式列表的映射通过matched_nodes_累积最终暴露给 Python 回调的node_map参数。九、速查与最佳实践9.1 常用组合速查需求写法匹配任意卷积is_op(nn.conv2d)(wildcard(), wildcard())卷积后必有 ReLUis_op(nn.relu)(is_op(nn.conv2d)(wildcard(), wildcard()))卷积后可有可无 ReLUis_op(nn.conv2d)(wildcard(), wildcard()).optional(is_op(nn.relu))ReLU 或 LeakyReLUis_op(nn.relu)(wildcard()) \| is_op(nn.leaky_relu)(wildcard())权重为常量的 Denseis_op(nn.dense)(wildcard(), is_constant())限定数据布局is_op(nn.conv2d)(wildcard(), wildcard()).has_attr({data_layout: NCHW})限定算子类型标记is_op(nn.dense).has_attr({TOpPattern: K_ELEMWISE})9.2 实践建议先用match验证模式在写重写/划分逻辑前用pattern.match(expr)快速验证模式是否符合预期拓扑。善用is_var而非wildcard划分参数partition以模式中的VarPattern为依据生成被划分函数的参数需要哪些输入应成为函数参数时用is_var精确控制。属性约束注意类型has_attr只支持 Op 属性数组属性如kernel_size在 C 端会做递归展开比较可放心使用嵌套结构。重写时注意rewrite_once与require_type需要类型信息做判断时设require_typeTrue只想改第一处匹配时设rewrite_onceTrue。node_map用于获取匹配子表达式回调中通过node_map[pattern]取出命中的原始表达式例如拿到卷积的输入/权重做进一步分析。十、参考资料与延伸阅读API 参考入口docs/reference/api/python/relay/dataflow_pattern.rstPython 实现本文主体python/tvm/relay/dataflow_pattern/init.pyFFI 绑定python/tvm/relay/dataflow_pattern/_ffi.pyC 模式节点定义include/tvm/relay/dataflow_pattern.h匹配引擎src/relay/ir/dataflow_matcher.cc 与 src/relay/ir/dataflow_matcher_impl.h模式树访问器src/relay/ir/dataflow_pattern_functor.cc单元测试tests/python/relay/test_dataflow_pattern.py覆盖匹配、约束、optional、dominates、rewrite、partition 全场景C 测试tests/cpp/dataflow_pattern_test.ccRelax 侧同类实现供对照include/tvm/relax/dataflow_pattern.h 与 tests/python/relax/test_dataflow_pattern.py赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐TVM Relay DataFlow Pattern 模式匹配语言从图重写到算子融合的声明式框架TVM Relay DataFlow Pattern 模式匹配语言从图重写到算子融合的声明式框架 本文围绕 TVM 中 Relay 数据流图的模式匹配语言D编译器深度学习模型优化Apache TVM Relax 数据流模式语言DPL完全指南从图匹配到算子融合与后端分发Apache TVM Relax 数据流模式语言DPL完全指南从图匹配到算子融合与后端分发 本指南围绕 Apache TVM 中 Relax 前端内置的数模型编译深度学习推理引擎Apache TVM Relax 数据流模式语言DPL实战指南图模式匹配与自动改写Apache TVM Relax 数据流模式语言DPL实战指南图模式匹配与自动改写 导读 tvm.relax.dpl Dataflow Pattern模型编译深度学习推理引擎创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考