ops-broadcast：Tensor Shape 的自动扩展机制

摘要：广播(Broadcast)机制允许不同形状的张量进行运算，通过虚拟扩展而非物理复制来提升效率。昇腾NPU通过Vector Unit实现广播计算，小张量缓存在L1中复用，避免显存浪费。图编译层(GE)会优化广播操作，如融合广播与后续算子、调整Tensor布局等。当小张量超过L1容量时，ops-broadcast会调整分块策略以维持性能。该机制显著减少了内存占用和带宽消耗，是AI计算中的重要优化

2501_93573368

6人浏览 · 2026-05-22 18:10:08

2501_93573368 · 2026-05-22 18:10:08 发布

写 AI 代码时经常遇到这种情况：a = [4, 256, 256] + [256, 256]。两个 Tensor 的形状不同但可以相加——PyTorch 和 NumPy 自动把小 Tensor 扩展到跟大 Tensor 一样的形状。这个机制就是 Broadcast。

CANN 的 ops-broadcast 仓库实现了 NPU 上的广播算子。Broadcast 本身不改变数据——它只是让形状不匹配的 Tensor 在执行计算时能被"虚拟地"扩展到相同大小。

Broadcast 为什么存在

如果不支持 Broadcast，[4,256,256] + [256,256] 需要手动把 [256,256] 扩展到 [4,256,256]——在内存中复制 4 份。这浪费了 4 倍显存。

Broadcast 的做法是：计算时让 AI Core 知道"第二个 Tensor 在第一个维度上不够时就在该维度上循环使用"。不需要实际复制数据——只是遍历地址时多了一次维度映射。

Tensor Shape 如何自动扩展

Broadcast 的规则：

从最后一个维度开始对齐
如果某个维度的大小是 1 或不匹配，沿该维度复制
如果某个维度不存在，虚拟插入大小为 1 的维度

[4,256,256] + [256,256] 的扩展过程：

A.shape = [4, 256, 256]
B.shape = [   256, 256]  → 虚拟扩展为 [1, 256, 256]
                            → Broadcast 为 [4, 256, 256]

实际执行时不会创建 [1,256,256] 的物理 Tensor——Vector Unit 在遍历 A 的每个 batch 时，重复使用 B 的同一份数据。

昇腾NPU如何处理广播计算

Broadcast 在 Vector Unit 上执行。对于 [4,256,256] + [256,256]：

// Broadcast Add 的 Vector 执行（简化）
__vector__ void broadcast_add_kernel(...) {
    for (int batch = 0; batch < 4; batch++) {
        // B 的地址指针不变——循环使用同一份数据
        for (int j = 0; j < 256 * 256; j += 128) {
            float16 a_vec[128] = load_gm(A + batch * 65536 + j);
            float16 b_vec[128] = load_gm(B + j);
            float16 result[128] = a_vec + b_vec;
            store_gm(output + batch * 65536 + j, result);
        }
    }
}

B 只从 DDR 读一次（在 L1 中缓存），4 个 batch 的计算都复用 L1 上的 B 数据。如果 B 太大无法全部放到 L1，ops-broadcast 会分块搬运——每次搬一部分到 L1，算完再搬下一部分。

图编译层如何优化 Broadcast

GE 在图编译阶段对 Broadcast 做优化：

Broadcast Elimination。 如果广播的源 Tensor 在后续算子中会被多次使用，GE 把广播操作融合到消费算子中——不需要独立的 Broadcast Kernel。

Layout 适配。 Broadcast 在 NZ 格式 Tensor 上的性能比 ND 差——NZ 的分块结构让 Broadcast 的地址映射更复杂。GE 在发现 Broadcast 算子的输入是 NZ 格式时，会在必要时插入 ND 转换。

GE 的 Broadcast 优化实例

GE 在编译时对 Broadcast 的优化举例：

输入：[4, 4096, 4096] + [1, 4096, 4096]。GE 的分析结果：第二个 Tensor 在 batch 维度是 1——需要广播。但如果 GE 发现广播后的结果只被一个后续算子的 batch 维度使用，它把广播融合到后续算子中——不需要显式的 Broadcast Kernel。

// GE 优化前
Z = broadcast_add(A, B)  // 独立 Broadcast
O = gemm(Z, W)            // GEMM

// GE 优化后——广播的代价被 GEMM 的 Tile 循环吸收
O = broadcast_gemm(A, B, W)  // Broadcast + GEMM 融合

融合后缺少了一次中间 Z 的 DDR 读写，搬运量减少 [4, 4096, 4096] 约 128MB。

Vector Unit 处理广播的性能

Broadcast 在 Vector Unit 上执行时，小 Tensor 在 L1 中缓存——大 Tensor 的每个分块跟 L1 中的小 Tensor 分块配对计算。如果小 Tensor 太大无法完全缓存到 L1，则每次重新从 DDR 读取——Broadcast 的带宽利用率下降到 30% 以下。ops-broadcast 在检测到小 Tensor 超过 L1 容量时会用更大的分块粒度——平均每个分块在 L1 中复用更多次。

参考仓库

ops-broadcast 广播算子库

Tensor Layout 优化指南