catlass 是 Cube 单元的算子模板库，atvc 是 Vector 单元的算子模板库——atvoss 再下一层：提供可复用的 Vector 子程序（例程），atvc 的 LayerNorm、Softmax、Dropout 等模板底层都在调用 atvoss 的标准化子程序。

atvoss 的子程序分类

atvoss/
├─ math/          # 数学子程序
│   ├─ exp, log, sqrt, rsqrt, reciprocal
│   ├─ sin, cos, tan, sincos
│   └─ erf, gelu, silu（激活函数子程序）
│
├─ reduction/     # 规约子程序
│   ├─ sum, mean, max, min
│   ├─ argmax, argmin
│   └─ softmax_reduce（online softmax 用）
│
├─ permutation/   # 数据排列子程序
│   ├─ transpose2d, transpose3d
│   ├─ gather, scatter
│   └─ shuffle, pack/unpack（INT4/INT8）
│
├─ conversion/    # 类型转换子程序
│   ├─ fp32→fp16, fp16→fp32
│   ├─ bf16→fp32, fp32→bf16
│   └─ int8→fp32, fp32→int4
│
└─ misc/          # 辅助子程序
    ├─ warp_reduce（warp 内部规约）
    ├─ lane_shift（lane 间数据交换）
    └─ prefetch（预取）

每个子程序是一个独立的 Vector 指令序列——没有状态，输入→计算→输出。上层算子模板（atvc）像搭积木一样组合这些子程序。

数学子程序：高效 exp 实现

RSoftmax 和 GELU 都需要 exp 计算——这是 Vector 单元上最频繁调用的子程序之一。直接调用 expf()（C 标准库）太慢——FP32 浮点 exp 需要费大半版的泰勒展开加上浮点数位运算。

atvoss 的 exp 用查表 + 线性插值实现：

// atvoss/math/exp_fp16.cpp

// FP16 exp 的查表实现
// 数值范围：exp(x) for x in (-8, 8)
// 查表大小：256 个条目（2KB），完全塞进 L1
// 精度：< 0.2% 相对误差（对比 FP32 exp）

__aicore__ void VectorExpFp16(
    LocalTensor<half>& output,  // [N]
    LocalTensor<half>& input,   // [N]
    int N
) {
    // 预加载的 exp 查找表（编译期计算好的）
    // table[i] = expf(-8.0f + i * 16.0f / 255.0f)  // 范围 [-8, 8]
    static const half exp_table[256] = {
        // 编译期计算，存储为 FP16
    };

    for (int i = 0; i < N; i += 256) {
        // 256 个 lane 各处理一个元素
        // 第一步：把输入 clamp 到 [-8, 8]
        half x = input[i + __lane_id__];
        x = max(x, half(-8.0f));
        x = min(x, half(8.0f));

        // 第二步：查表索引
        // idx = (x + 8) * 255 / 16
        float idx_float = (float(x) + 8.0f) * 255.0f / 16.0f;
        int idx_lo = int(idx_float);
        int idx_hi = idx_lo + 1;
        float frac = idx_float - float(idx_lo);

        // 第三步：线性插值
        half v_lo = exp_table[idx_lo];
        half v_hi = exp_table[idx_hi];
        half result = v_lo + half(frac) * (v_hi - v_lo);

        output[i + __lane_id__] = result;
    }
}

关键优化点：256 个 lane 并行处理值，全在 L1 缓存内——查表延迟 1 cycle，线性插值 1 cycle。对比 FP32 exp（~20 cycles via hardware），这是 10× 的吞吐量提升。

规约子程序：Warp Reduce

Softmax 的 todenominator 需要对 exp 值求全组和——这是规约操作。atvoss 的 Warp Reduce 用 butterfly 规约（log-level 并行的共享）：

// atvoss/reduction/warp_reduce.cpp

// Warp 内 32 个 Lane 的最大值值的归约
// 使用 butterfly 模式（5 次 shuffle，并行度逐步翻倍）
__aicore__ float WarpReduceMax(float val) {
    // 第 1 步：lane 间隔 16
    float peer = __lane_shuffle_xor(val, 16);
    val = max(val, peer);

    // 第 2 步：lane 间隔 8
    peer = __lane_shuffle_xor(val, 8);
    val = max(val, peer);

    // 第 3 步：间隔 4
    peer = __lane_shuffle_xor(val, 4);
    val = max(val, peer);

    // 第 4 步：间隔 2
    peer = __lane_shuffle_xor(val, 2);
    val = max(val, peer);

    // 第 5 步：间隔 1（最终结果广播到所有 lane）
    peer = __lane_shuffle_xor(val, 1);
    val = max(val, peer);

    return val;
}

// Warp 内求和（同理，5 次 shuffle）
__aicore__ float WarpReduceSum(float val) {
    float peer;
    peer = __lane_shuffle_xor(val, 16); val += peer;
    peer = __lane_shuffle_xor(val, 8);  val += peer;
    peer = __lane_shuffle_xor(val, 4);  val += peer;
    peer = __lane_shuffle_xor(val, 2);  val += peer;
    peer = __lane_shuffle_xor(val, 1);  val += peer;
    return val;
}

__lane_shuffle_xor 是 NPU 专用的 lane 间数据交换指令——两个 lane 交换数据，延迟 < 4 cycles（直接通过 Cross-Lane 交换网络，不走 HBM）。butterfly 归约 5 次 shuffle = 20 cycles——比从 HBM 逐个累加快 30×。

上层的使用：LayerNorm 内部调用 atvoss

看 atvc 的 LayerNorm 如何组合 atvoss 子程序：

