封装函数+TVM FFI 层注册

6 阅读1分钟

/* relax.broadcast_to */ 广播

Expr broadcast_to(Expr x, Expr shape) { 封装函数(函数名为broadcast_to)并输入参数x和shape 另外这两个函数的类型均为Expr

static const Op& op = Op::Get("relax.broadcast_to"); 定义一个Op类的引用常量

return Call(op, {std::move(x), std::move(shape)}, Attrs(), {}); 返回call类型,该类型里面会输出op,并讲输入的参数传入左值,Attrs(), {}:传两个空值给Call();

}

封装函数

TVM_FFI_STATIC_INIT_BLOCK() { // 1. TVM_FFI_STATIC_INIT_BLOCK():TVM框架自定义的静态初始化宏 // 作用:包裹的代码会在程序启动/模块加载时自动执行(无需手动调用),且仅执行一次,用于完成FFI相关初始化

namespace refl = tvm::ffi::reflection; // 2. 命名空间别名定义(C++基础语法)

// 含义:将冗长的命名空间 tvm::ffi::reflection 简写为 refl,后续代码中可通过refl替代原命名空间,简化代码书写

// 补充:tvm::ffi::reflection 是TVM框架中负责FFI(外部函数接口)反射功能的命名空间,包含各类注册、反射相关工具。

refl::GlobalDef().def("relax.op.broadcast_to", broadcast_to); }

TVM FFI 层的函数注册代码