昇腾CANN ascend-boost-comm:M×N 算子复用是怎么做到的
CANN 生态里 50 多个仓库每个仓库有十几到几十个算子。这些算子之间存在大量公共功能内存搬运算子需要数据切分、通信算子需要拓扑发现、融合算子需要 shape 推导。如果每个仓库各自实现一遍代码膨胀的同时任何一个公共功能的 bug 修复或性能优化需要推开几十个仓库的 PR。ascend-boost-comm 的设计目标就是把这个问题变成 M×N 的复用M 个上层算子仓库通过 N 个公共模块共享实现而不是 M×N 的各自重复。它不是通信库——名字里带 comm 容易误读。ascend-boost-comm 是算子公共平台提供的是中间件性质的基础能力数据切片、拓扑感知、生命周期管理、跨算子状态共享。为什么需要中间件层看一个实际例子。CANN 里有三个仓库都需要对输入张量做二维切分ops-math 的 Reduction 算子需要沿 dim 切分ops-nn 的 MatMul 需要按 M×K 分块ops-transformer 的 FlashAttention 需要按 Br×Bc 分块没有 ascend-boost-comm 时三个仓库各自写切分逻辑// ops-math/reduce/tiling.cpp —— 自己的切分代码// ops-nn/matmul/tiling.cpp —— 差不多的切分代码换了个名字// ops-transformer/flash_attn/tiling.cpp —— 还是差不多的切分代码有 ascend-boost-comm 时三个仓库共用统一的切分框架// ascend-boost-comm/tiling/tiling_framework.h// 所有算子仓库共用此接口templateintDIMstructTilingStrategy{intnum_blocks[DIM];// 每维切多少块intblock_size[DIM];// 每块的大小intremainder_block[DIM];// 尾块的策略对齐/不对齐staticTilingStrategycompute(constShapeDIMtotal_shape,// 总shapeconstShapeDIMmax_block,// 单块最大容量L1约束TilingPolicy policy// 切分策略){TilingStrategy result;for(intd0;dDIM;d){// 按 L1 容量约束计算最优分块intmax_elementsmax_block[d];inttotaltotal_shape[d];result.num_blocks[d](totalmax_elements-1)/max_elements;result.block_size[d]max_elements;// 尾块处理可以选择对齐到 16Cube 约束或保持原始大小intlast_sizetotal-(result.num_blocks[d]-1)*max_elements;if(policyALIGN_TO_CUBElast_size%16!0){result.remainder_block[d](last_size15)/16*16;}else{result.remainder_block[d]last_size;}}returnresult;}};三个仓库的代码变成// 三个仓库各自只用一行调用autotilingTilingStrategy2::compute({M,N},// 矩阵尺寸{MAX_M_TILE,MAX_N_TILE},// L1 容量上限ALIGN_TO_CUBE// 对齐策略);修复一个切分 bug升级 ascend-boost-comm 里的 TilingStrategy 就行——所有 50 个仓库自动受益。五大公共模块一、数据切片引擎除了前面的 shape 维切分还处理数据布局转换。算子从 NCHW 换到 NHWC、从 RowMajor 换到 ColMajor 的跨步映射全部由切片引擎提供// 数据布局转换 分块一次调用#includeascend-boost-comm/tiling/data_slice.hautosliceDataSlice::builder().shape({BATCH,CHANNEL,HEIGHT,WIDTH}).layout(LAYOUT_NCHW)// 输入格式.target_layout(LAYOUT_NHWC)// 输出格式Cube 友好.max_block_size(L1_CAPACITY)// L1 容量约束.build();// slice 自动生成最优的切分计划——// 包含了 layout 转换所需的 stride 映射二、拓扑发现服务分布式算子AllReduce、AllGather 等需要知道 NPU 之间的物理拓扑来选最优算法。ascend-boost-comm 提供统一的拓扑发现// 任何算子仓库都可以调拓扑发现#includeascend-boost-comm/topology/topo_discovery.hTopologyGraph topoTopologyDiscovery::get_instance()-discover();// 判断任意两张 NPU 之间走什么链路for(inti0;inum_npus;i){for(intji1;jnum_npus;j){autopathtopo.shortest_path(i,j);// path.type: NVLink / RoCE / PCIe// path.bandwidth: 链路有效带宽GB/s// path.latency: 链路延迟μs}}// 基于拓扑选算法if(topo.is_nvlink_full_mesh()){returnALG_HALVING_DOUBLING;}elseif(topo.is_ring()){returnALG_RING;}else{returnALG_NAIVE;}这个接口被 hcomm、hccl、asc-comm 三个通信层共用。拓扑发现逻辑只在 ascend-boost-comm 里维护一份改动了 NPU 拓扑描述数据结构后三个通信层自动同步。三、算子生命周期管理CANN 算子从注册到执行有完整的生命周期注册 → InferShape → Tiling → 内存分配 → Kernel Dispatch → 执行 → 内存释放。ascend-boost-comm 管理这个生命周期让每个算子只关注「计算逻辑」部分// 算子生命周期——ascend-boost-comm 统一管理#includeascend-boost-comm/lifecycle/op_lifecycle.h// 算子开发者只需要实现 OpInterfaceclassMyAddOp:publicOpInterface{ShapeInferShape(constvectorShapeinputs)override{...}KernelTypeDispatchKernel(constOpConfigconfig)override{...}StatusExecute(constvectorTensorinputs,Tensoroutput)override{...}};// ascend-boost-comm 管剩下的所有事// - 内存预分配从内存池复用// - workspace 管理// - 异步执行流绑定// - 执行完成的同步点autolifecycleOpLifecycle::createMyAddOp();lifecycle-infer_shape(inputs);lifecycle-allocate_memory();lifecycle-dispatch_kernel();lifecycle-execute();lifecycle-free_memory();四、跨算子状态共享某些状态需要跨多个算子共享——比如混合精度训练的 loss scale 因子、推理的 KV Cache 块池。ascend-boost-comm 提供了一个分布式状态管理器// 跨算子全局状态#includeascend-boost-comm/state/global_state.h// 设置全局状态任意算子可读写GlobalState::set(amp_loss_scale,65536.0f);GlobalState::set(kv_cache_block_pool,pool_ptr);// op-nn 的 LayerNorm 读 loss_scalefloatscaleGlobalState::getfloat(amp_loss_scale);// op-transformer 的 Attention 读 KV Cache 池auto*poolGlobalState::getvoid*(kv_cache_block_pool);状态管理器解决了「全局配置项到处传参数」的问题——loss_scale 只需要在 AMP 初始化时设一次后续所有算子的梯度缩放自动感知。五、调试与诊断算子出问题时快速定位是哪个阶段出的错。ascend-boost-comm 内建了分阶段的 profiling 和诊断// 分阶段 profiling// ascend-boost-comm 在生命周期每个阶段自动插桩#includeascend-boost-comm/debug/profiler.hOpProfilerprofiler(MatMulV2);profiler.enable_trace();// 开启全生命周期跟踪// 执行后输出// [MatMulV2] InferShape: 0.12ms// [MatMulV2] Tiling: 0.05ms// [MatMulV2] AllocMem: 0.23ms ← 瓶颈内存分配慢了// [MatMulV2] Dispatch: 0.01ms// [MatMulV2] Execute: 2.34ms// [MatMulV2] FreeMem: 0.08msProfiling 是分阶段自动注入的不需要算子开发者手动加计时器。依赖关系全景ascend-boost-comm 在 CANN 依赖链中的位置opbase基础组件Tensor、DataType ↓ ascend-boost-comm公共平台Tiling、Topology、Lifecycle、State、Debug ↓ ├─ ops-math / ops-nn / ops-blas / ops-cv ...核心算子仓库 ├─ hccl集合通信库——用 Topology 做算法选择 ├─ hcomm高层通信原语——用 Topology Lifecycle └─ ge图编译器——用 Lifecycle 管理算子执行流每个上层仓库通过 ascend-boost-comm 的模块各取所需。hccl 可能只用 Topology 模块ops-nn 用了 Tiling Lifecycle Debug 三个模块——但代码是同一套维护也是同一套。M×N 复用不是新鲜的架构概念——操作系统的内核模块、浏览器的渲染引擎、游戏引擎的 ECS 框架——在各自领域用了几十年。但算子生态里的 M×N 复用在 CANN 开源之前从未被系统性解决。大多数框架的做法是让每个算子仓库自己维护一套 tiling/topology/lifecycle 代码靠 code review 保持一致性。ascend-boost-comm 把这条路径反过来了——先建公共层再在上面长
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.coloradmin.cn/o/2630647.html
如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!