写一个高性能的 Vector 算子，不是每次从头写 pipeline。像 GEMM 里的分块加载、Softmax 里的 warp reduce、LayerNorm 里的均方差计算——这些是"子程序"：小粒度的、可复用的计算片段。手写 20 个算子，其中 15 个用到了同一套 reduce 逻辑，改了其中一个 bug，剩下 14 个还得逐个改。

atvoss（Ascend Template Vector Operator Subroutine Set）就是来解决这个问题的：把高频子程序抽成模板，算子开发时组合调用，而不是复制粘贴。

atvoss 和 atvc 的关系：atvc 是完整算子的参考实现模板（一个 MatMul 怎么调 Cube、一个 Add 怎么配 Vector），atvoss 是更细一层的子程序模板（一个 warp reduce 怎么写、一个 L1 分块怎么配 stride）。atvc 调 atvoss，不是替代关系。

子程序分类

atvoss/
├── memory/          # 搬运子程序
│   ├── tiling_2d   ← 2D 分块 + stride 计算
│   ├── load_store  ← 对齐搬运（burst / scattered）
│   └── prefetch    ← 预取 + 双缓冲管理
├── compute/         # 计算子程序
│   ├── reduce      ← warp reduce / block reduce (sum/max/min)
│   ├── gemm_micro  ← GEMM 微内核 (k-block matmul on Vector)
│   ├── softmax     ← online softmax (max diff + exp + normalize)
│   ├── layernorm   ← mean/std + affine
│   └── activation  ← gelu/swiglu/relu fused
├── transform/       # 变换子程序
│   ├── transpose   ← in-register transpose
│   ├── shuffle     ← warp-level data exchange
│   └── pack_unpack ← FP16 ↔ FP32 ↔ BF16 conversion
└── schedule/        # 调度子程序
    ├── pipeline    ← 搬运/计算流水线编排
    └── sync        ← barrier / fence 插入策略

warp reduce——最常见的子程序复用

// atvoss/compute/reduce/warp_reduce.h
//
// warp reduce: 32 个 Vector 单元的并行归约
// 替代手写 for-loop sum，一次调用完成 32→1 的归约

#pragma once
#include "vector_common.h"

namespace atvoss {
namespace reduce {

/**
 * Warp-level sum reduction
 *
 * 输入: val[32] 分布在 warp 的 32 个 lane 上
 * 输出: 每个 lane 都得到 sum(val[0..31])
 *
 * 硬件利用: shuffle_down 指令实现 butterfly reduction
 * 复杂度: O(log N) = 5 步 (N=32)
 */
template<typename T>
__aicore__ inline T WarpReduceSum(T val) {
    // Butterfly reduction with shuffle_down
    // Step 1: offset=16, 2: offset=8, 3: offset=4, 4: offset=2, 5: offset=1
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += ShuffleDown(val, offset);
    }
    return val;
}

/**
 * Warp-level max reduction
 * 同上，操作符替换为 max
 */
template<typename T>
__aicore__ inline T WarpReduceMax(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        T other = ShuffleDown(val, offset);
        val = (val > other) ? val : other;
    }
    return val;
}

/**
 * Block-level sum (多 warp 归约)
 *
 * 分层策略:
 * 1. 每个 warp 内做 WarpReduceSum → warp_sum
 * 2. warp 0 收集所有 warp_sum → 再做一次 warp reduce
 *
 * @param val      当前 lane 的值
 * @param shared   共享内存缓存 warp 结果
 * @param warp_id  当前 warp ID
 * @param num_warps 总 warp 数
 */
