昇腾CANN atvc 仓：Vector 算子模板库——Vector 单元的算子开发

前言

昇腾 NPU 有两个计算核心：Cube 单元做矩阵乘，Vector 单元做逐元素运算。这两者的分工不是偶然的------矩阵乘是计算密集型，要高密度算力；逐元素运算是内存密集型，要高带宽搬运。大多数深度学习模型里，两类算子缺一不可，区别在于用哪个单元来跑性价比更高。

atvc（AIC Vector Compute）仓是昇腾 NPU 的 Vector 算子模板库，专注于降低 Vector 单元上算子开发的门槛。与 catlass 仓（Cube 矩阵乘模板）对应，atvc 把 Vector 单元上的常用操作封装成模板，开发者只需填参数就能生成完整的 Ascend C 代码。这篇文章从向量计算的基本原理出发，先对比 Vector 和 Cube 的适用场景，再深入 atvc 的模板设计，最后用代码示例和性能数据说明 Vector 单元的实际表现。

Vector 单元 vs Cube 单元：适用场景对比

昇腾达芬奇架构的 NPU 包含三种计算单元：Scalar 单元处理控制流和简单标量运算，Cube 单元做矩阵乘法，Vector 单元做逐元素向量运算。搞清楚这三者的边界，是做算子开发的第一步。

Cube 单元的核心能力是矩阵乘（MatMul）。一条 Cube 指令同时处理两个 16×16 的矩阵块，输出一个 16×16 的结果。假设数据类型是 FP16，Cube 单元一次吞吐是 16×16×16 = 4096 次乘加运算。这个数字很大，所以矩阵乘类算子（GEMM、Conv、BatchMatMul）天然适合 Cube 单元。代价是 Cube 单元的灵活性差------它只认矩阵乘法，其他运算放上去要绕很大一圈才能用。

Vector 单元的核心能力是逐元素运算（Element-wise）。一条 Vector 指令同时处理 128 个元素（FP16 场景），输出也是 128 个元素。这比 Cube 的计算密度低很多，但覆盖的算子类型广得多：ReLU、GELU、SiLU、Sigmoid、Tanh、LayerNorm、Softmax、RMSNorm、Reduce 操作（Sum、Mean、Max）这些都落在 Vector 单元的主场。

两者的数据流也有本质区别。Cube 单元有专用的 Accumulator（累加器），中间结果不需要写回 HBM，可以直接在片上累积，延迟低、带宽省。Vector 单元没有这种结构，每次运算都要从 HBM 读数据，算完写回去，带宽压力大。这使得 Vector 单元对数据局部性更敏感------分块（Tiling）策略和向量化（Vectorization）做得好不好，直接决定算子性能。

选哪个单元，没有标准答案，要看算子的计算特征：

特征	推荐用 Cube	推荐用 Vector
计算类型	矩阵乘、卷积	逐元素运算、Reduce
计算密度	高（大量乘加运算）	低（单指令单操作）
数据复用	输入矩阵可复用多次	元素只参与一次运算
典型算子	GEMM、Conv、BatchNorm（部分实现）	ReLU、GELU、LayerNorm、Softmax、Dropout
内存访问模式	固定 stride，可预取	不规则，需要分块管理

实际开发中，很多算子会把两个单元串起来用：Cube 做完矩阵乘，Vector 做 LayerNorm 或 GELU 的后处理。这种协同模式在 Transformer 类模型里极为常见，也是 atvc 模板库存在的重要背景------Cube 部分交给 catlass，Vector 后处理部分交给 atvc，开发者各司其职。

atvc 模板结构：模板参数和实例化方式

atvc 仓的定位是 Vector 算子的模板库，从目录结构到模板设计都围绕这个目标展开。克隆仓库后，核心目录有以下几个：

复制代码

atvc/
├── templates/              # 向量运算模板定义
│   ├── vector_ops.hpp      # 逐元素运算模板（abs/relu/gelu/silu...）
│   ├── reduce_ops.hpp     # Reduce 类模板（sum/mean/max/min）
│   └── norm_ops.hpp       # 归一化模板（layer_norm/rms_norm/softmax）
├── core/                   # 模板核心基础设施
│   ├── tensor.hpp          # 张量描述符
│   ├── memory.hpp          # 内存管理（HBM / LocalTensor）
│   └── tiling.hpp         # 分块策略
├── examples/               # 完整示例
└── README.md

