Relay 算子调用流程
TVM 0.8 中 Relay 算子调用全流程解析
在 TVM 0.8 版本中，从 Relay 前端到 TOPI 计算实现的完整调用流程涉及算子定义、计算逻辑绑定和策略管理三个关键层次。以下是详细总结：
1. 算子定义层 (Relay OP)
路径: src/relay/op/nn/
作用: 声明算子的接口和基础属性
关键组件
// 以卷积算子为例 (src/relay/op/nn/convolution.cc)
struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {Array<IndexExpr> strides;Array<IndexExpr> padding;String data_layout; // e.g. "NCHW"// ...
};// 算子注册
TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d").set_body_typed(MakeConv<Conv2DAttrs>).set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout);
核心功能
- 定义算子的参数结构（如
Conv2DAttrs） - 注册到 Relay 的全局符号表（
relay.op.nn._make.conv2d） - 声明布局推断方法（
FInferCorrectLayout）
2. 计算逻辑层 (TOPI)
路径:
- 接口声明：
include/tvm/topi/nn.h - 实现：
src/topi/nn/convolution.cc - 硬件优化：
src/topi/cuda/,src/topi/x86/
计算绑定机制
(1) 传统方式：FTVMCompute
// 在算子注册时直接绑定TOP实现
.set_attr<FTVMCompute>("FTVMCompute", [](const Attrs& attrs, const Array<Tensor>& inputs, const Type& out_type) {return topi::nn::conv2d(inputs[0], inputs[1], attrs.strides, attrs.padding);});
(2) 现代方式：Strategy 策略
Python 策略注册 (python/tvm/relay/op/strategy/cuda.py)：
@conv2d_strategy.register("cuda")
def conv2d_strategy_cuda(attrs, inputs, out_type, target):return topi.cuda.conv2d_nchw_strategy(attrs, inputs, out_type, target)
3. 策略管理层
路径: python/tvm/relay/op/strategy/
作用: 动态选择最优计算实现
策略选择流程
策略优先级
- 优先查找注册的
strategy函数 - 若未找到，回退到
FTVMCompute - 最终触发
topi::generic默认实现
4. 完整调用链示例（以Conv2D为例）
步骤 1: Relay 前端解析
# Python 前端调用
x = relay.nn.conv2d(data, weight, strides=(1,1), padding=(0,0))
步骤 2: C++ 算子注册层
// src/relay/op/nn/convolution.cc
TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d").set_body_typed(MakeConv<Conv2DAttrs>);
步骤 3: 策略选择
# python/tvm/relay/op/strategy/cuda.py
@conv2d_strategy.register("cuda")
def conv2d_strategy_cuda(...):return topi.cuda.conv2d_nchw_strategy(...)
步骤 4: TOPI 计算实现
// src/topi/cuda/conv2d.cu
Tensor conv2d_nchw_cuda(..., Target target) {if (target.libs.count("cudnn")) {return cudnn_conv2d(...); // 使用cuDNN} else {return generic_gpu_conv2d(...); // 通用CUDA核函数}
}
步骤 5: 生成计算图
最终生成的计算表达式示例：
// 输出的Relay IR
%1 = nn.conv2d(%data, %weight, strides=[1,1], padding=[0,0]) /* target=cuda */;
5. TVM 0.8 的特性说明
(1) 与新版TVM的区别
- 策略机制：0.8 版本策略管理已存在，但接口较新版本简单
- 混合绑定：同时支持
FTVMCompute和Strategy，但推荐使用 Strategy
(2) 关键调试方法
# 查看算子绑定的策略
print(tvm.relay.op.get("nn.conv2d").get_attr("FTVMStrategy"))# 打印最终调度
print(tvm.lower(mod, target="cuda"))
总结：TVM 0.8 的算子调用架构
| 层级 | 组件 | 责任 |
|---|---|---|
| 定义层 | src/relay/op/nn/ | 算子接口和属性定义 |
| 计算层 | src/topi/ + include/tvm/topi | 数学实现和硬件优化 |
| 策略层 | python/tvm/relay/op/strategy/ | 动态选择最优实现路径 |
这种分层设计实现了：
- 前端稳定性：Relay 接口与底层实现解耦
- 硬件可扩展性：轻松添加新后端的优化
- 运行时灵活性：根据 Target 动态选择最优策略