template<typename T>
__aicore__ inline T BlockReduceSum(
    T val,
    T* shared,       // [num_warps] 每个 warp 的结果暂存
    int warp_id,
    int num_warps
) {
    // Phase 1: warp-level
    T warp_sum = WarpReduceSum(val);

    // Phase 2: 只有 lane 0 写 warp 结果到 shared memory
    if (GetLaneId() == 0) {
        shared[warp_id] = warp_sum;
    }
    __sync_warp();

    // Phase 3: warp 0 再做一次 reduce（读 shared 到寄存器，再 reduce）
    T block_sum = (warp_id == 0) ? shared[GetLaneId()] : T(0);
    if (warp_id == 0 && GetLaneId() < num_warps) {
        block_sum = shared[GetLaneId()];
    }
    if (warp_id == 0) {
        block_sum = WarpReduceSum(block_sum);
    }

    // 广播结果给所有 warp（通过 shared 或 shuffle broadcast）
    if (warp_id == 0 && GetLaneId() == 0) {
        shared[0] = block_sum;
    }
    __sync_warp();

    return shared[0];
}

} // namespace reduce
} // namespace atvoss

online softmax——数值稳定的分块实现

// atvoss/compute/softmax/online_softmax.h
//
// online softmax: 分块计算 softmax，每块更新 max 和 sum
// 关键: 用指数修正因子避免全量重算

#pragma once
#include "vector_common.h"

namespace atvoss {
namespace compute {

/**
 * Online Softmax State
 *
 * 维护两个状态变量:
 * - m: 当前已知的全局最大值
 * - s: 当前已知的 exp 累加和（已用历史 max 标定过）
 *
 * 每次新来一块数据 x[0..N-1]:
 *   1. m_new = max(m, max(x))
 *   2. s = s * exp(m - m_new) + sum(exp(x - m_new))
 *   3. m = m_new
 *
 * 最终: softmax(x_i) = exp(x_i - m) / s
 */
struct OnlineSoftmaxState {
    float m;  // running max
    float s;  // running sum of exp(x - m)

    __aicore__ void init() {
        m = -INFINITY;
        s = 0.0f;
    }

    /**
     * 更新状态: 混合新的数据块
     * @param data   新数据块的指针
     * @param count  数据块大小
     */
    __aicore__ void update(const float* data, int count) {
        // 找当前块的最大值
        float local_max = data[0];
        #pragma unroll
        for (int i = 1; i < count; ++i) {
            local_max = (data[i] > local_max) ? data[i] : local_max;
        }

        float new_max = (local_max > m) ? local_max : m;

        // 修正历史 sum: 乘以 exp(m - new_max)
        // exp(m - new_max) ≤ 1.0: 修正因子总是 ≤ 1（数值安全）
        float correction = expf(m - new_max);
        s = s * correction;

        // 累加新块的 exp(x - new_max)
        float local_sum = 0.0f;
        #pragma unroll
        for (int i = 0; i < count; ++i) {
            local_sum += expf(data[i] - new_max);
        }
        s += local_sum;

        m = new_max;
    }

    /**
     * 从状态恢复 softmax 结果
     * 调用者在所有 update 完成后，对每个原始值调用此函数
     */
    __aicore__ float normalize(float x) const {
        return expf(x - m) / s;
    }
};

/**
 * 全量 softmax 便捷封装
 *
 * 分块大小 = SRAM 缓冲区大小 / sizeof(float)
 * 自动分块 + 两遍计算: pass1 更新 state, pass2 写输出
 */
template<int BLOCK_SIZE = 256>
__aicore__ void Softmax(
    const float* input,
    float* output,
    int N
) {
    OnlineSoftmaxState state;
    state.init();

    // Pass 1: 分块更新 online softmax state
    for (int offset = 0; offset < N; offset += BLOCK_SIZE) {
        int count = (offset + BLOCK_SIZE <= N) ? BLOCK_SIZE : N - offset;
        state.update(input + offset, count);
    }

    // Pass 2: 用最终 state 归一化每个元素
    #pragma unroll
    for (int i = 0; i < N; ++i) {
        output[i] = state.normalize(input[i]);
    }
}

} // namespace compute
} // namespace atvoss

2D 分块——搬运子程序

// atvoss/memory/tiling_2d.h
//
// 2D 分块: 自动计算最优分块 + stride 对齐 + 剩余处理

