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

写 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__voidbroadcast_add_kernel(...){for(intbatch0;batch4;batch){// B 的地址指针不变——循环使用同一份数据for(intj0;j256*256;j128){float16 a_vec[128]load_gm(Abatch*65536j);float16 b_vec[128]load_gm(Bj);float16 result[128]a_vecb_vec;store_gm(outputbatch*65536j,result);}}}B 只从 DDR 读一次在 L1 中缓存4 个 batch 的计算都复用 L1 上的 B 数据。如果 B 太大无法全部放到 L1ops-broadcast 会分块搬运——每次搬一部分到 L1算完再搬下一部分。图编译层如何优化 BroadcastGE 在图编译阶段对 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 优化前Zbroadcast_add(A,B)// 独立 BroadcastOgemm(Z,W)// GEMM// GE 优化后——广播的代价被 GEMM 的 Tile 循环吸收Obroadcast_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 优化指南