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
// 所有算子仓库共用此接口

template <int DIM>
struct TilingStrategy {
    int num_blocks[DIM];         // 每维切多少块
    int block_size[DIM];         // 每块的大小
    int remainder_block[DIM];    // 尾块的策略(对齐/不对齐)
    
    static TilingStrategy compute(
        const Shape<DIM>& total_shape,  // 总shape
        const Shape<DIM>& max_block,    // 单块最大容量(L1约束)
        TilingPolicy policy             // 切分策略
    ) {
        TilingStrategy result;
        for (int d = 0; d < DIM; d++) {
            // 按 L1 容量约束计算最优分块
            int max_elements = max_block[d];
            int total = total_shape[d];
            
            result.num_blocks[d] = (total + max_elements - 1) / max_elements;
            result.block_size[d] = max_elements;
            
            // 尾块处理:可以选择对齐到 16(Cube 约束)或保持原始大小
            int last_size = total - (result.num_blocks[d] - 1) * max_elements;
            if (policy == ALIGN_TO_CUBE && last_size % 16 != 0) {
                result.remainder_block[d] = (last_size + 15) / 16 * 16;
            } else {
                result.remainder_block[d] = last_size;
            }
        }
        return result;
    }
};

三个仓库的代码变成:

// 三个仓库各自只用一行调用
auto tiling = TilingStrategy<2>::compute(
    {M, N},           // 矩阵尺寸
    {MAX_M_TILE, MAX_N_TILE},  // L1 容量上限
    ALIGN_TO_CUBE      // 对齐策略
);

修复一个切分 bug,升级 ascend-boost-comm 里的 TilingStrategy 就行——所有 50 个仓库自动受益。

五大公共模块

一、数据切片引擎

除了前面的 shape 维切分,还处理数据布局转换。算子从 NCHW 换到 NHWC、从 RowMajor 换到 ColMajor 的跨步映射,全部由切片引擎提供:

// 数据布局转换 + 分块:一次调用
#include "ascend-boost-comm/tiling/data_slice.h"

auto slice = DataSlice::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 提供统一的拓扑发现:

// 任何算子仓库都可以调拓扑发现
#include "ascend-boost-comm/topology/topo_discovery.h"

TopologyGraph topo = TopologyDiscovery::get_instance()->discover();

// 判断任意两张 NPU 之间走什么链路
for (int i = 0; i < num_npus; i++) {
    for (int j = i + 1; j < num_npus; j++) {
        auto path = topo.shortest_path(i, j);
        // path.type: NVLink / RoCE / PCIe
        // path.bandwidth: 链路有效带宽(GB/s)
        // path.latency: 链路延迟(μs)
    }
}

// 基于拓扑选算法
if (topo.is_nvlink_full_mesh()) {
    return ALG_HALVING_DOUBLING;
} else if (topo.is_ring()) {
    return ALG_RING;
} else {
    return ALG_NAIVE;
}

这个接口被 hcomm、hccl、asc-comm 三个通信层共用。拓扑发现逻辑只在 ascend-boost-comm 里维护一份,改动了 NPU 拓扑描述数据结构后,三个通信层自动同步。

三、算子生命周期管理

CANN 算子从注册到执行有完整的生命周期:注册 → InferShape → Tiling → 内存分配 → Kernel Dispatch → 执行 → 内存释放。ascend-boost-comm 管理这个生命周期,让每个算子只关注「计算逻辑」部分:

// 算子生命周期——ascend-boost-comm 统一管理
#include "ascend-boost-comm/lifecycle/op_lifecycle.h"

// 算子开发者只需要实现 OpInterface
class MyAddOp : public OpInterface {
    Shape InferShape(const vector<Shape>& inputs) override { ... }
    KernelType DispatchKernel(const OpConfig& config) override { ... }
    Status Execute(const vector<Tensor>& inputs, Tensor& output) override { ... }
};

// ascend-boost-comm 管剩下的所有事:
// - 内存预分配(从内存池复用)
// - workspace 管理
// - 异步执行流绑定
// - 执行完成的同步点
auto lifecycle = OpLifecycle::create<MyAddOp>();
lifecycle->infer_shape(inputs);
lifecycle->allocate_memory();
lifecycle->dispatch_kernel();
lifecycle->execute();
lifecycle->free_memory();

四、跨算子状态共享

某些状态需要跨多个算子共享——比如混合精度训练的 loss scale 因子、推理的 KV Cache 块池。ascend-boost-comm 提供了一个分布式状态管理器:

// 跨算子全局状态
#include "ascend-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_scale
float scale = GlobalState::get<float>("amp_loss_scale");

// op-transformer 的 Attention 读 KV Cache 池
auto* pool = GlobalState::get<void*>("kv_cache_block_pool");

状态管理器解决了「全局配置项到处传参数」的问题——loss_scale 只需要在 AMP 初始化时设一次,后续所有算子的梯度缩放自动感知。

五、调试与诊断

算子出问题时,快速定位是哪个阶段出的错。ascend-boost-comm 内建了分阶段的 profiling 和诊断:

// 分阶段 profiling
// ascend-boost-comm 在生命周期每个阶段自动插桩

#include "ascend-boost-comm/debug/profiler.h"

OpProfiler profiler("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.08ms

Profiling 是分阶段自动注入的,不需要算子开发者手动加计时器。

依赖关系全景

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 把这条路径反过来了——先建公共层,再在上面长

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