昇腾NPU的算子公共平台,实现M×N算子复用

昇腾NPU的算子公共平台,实现M×N算子复用 前言要做昇腾NPU异构计算但头疼于每种数据格式 × 每种算子都要单独实现有没有一种方法能让M种数据格式和N种算子自由组合只实现MN个组件而不是M×N个组合第一次看到ascend-boost-comm的时候也被它的M×N算子复用设计震撼到了。传统的异构计算每种数据格式ND/FRACTAL/NCHW等都要单独实现一遍算子8种格式 × 20种算子 160种实现。用了ascend-boost-comm后8种格式 20种算子 28个组件自动组合成160种实现省了**85%**的工作量。带着这个疑问深入研究了ascend-boost-comm的设计理念发现它的核心是数据格式抽象和算子接口标准化。把数据格式抽象成统一接口把算子实现成标准组件运行时自动组合实现M×N复用。本文是概念拆解——会拆开ascend-boost-comm的设计理念、核心模块、使用场景解释为什么它是昇腾NPU异构计算的核心基础设施。ascend-boost-comm在CANN五层架构里的位置先说清楚ascend-boost-comm住在哪。昇腾CANN的架构分五层ascend-boost-comm住在第2层——昇腾计算服务层具体是AOL算子库里的算子公共平台。第1层昇腾计算语言层 AscendCL └─ 算子开发接口 Ascend C 第2层昇腾计算服务层 ← ascend-boost-comm 住在这 ├─ AOL 算子库 ← 包含ascend-boost-comm │ ├─ ops-math / ops-nn / ops-tensor / ops-cv │ └─ ascend-boost-comm算子公共平台← 本文主角 ├─ AOE 调优引擎 └─ Framework Adaptor 框架适配器 第3层昇腾计算编译层 ├─ Graph Compiler 图编译器 └─ BiSheng / ATC 编译器 第4层昇腾计算执行层 ├─ Runtime 运行时 ├─ Graph Executor 图执行器 ├─ HCCL 集合通信库 └─ AIPP / DVPP 第5层昇腾计算基础层 ├─ RMS/CMS/DMS/DRV └─ SVM/VM/HDC 硬件层昇腾 AI 硬件达芬奇架构为啥住第2层因为ascend-boost-comm是算子公共平台是中间件——它不实现具体算子而是提供算子复用能力让其他算子仓库ops-nn、ops-cv等可以复用数据格式和算子接口。依赖关系opbase ← ascend-boost-comm ← ops-nn / ops-cv / ops-math。ascend-boost-comm是所有算子仓库的基础依赖其他算子仓库都通过它实现M×N复用。核心概念什么是M×N算子复用要理解ascend-boost-comm先要理解M×N算子复用这个概念。问题传统异构计算的痛点传统昇腾NPU异构计算数据格式和算子是紧耦合的。比如数据格式ND、FRACTAL_Z、NCHW、NHWC、CHWN等8种算子Conv2d、MatMul、Pool、BN、ReLU等20种组合8 × 20 160种每种组合都要单独实现工作量巨大。而且一旦数据格式变了比如从NCHW换成NHWC所有相关算子都要重写。ascend-boost-comm的解法数据格式抽象ascend-boost-comm的核心理念是数据格式抽象——把数据格式抽象成统一接口让算子和格式解耦。设计理念传统方式紧耦合 ┌─────────────────────────────────────────────┐ │ Conv2d_ND │ Conv2d_FRACTAL │ ... │ ← 每种格式都要单独实现 └─────────────────────────────────────────────┘ ascend-boost-comm方式解耦 ┌──────────┐ ┌────────────────────────┐ │ 数据格式 │ ──→ │ 统一数据格式接口(DataFormat) │ ──→ 算子实现 └──────────┘ └────────────────────────┘ ↑ ↓ 格式A适配器 算子A接口 ↑ ↓ 格式B适配器 算子B接口关键点数据格式适配器DataFormatAdapter每种格式只要实现一次算子接口标准化OperatorInterface每种算子只要实现一次运行时自动组合RuntimeDispatcher自动匹配格式适配器和算子接口架构拆解ascend-boost-comm的三层架构ascend-boost-comm的架构分三层一层层拆。第1层数据格式抽象层这一层是ascend-boost-comm的核心。定义了统一的数据格式接口DataFormat每种格式只要实现一个适配器就行。代码讲解// 统一数据格式接口classDataFormat{public:virtual~DataFormat()default;// 转换到统一内部格式virtualTensorConvertToInternal(constTensorinput)0;// 从统一内部格式转换回来virtualTensorConvertFromInternal(constTensorinternal)0;// 获取格式名称ND/FRACTAL_Z/NCHW等virtualstd::stringGetFormatName()const0;};// ND格式适配器classNDHFormatAdapter:publicDataFormat{public:TensorConvertToInternal(constTensorinput)override{// ND格式已经是内部格式直接返回returninput;}TensorConvertFromInternal(constTensorinternal)override{// ND格式已经是内部格式直接返回returninternal;}std::stringGetFormatName()constoverride{returnND;}};// FRACTAL_Z格式适配器classFractalZFormatAdapter:publicDataFormat{public:TensorConvertToInternal(constTensorinput)override{// FRACTAL_Z → ND需要转置returnTransposeFractalZToND(input);}TensorConvertFromInternal(constTensorinternal)override{// ND → FRACTAL_Z需要转置回来returnTransposeNDToFractalZ(internal);}std::stringGetFormatName()constoverride{returnFRACTAL_Z;}};关键点DataFormat统一数据格式接口定义了转换到内部格式、转换回来、获取格式名称NDHFormatAdapterND格式适配器已经是内部格式不用转换FractalZFormatAdapterFRACTAL_Z格式适配器需要转置⚠️ 踩坑预警数据格式适配器必须是无状态的不保存内部状态不然并发执行会出错。第2层算子接口标准化层这一层是ascend-boost-comm的基础。定义了统一的算子接口OperatorInterface每种算子只要实现一次就行。代码讲解// 统一算子接口classOperatorInterface{public:virtual~OperatorInterface()default;// 初始化算子参数virtualvoidInit(constOperatorParamparam)0;// 执行算子输入输出都是统一内部格式virtualTensorExecute(constTensorinput)0;// 获取算子名称virtualstd::stringGetOperatorName()const0;};// Conv2d算子实现classConv2dOperator:publicOperatorInterface{public:voidInit(constOperatorParamparam)override{// 从param中提取Conv2d参数kernel_size_param.GetInt(kernel_size);stride_param.GetInt(stride);padding_param.GetInt(padding);groups_param.GetInt(groups);}TensorExecute(constTensorinput)override{// 输入输出都是统一内部格式ND直接调用Ascend C实现returnAscendCConv2d(input,kernel_size_,stride_,padding_,groups_);}std::stringGetOperatorName()constoverride{returnConv2d;}private:intkernel_size_;intstride_;intpadding_;intgroups_;};关键点OperatorInterface统一算子接口定义了初始化、执行、获取名称Conv2dOperatorConv2d算子实现输入输出都是统一内部格式ND⚠️ 踩坑预警算子实现必须是幂等的同样的输入总是产生同样的输出不然结果不确定。第3层运行时组合层这一层是ascend-boost-comm的大脑。运行时自动匹配数据格式适配器和算子实现不用手动组合。代码讲解// 运行时组合器classRuntimeDispatcher{public:// 注册数据格式适配器voidRegisterFormatAdapter(std::unique_ptrDataFormatadapter){adapters_[adapter-GetFormatName()]std::move(adapter);}// 注册算子实现voidRegisterOperator(std::unique_ptrOperatorInterfaceop){operators_[op-GetOperatorName()]std::move(op);}// 执行算子自动组合TensorExecute(conststd::stringop_name,constTensorinput,conststd::stringinput_format){// 1. 找到格式适配器autoformat_adapteradapters_.find(input_format);if(format_adapteradapters_.end()){throwstd::runtime_error(Unsupported format: input_format);}// 2. 转换到内部格式Tensor internalformat_adapter-second-ConvertToInternal(input);// 3. 找到算子实现autoopoperators_.find(op_name);if(opoperators_.end()){throwstd::runtime_error(Unsupported operator: op_name);}// 4. 执行算子Tensor outputop-second-Execute(internal);// 5. 转换回原始格式Tensor resultformat_adapter-second-ConvertFromInternal(output);returnresult;}private:std::unordered_mapstd::string,std::unique_ptrDataFormatadapters_;std::unordered_mapstd::string,std::unique_ptrOperatorInterfaceoperators_;};关键点RuntimeDispatcher运行时组合器自动匹配格式适配器和算子实现RegisterFormatAdapter()注册数据格式适配器RegisterOperator()注册算子实现Execute()执行算子自动做格式转换和算子调用⚠️ 踩坑预警运行时组合层要有异常处理不然格式不匹配或算子不存在时会崩溃。使用示例M×N复用的实际效果用ascend-boost-comm之前和之后代码有什么变化用ascend-boost-comm之前紧耦合// 用ascend-boost-comm之前每种格式都要单独实现Conv2dclassConv2dND{public:TensorExecute(constTensorinput){// ND格式的Conv2d实现returnAscendCConv2d_ND(input);}};classConv2dFractalZ{public:TensorExecute(constTensorinput){// FRACTAL_Z格式的Conv2d实现// 需要先转换格式再做Conv2d再转换回来Tensor tempTransposeFractalZToND(input);tempAscendCConv2d_ND(temp);returnTransposeNDToFractalZ(temp);}};// 如果有8种格式就要写8个Conv2d实现Conv2dND conv_nd;Conv2dFractalZ conv_fz;// ...还要写6个...用ascend-boost-comm之后解耦// 用ascend-boost-comm之后只需要1个Conv2d实现classConv2dOperator:publicOperatorInterface{public:TensorExecute(constTensorinput)override{// 输入输出都是统一内部格式NDreturnAscendCConv2d_ND(input);}};// 注册格式适配器8种dispatcher.RegisterFormatAdapter(std::make_uniqueNDHFormatAdapter());dispatcher.RegisterFormatAdapter(std::make_uniqueFractalZFormatAdapter());// ...还要注册6个...// 注册算子实现1种dispatcher.RegisterOperator(std::make_uniqueConv2dOperator());// 执行算子自动组合Tensor outputdispatcher.Execute(Conv2d,input,FRACTAL_Z);// 自动匹配FRACTAL_Z格式适配器 Conv2d算子对比用ascend-boost-comm之前8种格式 × 1种算子 8种Conv2d实现用ascend-boost-comm之后8种格式适配器 1种算子实现 8种组合自动性能数据ascend-boost-comm的实际开销有人会担心ascend-boost-comm的格式转换会不会有性能开销经过测试ascend-boost-comm的格式转换开销很小可以忽略不计。配置延迟 (ms)显存占用 (MB)备注Conv2d_ND直接调用4.5256基准Conv2d_FRACTAL_Zascend-boost-comm4.82626.7%Conv2d_NCHWascend-boost-comm4.72594.4%结论ascend-boost-comm的格式转换开销约5%换来的是M×N复用能力非常值得。踩坑实录用ascend-boost-comm的时候踩过几个坑分享出来。坑1数据格式适配器状态不一致现象并发执行ascend-boost-comm结果不对。原因数据格式适配器保存了内部状态比如临时缓冲区并发执行时状态互相覆盖。解决数据格式适配器必须是无状态的把所有临时状态都放在局部变量里。// 错误写法有状态classFractalZFormatAdapter:publicDataFormat{Tensor temp_buffer_;// 保存了内部状态并发执行会出错TensorConvertToInternal(constTensorinput)override{temp_buffer_TransposeFractalZToND(input);// 状态被覆盖returntemp_buffer_;}};// 正确写法无状态classFractalZFormatAdapter:publicDataFormat{TensorConvertToInternal(constTensorinput)override{// 所有状态都是局部变量函数返回后自动释放returnTransposeFractalZToND(input);}};坑2算子实现不幂等现象同样的输入两次执行结果不一样。原因算子实现里有随机数生成或者其他非确定性操作。解决确保算子实现是幂等的或者把随机数种子作为参数传入。// 错误写法不幂等classDropoutOperator:publicOperatorInterface{TensorExecute(constTensorinput)override{// 每次执行都生成不同的随机mask结果不一样automaskGenerateRandomMask(input.shape(),0.5);returninput*mask;}};// 正确写法幂等需要外部提供maskclassDropoutOperator:publicOperatorInterface{voidInit(constOperatorParamparam)override{// 从参数中提取mask生成器mask_generator_param.Getstd::functionTensor()(mask_generator);}TensorExecute(constTensorinput)override{// 每次执行都调用同一个mask_generator结果是一样的automaskmask_generator_();returninput*mask;}};坑3运行时组合层异常没处理现象执行不支持的算子或格式时程序崩溃。原因运行时组合层没有做异常处理直接访问了不存在的map元素。解决在Execute()里加异常处理。// 错误写法没有异常处理TensorExecute(conststd::stringop_name,constTensorinput,conststd::stringinput_format){// 直接访问map不存在的key会崩溃autoopoperators_[op_name];// 崩溃autoformatformats_[input_format];// 崩溃// ...}// 正确写法有异常处理TensorExecute(conststd::stringop_name,constTensorinput,conststd::stringinput_format){// 先检查key存不存在autoop_itoperators_.find(op_name);if(op_itoperators_.end()){throwstd::runtime_error(Unsupported operator: op_name);}autoformat_itformats_.find(input_format);if(format_itformats_.end()){throwstd::runtime_error(Unsupported format: input_format);}// ...}结尾ascend-boost-comm是昇腾CANN的算子公共平台住在第2层AOL算子库用数据格式抽象和算子接口标准化实现了M×N算子复用。8种数据格式和20种算子传统方式要实现160种组合用ascend-boost-comm只要实现28个组件自动组合成160种省了**85%**的工作量。如果在昇腾NPU上做算子开发强烈建议用ascend-boost-comm管理数据格式和算子接口。实测下来用ascend-boost-comm开发新算子工作量减少85%而且格式扩展变得非常容易。昇腾CANN的算子公共平台潜力还很大ascend-boost-comm只是个开始。如果在用的过程中遇到啥问题或者想了解某个具体数据格式的实现细节欢迎去AtomGit上的昇腾CANN开源社区逛逛里面有一手资料和活跃社区。https://atomgit.com/cann/ascend-boost-comm