从 Ascend C 到 MindSpore:自定义算子端到端集成实战
MindSpore 支持通过Custom OP机制集成第三方算子。↓↓↓:Host 侧调度逻辑;(GPU)或xxx.cpp(Ascend C);:算子定义(输入输出规格);:Python 接口。本文打通了的完整链路,使自定义算子真正可用、可训练、可部署。这套方法论适用于任何 Ascend C 算子(GEMM、Conv、LayerNorm 等),是昇腾生态开发者的必备技能。最佳实践算子命名遵循;输入
引言
在昇腾 AI 生态中,Ascend C 赋予开发者极致的底层控制能力,而 MindSpore 作为华为推出的全场景 AI 框架,则提供了简洁高效的模型开发接口。然而,许多开发者在尝试将两者结合时,常陷入“算子能单独运行,却无法在 MindSpore 中调用”的困境。
本文将提供一套 工业级、可复现、带调试技巧 的端到端集成方案,以 VectorAdd 算子为例,完整演示如何:
- 在 MindSpore 源码中注册自定义 Ascend C 算子;
- 实现 Host 侧调度逻辑与 Python 接口;
- 支持前向/反向传播;
- 编写单元测试与性能验证;
- 定位典型编译/运行错误。
适用环境:
- MindSpore 2.3.0 或 2.4.0(源码编译)
- CANN Toolkit 7.0.RC1 或更高
- 昇腾 910/310 设备或 Atlas 800 服务器
前置知识:熟悉 C++、Python、基本 CMake,了解 MindSpore 张量模型。
一、整体架构与数据流
在 MindSpore 中,自定义 Ascend 算子的执行流程如下:
[Python Layer]
↓ (调用 ops.vector_add)
[MindSpore Frontend: Primitive + InferShape]
↓
[MindSpore Backend: AscendKernelMod::Launch()]
↓ (调用 aclrtLaunchKernel)
[CANN Runtime → 加载 .o 算子二进制]
↓
[Ascend NPU: 执行 Ascend C Kernel]
↓
[结果写回 Device Memory → 返回 Python]
关键组件说明:
| 组件 | 作用 | 文件位置 |
|---|---|---|
Primitive |
Python 算子定义 | python/mindspore/ops/operations/xxx.py |
InferImpl |
Shape/Dtype 推导 | op_def/xxx_op.cc |
KernelMod |
Host 侧调度入口 | plugin/device/ascend/kernel/xxx_kernel.cc |
Ascend C Kernel |
NPU 上执行的计算逻辑 | custom/op_kernel/xxx.cpp |
.o 文件 |
编译后的算子二进制 | 由 aic 编译器生成 |
二、工程目录结构搭建
假设你已克隆 MindSpore 源码(git clone https://gitee.com/mindspore/mindspore.git),我们将新建以下文件:
mindspore/
├── custom/ # 【新增】自定义算子根目录
│ └── vector_add/
│ ├── op_kernel/
│ │ └── vector_add.cpp # Ascend C 核心实现
│ └── build.sh # Ascend C 编译脚本
├── mindspore/
│ └── ops/
│ └── operations/
│ └── math_ops.py # Python 接口(追加)
├── op_def/
│ └── vector_add_op.cc # 算子注册与推导
├── plugin/
│ └── device/
│ └── ascend/
│ └── kernel/
│ └── vector_add_kernel.cc # Host 调度逻辑
└── tests/
└── ut/
└── python/
└── ops/
└── test_vector_add.py # 单元测试
注意:生产项目建议将
custom/目录独立为子模块,便于版本管理。
三、Step 1:Ascend C 算子实现(vector_add.cpp)
我们复用并增强前文的 VectorAdd 实现,增加对 非对齐长度 的处理:
// custom/vector_add/op_kernel/vector_add.cpp
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BLOCK_SIZE = 16; // 向量化单位
extern "C" __global__ __aicore__ void VectorAddCustom(
uint32_t coreId,
void* x1_gm,
void* x2_gm,
void* y_gm,
uint32_t total_elem) {
KernelHandle handle;
handle.Init();
// 分配工作 Core
uint32_t core_num = GetCoreNum();
if (coreId >= core_num) return;
// 计算当前 Core 负责的元素范围
uint32_t per_core = (total_elem + core_num - 1) / core_num;
uint32_t start = coreId * per_core;
uint32_t end = min(start + per_core, total_elem);
if (start >= total_elem) return;
uint32_t process_elem = end - start;
uint32_t align_elem = ((process_elem + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
// 分配 UB
Queue<QuePosition::QueSram> sram_queue;
sram_queue.Init();
LocalTensor<half> x1_ub = AllocTensor<half>(sram_queue, {align_elem});
LocalTensor<half> x2_ub = AllocTensor<half>(sram_queue, {align_elem});
LocalTensor<half> y_ub = AllocTensor<half>(sram_queue, {align_elem});
// 搬运输入(自动 padding 尾部)
GlobalTensor<half> x1_gm_tensor(reinterpret_cast<half*>(x1_gm) + start, {process_elem});
GlobalTensor<half> x2_gm_tensor(reinterpret_cast<half*>(x2_gm) + start, {process_elem});
DataCopy(x1_ub, x1_gm_tensor, process_elem);
DataCopy(x2_ub, x2_gm_tensor, process_elem);
// 补零尾部(保证向量化安全)
if (process_elem < align_elem) {
for (uint32_t i = process_elem; i < align_elem; i++) {
x1_ub.SetValue(i, 0.0_h);
x2_ub.SetValue(i, 0.0_h);
}
}
// 向量化加法
Add(y_ub, x1_ub, x2_ub, align_elem);
// 写回有效部分
GlobalTensor<half> y_gm_tensor(reinterpret_cast<half*>(y_gm) + start, {process_elem});
DataCopy(y_gm_tensor, y_ub, process_elem);
Pipe::SyncAll();
FreeTensor(x1_ub); FreeTensor(x2_ub); FreeTensor(y_ub);
}
关键改进:
- 自动处理
total_elem % 16 != 0的情况;- 使用
GetCoreNum()动态获取 Core 数量,提升可移植性。
四、Step 2:编译 Ascend C 算子(build.sh)
创建独立编译脚本,生成 .o 文件供 MindSpore 加载:
#!/bin/bash
# custom/vector_add/build.sh
set -e
source /usr/local/Ascend/ascend-toolkit/set_env.sh
KERNEL_NAME="vector_add_custom"
SRC_DIR="$(dirname $0)/op_kernel"
BUILD_DIR="${SRC_DIR}/build"
mkdir -p ${BUILD_DIR}
# 编译 Ascend C
aic -e aic-vec-intrinsic-check=off \
-c ${SRC_DIR}/vector_add.cpp \
-o ${BUILD_DIR}/${KERNEL_NAME}.o \
--host-os linux \
--host-arch x86_64
echo "✅ Ascend C kernel compiled to: ${BUILD_DIR}/${KERNEL_NAME}.o"
执行后生成 vector_add_custom.o,后续需将其路径告知 MindSpore。
五、Step 3:Host 侧 Kernel 实现(vector_add_kernel.cc)
这是 MindSpore 与 Ascend C 的桥梁:
// plugin/device/ascend/kernel/vector_add_kernel.cc
#include "plugin/device/ascend/kernel/ascend_kernel_mod.h"
#include "acl/acl_rt.h"
#include "acl/acl_op_compiler.h"
#include "include/common/utils/utils.h"
namespace mindspore {
namespace kernel {
class VectorAddAscendKernel : public AscendKernelMod {
public:
VectorAddAscendKernel() = default;
~VectorAddAscendKernel() override = default;
bool Init(const BaseOperatorPtr &base_operator,
const std::vector<KernelTensorPtr> &inputs,
const std::vector<KernelTensorPtr> &outputs) override {
// 可选:校验输入输出
return true;
}
bool Launch(const std::vector<AddressPtr> &inputs,
const std::vector<AddressPtr> &,
const std::vector<AddressPtr> &outputs,
void *stream_ptr) override {
auto x1 = GetDeviceAddress<half>(inputs, 0);
auto x2 = GetDeviceAddress<half>(inputs, 1);
auto y = GetDeviceAddress<half>(outputs, 0);
size_t elem_num = inputs[0]->size / sizeof(half);
// 获取算子二进制路径(可硬编码或通过 env 传入)
const char *kernel_path = "/path/to/mindspore/custom/vector_add/op_kernel/build/vector_add_custom.o";
// 构造内核参数
void *args[5];
args[0] = &block_idx_; // coreId(通常设为 0,由 Runtime 分配)
args[1] = &x1;
args[2] = &x2;
args[3] = &y;
args[4] = &elem_num;
// 启动 Kernel
auto ret = aclrtLaunchKernel(kernel_path,
"VectorAddCustom", // 函数名需与 __global__ 一致
1, 1, 1, // gridDim = blockDim = 1(单核模式)
args, 5 * sizeof(void*),
nullptr, 0,
reinterpret_cast<aclrtStream>(stream_ptr));
if (ret != ACL_SUCCESS) {
MS_LOG(ERROR) << "aclrtLaunchKernel failed, ret=" << ret;
return false;
}
return true;
}
private:
uint32_t block_idx_ = 0;
};
// 注册 Kernel
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeCpuKernelMod, VectorAdd,
[]() { return std::make_shared<VectorAddAscendKernel>(); });
} // namespace kernel
} // namespace mindspore
重要说明:
aclrtLaunchKernel是 CANN 提供的底层启动接口;- 算子路径建议通过环境变量配置,避免硬编码;
gridDim=1表示使用单个 Block,多 Core 调度由 Ascend C 内部coreId处理。
六、Step 4:算子定义与 Shape 推导(vector_add_op.cc)
// op_def/vector_add_op.cc
#include "abstract/abstract_value.h"
#include "ops/vector_add.h"
#include "utils/check_convert_utils.h"
namespace mindspore {
namespace ops {
abstract::AbstractBasePtr VectorAddInfer(const abstract::AnalysisEnginePtr &,
const PrimitivePtr &primitive,
const std::vector<abstract::AbstractBasePtr> &input_args) {
// 校验输入数量
CheckArgsSize(primitive->name(), input_args, 2);
// 获取输入 shape 和 dtype
auto x1_shape = input_args[0]->BuildShape();
auto x1_type = input_args[0]->BuildType();
auto x2_shape = input_args[1]->BuildShape();
auto x2_type = input_args[1]->BuildType();
// 校验 shape 是否可广播(简化:要求完全一致)
if (!(*x1_shape == *x2_shape)) {
MS_EXCEPTION(ValueError) << "Input shapes must be equal.";
}
if (!(*x1_type == *x2_type)) {
MS_EXCEPTION(TypeError) << "Input dtypes must be equal.";
}
// 输出 shape 与 dtype 同输入
return abstract::MakeAbstract(x1_shape, x1_type);
}
REGISTER_PRIMITIVE_OP_INFER_IMPL(VectorAdd, prim::kPrimVectorAdd, VectorAddInfer, false);
} // namespace ops
} // namespace mindspore
同时需在 ops/CMakeLists.txt 中添加:
target_sources(mindspore_op_obj PRIVATE
${CMAKE_CURRENT_LIST_DIR}/vector_add_op.cc
)
七、Step 5:Python 接口封装
# mindspore/ops/operations/math_ops.py (追加到文件末尾)
class VectorAdd(Primitive):
r"""
Computes element-wise addition of two input tensors.
Inputs:
- **x1** (Tensor) - First input tensor.
- **x2** (Tensor) - Second input tensor.
Outputs:
Tensor, has the same shape and type as the inputs.
Supported Platforms:
``Ascend``
"""
@prim_attr_register
def __init__(self):
self.init_prim_io_names(inputs=['x1', 'x2'], outputs=['y'])
def vector_add(x1, x2):
r"""
Alias for `VectorAdd`.
"""
return VectorAdd()(x1, x2)
八、Step 6:编译 MindSpore 并集成算子
8.1 修改主 CMakeLists.txt
在 plugin/device/ascend/kernel/CMakeLists.txt 中添加:
add_library(vector_add_kernel SHARED vector_add_kernel.cc)
target_link_libraries(vector_add_kernel ${MS_ASCEND_LIBS})
并在 plugin/device/ascend/CMakeLists.txt 中引入:
add_subdirectory(kernel)
8.2 编译整个项目
cd mindspore
bash build.sh -S on -A x86_64 -j16
编译成功后,生成 libvector_add_kernel.so,MindSpore 会自动加载。
九、Step 7:编写单元测试(test_vector_add.py)
# tests/ut/python/ops/test_vector_add.py
import numpy as np
import pytest
from mindspore import Tensor, context
from mindspore.ops.operations.math_ops import vector_add
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
@pytest.mark.level0
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
def test_vector_add_basic():
"""Test basic functionality."""
x1 = Tensor(np.array([1.0, 2.0, 3.0], dtype=np.float16))
x2 = Tensor(np.array([4.0, 5.0, 6.0], dtype=np.float16))
expect = np.array([5.0, 7.0, 9.0], dtype=np.float16)
output = vector_add(x1, x2).asnumpy()
assert np.allclose(output, expect, atol=1e-3)
@pytest.mark.level1
def test_vector_add_large():
"""Test large tensor."""
n = 1024 * 1024
x1 = Tensor(np.random.randn(n).astype(np.float16))
x2 = Tensor(np.random.randn(n).astype(np.float16))
output = vector_add(x1, x2)
assert output.shape == (n,)
运行测试:
pytest tests/ut/python/ops/test_vector_add.py -v
十、支持反向传播(训练场景)
若用于训练,需注册梯度函数:
# mindspore/ops/composite/basic_grad.py (追加)
from mindspore.ops.operations.math_ops import VectorAdd
@bprop_getters.register(VectorAdd)
def get_bprop_vector_add(self):
def bprop(x1, x2, out, dout):
return dout, dout # ∂L/∂x1 = ∂L/∂y, ∂L/∂x2 = ∂L/∂y
return bprop
验证训练:
import mindspore.nn as nn
class Net(nn.Cell):
def construct(self, x1, x2):
return vector_add(x1, x2)
net = Net()
grad_fn = ms.value_and_grad(net, grad_position=(0, 1))
x1, x2 = Tensor([1.0]), Tensor([2.0])
loss, grads = grad_fn(x1, x2)
print(grads) # 应输出 (Tensor([1.]), Tensor([1.]))
十一、性能分析与优化建议
11.1 使用 Profiler
from mindspore.profiler import Profiler
profiler = Profiler()
y = vector_add(x1, x2)
profiler.analyse()
查看 PROFILING 目录下的 Timeline,确认:
- 算子是否被正确调度;
- 是否存在 Host-Device 同步瓶颈。
11.2 优化建议
| 问题 | 优化手段 |
|---|---|
| 小 tensor 性能差 | 合并多个小算子(Kernel Fusion) |
| 启动开销大 | 使用 acl.op.load 预加载 .o 文件 |
| 内存拷贝多 | 确保输入输出为 Device Tensor |
十二、常见错误与排查表
| 错误现象 | 可能原因 | 解决方案 |
|---|---|---|
aclrtLaunchKernel failed, ret=507004 |
算子路径错误或 .o 不存在 | 检查 kernel_path,确认文件权限 |
Shape mismatch |
InferShape 未正确实现 | 在 VectorAddInfer 中打印 shape 调试 |
| 算子不执行 | 未注册 Kernel | 检查 MS_KERNEL_FACTORY_REG 是否拼写正确 |
| 结果全零 | UB 未初始化或搬运 size 错 | 在 Ascend C 中加 ASSERT 校验 |
| 编译失败 | CANN 版本不匹配 | 确认 aic 与 MindSpore CANN 依赖一致 |
十三、总结与最佳实践
本文提供了一套 完整、可落地 的 Ascend C 算子集成到 MindSpore 的方案。关键成功要素包括:
- 清晰的分层设计:Ascend C / Host Kernel / Python 接口职责分离;
- 严格的 Shape/Dtype 校验:避免运行时崩溃;
- 完善的测试覆盖:从小 tensor 到大 tensor,从前向到反向;
- 性能可观测:通过 Profiler 验证优化效果。
企业级建议:
- 将自定义算子打包为 MindSpore 插件(.whl),避免修改主干;
- 使用 CI/CD 自动化测试(如 Jenkins + Atlas 设备池);
- 文档化算子规格(输入范围、精度误差、性能基线)。
通过本文方法,您可将任何 Ascend C 算子(GEMM、Conv、Softmax 等)无缝集成到 MindSpore,构建高性能、定制化的 AI 模型。
2025年昇腾CANN训练营第二季,基于CANN开源开放全场景,推出0基础入门系列、码力全开特辑、开发者案例等专题课程,助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证,即可领取精美证书,完成社区任务更有机会赢取华为手机,平板、开发板等大奖。
报名链接:https://www.hiascend.com/developer/activities/cann20252
昇腾计算产业是基于昇腾系列(HUAWEI Ascend)处理器和基础软件构建的全栈 AI计算基础设施、行业应用及服务,https://devpress.csdn.net/organization/setting/general/146749包括昇腾系列处理器、系列硬件、CANN、AI计算框架、应用使能、开发工具链、管理运维工具、行业应用及服务等全产业链
更多推荐


所有评论(0)