单算子使用

使用流程

1. 构造Operation参数。

   atb::infer::ElewiseParam opParam = { ... };

2. 调用CreateOperation函数返回Operation对象。

   atb::Operation *operation = nullptr;
   atb::Status st = atb::CreateOperation(opParam, &operation);

3. 调用Operation对象的Setup函数，返回workspace大小。其中，variantPack中传入的每个输入Tensor要求大于0且不超过256GB。

   atb::VariantPack variantPack = { ... };
   uint64_t workspaceSize = 0;
   operation->Setup(variantPack, workspaceSize);

4. 根据workspace（workspace包含中间过程输出的Tensor）大小，申请NPU内存。

   void* workspace = nullptr;
   aclrtMalloc(&workspace, workspaceSize);

5. 创建Context并设置Stream。

   atb::Context *context = nullptr;
   int ret = atb::CreateContext(&context);
   void *stream = nullptr;
   ret = aclrtCreateStream(&stream);
   context->SetExecuteStream(stream);

6. 调用Operation对象的Execute函数。

   operation->Execute(variantPack, workspace, workspaceSize, context);

7. 销毁Context、Stream、workspace和Operation对象。

   int ret = DestroyContext(context);
   ret = aclrtDestroyStream(stream);
   ret = aclrtFree(workspace);
   st = DestroyOperation(operation);

典型样例说明

以add Operation举例。

1. 构造Operation参数，add算子使用的是ElewiseParam，此处对此进行初始化。

   atb::infer::ElewiseParam addParam;
   addParam.elewiseType =  atb::infer::ElewiseParam::ELEWISE_ADD; // 此处指定ElewiseParam的elewiseType为ELEWISE_ADD

2. 调用CreateOperation函数返回Operation对象。

   atb::Operation *addOp = nullptr;                         // 声明一个Operation类对象指针addOp
   atb::Status st = atb::CreateOperation(addParam, &addOp); // 传入步骤一中的参数和Operation类对象指针引用

3. 调用Operation对象的Setup函数，返回workspace大小。

   atb::VariantPack variantPack;                  // 此处先声明VariantPack类对象，其中包含Tensor的信息，此处省略对其具体的构造过程
   uint64_t workspaceSize = 0;
   st = addOp->Setup(variantPack, workspaceSize); // 传入variantPack和workspaceSize参数，setup函数中根据实际variantPack所需空间计算出workspaceSize大小并返回

4. 根据workspace大小，申请NPU内存。

   void* workspace = nullptr;
   aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); // AscendCL的NPU内存分配函数，根据步骤三中计算出来的workspaceSize以ACL_MEM_MALLOC_HUGE_FIRST方式对workspace分配NPU内存空间

5. 创建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

6. 调用Operation对象的Execute函数。

   addOp>Execute(variantPack, workspace, workspaceSize, context); // 进行具体的addOp的执行操作

7. 销毁Context、Stream、workspace和Operation对象。

   int ret = DestroyContext(context); // 销毁context
   ret = aclrtDestroyStream(stream); // 销毁stream
   ret = aclrtFree(workspace); // 销毁workspace
   st = DestroyOperation(operation); // 销毁operation