CANN 组织链接： https://atomgit.com/cann
ops-nn 仓库链接： https://atomgit.com/cann/ops-nn

---

循环神经网络（RNN）及其变体（LSTM, GRU）因其天然的时序依赖性，在并行计算架构上一直面临巨大的性能挑战。ops-nn 库中的 DynamicRNNV2 算子并非简单的逻辑移植，而是针对 NPU 异构计算架构进行的深度定制。它通过算子融合、极致的片上内存管理以及混合精度量化策略，成功解决了序列计算中的“内存墙”问题，实现了高吞吐量的时序推理。

1. 序列依赖与并行计算的架构冲突

RNN 的核心数学特征是 $H_t=f(X_t,H_{t-1})$。这种强依赖关系意味着在计算完成 $t-1$ 时刻之前，$t$ 时刻的计算无法启动。这与 GPU/NPU 这种依赖大规模数据并行的硬件架构存在天然的冲突。

1.1 细粒度调度的开销

在传统的深度学习框架实现中，一个 LSTM 单元（Cell）可能被拆分为多个独立的算子（如 Split, MatMul, Add, Sigmoid, Tanh, Mul）。

• 内核启动开销：对于长度为 $T$ 的序列，框架需要启动 $T\times N_{ops}$ 个内核。当 $T$ 很大时，CPU 下发任务的开销将超过 GPU/NPU 的计算时间。
• 显存带宽浪费：每个微小算子都会产生中间结果的读写（Global Memory R/W），导致带宽被大量无效的数据搬运占据。

1.2 算子融合（Operator Fusion）策略

DynamicRNNV2 采用了“大核”设计思路，将整个时间步的计算逻辑封装在一个 Kernel 中。

• 指令级流水线：在 Kernel 内部，利用指令并行性，将矩阵乘（Cube Unit）与激活函数（Vector Unit）进行流水线编排。
• 一次启动：无论序列长度如何，整个 RNN 层仅需一次 Kernel Launch，彻底消除了框架层面的调度开销。

2. 极致的内存管理：State Retention 机制

在 NPU 架构中，数据搬运的能耗和延迟远高于计算。DynamicRNNV2 的核心优化在于如何处理隐藏状态（Hidden State, $H$）和细胞状态（Cell State, $C$）。

2.1 状态数据的片上驻留

为了避免反复访问慢速的全局内存（HBM），算子利用片上 Unified Buffer (UB) 实现了状态驻留。

• 初始化加载：仅在 $t=0$ 时，从 HBM 读取初始状态 $H_0,C_0$。
• 闭环更新：在处理 $t\to t+1$ 的过程中，计算出的 $H_t,C_t$ 直接写入 UB 的特定区域，供下一个时间步的指令直接读取。
• 最终回写：仅在序列结束或 UB 空间不足时，才将必要的状态序列回写至 HBM。

2.2 权重矩阵的预处理与常驻

LSTM 的权重通常较大，DynamicRNNV2 采用了权重拆分与复用策略：

• 权重拼接：将 $W_{ii},W_{if},W_{ig},W_{io}$ 拼接为一个大矩阵，一次性完成所有门的线性变换，提高矩阵乘单元（Cube）的利用率。
• L1 Cache 驻留：利用 NPU 的 L1 缓存特性，尽可能让权重数据在 L1 中保持“温热”状态，减少从 L2 或 HBM 重新加载的次数。

3. 动态序列长度（Dynamic Shape）的运行时适配

实际业务（如语音识别或 NLP）中的 Batch 数据通常包含不同长度的序列。DynamicRNNV2 摒弃了低效的 Padding 填充计算，实现了真正的动态执行。

3.1 基于序列长度的控制流

算子接收一个辅助输入 seq_length，记录 Batch 中每个样本的真实长度。

• 条件执行：在 Kernel 内部循环中，针对 Batch 维度的每一个样本，单独判断当前时间步 $t$ 是否小于其 seq_length。
• 无效计算屏蔽：对于已经结束的序列，算子会自动跳过后续的计算指令，或者利用 Predicate 寄存器屏蔽写操作，防止脏数据污染结果。

3.2 尾部（Tail）数据的 Tiling 处理

当序列长度无法被硬件的 Block 大小整除时，边界处理变得至关重要。

• 对齐填充：在片上 UB 中，算子会自动对尾部数据进行零填充以满足 32 字节对齐要求。
• Mask 机制：在写回结果时，使用向量掩码（Vector Mask）指令，仅将有效的数据位写回 Global Memory，确保内存安全。

4. 异构计算单元协同：Cube 与 Vector 的这种平衡

NPU 拥有专门处理矩阵运算的 Cube Unit 和处理向量运算的 Vector Unit。DynamicRNNV2 必须在这两者之间进行精细的任务分配。

4.1 线性变换与非线性激活的分离

LSTM 的计算可以清晰地划分为两部分：