#pragma once
#include "vector_common.h"

namespace atvoss {
namespace memory {

/**
 * Tiling2D 配置
 *
 * 自动计算 2D 矩阵分块的最优参数:
 * - 对齐到 32B burst 边界
 * - 避免 stride=2^n（银行冲突）
 * - 尽量填满 L1 缓存
 */
struct Tiling2DConfig {
    int M;           // 原始行数
    int K;           // 原始列数
    int lda;         // leading dimension (stride)
    int tile_m;      // 行分块大小
    int tile_k;      // 列分块大小
    int num_m_tiles; // 行方向分块数
    int num_k_tiles; // 列方向分块数
    int last_m;      // 最后一行块的实际大小
    int last_k;      // 最后一列块的实际大小
    int lda_tile;    // 分块后的 local leading dimension

    /**
     * 构造: 自动推导最优分块
     *
     * @param M      矩阵行数
     * @param K      矩阵列数
     * @param lda    原始 leading dimension
     * @param elem_size  元素字节数 (2=FP16, 4=FP32)
     * @param l1_bytes   L1 缓存可用字节数
     */
    __aicore__ void configure(
        int M, int K, int lda, int elem_size, int l1_bytes
    ) {
        this->M = M;
        this->K = K;
        this->lda = lda;

        // 目标: tile_m * tile_k * elem_size ≈ l1_bytes * 0.8 (留 20% 给流水)
        int max_elements = (l1_bytes * 8 / 10) / elem_size;

        // tile_k: 对齐到 32B (burst=32B, FP16→16elem, FP32→8elem)
        int align = (elem_size == 2) ? 16 : 8;
        tile_k = ((K > max_elements / 32) ? max_elements / 32 : K);
        tile_k = (tile_k / align) * align;
        if (tile_k == 0) tile_k = align;

        // 避免 2^n stride 银行冲突: 如果 tile_k 是 2 的幂 → +1
        if ((tile_k & (tile_k - 1)) == 0) {
            tile_k += align;  // 加一个对齐单位
        }

        // tile_m: 用剩余空间
        tile_m = max_elements / tile_k;
        if (tile_m > M) tile_m = M;
        if (tile_m < 1) tile_m = 1;

        // 分块数
        num_m_tiles = (M + tile_m - 1) / tile_m;
        num_k_tiles = (K + tile_k - 1) / tile_k;

        // 尾块大小
        last_m = M - (num_m_tiles - 1) * tile_m;
        last_k = K - (num_k_tiles - 1) * tile_k;

        // Local leading dimension: 对齐到 burst 边界
        lda_tile = (tile_k * elem_size / 32) * 32 / elem_size;
    }

    /**
     * 获取分块 (mi, ki) 对应的全局偏移 + 局部大小
     */
    __aicore__ void get_tile_info(
        int mi, int ki,
        int& global_offset,  // 分块起始在全局矩阵中的偏移
        int& local_m,         // 分块实际行数（含尾块处理）
        int& local_k          // 分块实际列数
    ) const {
        local_m = (mi == num_m_tiles - 1) ? last_m : tile_m;
        local_k = (ki == num_k_tiles - 1) ? last_k : tile_k;

        int row_start = mi * tile_m;
        int col_start = ki * tile_k;
        global_offset = row_start * lda + col_start;
    }

    /**
     * 搬运一个分块: global → L1
     *
     * 自动处理 burst 对齐 + 尾块截断 + stride 适配
     */
    __aicore__ void load_tile(
        const float* global,   // HBM 地址
        float* l1_buf,         // L1 缓冲区
        int mi, int ki         // 分块坐标
    ) const {
        int global_offset, local_m, local_k;
        get_tile_info(mi, ki, global_offset, local_m, local_k);

        // 逐行搬运（每行对齐到 burst）
        for (int r = 0; r < local_m; ++r) {
            const float* src = global + global_offset + r * lda;
            float* dst = l1_buf + r * lda_tile;

            // 对齐 copy
            int burst_count = (local_k * sizeof(float)) / 32;
            int remainder = (local_k * sizeof(float)) % 32;

            // Burst copy
            // 实际使用 Vector 的 DataCopy 宏，这里简化
            for (int b = 0; b < burst_count; ++b) {
                // VecDataCopy(src + b*8, dst + b*8, 32);  // 32B burst
            }

            // 尾数逐元素搬运
            if (remainder > 0) {
                for (int e = burst_count * 8; e < local_k; ++e) {
                    dst[e] = src[e];
                }
            }
        }
    }
};

} // namespace memory
} // namespace atvoss