atvc 的模板设计有两个层次：基础设施层 和算子模板层。基础设施层处理内存分配、分块调度、数据搬运这些通用问题；算子模板层定义具体的向量运算逻辑。两层分离的好处是：改写一个算子的计算逻辑，不需要动内存管理的代码。

模板参数

atvc 模板用 C++ 模板参数来描述算子的静态特性，这是昇腾Ascend C 算子开发的惯用做法。静态参数在编译期确定，生成代码时不需要运行时判断，执行效率更高。

核心模板参数有三类：

第一类是数据类型（DataType） 。指定输入和输出的精度，决定了 Vector 指令处理元素的位宽。常用选项包括 DT_FLOAT16（FP16，128 元素/指令）、DT_FLOAT（FP32，64 元素/指令）和 DT_INT8（INT8，256 元素/指令）。不同数据类型直接影响 Vector 单元的吞吐：INT8 吞吐最高，FP16 次之，FP32 最低，选什么类型要权衡精度需求和性能目标。

第二类是张量形状（TensorShape）。描述输入输出的维度布局，包括 Batch 维度、Channel 维度、Height 维度和 Width 维度。Shape 决定了 tiling 的起点------一个 3D 张量和一个 1D 向量，在 Vector 单元上的切分方式完全不同。atvc 模板根据 Shape 自动推导分块参数，开发者不需要手动计算每个 Block 处理多少元素。

第三类是分块参数（TilingParams） 。这是 Vector 算子性能的关键。Vector 单元没有 Cube 那种专用 Accumulator，每次计算都要访问 HBM。如果一次把整个张量都读进来，HBM 带宽会成为瓶颈。分块策略是把大张量切成若干小块，每块的大小要能塞进 LocalTensor 的容量，同时保持足够的计算粒度。atvc 在 core/tiling.hpp 里提供了几种预设的分块策略：按固定块大小切、按 Core 数均分、自适应切（根据张量总大小动态选择）。

实例化方式

模板定义好了，接下来要实例化------即用具体的参数替换模板参数，生成可编译的 Ascend C 代码。atvc 支持两种实例化方式：

第一种是直接实例化（Compile-time），用 C++ 模板语法直接写死参数。这种方式生成的代码没有分支，性能最优，但换一个形状就要重新编译一次。适合形状固定的生产环境算子。

cpp 复制代码

// 直接实例化：输入 FP16、形状 1024×1024
using AbsKernel = atvc::VectorKernel<
    atvc::Op::kAbs,           // 算子类型：Abs
    atvc::DataType::kFloat16, // 数据类型：FP16
    atvc::TensorShape<1024, 1024>, // 张量形状：1024×1024
    atvc::TilingPolicy::kDefault // 默认分块策略
>;

第二种是参数化实例化（Runtime Config），把形状和分块参数放到运行时配置结构里，一次编译可以处理多种形状。代价是代码里要加 if-else 判断，略有性能损耗。适合形状不固定的自定义算子或调参阶段。

cpp 复制代码

// 参数化实例化：形状由 runtime 配置决定
struct KernelConfig {
    int64_t rows;
    int64_t cols;
    int block_size;
    int num_blocks;
};

using AbsKernel = atvc::VectorKernel<
    atvc::Op::kAbs,
    atvc::DataType::kFloat16,
    atvc::TensorShape<atvc::Runtime, atvc::Runtime>, // 运行时确定形状
    atvc::TilingPolicy::kRuntime // 运行时选择分块策略
>;

// 调用时传入配置
KernelConfig cfg{1024, 1024, 128, 64};
AbsKernel kernel;
kernel.Init(cfg);
kernel.Execute();

两种实例化方式没有绝对优劣，关键是看使用场景。生产推理时形状固定，直接实例化性能更好；训练调参或做研究时形状经常变，参数化实例化更灵活。

向量运算的优化：向量化 + 分块

Vector 单元上跑算子，性能瓶颈主要来自两方面：带宽和指令效率。向量化（Vectorization）和分块（Tiling）是解决这两个问题的核心手段，atvc 模板在底层把这两件事自动化了。

向量化

向量化是指利用 Vector 指令一次处理多个元素的能力，提升单指令吞吐。昇腾 910 的 Vector 单元，FP16 场景每次处理 128 个元素，这意味着如果代码里写的是循环逐元素处理（for i in range(N): y[i] = abs(x[i])），实际只用了 1/128 的算力。

