📚 课程目录第2章 NPU目标后端注册

2. NPU目标后端注册

任务说明

本案例目标是在TVM编译框架中注册一个名为example_npu的自定义NPU后端，并实现最基本的代码生成接口。完成注册后，可通过TVM的Target API创建该后端目标，并调用代码生成入口生成一个空模块，为后续扩展NPU算子lowering与代码生成奠定基础。

实现方案

TVM支持通过TVM_REGISTER_TARGET_KIND宏注册自定义后端。注册时需要指定目标名称、可选的FTVMTVMParseTarget解析函数，以及FTVMCodeGen代码生成函数。本案例实现一个最小可用的NPU后端：解析函数负责读取mcpu属性，代码生成函数接收IRModule并返回一个仅包含空runtime::Module的生成结果，用于验证后端注册与调用链路通畅。

代码实现

创建文件example_npu_target.cc，内容如下：

#include <tvm/runtime/module.h>
#include <tvm/runtime/registry.h>
#include <tvm/target/target.h>
#include <tvm/ir/module.h>
#include <tvm/runtime/packed_func.h>

namespace tvm {
namespace codegen {

// 空NPU模块，仅用于验证代码生成接口可被调用
class ExampleNPUModule : public runtime::ModuleNode {
 public:
  const char* type_key() const final { return "example_npu"; }
  // 不实现任何函数，仅作为占位
};

// 代码生成入口：接收IRModule，返回空模块
runtime::Module ExampleNPUCodeGen(const IRModule& mod) {
  // 此处仅验证接口调用，不执行实际编译
  auto n = make_object<ExampleNPUModule>();
  return runtime::Module(n);
}

// 目标解析函数：读取mcpu属性，返回配置Map
TargetKind::AttrMap<ObjectRef> ExampleNPUParseTarget(const Target& target) {
  TargetKind::AttrMap<ObjectRef> config;
  if (target->kind->name == "example_npu") {
    // 从target属性中读取mcpu，若未提供则使用默认值
    auto it = target->attrs.find("mcpu");
    if (it != target->attrs.end()) {
      config.Set("mcpu", it->second);
    } else {
      config.Set("mcpu", String("npu_v1"));
    }
  }
  return config;
}

TVM_REGISTER_GLOBAL("target.example_npu")
    .set_body([](Target target) {
      return ExampleNPUParseTarget(target);
    });

TVM_REGISTER_TARGET_KIND("example_npu", kDLCPU)
    .set_parser(ExampleNPUParseTarget)
    .set_attr<FTVMCodeGen>("codegen", ExampleNPUCodeGen);

}  // namespace codegen
}  // namespace tvm

编译与链接

将上述文件编译为动态库，或直接加入TVM源码树中重新编译。若采用独立编译方式，使用以下命令：

g++ -std=c++17 -shared -fPIC -o libexample_npu.so example_npu_target.cc \
    -I${TVM_HOME}/include -I${TVM_HOME}/3rdparty/dlpack/include \
    -I${TVM_HOME}/3rdparty/dmlc-core/include \
    -L${TVM_HOME}/build -ltvm_runtime

其中${TVM_HOME}为TVM源码根目录。编译前需确保TVM已构建完成，且build/目录下存在libtvm_runtime.so。

验证方法

编写Python脚本test_npu_target.py验证后端注册是否成功：

import tvm
from tvm import target

# 加载自定义后端动态库
tvm.load_lib("libexample_npu.so")

# 创建NPU目标
tgt = target.Target("example_npu -mcpu=npu_v2")
print("Target kind:", tgt.kind.name)
print("mcpu attr:", tgt.attrs.get("mcpu", ""))

# 调用代码生成接口
mod = tvm.IRModule({})
codegen_func = tvm.get_global_func("target.example_npu.codegen")
if codegen_func:
    result = codegen_func(mod)
    print("Codegen result type:", result.type_key)
else:
    print("Codegen function not found")

运行验证：

export TVM_HOME=/path/to/tvm
export PYTHONPATH=${TVM_HOME}/python:${PYTHONPATH}
export LD_LIBRARY_PATH=${TVM_HOME}/build:${LD_LIBRARY_PATH}
python test_npu_target.py

判定标准

满足以下条件即判定案例通过：

脚本输出Target kind: example_npu，说明目标注册成功；
输出mcpu attr: npu_v2，说明解析函数正确读取了mcpu属性；
输出Codegen result type: example_npu，说明代码生成接口被成功调用并返回了自定义模块。

关键说明

本案例仅实现后端注册与代码生成接口的占位调用，不包含实际指令生成逻辑。代码生成函数接收IRModule后直接返回空模块，用于验证从Python侧创建目标到C++侧代码生成入口的完整链路。后续案例将在此基础上扩展算子lowering与NPU指令发射功能。