1. Relay算子注冊 (C++層)
(a) 算子屬性注冊
路徑: src/relay/op/nn/nn.cc
RELAY_REGISTER_OP("hardswish").set_num_inputs(1).add_argument("data", "Tensor", "Input tensor.").set_support_level(3).add_type_rel("Identity", Identity);
(b) 調用節點構造
路徑: src/relay/op/nn/activation.cc
TVM_REGISTER_GLOBAL("relay.op._make.hardswish").set_body_typed([](Expr data) {static const Op& op = Op::Get("hardswish");return Call(op, {data}, Attrs(), {});});
2. TOPI計算實現 (C++層)
? TOPI注冊入口
路徑: src/topi/elemwise.cc
TVM_REGISTER_GLOBAL("topi.hardswish").set_body([](TVMArgs args, TVMRetValue* rv) {*rv = hardswish(args[0]);});
(d) 數學內核實現
路徑: include/tvm/topi/nn.h
inline Tensor hardswish(const Tensor& x, std::string name = "T_hardswish") {auto three = make_const(x->dtype, 3);auto six = make_const(x->dtype, 6);return compute(x->shape,[&](const Array<Var>& i) {return x(i) * max(min(x(i) + three, six), 0) / six;},name, kElementWise);
}
3. Python接口層
(e) Relay Python API
路徑: python/tvm/relay/op/nn/_nn.py
def hardswish(data):return _make.hardswish(data)
(f) TOPI通用接口
路徑: python/tvm/topi/nn.py
@tvm.target.generic_func
def hardswish(x):return cpp.hardswish(x)
4. 計算調度注冊
(g) Compute注冊
路徑: python/tvm/relay/op/strategy/generic.py
@register_compute("hardswish")
def hardswish_compute(attrs, inputs, out_type):return [topi.hardswish(inputs[0])]
(h) 調度策略
路徑: `python/tvm/relay/op/op.py**
register_broadcast_schedule("hardswish")
register_shape_func("hardswish", False, elemwise_shape_func)
5. 硬件專用實現
(i) NPU支持聲明
路徑: `src/relay/backend/contrib/npu/src/op_map.cc**
const std::vector<std::string> _NPU_OP = {...,"hardswish" // 添加算子名
};
(j) NPU內核實現
路徑: `python/tvm/relay/backend/contrib/npu/ops.py**
def custom_hardswish(x):x1 = custom_add(x, te.extern_scalar_value(3.0))x2 = custom_relu(x1)return npu_hardwish(x2, ...)
(k) NPU策略注冊
路徑: `python/tvm/relay/op/strategy/npu.py**
@hardswish.register("npu")
def hardswish_npu(x):return npu_api.custom_hardswish(x)
6. 前端框架對接
(l) PyTorch轉換器
路徑: `python/tvm/relay/frontend/pytorch.py**
def _hardswish():def _impl(inputs, input_types):return _op.hardswish(inputs[0])return _impl
關鍵文件路徑總結
功能模塊 | 關鍵路徑 |
---|---|
Relay核心注冊 | src/relay/op/nn/{nn.cc, activation.cc} |
TOPI計算 | {include,src}/topi/{nn.h, elemwise.cc} |
Python接口 | python/tvm/{relay/op/nn/_nn.py, topi/nn.py} |
策略注冊 | python/tvm/relay/op/strategy/{generic.py, npu.py} |
硬件后端 | src/relay/backend/contrib/npu/ |
前端對接 | python/tvm/relay/frontend/pytorch.py |
開發流程示意圖
通過這種清晰的路徑劃分,TVM實現了:
- 模塊化開發:各層級代碼物理隔離
- 可擴展性:新增硬件只需在對應目錄添加實現
- 維護性:相關功能的代碼集中存放