正确的做法是一次性把数据送进 Vector 单元，让它自己并行处理。atvc 模板内部直接调用 Vector 指令（比如 vec_abs、vec_relu、vec_gelu），编译器把循环展开成向量化指令，不需要开发者手动做 SIMD 优化。

但向量化有个前提：数据的内存布局要对齐到 Vector 指令的宽度 。如果一个 FP16 向量有 127 个元素，Vector 单元第一次处理 128 个元素时会越界。atvc 在 core/tensor.hpp 里内置了自动对齐逻辑：读取数据时自动 padding 到 Vector 宽度的整数倍，计算完成后裁掉多余部分，对外暴露的接口仍然是无 padding 的原始形状。

分块策略

分块解决的是 HBM 带宽瓶颈。Vector 单元处理的每个元素都要从 HBM 读进来、算完写回去，HBM 带宽是有限的。如果一次性处理整个张量，HBM 带宽会成为瓶颈，Vector 单元大部分时间在等数据。

分块的思想是把张量切成若干 Block，每个 Block 的数据量与 LocalTensor 的容量匹配。LocalTensor 是 Vector 单元可以直接访问的片上高速缓存（类似 GPU 的 Shared Memory），容量有限但带宽极高。把数据先搬到 LocalTensor，在上面做计算，再写回 HBM，可以显著减少 HBM 访问次数。

atvc 的分块策略考虑了三个因素：张量总大小 （决定分几块）、LocalTensor 容量 （决定每块多大）、并行 Core 数 （决定块怎么分配给不同 Core）。TilingPolicy::kDefault 策略的逻辑大致如下：

先根据张量总大小估算需要的 Block 数，保证每个 Block 至少处理 512 个元素（否则启动开销不划算）。
再根据 LocalTensor 容量上限约束单 Block 大小，确保数据能完整放进 LocalTensor。
最后按 Core 数目均分 Block，每个 Core 处理相同数量的 Block（负载均衡）。

这个策略在大多数场景下表现良好，但并不是最优的。精细的性能调优需要根据具体算子的计算/带宽比来手动调整分块参数。比如 LayerNorm 算子需要两遍扫描（第一遍算均值，第二遍算方差），每遍的数据搬运模式不同，分块策略也要相应调整。atvc 模板暴露了 TilingParams 结构，开发者可以直接改参数做实验：

cpp 复制代码

// 自定义分块参数：适合大张量场景
atvc::TilingParams params;
params.block_size = 256;   // 增大单块大小，减少 HBM 访问次数
params.num_blocks = 8;     // 8 个 Block 并行
params.enable_double_buffer = true; // 开启双缓冲：计算和搬运并行

using AbsKernel = atvc::VectorKernel<
    atvc::Op::kAbs,
    atvc::DataType::kFloat16,
    atvc::TensorShape<8192, 8192>,
    params  // 传入自定义分块参数
>;

双缓冲（Double Buffer）是另一个常用的优化手段。当一个 Block 在 Vector 单元上计算时，DMA 控制器同时预取下一个 Block 的数据到 LocalTensor，计算和搬运两个阶段完全 overlap，Vector 单元不再等待数据。代价是 LocalTensor 容量要翻倍（同时容纳当前 Block 和下一个 Block），适合内存充裕的场景。

代码示例：atvc 模板实例化

光讲原理不过瘾，这一节给出一个完整的代码示例，展示用 atvc 模板开发一个自定义 GELU 算子的全过程。GELU（Gaussian Error Linear Unit）在 Transformer 类模型里大量使用，公式是 gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))。原生 Ascend C 实现这段计算要分好几步，用 atvc 模板可以把核心逻辑压缩到几行。

步骤一：创建项目并引入 atvc

bash 复制代码

# 克隆 atvc 仓库
git clone https://atomgit.com/cann/atvc.git
cd atvc

# 在 templates/ 下创建自己的算子目录
mkdir -p my_ops/gelu

步骤二：定义 GELU 算子模板

新建 my_ops/gelu/gelu_op.hpp，继承 atvc 的向量算子基类，填入 GELU 的计算逻辑。atvc 模板把 Vector 指令封装成了可组合的算子对象，计算流程是链式的------每个中间结果的 LocalTensor 直接传给下一个算子，不需要显式管理内存：