// atvc/layernorm.cpp —— 内部调用 atvoss 子程序

__aicore__ void LayerNorm(
    LocalTensor<float>& out,
    LocalTensor<float>& inp,
    LocalTensor<float>& gamma,
    LocalTensor<float>& beta,
    float eps,
    int N
) {
    // Step 1：Warp Reduce 求均值（atvoss::warp_reduce_sum）
    float warp_sum = 0.0f;
    for (int i = 0; i < N; i += 32) {
        float val = inp[__lane_id__ + i];
        warp_sum = atvoss::WarpReduceSum(val);
    }
    float mean = warp_sum / float(N);

    // Step 2：Warp Reduce 求方差（复用 atvoss::warp_reduce_sum）
    float warp_var = 0.0f;
    for (int i = 0; i < N; i += 32) {
        float diff = inp[__lane_id__ + i] - mean;
        warp_var = atvoss::WarpReduceSum(diff * diff);
    }
    float var = warp_var / float(N);
    float inv_std = atvoss::rsqrt(var + eps);  // atvoss/math/rsqrt

    // Step 3：归一化 + 缩放 + 偏移（atvoss::fma）
    for (int i = 0; i < N; i += 256) {
        float x = inp[i + __lane_id__];
        out[i + __lane_id__] = atvoss::fma(
            (x - mean) * inv_std,   // 归一化
            gamma[__lane_id__],     // 缩放
            beta[__lane_id__]       // 偏移
        );
    }
}

atvoss 的三层抽象：

• Level 0：原始 PTO 指令（MMA, LOAD, STore…）
• Level 1：atvoss 子程序（WarpReduceSum, exp_fp16, rsqrt…）
• Level 2：atvc 模板（LayerNorm, Softmax, Dropout…）

每层的累计复杂度被下层封装——atvc 的 LayerNorm 只需组合几个 atvoss 子程序，atvoss 的子程序内部是高度优化的 PTO 指令序列。

踩坑一：FP16 EXP 查表的边界条件

FP16 的 exp 在查表时键范围 [-8, 8]。但 FP16 最大值是 65504——如果传入了 exp(x) x>11.1 作为输入，表查不到对应的值，返回 0.0。

错误：没有 clamp 输入就查表。

// 查表索引 = (x + 8) * 255 / 16
// x = 20.0 → idx = (28 * 255 / 16) = 446 → 越界
// 访问 exp_table[446] → 读到了未初始化的 L1 cache 区域
// 返回随机值 → Softmax 输出全是 NaN

正确：查表前 clamp 输入到 [-8, 8]。

float x = max(min(x, 8.0f), -8.0f);  // Clamp 到有效表范围
int idx = int((x + 8.0f) * 255.0f / 16.0f + 0.5f);  // 四舍五入

踩坑二：Warp Reduce 假设 All Lanes Active

Warp Reduce 的 5 次 butterfly shuffle 假设所有 32 个 lane 都活跃。但如果输入 N 不能被 32 整除——最后一个 warp 的部分 lane 是 inactive 的——这些 lane 对应的寄存器是未定义值。Warp Reduce 把它们也归进去了，导致规约结果错误。

错误：

// N=100, 4 个 warp (128 lane)
// 最后一个 warp 有 28 个 active lane，4 个 inactive
// WarpReduceSum 归约了 32 个值 → 4 个是垃圾值
float total = 0.0f;
for (int i = 0; i < N; i += 32) {
    float x = (i + __lane_id__ < N) ? inp[i + __lane_id__] : 0.0f;
    total += WarpReduceSum(x);  // ← 问题：最后一个 warp 归约了 0 + 垃圾值
}

正确：在 Warp Reduce 前把 inactive lane 的值置零。

float x = (i + __lane_id__ < N) ? inp[i + __lane_id__] : 0.0f;
total += WarpReduceSum(x);  // inactive lane 贡献 0 → 不影响结果

踩坑三：子程序的 L1 寄存器污染

atvoss 的子程序内部会使用 L1 寄存器（Vector 单元的本地缓存）。多个 atvoss 子程序串联时，如果没有显式管理寄存器生命周期，后一个子程序可能会覆盖前一个子程序的中间结果。

错误：

// 错误：假设 atvoss 子程序不污染寄存器
float a = atvoss::exp(x);       // exp 内部用了 L1 registers 0-31
float b = atvoss::rsqrt(y);     // rsqrt 内部也用了 L1 registers 0-31 ❌
// a 的中间结果被 rsqrt 覆盖 → a 是垃圾值

正确：用 atvoss::PreserveRegs 显式保护中间值。

float a;
{
    auto preserve = atvoss::PreserveRegs(0, 31);  // 保护 regs 0-31
    a = atvoss::exp(x);
}  // preserve 析构，中间值 a 已经写入安全区域

float b = atvoss::rsqrt(y);  // 可以安全使用 regs 0-31

atvoss 内部其实没有这么复杂的手动寄存器管理——编译器的寄存器分配器会自动处理。但当一个 kernel 调用 10+ 个 atvoss 子程序时，编译器可能高估寄存器压力，导致部分中间值被 spilling 到 HBM——性能从 50ms 跌到 200ms。这时候需要减少 kernel 内的子程序调用次数（拆成多个 pass）。

---

atvoss 是「搭建积木的工厂」——atvc 的算子模板从 atvoss 拿现成的标准化子程序，不用自己重写 exp/rsqrt/warp_reduce。atvoss 本身的性能特征决定了一个 Vector 算子的天花板——倒数、开方、规约、类型转换——这些基础子程序的性能是 L2 算子性能的基石。