LayerNorm——调用子程序组合

// atvoss/compute/layernorm/layernorm.h
//
// LayerNorm 实现: 调用 atvoss 子程序组合
// 依赖: BlockReduceSum (warp + block reduce) + Online Softmax pattern
//
// LayerNorm(x) = (x - mean) / sqrt(var + eps) * gamma + beta
// mean = sum(x) / N
// var = sum((x - mean)^2) / N

#pragma once
#include "atvoss/compute/reduce/warp_reduce.h"
#include "atvoss/memory/tiling_2d.h"
#include "vector_common.h"

namespace atvoss {
namespace compute {

template<int HIDDEN_SIZE = 4096>
__aicore__ void LayerNorm(
    const float* input,       // [HIDDEN_SIZE]
    float* output,            // [HIDDEN_SIZE]
    const float* gamma,       // [HIDDEN_SIZE]
    const float* beta,        // [HIDDEN_SIZE]
    float eps = 1e-5f
) {
    using namespace atvoss::reduce;

    // Step 1: 计算 mean = sum(x) / HIDDEN_SIZE
    float sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < HIDDEN_SIZE; ++i) {
        sum += input[i];
    }
    // 如果 HIDDEN_SIZE > warp_size: 需要 block reduce
    // 这里简化: 假设单 warp 内可完成

    float mean = sum / float(HIDDEN_SIZE);

    // Step 2: 计算 variance
    float var_sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < HIDDEN_SIZE; ++i) {
        float diff = input[i] - mean;
        var_sum += diff * diff;
    }
    float inv_std = 1.0f / sqrtf(var_sum / float(HIDDEN_SIZE) + eps);

    // Step 3: normalize + affine
    #pragma unroll
    for (int i = 0; i < HIDDEN_SIZE; ++i) {
        output[i] = (input[i] - mean) * inv_std * gamma[i] + beta[i];
    }
}

/**
 * RMSNorm: 简化版 LayerNorm，不加 mean centered
 *
 * RMSNorm(x) = x / RMS(x) * gamma
 * RMS(x) = sqrt(sum(x^2) / N + eps)
 *
 * 相比 LayerNorm: 省一次 mean 计算 + 一次减法
 * 大模型 (LLaMA, Qwen) 里已基本替代 LayerNorm
 */
template<int HIDDEN_SIZE = 4096>
__aicore__ void RMSNorm(
    const float* input,
    float* output,
    const float* gamma,
    float eps = 1e-6f
) {
    // RMS 计算
    float sq_sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < HIDDEN_SIZE; ++i) {
        sq_sum += input[i] * input[i];
    }

    float rms = sqrtf(sq_sum / float(HIDDEN_SIZE) + eps);
    float inv_rms = 1.0f / rms;

    // normalize
    #pragma unroll
    for (int i = 0; i < HIDDEN_SIZE; ++i) {
        output[i] = input[i] * inv_rms * gamma[i];
    }
}

} // namespace compute
} // namespace atvoss

GEMM 微内核——Vector 单元的矩阵乘子程序

// atvoss/compute/gemm_micro/gemm_micro.h
//
// GEMM 微内核: 在 Vector 单元上用 outer product 做矩阵乘
// C[m][n] += sum_k A[m][k] * B[k][n]
//
// 适用场景: K 维度较小（≤512）的矩阵乘，不适合 Cube 单元时退化为 Vector

