文档
注册

单算子使用

使用流程

图1 单算子使用流程
  1. 构造Operation参数。
    atb::infer::ElewiseParam opParam = { ... };
  2. 调用CreateOp函数返回Operation对象。
    atb::Operation *operation = nullptr;
    atb::Status st = atb::CreateOp(opParam, &operation);
  3. 调用Operation对象的Setup函数,返回workspace大小。
    atb::VariantPack variantPack = { ... };
    uint64_t workspaceSize = 0;
    operation->Setup(variantPack, workspaceSize);
  4. 根据workspace大小,申请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。
    int ret = DestroyContext(context);
    ret = aclrtDestroyStream(stream);

典型样例说明

以add Operation举例。

  1. 构造Operation参数,add算子使用的是ElewiseParam,此处对此进行初始化。
    atb::infer::ElewiseParam addParam;
    addParam.elewiseType =  atb::infer::ElewiseParam::ELEWISE_ADD; // 此处指定ElewiseParam的elewiseType为ELEWISE_ADD
  2. 调用CreateOp函数返回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。
    int ret = DestroyContext(context); // 销毁context
    ret = aclrtDestroyStream(stream); // 销毁stream
搜索结果
找到“0”个结果

当前产品无相关内容

未找到相关内容,请尝试其他搜索词