【昇腾CANN训练营·行业篇】科学计算之光：基于 AI Core 的高性能 FFT (快速傅里叶变换) 算子开发

训练营简介 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 的边界正在无限延展。