#pragma once
#include "vector_common.h"

namespace atvoss {
namespace compute {

/**
 * GEMM micro-kernel on Vector unit
 *
 * Outer product 策略:
 *   for k in 0..K-1:
 *     a_col = A[:, k]    // [M]
 *     b_row = B[k, :]    // [N]
 *     C[:, :] += outer(a_col, b_row)  // [M, N] += [M, 1] × [1, N]
 *
 * M, N 是固定的小维度（如 16×16），K 是收缩维度
 */
template<int M_BLOCK = 16, int N_BLOCK = 16, int K_UNROLL = 8>
struct GemmMicroKernel {

    /**
     * 执行一个 M×N 块的矩阵乘
     *
     * @param A    [M][K], row-major, leading dimension = lda
     * @param B    [K][N], row-major, leading dimension = ldb
     * @param C    [M][N], row-major, leading dimension = ldc (accumulate)
     * @param K    收缩维度
     */
    __aicore__ static void compute(
        const float* A, int lda,
        const float* B, int ldb,
        float* C, int ldc,
        int K
    ) {
        // 寄存器分块: 预加载一行 A + 一列 B
        float a_reg[M_BLOCK];   // A 的一列（M 个元素）
        float b_reg[N_BLOCK];   // B 的一行（N 个元素）

        // C 全在寄存器: 16×16 = 256 floats
        float c_reg[M_BLOCK * N_BLOCK] = {0};

        // 主循环: k 维度收缩
        for (int k = 0; k < K; k += K_UNROLL) {
            int k_end = (k + K_UNROLL <= K) ? k + K_UNROLL : K;

            for (int kk = k; kk < k_end; ++kk) {
                // 加载 A[:, kk]
                #pragma unroll
                for (int m = 0; m < M_BLOCK; ++m) {
                    a_reg[m] = A[m * lda + kk];
                }

                // 加载 B[kk, :]
                #pragma unroll
                for (int n = 0; n < N_BLOCK; ++n) {
                    b_reg[n] = B[kk * ldb + n];
                }

                // Outer product: C[m][n] += a_reg[m] * b_reg[n]
                #pragma unroll
                for (int m = 0; m < M_BLOCK; ++m) {
                    float a_val = a_reg[m];
                    #pragma unroll
                    for (int n = 0; n < N_BLOCK; ++n) {
                        c_reg[m * N_BLOCK + n] += a_val * b_reg[n];
                    }
                }
            }
        }

        // 写回 C（accumulate）
        #pragma unroll
        for (int m = 0; m < M_BLOCK; ++m) {
            #pragma unroll
            for (int n = 0; n < N_BLOCK; ++n) {
                C[m * ldc + n] += c_reg[m * N_BLOCK + n];
            }
        }
    }
};

} // namespace compute
} // namespace atvoss

流水线调度——搬运和计算重叠

// atvoss/schedule/pipeline.h
//
// 流水线调度: 双缓冲 + 搬运/计算 overlap
// 模式: 当前块计算 | 下一块预取

#pragma once
#include "vector_common.h"

namespace atvoss {
namespace schedule {

/**
 * 双缓冲流水线
 *
 * 模式:
 *   slot 0: 计算块 i   | 预取块 i+2  | 计算块 i+2 | ...
 *   slot 1: 预取块 i+1  | 计算块 i+1  | 预取块 i+3 | ...
 *
 * 时间线:
 *   T0: 搬块0到slot0 (Wait)
 *   T1: 搬块1到slot1 | 算slot0
 *   T2: 算slot1 | 搬块2到slot0
 *   T3: 算slot2(slot0) | 搬块3到slot1
 *   ...
 *
 * 搬运和计算只在第一步串行，之后完全重叠
 *
 * @param num_blocks  总块数
 * @param load_fn     (slot_idx, block_idx) → 搬运函数
 * @param compute_fn  (slot_idx, block_idx) → 计算函数
 */
template<typename LoadFn, typename ComputeFn>
__aicore__ void DoubleBufferedPipeline(
    int num_blocks,
    LoadFn load_fn,
    ComputeFn compute_fn
) {
    if (num_blocks <= 0) return;

    // 预加载块 0
    load_fn(0, 0);
    __sync();

    int current_slot = 0;
    int next_slot = 1;

    for (int i = 0; i < num_blocks; ++i) {
        // 预取下一块（如果存在）
        if (i + 1 < num_blocks) {
            load_fn(next_slot, i + 1);
        }

        // 计算当前块（和预取并行）
        compute_fn(current_slot, i);
        __sync();

        // 交换 slot
        int tmp = current_slot;
        current_slot = next_slot;
        next_slot = tmp;
    }
}

} // namespace schedule
} // namespace atvoss

