训练营简介 2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。

报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro

前言

在深度学习爆发之前，FFT 曾被誉为“20 世纪最重要的算法”。 它不仅用于信号处理（MP3, JPEG），现在更是 AI for Science 的基石。例如，在求解偏微分方程（PDE）时，利用 Spectral Method（谱方法），我们可以将复杂的卷积运算转化为频域的简单乘法：

$$f * g = \mathcal{F}^{-1}(\mathcal{F}(f) \cdot \mathcal{F}(g))$$

FFT 的核心挑战在于：

1. 复数运算：AI Core 的 Vector 单元原生支持 FP16/FP32，但没有原生的 complex 数据类型，需要我们用两个 float 模拟。

2. 蝶形运算 (Butterfly)：数据依赖关系复杂，每一级计算都需要跨步访问。

3. 非连续访存：标准的 Cooley-Tukey 算法需要 Bit-Reversal（位反转） 排列（如 001 -> 100），这对擅长连续读写的 MTE 单元来说是极其低效的。

本期文章，我们将避开低效的位反转，采用更适合 SIMD 硬件的 Stockham 算法，在 Ascend C 上实现高效 FFT。

一、 核心图解：时域与频域的棱镜

FFT 就像是一个数学棱镜，把混合在一起的波形（时域），拆解成一个个纯净的频率（频域）。

二、 算法原理：为什么选择 Stockham？

在 NPU 上写 FFT，算法的选择比代码优化更重要。

• Cooley-Tukey：教科书里的标准算法。最大的缺点是In-place（原地）计算需要Bit-Reversal重排。这意味着你需要把内存地址 0, 1, 2, 3 变成 0, 2, 1, 3，这种随机 Scatter 操作会把 MTE 带宽打至谷底。

• Stockham：虽然需要两倍的存储空间（Ping-Pong Buffer），但它不需要位反转。每一级迭代（Stage）都是连续读、连续写。

结论：用空间换时间，Stockham 完美契合 Ascend C 的流水线架构。

蝶形运算 (Butterfly Unit)： 基本单元是计算两个复数 $A$ 和 $B$：

$$X = A + W \cdot B$$$$Y = A - W \cdot B$$

其中 $W$ 是旋转因子（Twiddle Factor，复数）。

三、 实战：Ascend C 实现 FFT

我们假设输入是两个 Tensor：real 和 imag（实部和虚部）。

3.1 Kernel 类定义与双缓冲设计

为了实现 Stockham 算法的 Ping-Pong 迭代，我们需要在 UB 中申请两组 Buffer。

class KernelFFT {
public:
    __aicore__ inline void Init(GM_ADDR real, GM_ADDR imag, GM_ADDR out_real, GM_ADDR out_imag, 
                                uint32_t n, GM_ADDR twiddles) {
        this->n = n;
        // Init GlobalTensors...
        
        // 申请 Ping-Pong Buffer (UB)
        // 假设 N=1024，足够放入 UB
        // 如果 N 很大，需要做分块 FFT (Block FFT)
        pipe.InitBuffer(pingQ, 2, n * sizeof(float)); // 存实部和虚部
        pipe.InitBuffer(pongQ, 2, n * sizeof(float));
        
        // Twiddle Factors (旋转因子)
        // 最好在 Host 侧算好传入，不要在 Kernel 里算 Cos/Sin
        this->twiddlesGm.SetGlobalBuffer((__gm__ float*)twiddles);
    }
    
    __aicore__ inline void Process() {
        // 1. CopyIn
        LocalTensor<float> pingReal = pingQ.AllocTensor<float>();
        LocalTensor<float> pingImag = pingQ.AllocTensor<float>();
        DataCopy(pingReal, xRealGm, n);
        DataCopy(pingImag, xImagGm, n);
        
        // 2. FFT Main Loop (log2(N) stages)
        ComputeFFT(pingReal, pingImag);
        
        // 3. CopyOut
        // 最终结果可能在 Ping 也可能在 Pong，取决于 log2(N) 的奇偶性
    }
};

3.2 复数乘法 (Complex Mul)

AI Core Vector 单元没有 ComplexMul 指令，我们需要用实数指令组合。 公式：

$$(a + bi)(c + di) = (ac - bd) + (ad + bc)i$$