• Cube 密集计算：$X_t\times W_{ih}$ 和 $H_{t-1}\times W_{hh}$ 占据了 90% 以上的计算量，这部分被完全映射到 Cube Unit 上执行，享受高算力。
• Vector 复杂计算：门控信号的 Sigmoid、Tanh 激活以及 $C_t$ 的逐元素加法、乘法，被映射到 Vector Unit。

4.2 依赖关系的流水线掩盖

为了防止 Vector Unit 等待 Cube Unit 的结果（或是反之），算子内部采用了双缓冲（Double Buffering）技术。
当 Cube Unit 正在计算 $t+1$ 时刻的矩阵乘时，Vector Unit 并行地处理 $t$ 时刻的激活函数计算。这种深度流水线设计最大化了硬件单元的并发度。

5. 边缘侧加速：INT8 量化与混合精度计算

在推理场景下，DynamicRNNV2 支持 INT8 量化，这不仅减少了模型体积，更大幅提升了计算吞吐量。

5.1 矩阵乘法的 INT8 路径

量化主要针对计算密集型的矩阵乘法：

• 输入量化：输入 $X_t$ 和隐藏状态 $H_{t-1}$ 被量化为 INT8 格式。
• 高吞吐计算：Cube Unit 执行 INT8 乘法，并将结果累加到 INT32 寄存器中，防止溢出。

5.2 激活函数的 FP16 精度恢复

由于 Sigmoid 和 Tanh 函数对输入范围非常敏感，直接使用 INT32 结果进行查表或近似计算会导致精度严重下降。

• Dequantization：将 INT32 的累加结果乘以量化比例因子（Scale），恢复为 FP16 浮点数。
• 高精度激活：在 FP16 精度下执行非线性激活函数。
• Requantization：将激活后的 $H_t$ 再次量化为 INT8，以便用于下一时刻的输入。

6. 数据排布（Layout）与 DMA 搬运优化

NPU 的硬件架构对数据在内存中的物理排布非常敏感。DynamicRNNV2 在数据加载和存储阶段进行了专门的格式优化。

6.1 分形格式（Fractal Format）的转换

Cube Unit 并不擅长处理连续的行主序（Row-Major）数据，它更喜欢分形块格式（如 $16\times 16$ 的小块）。

• ND 到 5HD：算子在 DMA 搬运阶段，利用数据搬运单元（MTE）的“随路转换”能力，在不增加额外指令周期的情况下，将输入的 ND 格式转换为 NPU 内部的私有格式（如 NC1HWC0 或 Fractal_NZ）。
• 减少 Bank Conflict：这种特殊的排布方式还能有效减少片上 Bank 冲突，提升 L1 Cache 的读取带宽。

6.2 算子内部逻辑示意

以下代码展示了 DynamicRNNV2 核心循环在逻辑层面的实现结构，体现了状态驻留和流水线并行的思想：

template <typename T>
void DynamicRNNV2_Kernel(
    const T* input_x, 
    const T* weights, 
    T* output_h, 
    const int32_t* seq_lengths,
    int32_t max_time_step, 
    int32_t batch_size
) {
    // 1. 初始化片上缓存 (Unified Buffer)
    // h_state_ub 和 c_state_ub 常驻片上，避免重复 HBM 读写
    LocalTensor<T> h_state_ub = InitBuffer(batch_size, hidden_size);
    LocalTensor<T> c_state_ub = InitBuffer(batch_size, hidden_size);
  
    // 加载初始状态
    LoadToUB(h_state_ub, initial_h_gm);
    LoadToUB(c_state_ub, initial_c_gm);

    // 2. 时间步循环 (Time-Step Loop)
    for (int t = 0; t < max_time_step; ++t) {
      
        // 2.1 批量检查：跳过已结束的序列 (Dynamic Shape Support)
        // 使用 mask 标记当前 batch 中哪些样本依然有效
        auto valid_mask = Compare(seq_lengths, t); 
        if (AllSamplesFinished(valid_mask)) break;

        // 2.2 加载当前时刻输入 X_t
        LocalTensor<T> x_t_ub = LoadInputStep(input_x, t);

        // 2.3 密集计算 (Cube Unit Task)
        // 并行计算 [Xt * W_ih] 和 [Ht-1 * W_hh]
        // 结果累加到 gates_buffer
        LocalTensor<T> gates = MatMul_Cube(x_t_ub, weights_x);
        gates += MatMul_Cube(h_state_ub, weights_h); // 使用片上驻留的 Ht-1

        // 2.4 激活与状态更新 (Vector Unit Task)
        // 包含 Sigmoid, Tanh, Element-wise Mul/Add
        // 这一步与下一步的 Cube 计算可以形成流水线
        VectorCompute_LSTM_Cell(gates, c_state_ub, h_state_ub);

        // 2.5 结果输出
        // 仅将有效的 H_t 写回 Global Memory 用于输出序列
        StoreToGM(output_h, h_state_ub, t, valid_mask);
    }
}

通过上述六大维度的深度优化，DynamicRNNV2 将原本难以并行化的串行算法，成功适配到了以吞吐量见长的 NPU 架构上，成为了 ops-nn 库中处理时序数据的标杆算子。