算子开发者怎么用

// 自定义 LayerNorm 算子，组合 atvoss 子程序
//
// 依赖: atvoss::reduce::WarpReduceSum + atvoss::reduce::BlockReduceSum

#include "atvoss/compute/reduce/warp_reduce.h"
#include "atvoss/compute/layernorm/layernorm.h"

__aicore__ void MyCustomNormalize(
    const float* input, float* output,
    const float* gamma, const float* beta, int hidden_size
) {
    // 方案 A: 直接用 atvoss 封装（小 hidden_size）
    if (hidden_size <= 4096) {
        atvoss::compute::LayerNorm<4096>(input, output, gamma, beta);
        return;
    }

    // 方案 B: 大 hidden_size → 调子程序手动组装
    float sum = 0.0f;
    for (int i = 0; i < hidden_size; ++i) sum += input[i];

    // 多 warp block reduce
    float global_sum = atvoss::reduce::BlockReduceSum(
        sum, shared_buf, GetWarpId(), num_warps
    );

    float mean = global_sum / hidden_size;

    // ... 后续逻辑
}

踩坑：warp reduce 的 divergence——warp 不均匀时部分结果归零

// ❌ 直接 warp reduce: 最后几个 warp 的 lane 数不足 32
// warp 7 只有 4 个有效 lane → shuffle_down(offset=16,8,4) 读到 0
// → 结果被 0 污染

// ✅ 先掩码无效 lane 为 identity 值
// sum reduce: identity = 0 (不影响加法)
// max reduce: identity = -INF

template<typename T>
__aicore__ T MaskedWarpReduceSum(T val, int active_lanes) {
    // 掩码: 无效 lane 设 0
    if (GetLaneId() >= active_lanes) {
        val = T(0);
    }
    return WarpReduceSum(val);
}

踩坑：GEMM 微内核 outer product 寄存器溢出——16×16 的 float 就 256 个

// ❌ 设 M=N=32，寄存器需要 32×32=1024 floats → 4096 bytes
// Vector 寄存器文件不够 → 编译器 spill 到 L1 → 慢 10×
//
// ✅ 寄存器分块上限: FP32 时 M×N ≤ 256 (16×16)
//     FP16 时 M×N ≤ 512 (16×32)
//     超过→用 L1 做累加缓存，而不是寄存器

atvoss 把 Vector 算子的高频子程序抽成可复用模板：warp reduce（shuffle_down butterfly 5 步 32→1）→ block reduce（warp reduce + shared 暂存 + 二次 reduce）→ online softmax（max+sum 双状态增量更新，修正因子 exp(m-m_new)≤1 数值安全）→ LayerNorm/RMSNorm（调用 block reduce 算 mean/var + affine）→ GEMM 微内核（outer product，M×N≤256 FP32）→ 双缓冲流水线（搬运/计算 overlap，第二步起完全并行）。踩坑：warp lane 不足 32 时 shuffle_down 读到 0 污染→identity value 掩码、GEMM 微内核寄存器 M×N>256 溢出→降维或退 L1 累加。这套子程序库让算子开发从"每次重写 reduce"变成"include + 调用"，改一个 bug 所有算子自动受益。