昇腾社区首页
中文
注册

自定义插件

图1 ATB算子调用流程

使用流程

  • 以下代码为简单示例,未对接口函数进行返回值校验,实际使用时需校验返回值。
  • 在通过加速库进行开发时,请使用张量作为输入。
  1. 包含ACL与加速库接口头文件。
    #include <acl/acl.h>
    #include <atb/atb_infer.h>
  2. 配置deviceId。
    根据需求设置deviceId。
    int deviceId = 0;
    aclError status = aclrtSetDevice(deviceId);
  3. 创建自定义插件算子(PluginOperation)对象实例。

    在使用插件机制时,用户需要自行管理因编写的插件代码导致的安全性或系统不可控行为。请确保编写的代码可靠并遵循相关安全规范和最佳实践。

    1. 开发算子

      使用Ascend C或其他方式实现自定义算子,并在加速库中实现对应的算子类,如CustomizeOperation。

    2. 创建算子对象实例
      CustomizeOperation *op = new CustomizeOperation("CustomizeOperation");
  4. 创建VariantPack。
    VariantPack中包含输入和输出tensor列表。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(a); // 获取Tensor内存大小
    status = aclrtMalloc(&a.deviceData, a.dataSize, ACL_MEM_MALLOC_HUGE_FIRST); // 申请device内存
    // 按上述方法构造所有输入和输出tensor,存入VariantPack
    atb::VariantPack variantPack;
    variantPack.inTensors = { a, ... };
    variantPack.outTensors = { output, ... };
  5. 创建Context,配置Stream。
    Context主要负责对NPU中使用的Stream进行管理。
    atb::Context *context = nullptr;
    st = atb::CreateContext(&context);
    aclrtStream stream = nullptr;
    status = aclrtCreateStream(&stream);
    context->SetExecuteStream(stream);
  6. 调用Setup接口,计算workspace大小。
    uint64_t workspaceSize = 0;
    st = op->Setup(variantPack, workspaceSize, context);
  7. 根据workspace大小申请NPU内存。
    void *workspace = nullptr;
    status = aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);

    当workspace大小为0时,无需执行该步骤,否则会报错。

  8. 调用Execute接口,执行算子。
    st = op->Execute(variantPack, (uint8_t *)workspace, workspaceSize, context);
  9. 销毁创建的对象,释放内存。
    status = aclrtDestroyStream(stream); // 销毁stream
    status = aclrtFree(workspace);       // 销毁workspace
    st = atb::DestroyOperation(op);      // 销毁op对象
    st = atb::DestroyContext(context);   // 销毁context
    // 下面代码为释放Tensor的示例代码,实际使用时需释放VariantPack中的所有Tensor
    status = aclrtFree(tensor.deviceData);
    tensor.deviceData = nullptr;
    tensor.dataSize = 0;
若想执行上述demo,按以下步骤进行操作。
# g++编译demo工程,demo.cpp为demo对应的源码文件
g++ -I "${ATB_HOME_PATH}/include" -I "${ASCEND_HOME_PATH}/include" -L "${ATB_HOME_PATH}/lib" -L "${ASCEND_HOME_PATH}/lib64" demo.cpp -l atb -l ascendcl -o demo
./demo # 运行可执行文件

当abi=0时,需在g++命令中添加编译选项-D_GLIBCXX_USE_CXX11_ABI=0。

ATB加速库需配置abi参数时,请参考如下两种方式配置:

  • 自动配置:执行set_env.sh脚本时,若不加任何参数,且已检测到PyTorch环境时会自动调用torch.compiled_with_cxx11_abi()接口,自动选择PyTorch编译时abi参数作为ATB的abi参数,如果没有检测到PyTorch环境则默认配置abi=1
  • 手动配置:执行set_env.sh时,支持用户通过--cxx_abi=1--cxx_abi=0参数指定ATB的abi参数,例如source set_env.sh --cxx_abi=1