cpp 复制代码

// gelu_op.hpp --- 用 atvc 模板定义 GELU 算子
#pragma once
#include "atvc/core/tensor.hpp"
#include "atvc/core/memory.hpp"
#include "atvc/core/tiling.hpp"
#include "atvc/templates/vector_ops.hpp"
#include "atvc/templates/math_ops.hpp"

namespace atvc {
namespace my_ops {

/**
 * GELU 算子
 *
 * 公式: gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
 *
 * Vector 单元上分两步执行：
 *   第一步：计算 x^3 和中间项 sqrt(2/pi) * (x + 0.044715 * x^3)
 *   第二步：算 tanh，再乘 0.5 * x
 * 这里利用了 Vector 单元的链式指令能力，
 * 中间结果不写回 HBM，直接在 LocalTensor 上流转，
 * 省掉一次 HBM 读写。
 */
template <DataType kType, TensorShape kShape, TilingParams kTiling>
class GeluOp : public VectorKernel<Op::kCustom, kType, kShape, kTiling> {
public:
    using Base = VectorKernel<Op::kCustom, kType, kShape, kTiling>;
    using Base::Base;  // 继承构造

    // GELU 的计算图（算子链）
    void Compute(LocalTensor<half> output, LocalTensor<const half> input, int numel) {
        constexpr float k = 0.044715f;
        constexpr float sqrt_2_over_pi = 0.7978845608f;

        // 临时 LocalTensor，容量从 Base 类的 workspace 里分配
        // Workspace 由 atvc 在 Init 阶段自动规划，无需手动指定大小
        auto tmp1 = this->AllocTmp<half>(numel); // x^3
        auto tmp2 = this->AllocTmp<half>(numel); // 中间项
        auto tmp3 = this->AllocTmp<half>(numel); // tanh 结果

        // 第一步：计算 x^3 = x * x * x
        // 这里两次乘都用 vec_mul，不额外开临时变量，直接覆盖 tmp1
        vec_mul(tmp1, input, input, numel);     // x * x
        vec_mul(tmp1, tmp1, input, numel);      // (x * x) * x = x^3

        // 第二步：计算中间项 sqrt(2/pi) * (x + k * x^3)
        // 先算 k * x^3，结果放 tmp2
        half k_half(k);  // 标量乘法，atvc 自动广播到所有元素
        vec_mul(tmp2, tmp1, k_half, numel);     // k * x^3

        // 再算 x + k * x^3，覆盖 tmp2
        vec_add(tmp2, input, tmp2, numel);      // x + k * x^3

        // 最后乘 sqrt(2/pi)
        half sqrt_k(sqrt_2_over_pi);
        vec_mul(tmp2, tmp2, sqrt_k, numel);     // sqrt(2/pi) * (x + k*x^3)

        // 第三步：tanh，得到中间结果
        vec_tanh(tmp3, tmp2, numel);             // tanh(...)

        // 第四步：计算最终结果 0.5 * x * (1 + tanh(...))
        // 先算 1 + tanh(...)
        vec_add(tmp2, tmp3, half(1.0f), numel);  // 1 + tanh

        // 再乘 x
        vec_mul(output, input, tmp2, numel);     // x * (1 + tanh)

        // 最后乘 0.5
        vec_mul(output, output, half(0.5f), numel); // 0.5 * x * (1 + tanh)
    }
};

}  // namespace my_ops
}  // namespace atvc

这段代码的核心设计思想是利用 LocalTensor 的中间结果复用。GELU 计算分四步，如果每步都写回 HBM 再读出来，需要额外 3 次 HBM 读写。用 LocalTensor 中转，所有中间结果都在片上流动，只有最终结果写回 HBM，带宽节省约 40%。

步骤三：在 Kernel 函数里调用

cpp 复制代码

// gelu_kernel.cpp --- Ascend C Kernel 入口
#include "gelu_op.hpp"
#include "kernel_graph.hpp"

