CustomOpBuilder 通过 AclnnOpRunner 接入 ACLNN 算子

概述

CANN中的AOL算子加速库提供了大量深度优化、硬件亲和的高性能算子。若MindSpore尚未封装aclnn算子的Python接口，或你已基于Ascend C自研开发了算子，只需CustomOpBuilder + AclnnOpRunner 即可在 动态图（PyNative）模式下无感接入，无需关心内存、stream、workspace等底层细节。

aclnn算子的调用方式，通常是基于“两段式”接口，样式形如：

aclnnStatus aclxxXxxGetWorkspaceSize(const aclTensor * src, ..., aclTensor * out, ..., uint64_t * workspaceSize, aclOpExecutor ** executor);
aclnnStatus aclxxXxx(void * workspace, uint64_t workspaceSize, aclOpExecutor * executor, aclrtStream stream);

必须先调用第一段接口aclxxXxxGetWorkspaceSize，用于计算本次API调用过程中需要多少workspace内存，获取到计算所需的workspace size后，按照workspace size申请NPU内存，然后调用第二段接口aclxxXxx执行计算。

在基于CustomOpBuilder的自定义算子中，MindSpore提供了 PyboostRunner 方便用户在动态图接入自定义算子。为了简化调用流程，隐藏接口数据类型转换操作，MindSpore针对aclnn算子提供了统一的执行入口ms::pynative::AclnnOpRunner，其支持PyBoost多级流水和MindSpore框架的算子缓存能力，提高算子和网络执行效率。

本教程以 ArgMin为例，展示完整接入流程。完整代码可参考 MindSpore 代码仓。

安装 ACLNN 开发环境

CANN中算子

若算子为CANN包中算子，则用户无需额外配置环境，只需按照MindSpore安装指南配置MindSpore使用环境即可。

Ascend C自定义算子

若算子是用户基于Ascend C开发的自定义算子，则需将算子编译结果路径添加环境变量ASCEND_CSUTOM_OPP_PATH中，例如：

export ASCEND_CUSTOM_OPP_PATH={build_out_path}/build_out/_CPack_Package/Linux/External/custom_opp_euleros_aarch64.run/packages/vendors/{your_custom_name}:$ASCEND_CUSTOM_OPP_PATH

ArgMin算子接入

下面给出示例完整代码。

#include <set>
#include <optional>
#include "ms_extension/all.h"

namespace custom {

/* 1. 推导输出shape */
static ShapeVector InferArgMinShape(const ShapeVector &in_shape, int64_t dim, bool keep_dims) {
  const int64_t rank = static_cast<int64_t>(in_shape.size());
  if (rank == 0) {
    return in_shape;
  }

  int64_t axis = (dim < 0) ? (dim + rank) : dim;
  if (axis < 0 || axis >= rank) {
    MS_LOG(EXCEPTION) << "Infer shape failed";
  }

  ShapeVector out_shape;
  out_shape.reserve(keep_dims ? rank : rank - 1);

  for (int64_t i = 0; i < rank; ++i) {
    if (i == axis) {
      if (keep_dims) {
        out_shape.push_back(1);
      }
    } else {
      out_shape.push_back(in_shape[i]);
    }
  }

  return out_shape;
}

/* 2. 构造空输出Tensor */
ms::Tensor GenResultTensor(const ms::Tensor &t, int64_t dim, bool keep_dim, ms::TypeId type_id) {
  ShapeVector in_shape = t.shape();
  ShapeVector out_shape = InferArgMinShape(in_shape, dim, keep_dim);
  return ms::Tensor(type_id, out_shape);
}

/* 3. 算子入口：被Python直接调用 */
ms::Tensor npu_arg_min(const ms::Tensor &x, int64_t dim, bool keep_dim) {
  auto result = GenResultTensor(x, dim, keep_dim, ms::TypeId::kNumberTypeInt64);
  auto runner = std::make_shared<ms::pynative::AclnnOpRunner>("ArgMin");
  runner->SetLaunchFunc(LAUNCH_ACLNN_FUNC(aclnnArgMin, x, dim, keep_dim, result));
  runner->Run({x}, {result});
  return result;
}
}  // namespace custom

/* 4. PYBIND11接口定义 */
PYBIND11_MODULE(MS_EXTENSION_NAME, m) { m.def("npu_arg_min", PYBOOST_CALLER(1, custom::npu_arg_min)); }

1. 推导算子输出信息

auto y = GenResultTensor(x, axis, keep_dims);

这一步骤是基于算子的计算逻辑，使用shape和type创建输出Tensor。例如，aclnnArgMin根据 axis 与 keep_dims提前计算出输出shape和type并使用ms::Tensor(dtype, shape)构造空Tensor。该Tensor只分配元信息，不分配device内存，AclnnOpRunner::Run会在内部分配device内存。

2. 创建AclnnOpRunner

在基于CustomOpBuilder的自定义算子中，MindSpore提供了通用的接入自定义算子类PyboostRunner。针对aclnn算子，用户可直接使用AclnnOpRunner类创建对象。

auto runner = std::make_shared<ms::pynative::AclnnOpRunner>("ArgMin");

3. 调用接口执行算子

runner->SetLaunchFunc(LAUNCH_ACLNN_FUNC(aclnnArgMin, x, axis, keep_dims, y));
runner->Run({x}, {y});

在LAUNCH_ACLNN_FUNC中依次传入算子名字、输入、输出，并调用SetLaunchFunc方法将launch函数设置到runner中。调用Run方法，入参分别是算子ms::Tensor类型的输入和输出。

4. 通过pybind11将C++函数绑定一个Python函数

PYBIND11_MODULE(MS_EXTENSION_NAME, m) { m.def("npu_arg_min", PYBOOST_CALLER(1, custom::npu_arg_min)); }

npu_arg_min: 前端接口名称。 custom::npu_arg_min: 实际调用的后端接口。 PYBOOST_CALLER: 入参为输出个数和后端接口。

5. 使用CustomOpBuilder编译自定义算子

将上述C++代码保存成文件argmin.cpp，然后使用Python接口CustomOpBuilder编译。

import mindspore as ms
import numpy as np

my_ops = CustomOpBuilder("my_custom", 'argmin.cpp', backend="Ascend").load()
x = np.random.randn(2, 3, 4, 5).astype(np.float32)
output = my_ops.npu_arg_min(ms.Tensor(x), 0, False)