前言CANN 自带的算子库已经覆盖了大部分常用算子但总会遇到不支持的。这时候要用 Ascend C 自己写一个算子。Ascend C 是昇腾的算子编程语言语法类似 C但专门针对达芬奇架构做了优化。这篇文章讲怎么从零写一个矩阵乘法算子。开发环境准备安装算子开发包# CANN 工具包自带算子开发工具whichcanndev# 输出/usr/local/Ascend/ascend-toolkit/latest/bin/canndev# 如果没有安装 toolkitaptinstallascend-toolkit创建算子工程# 用 canndev 创建算子工程canndev create--projectmatmul_custom--output./workspace# 工程结构tree ./workspace/matmul_custom/# output:# matmul_custom/# ├── CMakeLists.txt# ├── framework/# │ └── tensor_add_impl.cpp # 算子实现# ├── op_proto/# │ └── tensor_add.h # 算子原型# ├── unittest/# │ └── test_tensor_add.py # 单元测试# └── CMakeLists.txt算子原型定义第一步是定义算子的原型输入、输出、属性。头文件定义// matmul_custom.h#ifndefMATMUL_CUSTOM_H#defineMATMUL_CUSTOM_H#includetiking_pub_api.hnamespaceacl{namespaceops{classMatMulCustom{public:MatMulCustom()default;~MatMulCustom()default;// 初始化算子graphStatusInit(constTensora,constTensorb,Tensorc,booltrans_afalse,booltrans_bfalse);// 推理接口graphStatusInferShape();// 数据类型推断graphStatusInferDataType();};}// namespace ops}// namespace acl#endif// MATMUL_CUSTOM_H实现文件// matmul_custom.cpp#includematmul_custom.hnamespaceacl{namespaceops{graphStatusMatMulCustom::Init(constTensora,constTensorb,Tensorc,booltrans_a,booltrans_b){// 设置输入(void)ge::Node::GetNodeFromTensor(a);(void)ge::Node::GetNodeFromTensor(b);// 设置属性SetAttr(trans_a,trans_a);SetAttr(trans_b,trans_b);returnGRAPH_SUCCESS;}graphStatusMatMulCustom::InferShape(){// 获取输入 shapestd::vectorint64_tshape_a;std::vectorint64_tshape_b;GetInputDesc(0).GetShape(shape_a);GetInputDesc(1).GetShape(shape_b);// 计算输出 shapebooltrans_aGetAttr(trans_a)-GetBool();booltrans_bGetAttr(trans_b)-GetBool();int64_tmtrans_a?shape_a[1]:shape_a[0];int64_tntrans_b?shape_b[0]:shape_b[1];std::vectorint64_tshape_c{m,n};GetOutputDesc(0).SetShape(shape_c);returnGRAPH_SUCCESS;}graphStatusMatMulCustom::InferDataType(){// 输出数据类型和输入 A 一致DataType dtypeGetInputDesc(0).GetDataType();GetOutputDesc(0).SetDataType(dtype);returnGRAPH_SUCCESS;}}// namespace ops}// namespace acl算子实现核心这是最关键的部分用 Ascend C 实现矩阵乘法。基本框架#includekernel_operator.hclassMatMulCustomKernel{public:__aicore__inlineMatMulCustomKernel(){}__aicore__voidInit(GM_ADDR a,GM_ADDR b,GM_ADDR c,int32_tM,int32_tN,int32_tK);__aicore__voidProcess();private:GlobalTensorfloataGm;// 输入 AGM 内存GlobalTensorfloatbGm;// 输入 BGM 内存GlobalTensorfloatcGm;// 输出 CGM 内存LocalTensorfloataLocal;// 输入 AUB 内存LocalTensorfloatbLocal;// 输入 BUB 内存LocalTensorfloatcLocal;// 输出 CUB 内存int32_tM,N,K;};externC__global__voidmatmul_custom(GM_ADDR a,GM_ADDR b,GM_ADDR c,int32_tM,int32_tN,int32_tK){MatMulCustomKernel op;op.Init(a,b,c,M,N,K);op.Process();}Init 函数__aicore__voidMatMulCustomKernel::Init(GM_ADDR a,GM_ADDR b,GM_ADDR c,int32_tM,int32_tN,int32_tK){// 设置 GM 地址aGm.SetGlobalBuffer((__gm__float*)a);bGm.SetGlobalBuffer((__gm__float*)b);cGm.SetGlobalBuffer((__gm__float*)c);// 保存参数this-MM;this-NN;this-KK;// 分配 UB 空间constexprint32_tBLOCK_SIZE32;// 一个 block 的大小aLocalGetUbBlockfloat(M*K/BLOCK_SIZE);bLocalGetUbBlockfloat(K*N/BLOCK_SIZE);cLocalGetUbBlockfloat(M*N/BLOCK_SIZE);}Process 函数核心计算__aicore__voidMatMulCustomKernel::Process(){// 1. 把数据从 GM 搬到 UBCopyIn(aGm,aLocal,M*K);CopyIn(bGm,bLocal,K*N);// 2. 矩阵乘法计算MatMul(cLocal,aLocal,bLocal,M,N,K);// 3. 把结果从 UB 写回 GMCopyOut(cLocal,cGm,M*N);}// 矩阵乘法实现voidMatMul(LocalTensorfloatc,LocalTensorfloata,LocalTensorfloatb,int32_tM,int32_tN,int32_tK){// 简化版本逐行逐列计算for(int32_ti0;iM;i){for(int32_tj0;jN;j){floatsum0.0f;for(int32_tk0;kK;k){suma[i*Kk]*b[k*Nj];}c[i*Nj]sum;}}}优化版本使用 Cube Unit#includelib_api.h// 使用 Cube Unit 做矩阵乘法voidMatMulCube(LocalTensorfloatc,LocalTensorfloata,LocalTensorfloatb,int32_tM,int32_tN,int32_tK){// 调用 Cube Unit 的 MatMul 接口matmul_t matmul_para;matmul_para.MM;matmul_para.NN;matmul_para.KK;matmul_para.is_trans_afalse;matmul_para.is_trans_bfalse;// 用 Cube Unit 计算MatMul(c,a,b,matmul_para);}编译算子CMake 配置# CMakeLists.txt cmake_minimum_required(VERSION 3.10) project(matmul_custom) # 设置 Ascend C 编译选项 set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS} -marchdavinci) set(CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS} -O2) # 添加算子实现文件 add_library(matmul_custom SHARED matmul_custom.cpp matmul_custom_kernel.cpp ) # 链接 Ascend C 库 target_link_libraries(matmul_custom ascendcl tbe_operator )编译命令# 创建编译目录mkdirbuildcdbuild# 配置cmake..# 编译make-j8# 输出libmatmul_custom.so在 PyTorch 里调用注册自定义算子importtorchimporttorch.npuimportctypes# 加载自定义算子库libctypes.CDLL(./libmatmul_custom.so)# 定义 Python 接口defmatmul_custom(a:torch.Tensor,b:torch.Tensor)-torch.Tensor:自定义矩阵乘法# 检查输入asserta.dim()2andb.dim()2asserta.shape[1]b.shape[0]M,Ka.shape K2,Nb.shape# 创建输出 tensorctorch.empty(M,N,devicenpu:0,dtypetorch.float32)# 调用自定义算子lib.matmul_custom(a.data_ptr(),b.data_ptr(),c.data_ptr(),M,N,K)returnc测试自定义算子# 创建测试数据atorch.randn(128,256,devicenpu:0)btorch.randn(256,64,devicenpu:0)# 用自定义算子计算c_custommatmul_custom(a,b)# 用 PyTorch 内置算子计算参考c_reftorch.matmul(a,b)# 对比精度cosine_simtorch.nn.functional.cosine_similarity(c_custom.flatten(),c_ref.flatten(),dim0)print(f余弦相似度:{cosine_sim.item():.6f})# 对比性能importtime# 预热for_inrange(10):matmul_custom(a,b)torch.npu.synchronize()starttime.time()for_inrange(100):matmul_custom(a,b)torch.npu.synchronize()endtime.time()print(f自定义算子延迟:{(end-start)/100*1000:.2f}ms)性能优化使用 Cube Unit达芬奇架构的 Cube Unit 是专门为矩阵乘法设计的硬件单元比用 Vector Unit 算快得多。// 使用 Cube Unit#includelib_api.hvoidMatMulOptimized(LocalTensorfloatc,LocalTensorfloata,LocalTensorfloatb,int32_tM,int32_tN,int32_tK){// 检查是否可以用 Cubeif(M16N16K16){// 用 Cube Unitmatmul_t params;params.MM;params.NN;params.KK;MatMul(c,a,b,params);}else{// 用小矩阵算法MatMulSmall(c,a,b,M,N,K);}}分块计算当矩阵太大UB 放不下时要分块计算。constexprint32_tTILE_SIZE32;// 块大小voidMatMulTiled(LocalTensorfloatc,GlobalTensorfloataGm,GlobalTensorfloatbGm,int32_tM,int32_tN,int32_tK){// 分块计算for(int32_ti0;iM;iTILE_SIZE){for(int32_tj0;jN;jTILE_SIZE){// 搬运当前块CopyInA(aGm,aLocal,i,min(iTILE_SIZE,M),K);CopyInB(bGm,bLocal,j,min(jTILE_SIZE,N),K);// 计算当前块MatMulTile(cLocal,aLocal,bLocal,min(TILE_SIZE,M-i),min(TILE_SIZE,N-j),K);// 写回当前块CopyOutC(cLocal,cGm,i,j,min(TILE_SIZE,M-i),min(TILE_SIZE,N-j));}}}调试技巧打印调试// 在算子实现里加打印#includecstdio__aicore__voidMatMulCustomKernel::Process(){// 打印输入参数printf(M%d, N%d, K%d\n,M,N,K);// 打印输入数据前 10 个for(inti0;imin(10,M*K);i){printf(a[%d]%f\n,i,aLocal.GetValue(i));}// 计算结果MatMul(cLocal,aLocal,bLocal,M,N,K);// 打印输出数据前 10 个for(inti0;imin(10,M*N);i){printf(c[%d]%f\n,i,cLocal.GetValue(i));}}用 Ascend CL 调试importacl# 初始化acl.init()acl.rt.set_device(0)# 运行算子调试模式acl.rt.set_op_execute_mode(debug)# 运行# ...运行算子# 查看日志# 日志在 /var/log/Ascend/ascend_toolkit/matmul_custom.log常见问题问题一编译报错error: matmul_t was not declared in this scope解决包含正确的头文件#includelib_api.h// 包含 matmul_t 的定义问题二运行时报错[ERROR] Kernel execute failed: out of memory解决减小块大小或者检查 UB 大小是否足够// 检查 UB 大小uint32_tub_sizeGetUbSize();printf(UB size: %u bytes\n,ub_size);// 减小块大小constexprint32_tTILE_SIZE16;// 从 32 改成 16问题三精度不达标余弦相似度: 0.97 (应该 0.99)解决检查计算精度可能需要用 FP32// 用 FP32 计算GlobalTensorfloataGm;// float FP32GlobalTensorfloatbGm;GlobalTensorfloatcGm;// 不要用 FP16// GlobalTensorhalf aGm; // half FP16精度可能不够参考资源Ascend C 编程指南: https://www.hiascend.com/document/detail/zh/CANN/算子开发最佳实践: https://www.hiascend.com/document/detail/zh/CANN/算子样例代码: https://atomgit.com/cann/samples达芬奇架构白皮书: https://www.hiascend.com/document/detail/zh/CANN/总结用 Ascend C 开发自定义算子流程是定义算子原型 → 实现算子逻辑 → 编译成.so→ 在 PyTorch 里调用。核心是要理解达芬奇架构的内存层次GM → UB → 计算单元以及 Cube Unit 和 Vector Unit 的适用场景。性能优化的关键是尽量用 Cube Unit 做矩阵乘法太大的矩阵要分块计算。调试可以用printf打印中间结果或者用 Ascend CL 的调试模式。精度问题通常是因为用了 FP16改成 FP32 一般能解决。
