单算子使用
使用流程
图1 单算子使用流程
- 构造Operation参数。
atb::infer::ElewiseParam opParam = { ... };
- 调用CreateOp函数返回Operation对象。
atb::Operation *operation = nullptr; atb::Status st = atb::CreateOp(opParam, &operation);
- 调用Operation对象的Setup函数,返回workspace大小。
atb::VariantPack variantPack = { ... }; uint64_t workspaceSize = 0; operation->Setup(variantPack, workspaceSize);
- 根据workspace大小,申请NPU内存。
void* workspace = nullptr; aclrtMalloc(&workspace, workspaceSize);
- 创建Context并设置Stream。
atb::Context *context = nullptr; int ret = atb::CreateContext(&context); void *stream = nullptr; ret = aclrtCreateStream(&stream); context->SetExecuteStream(stream);
- 调用Operation对象的Execute函数。
operation->Execute(variantPack, workspace, workspaceSize, context);
- 销毁Context和Stream。
int ret = DestroyContext(context); ret = aclrtDestroyStream(stream);
典型样例说明
以add Operation举例。
- 构造Operation参数,add算子使用的是ElewiseParam,此处对此进行初始化。
atb::infer::ElewiseParam addParam; addParam.elewiseType = atb::infer::ElewiseParam::ELEWISE_ADD; // 此处指定ElewiseParam的elewiseType为ELEWISE_ADD
- 调用CreateOp函数返回Operation对象。
atb::Operation *addOp = nullptr; // 声明一个Operation类对象指针addOp atb::Status st = atb::CreateOperation(addParam, &addOp); // 传入步骤一中的参数和Operation类对象指针引用
- 调用Operation对象的Setup函数,返回workspace大小。
atb::VariantPack variantPack; // 此处先声明VariantPack类对象,其中包含Tensor的信息,此处省略对其具体的构造过程 uint64_t workspaceSize = 0; st = addOp->Setup(variantPack, workspaceSize); // 传入variantPack和workspaceSize参数,setup函数中根据实际variantPack所需空间计算出workspaceSize大小并返回
- 根据workspace大小,申请NPU内存。
void* workspace = nullptr; aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); // AscendCL的NPU内存分配函数,根据步骤三中计算出来的workspaceSize以ACL_MEM_MALLOC_HUGE_FIRST方式对workspace分配NPU内存空间
- 创建Context并设置Stream。
atb::Context *context = nullptr; // Context主要负责对NPU中使用的Stream进行管理 int ret = atb::CreateContext(&context); void *stream = nullptr; ret = aclrtCreateStream(&stream); context->SetExecuteStream(stream); // 设置指定Stream为Context中的executeStream
- 调用Operation对象的Execute函数。
addOp>Execute(variantPack, workspace, workspaceSize, context); // 进行具体的addOp的执行操作
- 销毁Context和Stream。
int ret = DestroyContext(context); // 销毁context ret = aclrtDestroyStream(stream); // 销毁stream
父主题: 加速库使用指导