1. 为什么要在 Relay 里加算子从一次真实需求说起深度学习编译器这个领域Relay 作为 TVM 的高层中间表示承担着把前端框架的计算图翻译成可优化、可调度 IR 的核心职责。日常做模型部署的时候绝大多数算子 Relay 的算子库都已经覆盖了卷积、矩阵乘、各类激活函数、归一化、池化这些常规操作基本都能直接找到对应实现。但真正落到项目里总会遇到一些“标准库里没有”的情况——可能是某个论文里新提出的激活函数可能是业务侧自定义的融合逻辑也可能是某个冷门框架导出的算子 Relay 没有对应解析。这时候绕不过去的一件事就是自己给 Relay 添加一个新算子。这篇文章面向的读者是已经跑通过 TVM 基础流程、能看懂 Relay IR、知道relay.op大概长什么样但还没完整走过“注册一个新算子”全流程的人。我会把从算子定义、属性注册、类型关系推导、到 lowering 落地、再到 Python 侧调用和测试的整条链路拆开讲清楚。整个过程涉及 C 和 Python 两侧的配合涉及 TVM 的注册机制、类型系统、以及 Relay 的 pass 体系信息量不小但走通一遍之后你会发现它的套路其实非常固定。先说清楚一个前提给 Relay 加算子本质上是在做三件事。第一告诉 Relay 这个算子长什么样——它的名字、输入输出、有哪些属性参数第二告诉 Relay 这个算子的类型规则——给定输入类型输出类型怎么推第三告诉 Relay 这个算子最终怎么算——也就是 lowering 到更低层的实现。这三件事分别对应算子注册、类型关系TypeRelation、以及计算实现。很多人卡住不是因为某一步特别难而是因为这三件事分散在不同文件、不同注册宏里第一次做的时候容易顾此失彼。我个人的建议是动手之前先想清楚你的算子属于哪一类。如果它可以用现有 Relay 算子组合表达出来那优先考虑用relay.op里的组合方式实现而不是新增算子——组合方式不需要碰 C维护成本低得多。只有当组合方式表达不了、或者性能上必须下沉到 C 甚至要自定义 schedule 的时候才值得走完整的新增算子流程。这个判断很关键我见过不少人一上来就加算子结果发现用te.compute加几个现有 op 就能搞定白白增加了维护负担。下面我按实际开发顺序把整个流程拆成几个部分来讲。每一部分我都会说明“为什么这么做”而不只是“怎么做”因为 TVM 的注册机制设计得比较绕不理解意图的话照抄代码很容易在细节上翻车。2. 动手前的整体设计新算子该放在哪一层2.1 先判断组合实现还是新增算子在决定新增算子之前有一个必须做的判断这个算子能不能用 Relay 已有的算子拼出来。举个常见的例子假设你要实现一个swish激活函数它的定义是x * sigmoid(x)。这个完全可以用relay.multiply(x, relay.sigmoid(x))表达不需要新增任何算子。再比如某些自定义的归一化如果能拆成 reduce、subtract、divide 的组合也优先走组合路线。组合路线的优势非常明显不碰 C不需要重新编译 TVMPython 侧直接就能用而且能自动享受 Relay 现有的融合、常量折叠等优化。缺点是如果组合出来的子图比较碎可能会影响融合效果或者在某些后端上性能不理想。什么时候必须新增算子我总结了几种典型情况。第一种是算子语义无法用现有算子组合表达比如涉及特殊的内存布局或者不可分解的数学运算。第二种是性能敏感组合实现会产生大量中间张量而新增算子可以在一个 compute 里完成。第三种是需要自定义 schedule比如你想对这个算子做特殊的循环切分和向量化组合实现没法控制底层调度。第四种是前端框架导出的模型里带了这个算子你需要让 Relay 能解析它。判断清楚之后如果确定要新增接下来要决定算子的“层级”。Relay 里算子可以注册在几个不同的层次最上层是 Relay oprelay.op命名空间它直接面向计算图往下是 TOPITVM Operator Inventory提供具体的 compute 和 schedule 实现再往下是 TETensor Expression描述张量计算。一个完整的算子通常需要在 Relay 层注册 op在 TOPI 层提供 compute 实现必要时还要写 schedule。2.2 文件组织新增算子该改哪些文件TVM 的代码组织比较清晰新增一个算子通常涉及以下几个位置。C 侧算子注册一般放在src/relay/op/下的对应子目录里比如张量运算放tensor/神经网络相关放nn/。类型关系通常和算子注册放在一起或者单独放在op/op_common.h之类的公共头文件里。TOPI 的 compute 实现放在python/tvm/topi/下对应的模块里比如topi/nn/、topi/math/。Python 侧的封装放在python/tvm/relay/op/下。这里有个经验不要一上来就想着把算子加到最合适的位置先把流程跑通再考虑代码放哪。我见过有人花大量时间纠结目录结构结果算子还没跑起来。正确的做法是先在一个临时文件里把注册、类型关系、compute 都写出来编译通过、测试通过之后再按 TVM 的代码规范挪到合适的位置。另外要提醒的是TVM 的注册宏是全局的算子名字在整个 Relay 命名空间里必须唯一。命名的时候建议遵循现有算子的命名习惯比如nn.contrib_xxx、math.xxx这种带命名空间前缀的形式避免和已有算子冲突。如果你只是本地实验可以用一个不太可能冲突的前缀比如myop.xxx。2.3 注册机制的核心三个宏各管什么TVM 里新增算子最核心的是三个注册宏理解它们各自的作用是走通流程的关键。第一个是RELAY_REGISTER_OP它注册算子的“元信息”算子名字、参数描述、类型关系、以及一些属性。这个宏决定了 Relay 在解析和类型推导阶段怎么看待这个算子。它里面会指定set_attrs_type、add_argument、set_support_level、set_attr等。第二个是TVM_REGISTER_GLOBAL它注册一个全局函数通常用来把算子的 compute 实现暴露给 Relay 的 lowering 阶段。Relay 在把 op 转成 TE 的时候会通过这个全局函数名去调用对应的 compute。第三个是类型关系相关的注册比如RelayExpr TypeRel或者用TypeRelation的子类。它负责在类型推导阶段计算输出类型。对于大多数逐元素算子可以直接用现成的BroadcastRel或者ElemwiseRel对于形状有特殊变化的算子就需要自己写一个 TypeRelation。这三个宏的关系可以这样理解RELAY_REGISTER_OP是“身份证”告诉系统有这么个算子类型关系是“规则书”告诉系统输入输出类型怎么对应TVM_REGISTER_GLOBAL是“施工队”告诉系统实际怎么算。三者缺一不可而且名字要对应上否则会在不同阶段报错。3. 核心细节拆解算子注册与类型关系3.1 用 RELAY_REGISTER_OP 描述算子元信息先看算子注册。一个典型的注册代码长这样以我实际做过的一个自定义激活算子为例这里用虚构的myop.leaky_clamp代称RELAY_REGISTER_OP(myop.leaky_clamp) .describe(Rcode(Leaky clamp activation. Computes max(min(x, upper), lower) with a leaky slope for the lower part. )code TVM_ADD_FILELINE) .set_num_inputs(1) .add_argument(data, Tensor, The input tensor.) .set_support_level(10) .set_attrs_typeLeakyClampAttrs() .add_type_rel(LeakyClamp, LeakyClampRel());这里逐项说明。describe是给文档和调试用的描述写清楚算子语义方便别人包括未来的自己理解。set_num_inputs(1)声明输入个数注意这里只算数据输入属性参数不算在内。add_argument描述每个输入的名字和类型这个名字会出现在错误信息和调试输出里。set_support_level是 TVM 内部的一个标记数值越大表示支持程度越低自定义算子一般给个较大的值比如 10表示“非官方支持”。set_attrs_type指定属性结构体类型如果你的算子有额外参数比如上面的lower、upper、slope就需要定义一个 Attrs 结构体。add_type_rel绑定类型关系第一个参数是类型关系的名字第二个是实例。属性结构体通常这样定义struct LeakyClampAttrs : public tvm::AttrsNodeLeakyClampAttrs { double lower; double upper; double slope; TVM_DECLARE_ATTRS(LeakyClampAttrs, relay.attrs.LeakyClampAttrs) { TVM_ATTR_FIELD(lower).set_default(-1.0).describe(Lower bound.); TVM_ATTR_FIELD(upper).set_default(1.0).describe(Upper bound.); TVM_ATTR_FIELD(slope).set_default(0.01).describe(Leaky slope.); } };这个结构体继承AttrsNode用TVM_DECLARE_ATTRS声明字段每个字段用TVM_ATTR_FIELD注册并可以设默认值。属性结构体的好处是 Relay 会自动帮你做参数解析、类型检查、以及序列化。定义好之后Python 侧调用时传的关键字参数会自动映射到这些字段。这里有个容易踩的坑属性字段的类型要选对。TVM 支持int、double、bool、String、Array等类型但如果你用了不支持的复杂类型编译会报错。另外属性结构体的名字上面那个字符串relay.attrs.LeakyClampAttrs要全局唯一建议带上算子名前缀。3.2 类型关系输出类型到底怎么推类型关系是很多人第一次做时最容易迷糊的地方。它的职责是给定输入表达式的类型推导出输出表达式的类型。对于 Relay 来说类型主要是TensorType包含 shape 和 dtype。最简单的做法是复用现成的类型关系。如果你的算子是逐元素的输出 shape 和输入完全一致直接用ElemwiseRel就行。如果涉及广播比如两个不同 shape 的输入做运算用BroadcastRel。这两个现成的类型关系覆盖了相当一部分场景。但有些算子形状会变比如 reduce 类算子输出维度会减少或者像 reshape 这种输出 shape 由属性决定。这时候就需要自己写 TypeRelation。一个自定义 TypeRelation 的骨架class LeakyClampRel : public TypeRelation { public: LeakyClampRel() : TypeRelation(nullptr, LeakyClampRel) {} TVM_DLL static ArrayType LeakyClampRelFunc( const ArrayType types, int num_inputs, const Attrs attrs, const TypeReporter reporter) { ICHECK_EQ(types.size(), 2); const auto* data types[0].asTensorTypeNode(); if (data nullptr) return types; reporter-Assign(types[1], TensorType(data-shape,>tvm.script.register_object def leaky_clamp(data, lower-1.0, upper1.0, slope0.01): return _make.leaky_clamp(data, lower, upper, slope)这里的_make.leaky_clamp是 TVM 根据 C 注册自动生成的参数顺序要和 C 侧add_argument加属性字段的顺序对应。实际上 TVM 的 FFI 机制会自动把 Python 调用转成 C 的Call节点属性会打包成 Attrs。更规范的做法是用tvm.relay.op的注册装饰器把算子注册到relay.op命名空间下这样用户可以用relay.op.myop.leaky_clamp(...)调用。同时建议加上类型检查比如检查输入是不是 Tensor属性值是否在合理范围内。Python 侧还有一个重要工作是注册算子的“梯度”或者“策略”如果你的算子需要参与自动微分或者有特定的算子融合策略。梯度注册用register_gradient融合策略用register_fusion_pattern。这些不是必须的但如果你希望新算子能融入 Relay 的优化流程最好补上。4. 实操过程从 compute 到 lowering 的完整链路4.1 用 TOPI 写 compute 实现算子注册和类型关系搞定之后接下来是实际的计算实现。Relay 的 lowering 会把 op 转成 TE而 TE 的 compute 通常由 TOPI 提供。TOPI 的 compute 函数接收输入张量和属性返回一个 TE 的Tensor。继续用leaky_clamp举例TOPI 实现大概是这样def leaky_clamp(data, lower-1.0, upper1.0, slope0.01): def _compute(ins): x ins[0] return te.compute( x.shape, lambda *idx: te.max( te.min(x[idx], tvm.tir.const(upper, x.dtype)), tvm.tir.const(lower, x.dtype) ) tvm.tir.const(slope, x.dtype) * te.min( x[idx] - tvm.tir.const(lower, x.dtype), 0.0 ), nameleaky_clamp, tagelemwise ) return _compute这里有几个细节值得说。第一te.compute的第一个参数是输出 shape逐元素算子直接用输入 shape。第二lambda 里的idx是索引变量x[idx]取出对应位置的元素。第三常量要用tvm.tir.const包装并指定 dtype否则可能因为类型不匹配报错。第四tag参数很重要它标记了算子的类型比如elemwise、broadcast、reduceRelay 的融合 pass 会根据 tag 决定能不能把这个算子和相邻算子融合。第五name参数会出现在生成的代码里取个有意义的名字方便调试。写完 compute 之后需要用TVM_REGISTER_GLOBAL把它注册成全局函数函数名要和 Relay 侧 lowering 时查找的名字一致。通常的约定是topi.myop.leaky_clamp这种形式。Relay 在 lowering 时会根据算子的名字拼出这个全局函数名去调用。这里有个非常容易踩的坑compute 函数的参数顺序。它必须和 Relay 侧add_argument加属性字段的顺序完全一致否则运行时会因为参数错位产生莫名其妙的错误。我建议在写的时候把参数列表写下来对照不要凭记忆。4.2 注册 compute 与 lowering 对接compute 写好后注册代码通常长这样TVM_REGISTER_GLOBAL(topi.myop.leaky_clamp) .set_body([](TVMArgs args, TVMRetValue* rv) { *rv LeakyClampCompute(args[0], args[1], args[2], args[3]); });或者更常见的是在 Python 侧用tvm.register_func注册tvm.register_func(topi.myop.leaky_clamp) def _leaky_clamp_compute(data, lower, upper, slope): return leaky_clamp(data, lower, upper, slope)两种方式都可以C 侧注册性能稍好Python 侧注册更灵活、改起来不用重新编译。开发阶段我强烈建议先用 Python 注册因为改一行代码就能生效不用等 C 编译。等逻辑稳定了再考虑挪到 C。lowering 对接的关键是 Relay 怎么知道去调用哪个全局函数。Relay 的OpImplementation或者FTVMCompute属性会指定 compute 函数名。在RELAY_REGISTER_OP里可以用set_attrFTVMCompute(FTVMCompute, ...)来指定或者依赖默认的命名约定。默认约定是算子名加前缀比如myop.leaky_clamp对应topi.myop.leaky_clamp。如果你用了非默认命名一定要显式设置FTVMCompute。4.3 编译、测试与验证代码写完之后编译 TVM。如果你只改了 Python 侧不需要重新编译直接跑测试就行。如果改了 C 侧需要重新编译这个过程可能比较久建议用ninja增量编译。测试分几个层次。第一层是算子能不能被正确解析和类型推导写个简单的 Relay 程序用relay.Function构造调用跑relay.transform.InferType看输出类型对不对。第二层是数值正确性构造随机输入用新算子算一遍再用 NumPy 手动实现算一遍对比结果。第三层是端到端把算子放进一个完整模型编译到目标后端跑推理看结果。数值测试的代码大概这样import numpy as np import tvm from tvm import relay data relay.var(data, relay.TensorType((2, 3), float32)) out relay.op.myop.leaky_clamp(data, lower-1.0, upper1.0, slope0.01) func relay.Function([data], out) mod tvm.IRModule.from_expr(func) mod relay.transform.InferType()(mod) x np.random.randn(2, 3).astype(float32) expected np.maximum(np.minimum(x, 1.0), -1.0) 0.01 * np.minimum(x 1.0, 0.0) with tvm.transform.PassContext(opt_level3): lib relay.build(mod, targetllvm) runtime tvm.contrib.graph_executor.GraphModule(lib[default](tvm.cpu())) runtime.set_input(data, x) runtime.run() result runtime.get_output(0).numpy() np.testing.assert_allclose(result, expected, rtol1e-5)这个测试覆盖了类型推导、lowering、编译、执行全流程。如果这一步过了说明算子基本可用。如果没过错误信息通常会指向某个具体阶段按阶段排查即可。5. 常见问题与排查技巧实录5.1 编译期报错注册相关的典型问题新增算子时编译期报错主要集中在注册环节。最常见的是“算子未注册”或者“找不到类型关系”。前者通常是RELAY_REGISTER_OP的名字和调用时用的名字不一致注意 Relay 里算子名是大小写敏感的而且命名空间前缀要完全匹配。后者是add_type_rel绑定的名字和实际 TypeRelation 注册的名字不一致。另一个常见问题是属性结构体没注册。如果你定义了LeakyClampAttrs但忘了用TVM_REGISTER_NODE_TYPE注册编译时可能不报错但运行时解析属性会失败。注册代码是TVM_REGISTER_NODE_TYPE(LeakyClampAttrs);这个宏要放在属性结构体定义之后。我踩过一次坑属性结构体定义在头文件里注册宏忘了加结果算子能注册成功但一调用就崩排查了半天才发现是属性没注册。还有一种情况是set_num_inputs和实际add_argument的次数不匹配。set_num_inputs声明的是数据输入个数add_argument也要调用相同次数。如果声明了 2 个输入但只add_argument了 1 次类型推导时types数组长度会对不上报数组越界的错。5.2 运行期报错类型推导与 lowering 的坑运行期报错里类型推导相关的最多。典型症状是Check failed: reporter-Assign之类的错误通常是因为输出类型和输入类型不兼容或者Assign被调用了多次。检查你的 TypeRelation确保每个输出只Assign一次且 dtype 和 shape 都合理。lowering 阶段的报错通常是“找不到 compute 函数”。这多半是全局函数名和 Relay 期望的名字不一致。排查方法是打印出 Relay 期望的函数名或者直接在RELAY_REGISTER_OP里显式设置FTVMCompute属性。另外compute 函数返回的必须是 TE 的 Tensor 或者 Tensor 列表如果返回了别的东西lowering 会失败。数值不对但能跑通的情况通常是 compute 实现里的索引或者常量写错了。比如te.min和te.max的参数顺序搞反或者常量 dtype 和输入 dtype 不一致导致隐式类型转换。这类问题建议先用小 shape 的输入手动验证比如(1, 1)或者(2, 2)把中间结果打印出来对比。5.3 常见问题速查表问题现象可能原因排查方向算子未注册名字不匹配或注册宏未执行检查RELAY_REGISTER_OP名字与调用名类型关系找不到add_type_rel名字与注册名不一致核对 TypeRelation 注册字符串属性解析失败属性结构体未注册补TVM_REGISTER_NODE_TYPE数组越界set_num_inputs与add_argument次数不符核对输入个数声明找不到 compute全局函数名不匹配显式设置FTVMCompute或核对命名数值错误compute 索引或常量写错小 shape 手动验证中间结果融合失效tag 设置不当检查 compute 的 tag 参数5.4 独家避坑经验第一条经验开发阶段全部用 Python 注册 compute不要碰 C。C 编译一次动辄十几分钟改一个符号就要重编效率极低。等 Python 版本验证通过再考虑要不要挪到 C。实际上很多自定义算子用 Python 注册完全够用性能差异在大多数场景下可以忽略。第二条经验类型关系先用现成的实在不行再自己写。ElemwiseRel和BroadcastRel能覆盖大部分逐元素和广播场景自己写 TypeRelation 容易在边界情况上出错。如果非写不可先写一个最简版本只处理TensorType把nullptr和动态 shape 的情况留到后面再补。第三条经验测试要从最小用例开始。不要一上来就用真实模型的复杂输入先用(1, 1)或者(2, 3)这种小 shape把类型推导、lowering、数值都对一遍。小用例跑通了再逐步放大。我见过有人直接上大模型报错信息淹没在几百行日志里根本定位不到问题。第四条经验善用relay.transform.InferType和relay.transform.PrintIR。类型推导出问题时先跑InferType看报错位置lowering 出问题时用PrintIR把 lowering 前后的 IR 打出来对比能快速定位是哪一步出的问题。第五条经验属性字段的默认值要合理。如果某个属性没有默认值用户调用时必须显式传否则解析会失败。给一个合理的默认值能降低使用门槛也能避免因为漏传参数导致的诡异错误。6. 算子扩展后的优化与融合考量6.1 让新算子参与算子融合新算子加进去只是第一步能不能融入 Relay 的优化流程才是决定性能的关键。Relay 的融合 pass 会把相邻的算子合并成一个 kernel减少内存访问和 kernel 启动开销。新算子能不能被融合取决于它的 tag 和融合模式注册。compute 里的tag参数是融合 pass 的重要依据。elemwise标签的算子通常可以和前后算子自由融合broadcast次之reduce和injective有各自的融合规则。如果你希望新算子能融合tag 要设对。比如逐元素算子设elemwise涉及广播的设broadcast。除了 tag还可以用register_fusion_pattern注册更精细的融合模式。比如你希望某个算子和它前面的卷积融合可以写一个模式匹配规则。这部分相对进阶不是所有算子都需要但如果你的算子性能敏感值得花时间研究。6.2 自定义 schedule 的接入点如果新算子的 compute 用默认 schedule 性能不理想可以自定义 schedule。TOPI 的 compute 函数通常配套一个 schedule 函数用te.create_schedule创建调度然后对循环做切分、向量化、并行化等操作。自定义 schedule 的接入方式是在 Relay 的 lowering 配置里指定。TVM 的relay.build可以传schedule_cfg或者用OpStrategy机制把 compute 和 schedule 绑定。对于自定义算子最简单的做法是在 TOPI 里提供一个 schedule 函数然后在 lowering 时通过FTVMSchedule属性指定。schedule 的编写是另一个大话题这里不展开。但要提醒一点不要过早优化 schedule。先用默认 schedule 把功能跑通测出性能瓶颈再针对性地写 schedule。很多时候默认 schedule 已经够用自定义 schedule 的收益可能还不如换个更合适的 tag 让融合 pass 发挥作用。6.3 跨后端兼容性检查新算子写完之后如果你需要它在多个后端上都能跑要做兼容性检查。不同后端对算子的支持程度不一样比如某些后端不支持特定的数据类型或者不支持动态 shape。测试的时候至少覆盖 CPU 和你有需求的目标后端。如果某个后端不支持通常有两种处理方式。一种是提供 fallback 实现让不支持的部分回退到 CPU 或者用其他算子组合。另一种是在 lowering 阶段做后端特定的转换把新算子拆解成后端支持的算子。具体选哪种取决于你的性能要求和实现成本。还有一个容易被忽略的点是 dtype 支持。你的 compute 实现里如果用了tvm.tir.const包装常量要确保常量的 dtype 和输入 dtype 一致。如果输入可能是float16也可能是float32compute 要能处理两种 dtype不能写死。测试的时候把常见 dtype 都跑一遍避免上线后才发现某个 dtype 不支持。7. 我个人在实际操作中的几点体会走完整个流程我最深的体会是TVM 的算子注册机制虽然看起来繁琐但它的设计是有内在逻辑的。注册、类型关系、compute 这三层分别对应“是什么”“怎么推”“怎么算”把这三层想清楚代码就是水到渠成的事。真正难的不是写代码而是理解每一层为什么存在、什么时候需要自定义、什么时候可以复用。另一个体会是文档和社区示例很重要但不能全信。TVM 版本迭代快不同版本的 API 可能有差异网上找到的示例代码不一定适用于你的版本。最可靠的做法是直接看当前版本 TVM 源码里已有的算子实现照葫芦画瓢。比如你想加一个逐元素算子就找一个现有的逐元素算子把它的注册、类型关系、compute 都读一遍然后照着改。这比看任何教程都管用。最后一点不要怕报错。新增算子涉及的环节多第一次做几乎不可能一次通过。报错信息通常很具体指向某个文件某一行顺着排查基本都能找到原因。我自己的习惯是每改一小步就编译测试一次不要攒一大堆改动再一起测那样出了问题很难定位。小步快跑逐步验证是走通这个流程最稳的方式。如果你后续还想扩展可以考虑的方向包括给算子加梯度支持让它能参与训练、写自定义 schedule 提升性能、或者把算子贡献回 TVM 社区。每个方向都有各自的坑但核心流程和这篇文章讲的是一致的。把基础流程走通后面的扩展就是在这个框架上做加法。