在通过加速库进行开发时,请使用张量作为输入。
大部分算子功能可以由加速库提供的单Operation或者根据单Operation组图完成。如果一些功能通过单Operation或者组合Operation无法支持,用户可以通过PluginOperation机制,即插件机制,开发自定义的PluginOperation,满足对应的功能。
在使用插件机制时,用户需要自行管理因编写的插件代码导致的安全性或系统不可控行为。请确保编写的代码可靠并遵循相关安全规范和最佳实践。
以使用Ascend C创建Add算子为例。
void add_custom_do(uint32_t blockDim, void* l2ctrl, void* stream, uint8_t* x, uint8_t* y, uint8_t* z) { add_custom<<<blockDim, l2ctrl, stream>>>(x, y, z); }
根据Add算子的具体情况实现AddOperation并重载Operation类的成员函数,具体可参考以下实现。
AddOperation *addOp = new AddOperation("addOperation");
atb::Context *context = nullptr; // Context主要负责对NPU中使用的Stream进行管理 int ret = atb::CreateContext(&context); void *stream = nullptr; ret = aclrtCreateStream(&stream); context->SetExecuteStream(stream); // 设置指定Stream为Context中的executeStream
atb::VariantPack variantPack; // 此处先声明VariantPack类对象,其中包含Tensor的信息,此处省略对其具体的构造过程
uint64_t workspaceSize = 0;
addOp ->Setup(variantPack, workspaceSize, context); // 传入variantPack和workspaceSize参数,setup函数中根据实际variantPack所需空间计算出workspaceSize大小并返回
void* workspace = nullptr; aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); // AscendCL的NPU内存分配函数,根据步骤三中计算出来的workspaceSize以ACL_MEM_MALLOC_HUGE_FIRST方式对workspace分配NPU内存空间
addOp->Execute(variantPack, workspace, workspaceSize, context); // 进行具体的addOp的执行操作
int ret = DestroyContext(context); // 销毁context ret = aclrtDestroyStream(stream); // 销毁stream ret = aclrtFree(workspace); // 销毁workspace delete addOp; // 销毁addOp