__aicore__ inline void ComplexMul(LocalTensor<float>& resR, LocalTensor<float>& resI,
                                  LocalTensor<float>& aR, LocalTensor<float>& aI,
                                  LocalTensor<float>& bR, LocalTensor<float>& bI, 
                                  uint32_t len) {
    // 申请临时寄存器
    LocalTensor<float> tmp1 = tmpQueue.AllocTensor<float>();
    LocalTensor<float> tmp2 = tmpQueue.AllocTensor<float>();

    // Real Part: ac - bd
    Mul(tmp1, aR, bR, len);
    Mul(tmp2, aI, bI, len);
    Sub(resR, tmp1, tmp2, len);

    // Imag Part: ad + bc
    Mul(tmp1, aR, bI, len);
    Mul(tmp2, aI, bR, len);
    Add(resI, tmp1, tmp2, len);
    
    tmpQueue.FreeTensor(tmp1);
    tmpQueue.FreeTensor(tmp2);
}

3.3 Stockham 迭代逻辑

Stockham 算法在第 s 级迭代时，数据不仅需要蝶形运算，还需要进行特定的混洗（Shuffle）。 在 Vector 单元上实现 Shuffle 有两种思路：

1. DataCopy：利用 dstStride 和 srcStride 进行交织搬运。

2. Mask + Offset：利用 Vector 指令的掩码和源地址偏移。

这里展示一种基于 步长（Stride） 的逻辑视角：

__aicore__ inline void ComputeFFT(LocalTensor<float>& inReal, LocalTensor<float>& inImag) {
    // Ping-Pong 切换指针
    LocalTensor<float>* currR = &inReal;
    LocalTensor<float>* nextR = &pongReal; // 简化示意
    
    // 这是一个 Radix-2 循环
    for (int s = 1; s <= log2N; s++) {
        uint32_t half_width = 1 << (s - 1);
        
        // 我们需要加载 Twiddle Factor
        // 这里的 Twiddle 访问模式是重复的，利用 Repeat/Stride 机制
        
        // 核心蝶形运算：
        // butterfly_upper = A + W*B
        // butterfly_lower = A - W*B
        // Ascend C 优化：一次计算一组，而不是一个点
        
        // 这里的难点在于：如何并行地拿到 A 和 B？
        // 在 Stockham 中，A 和 B 在内存中是分块连续的
        
        // 假设我们已经通过地址偏移拿到了 A_vec, B_vec 和 W_vec
        ComplexMul(W_B_Real, W_B_Imag, W_Real, W_Imag, B_Real, B_Imag, len);
        
        Add(nextR_upper, A_Real, W_B_Real, len/2);
        Sub(nextR_lower, A_Real, W_B_Real, len/2);
        
        // Swap Ping/Pong
        swap(currR, nextR);
        PipeBarrier<PIPE_V>(); // 确保计算完成
    }
}

四、 性能优化的“频率魔法”

FFT 是典型的 Compute Bound（大量的乘加）和 Latency Bound（复杂的依赖）混合体。

4.1 预计算旋转因子 (Precomputed Twiddles)

绝对不要在 Kernel 里动态计算 Cos 和 Sin！这会极其浪费算力。 必须在 Host 侧算好，甚至针对每一级 Stage 的访问顺序进行重排（Reorder），使得 Kernel 内可以顺序读取 W，无需 Gather。

4.2 向量化复数指令

虽然我们用 Float 模拟了 Complex，但在最新的 Ascend 芯片（如 910B）中，可能存在针对特定模式优化的指令。此外，利用 Muls、Mads（乘加）指令可以合并部分乘法和加法，减少指令发射数。

4.3 2D FFT 的优化策略

在气象预测中，通常做 2D FFT。

• 行变换 (Row FFT)：直接在 UB 内做 1D FFT，利用 Vector 并行处理多行。

• 转置 (Transpose)：利用第 45 期学的转置技巧（MTE3 或 Cube 转置），将数据转置，使得列变成行。

• 列变换 (Col FFT)：再次做 1D FFT（此时在物理上是行）。

• 再转置：转回来。

这种 "Row-Transpose-Col" 模式比直接做列变换（Stride 访问）效率高得多，因为它最大限度地保证了访存的连续性。

五、 总结

FFT 是连接 AI 与物理世界的桥梁。

1. 数据结构：用两个 Float Tensor 模拟 Complex，利用结构体或指针数组管理 Ping-Pong。

2. 算法选择：Stockham 优于 Cooley-Tukey，因为它避免了位反转，更适合 NPU 的流式架构。

3. 应用：掌握了 FFT，你就能处理 PDE 求解（如 Navier-Stokes 方程）、语音降噪 等高端任务。

至此，我们的行业篇也画上了一个硬核的句号。从自动驾驶的点云，到科学计算的波形，Ascend C 的边界正在无限延展。