解析
TVM源码中经常会看到PackedFunc,PackedFunc是为了支持多语言(Python,C++等)设计的封装结构。它使得某个C++的函数可以在注册后由Python调用,反之亦然。定义在tvm/runtime/packed_func.h下。
TypedPackedFunc,顾名思义就是注明类型的封装函数。
A PackedFunc wrapper to provide typed function signature. It is backed by a PackedFunc internally. TypedPackedFunc enables compile time type checking. TypedPackedFunc works with the runtime system: - It can be passed as an argument of PackedFunc. - It can be assigned to TVMRetValue. - It can be directly converted to a type-erased PackedFunc. Developers should prefer TypedPackedFunc over PackedFunc in C++ code as it enables compile time checking. We can construct a TypedPackedFunc from a lambda function with the same signature. user defined lambda function. auto addone = [](int x)->int { return x + 1; }; We can directly convert lambda function to TypedPackedFunc TypedPackedFunc<int(int)> ftyped(addone); invoke the function. int y = ftyped(1); Can be directly converted to PackedFunc PackedFunc packed = ftype;
这个结构支持runtime时的类型检查。它可以包裹C++ lambda表达式,例如上面提到的addone。函数的参数类型设为了Int,返回值也是Int。需要理解的是,TypedPackedFunc<int(int)>
前面的类型指的是返回值类型,括号内的类型是参数类型。
另外还有一个例子:
printer/text_printer.cc
String AsText(const ObjectRef& node, bool show_meta_data,
runtime::TypedPackedFunc<String(ObjectRef)> annotate) {
Doc doc;
doc << "#[version = \"" << kSemVer << "\"]" << Doc::NewLine();
runtime::TypedPackedFunc<std::string(ObjectRef)> ftyped = nullptr;
if (annotate != nullptr) {
ftyped = runtime::TypedPackedFunc<std::string(ObjectRef)>(
[&annotate](const ObjectRef& expr) -> std::string { return annotate(expr); });
}
doc << TextPrinter(show_meta_data, ftyped).PrintFinal(node);
return doc.str();
}
这里ftyped在annotate不为Null时,将会被annotate函数封装。ftyped就变成了一个接受ObjectRef类型,返回一个string的函数。
在开发时,如果知道函数的类型,将优先考虑使用TypedPackedFunc,因为它支持runtime时的类型检查。