单算子使用

操作步骤

1. 构造Operation参数。

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

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

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

3. 创建Context并设置Stream。

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

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

   // 构造Tensor
   atb::Tensor a;
   a.desc.dtype = ACL_FLOAT16;    // 配置Tensor数据类型
   a.desc.format = ACL_FORMAT_ND; // 配置Tensor格式
   a.desc.shape.dimNum = 2;       // 配置Tensor维度数
   a.desc.shape.dims[0] = 3;      // 配置Tensor第0维大小
   a.desc.shape.dims[1] = 3;      // 配置Tensor第1维大小
   a.dataSize = Utils::GetTensorSize(tensor); // 获取Tensor内存大小
   ret = aclrtMalloc(&a.deviceData, a.dataSize, ACL_MEM_MALLOC_HUGE_FIRST); // 申请device内存
   // 按上述方法构造所有输入和输出tensor，存入VariantPack
   atb::VariantPack variantPack;
   variantPack.inTensors = { a, ... };
   variantPack.outTensors = { output, ... };
   uint64_t workspaceSize = 0;
   operation->Setup(variantPack, workspaceSize, context);

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

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

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

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

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

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

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

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