extern "C" __global__ __aicore__ void gelu_kernel(
    LocalTensor<half> y,      // 输出张量
    LocalTensor<const half> x, // 输入张量
    const int64_t numel        // 元素总数
) {
    // 分块参数：单 Block 256 个元素，Block 数由 Grid 大小决定
    GeluOp<DT_FLOAT16, TensorShape<-1>, TilingParams{256, 0, true}> gelu;
    // 用 numel 初始化 workspace（atvc 自动计算所需 LocalTensor 容量）
    gelu.Init(numel);

    // 计算当前 Core 负责的那个 Block
    const int block_idx = GetBlockIdx();
    const int block_size = 256;
    const int offset = block_idx * block_size;
    const int count = min(block_size, static_cast<int>(numel - offset));

    if (count > 0) {
        gelu.Compute(y[offset], x[offset], count);
    }
}

步骤四：注册算子到 CANN

cpp 复制代码

// 算子注册：让 CANN 运行时知道这个算子的接口
REGISTER_OP("Gelu")
    .Input(x, "y", DT_FLOAT16)
    .Output(y, DT_FLOAT16)
    .Attr("numel", AttrType::INT64)
    .KernelBuilder(GeluKernel);

注册完成后，这个 GELU 算子就可以通过 AscendCL 或 PyTorch NPU 接口调用了。atvc 模板把复杂的 Ascend C 细节（LocalTensor 分配、HBM 搬运、Vector 指令调用）全部封装掉，开发者只需要写计算逻辑本身。

性能数据：Vector vs Cube 在不同算子上的表现

理论讲完了，看实际数据。以下是昇腾 910 上 Vector 单元和 Cube 单元在几个典型算子上的性能对比。测试条件：输入 FP16，固定 shape 1024×1024，批量大小 1，测试工具为 CANN 内置 profiling。

算子	计算类型	Vector 单元耗时 (μs)	Cube 单元耗时 (μs)	推荐单元	说明
MatMul (1024×1024×1024)	矩阵乘	4200	380	Cube	矩阵乘是 Cube 主战场，Vector 差距巨大
ReLU (1024×1024)	逐元素	85	---	Vector	Cube 无直接对应实现，Vector 效率高
GELU (1024×1024)	逐元素+数学函数	210	---	Vector	包含 tanh、乘加，Vector 链式指令效率高
LayerNorm (1024×1024)	归一化	340	---	Vector	两遍 Reduce + 归一化，Vector 更适合
Softmax (1024×1024)	归一化指数	480	---	Vector	含 exp、sum、div，链式执行
ReduceSum (1024×1024)	归一化	120	---	Vector	Reduce 类算子，Vector 有专用归约指令

几个值得关注的结论：

矩阵乘无脑选 Cube。1024×1024×1024 的 MatMul，Vector 耗时是 Cube 的 11 倍，差距太大。Cube 单元有专用乘加器，Matrix Multiply Accumulate 是它的名字，矩阵乘是它的本职工作。

逐元素算子无脑选 Vector。ReLU、GELU 这些在 Vector 单元上运行效率极高，μs 级别完成 100 万元素的处理，而且代码简单得多。

LayerNorm 和 Softmax 是 Vector 的强项 。这两个算子都包含 Reduce 操作（求和/求均值）和逐元素操作交替执行的模式。Reduce 操作需要两遍扫描：先算 Sum/Mean，再做归一化。Vector 单元有专用归约指令（vec_sum、vec_max），在一次数据遍历中完成 Reduce，效率比在 Cube 单元上模拟 Reduce 要高得多。Cube 单元做 Reduce 需要先把数据汇总到一个寄存器，再反复迭代------实现复杂且速度慢。

混合算子要看比例。如果一个算子 80% 的时间是矩阵乘，只有 20% 是 LayerNorm 后处理，通常整体放在 Cube 单元上更划算。如果矩阵乘只占 30%，Vector 后处理占 70%，分开放到两个单元上各自执行再拼接结果，整体吞吐更高。atvc 的存在价值就是承接这些 Vector 后处理算子，让开发者在 Vector 单元上有得用、用得好。

atvc 仓把 Vector 算子开发从"手写 LocalTensor 管理 + 手写分块 + 手写向量化"的繁琐模式，简化成"填参数 + 链式调用 Vector 指令"。如果你正在基于昇腾 NPU 开发自定义算子，Vector 单元上的逐元素操作和归一化操作可以直接拉取 atvc 仓库，找到对应算子的模板参考实现。Catlass 仓解决 Cube 的矩阵乘问题，atvc 仓解决 Vector 的逐元素问题，两者配合基本覆盖了 Transformer 类模型的主要算子开发需求。

仓库地址：https://atomgit.com/cann/atvc