如何开发动态输入算子

动态输入算子是指算子的输入个数是动态的，例如AddN，将N个输入tensor累加到一起，输出一个tensor，输入tensor的个数是不固定的。动态输入算子的开发在构造和解析输入数据方面有差异：核函数的入参采用ListTensorDesc的结构存储输入数据信息，对应的，调用时需构造TensorList结构保存参数信息。下面基于kernel直调和工程化算子开发两种开发方式分别介绍具体开发流程。

下文仅列出代码片段，完整样例请参考动态输入算子样例（工程化算子开发）和动态输入算子样例（kernel直调）。

• kernel直调
  • 参考ListTensorDesc数据结构自行定义ListTensorDesc和TensorDesc结构体，并将实际的输入数据保存至ListTensorDesc结构中。示例如下:
    ptrOffset传入为ListTensorDesc首地址和数据指针首地址dataPtr之间的偏移量，tensorDesc中保存两个输入的tensor描述信息， dataPtr传入为保存输入数据的地址指针。

    1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20

    constexpr uint32_t SHAPE_DIM = 2; struct TensorDesc { uint32_t dim{SHAPE_DIM}; uint32_t index; uint64_t shape[SHAPE_DIM] = {8, 2048}; }; TensorDesc xDesc; xDesc.index = 0; TensorDesc yDesc; yDesc.index = 1; constexpr uint32_t TENSOR_DESC_NUM = 2; struct ListTensorDesc { uint64_t ptrOffset; TensorDesc tensorDesc[TENSOR_DESC_NUM]; uintptr_t dataPtr[TENSOR_DESC_NUM]; } inputDesc; ... inputDesc = {(1 + (1 + SHAPE_DIM) * TENSOR_DESC_NUM) * sizeof(uint64_t), {xDesc, yDesc}, {(uintptr_t)xDevice, (uintptr_t)yDevice}};

  • kernel侧调用时，直接传入ListTensorDesc表达的输入信息。示例如下：

    1 2 3 4 5 6

    void *inputDescInDevice = nullptr; CHECK_ACL(aclrtMalloc((void **)&inputDescInDevice, sizeof(ListTensorDesc), ACL_MEM_MALLOC_HUGE_FIRST)); CHECK_ACL(aclrtMemcpy(inputDescInDevice, sizeof(ListTensorDesc), &inputDesc, sizeof(ListTensorDesc), ACL_MEMCPY_HOST_TO_DEVICE)); ACLRT_LAUNCH_KERNEL(addn_custom)(blockDim, stream, inputDescInDevice, zDevice);

  • kernel侧算子实现，通过ListTensorDesc和TensorDesc提供的接口解析ListTensorDesc输入信息，并处理。示例如下：

    1 2 3 4 5 6 7

    uint64_t buf[SHAPE_DIM] = {0}; AscendC::TensorDesc<int32_t> tensorDesc; tensorDesc.SetShapeAddr(buf); listTensorDesc.GetDesc(tensorDesc, 0); uint64_t totalLength = tensorDesc.GetShape(0) * tensorDesc.GetShape(1); __gm__ uint8_t *x = listTensorDesc.GetDataPtr<__gm__ uint8_t>(0); __gm__ uint8_t *y = listTensorDesc.GetDataPtr<__gm__ uint8_t>(1);

• 工程化算子开发
  • 单算子调用时，构造List类型tensor并传入。

    使用aclCreateTensor创建tensor后，需调用aclCreateTensorList，将创建好的tensor组成List形式，如下所示。

    1	inputTensorList = aclCreateTensorList(inputTensor_.data(), inputTensor_.size());

    获取算子使用的workspace空间大小接口的入参，也需使用aclTensorList结构参数，用来计算workspace的大小，调用示例如下。

    1 2	// 获取算子使用的workspace空间大小 aclnnStatus aclnnAddNCustomGetWorkspaceSize(const aclTensorList *srcList, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor);

  • 算子原型定义中，输入数据的参数类型设置为动态，示例如下。

    1 2 3 4

    this->Input("srcList") .ParamType(DYNAMIC) .DataType({ge::DT_FLOAT16}) .Format({ge::FORMAT_ND});

  • host侧算子实现，获取动态输入信息的接口，需使用对应的动态接口。

    例如，Tiling函数和InferShape函数中，GetDynamicInputShape接口用于获取动态输入的shape信息，InferDataType函数中，GetDynamicInputDataType接口用于获取动态输入的数据类型，示例如下。

    1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16

    namespace ge { static graphStatus InferShape(gert::InferShapeContext *context) { const gert::Shape *x1_shape = context->GetDynamicInputShape(0, 0); gert::Shape *y_shape = context->GetOutputShape(0); *y_shape = *x1_shape; return GRAPH_SUCCESS; } static graphStatus InferDataType(gert::InferDataTypeContext *context) { const auto inputDataType = context->GetDynamicInputDataType(0, 0); context->SetOutputDataType(0, inputDataType); return ge::GRAPH_SUCCESS; } } // namespace ge

  • kernel侧算子实现，入参需传入动态结构的数据，并使用AscendC::ListTensorDesc结构做解析。

    核函数入参需传入动态结构的数据，例如GM_ADDR srcList，示例如下。

    1	extern "C" __global__ __aicore__ void addn_custom(GM_ADDR srcList, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling)

    对传入的参数srcList，需使用AscendC::ListTensorDesc结构做解析，得到每个tensor的具体信息，示例如下。

    1 2 3

    AscendC::ListTensorDesc keyListTensorDescInit((__gm__ void*)srcList); GM_ADDR x = (__gm__ uint8_t*)keyListTensorDescInit.GetDataPtr<__gm__ uint8_t>(0); GM_ADDR y = (__gm__ uint8_t*)keyListTensorDescInit.GetDataPtr<__gm__ uint8_t>(1);