当前位置：首页 > news >正文

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

news 2026/7/10 16:47:29

训练营简介 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 的核心挑战在于：

复数运算：AI Core 的 Vector 单元原生支持 FP16/FP32，但没有原生的complex数据类型，需要我们用两个float模拟。
蝶形运算 (Butterfly)：数据依赖关系复杂，每一级计算都需要跨步访问。
非连续访存：标准的 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 有两种思路：

DataCopy：利用dstStride和srcStride进行交织搬运。
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>(); // 确保计算完成 } }