CANN ops-fft FFT 算子——频域卷积加速原理
前言
大卷积核的卷积操作如果在空域逐点相乘再求和，其计算复杂度为O(H×W×K×H×K)，随着卷积核尺寸增大，计算量呈平方级增长。当卷积核大于 7×7 时，空域卷积的计算开销已经大到难以忽视。FFT（快速傅里叶变换）提供了一条绕过这条困境的路径：将时域卷积转换为频域乘法，从而把复杂度从 O(N²) 降至 O(N log N)。ops-fft 仓正是昇腾 NPU 上 FFT 算子的完整实现，本文剖析它的原理与用法。
1. FFT 加速卷积的原理
1.1 时域卷积 → 频域相乘
两个信号 f 和 g 的卷积，在数学上定义为：
(f ⊛ g)(n) = Σₘ f(m) · g(n − m)
这是一个滑动求和的过程。如果把 f 和 g 都补零（zero-padding）到相同长度 N，然后对两端做离散傅里叶变换（DFT），则卷积定理告诉我们：
DFT(f ⊛ g) = DFT(f) × DFT(g)
其中 × 表示频域复数的逐点相乘。也就是说，整个卷积过程可以拆解为三步：
- FFT(f)→ 变换到频域
- FFT(g)→ 变换到频域
- 逐点相乘→ 频域乘法
- IFFT→ 逆变换回时域
每一步 FFT/IFFT 的复杂度为 O(N log N)，频域乘法为 O(N)，整体复杂度 O(N log N)，远优于空域卷积的 O(N²)。
1.2 补零对齐与圆周卷积
实际实现中必须处理一个关键技术细节：圆周卷积 vs 线性卷积。FFT 给出的是圆周卷积结果，只有当输入序列补零到足够长度时，圆周卷积才等于线性卷积。
设输入信号长度为 L₁，卷积核长度为 L₂，则线性卷积结果长度为 L₁ + L₂ − 1。补零规则为：将两个输入都补零到长度N ≥ L₁ + L₂ − 1，且 N 必须是 2 的整数次幂（因为 FFT 算法通常要求长度为 2 的幂）。
对于二维图像卷积（H × W 输入，K × K 卷积核），补零规则扩展为：
- 高度方向补零到 H + K − 1，并向上取整到 2 的幂
- 宽度方向补零到 W + K − 1，并向上取整到 2 的幂
这就是为什么 FFT 卷积在小卷积核上不一定占优——补零引入的开销在小尺寸时可能抵消频域计算带来的收益。
1.3 何时 FFT 卷积更划算
设输入尺寸为 H × W，卷积核为 K × K。空域卷积计算量约为H × W × K²次乘加。FFT 卷积的计算量约为：
- 两次 FFT + 一次 IFFT ≈ 3 × O(N log N)，N 为补零后尺寸
- 一次频域乘法：O(N)
判断条件为：3N log N + N < H × W × K²
粗略估计，当 K > 7 时，FFT 卷积开始显现优势；K 越大，优势越显著。对于大核卷积（K ≥ 11）、depthwise 大核卷积或 stride 卷积场景，FFT 加速效果尤为明显。
2. 昇腾 NPU 上的 FFT 实现
2.1 混合基数算法：Radix-2 + Radix-4
FFT 算法的基础是蝶形运算（butterfly operation），核心思想是将长度为 N 的 DFT 分解为更小的子 DFT。经典的 Cooley-Tukey 算法以 2 为基数（Radix-2），每级将问题规模减半。
ops-fft 仓采用了Radix-2 + Radix-4 混合方案：
- Radix-2：适合长度 N = 2ⁿ 的任意场景，实现简单，控制流稳定
- Radix-4：当 N 能被 4 整除时，每次蝶形合并 4 个输入，复数乘法次数减少约 25%，适合大尺寸变换
具体策略是：根据输入数据长度选择最优分解路径。对于昇腾 NPU 的硬件特性（Ascend 芯片的矩阵乘单元），Radix-4 可以在同一批指令中利用更多 SIMD 并行度，从而获得更高吞吐量。
2.2 单精度与双精度支持
ops-fft 支持 FP32（单精度）和 FP64（双精度）两种数据类型。深度学习训练中常用 FP32，前向推理则可能使用 FP16/BF16（但 FFT 核心仍以 FP32/FP64 执行以保证精度）。
对于实数输入（图像、特征图都是实数），ops-fft 还支持half-format 优化：利用一个复数 FFT 同时处理两个实数序列，节省约 50% 的存储和计算资源。
2.3 内存布局与张量格式
昇腾 NPU 的张量格式为 NCHW（Batch, Channel, Height, Width）。ops-fft 内部将 2D FFT 按行优先逐通道处理，数据从 NCHW 格式转换为计算友好的线性布局，完成变换后再恢复原格式。
补零操作在原地（in-place）完成，不额外分配临时缓冲区，以降低内存峰值。
3. 适用场景
3.1 大卷积核（> 7×7）
这是 FFT 卷积的核心优势场景。典型应用包括：
- 空洞卷积（Dilated Convolution）：感受野大但权重数量相对少，FFT 可高效计算
- 大核可分离近似：用大核 1D FFT 分别处理行列方向，降低 2D 大核的计算开销
- GAN / Diffusion 模型中的大核上采样：上采样卷积层通常 K=3 或 K=4，但反卷积/转置卷积的"有效卷积核"实际更大
3.2 Stride 卷积
当 stride > 1 时，空域卷积的有效计算量下降，但控制流和数据重排开销不变。FFT 方法天然支持 stride——通过在频域直接对输入下采样后的 FFT 结果进行运算，可以绕过空域 stride 的冗余计算。
3.3 深度可分离卷积（Depthwise Convolutions）
在 depthwise 卷积中，每个通道独立卷积，FFT 的单通道计算可以完全并行。ops-fft 通过向量化多个单通道 FFT，在昇腾 NPU 的 Tensor Engine 上实现了极高的并行度。
3.4 不适合的场景
- 小卷积核（K ≤ 5）：补零开销 > FFT 收益，空域卷积更快
- 非方正输入：非均匀补零会降低 FFT 效率
- 极端 batch size：内存限制导致无法容纳足够大的补零缓冲区
4. 代码示例：FFT 卷积实操
以下代码演示了如何使用 ops-fft 仓在昇腾 NPU 上执行 FFT 卷积。假设已正确安装 CANN 环境并配置了 Ascend 芯片。
importtorchimporttorch.nnasnnfromaclnnimportinner_product# CANN 线性代数核心（示例导入）# 模拟 ops-fft 仓的 FFT 卷积封装classFFTConv2d(nn.Module):""" 基于 ops-fft 思路的 FFT 卷积实现。 将输入和卷积核都补零到 2 的幂次长度， 分别做 FFT，相乘后 IFFT，再裁剪到目标尺寸。 """def__init__(self,in_channels,out_channels,kernel_size,stride=1,padding=0):super().__init__()self.in_channels=in_channels self.out_channels=out_channels self.kernel_size=kernel_size self.stride=stride self.padding=padding# 权重 shape: (out, in, H, W)self.weight=nn.Parameter(torch.randn(out_channels,in_channels,kernel_size,kernel_size))def_next_power_of_2(self,n):"""向上取整到 2 的幂"""return1<<(n-1).bit_length()defforward(self,x):B,C,H,W=x.shape# 1. 应用空域 padding（如果有）ifself.padding>0:x=torch.nn.functional.pad(x,[self.padding]*4)# 2. 计算补零后尺寸H_pad=self._next_power_of_2(H+self.kernel_size-1)W_pad=self._next_power_of_2(W+self.kernel_size-1)# 3. 补零到 FFT 友好尺寸x_padded=torch.nn.functional.pad(x,[0,W_pad-x.shape[3],0,H_pad-x.shape[2]])k_padded=torch.nn.functional.pad(self.weight,[0,W_pad-self.kernel_size,0,H_pad-self.kernel_size])# 4. 转换为复数格式（FFT 需要复数张量）x_complex=x_padded.float().to(torch.complex64)k_complex=k_padded.float().to(torch.complex64)# 5. 执行 FFT（PyTorch 内置实现作为演示）X=torch.fft.fft2(x_complex)K=torch.fft.fft2(k_complex)# 6. 频域逐通道相乘Y=X*K# 7. 逆变换回时域y=torch.fft.ifft2(Y).real# 8. 裁剪到原始空间尺寸 + stride 处理y=y[:,:,:H,:W]ifself.stride>1:y=y[:,:,::self.stride,::self.stride]returny# 使用示例model=FFTConv2d(in_channels=64,out_channels=128,kernel_size=11,stride=1,padding=5)x=torch.randn(1,64,32,32)output=model(x)print(f"Output shape:{output.shape}")# torch.Size([1, 128, 32, 32])上述示例展示了 FFT 卷积的核心流程：补零 → FFT → 频域相乘 → IFFT → 裁剪。实际 ops-fft 仓中使用了昇腾 NPU 原生的 FFT kernel，通过aclnnFft系列 API 直接调用硬件 FFT 单元，性能远超 PyTorch 软件实现。
5. 性能对比：空域卷积 vs FFT 卷积
5.1 理论计算量对比
| 卷积核尺寸 | 空域卷积乘加次数 | FFT 卷积操作数（约） | 节省比例 |
|---|---|---|---|
| K = 3 | 9HWC | 3N log₂N + N | 负收益（补零开销大） |
| K = 7 | 49HWC | 3N log₂N + N | 基本持平 |
| K = 11 | 121HWC | 3N log₂N + N | ~30% 节省 |
| K = 15 | 225HWC | 3N log₂N + N | ~50% 节省 |
N 为补零后尺寸（向上取整到 2 的幂），H × W 为输入空间尺寸。K 越大，节省越显著。
5.2 实际性能考量
理论分析之外，实际部署还需考虑：
- 内存占用：FFT 需要 2N 大小的临时缓冲区，对于大图像（1024×1024）和大卷积核（K=15），N 接近 2048×2048，单次 FFT 显存开销可达数十 MB。
- 数据搬移开销：输入从 global memory 搬到 FFT 单元，再搬回，数据路径上的带宽消耗在异构架构中不可忽视。
- 批量并行度：batch size 较大时，FFT 单元可以被充分利用，但 batch size = 1 时，FFT 开销相对固定，收益降低。
ops-fft 仓在内存管理上做了精细优化：原地计算策略减少临时缓冲区、通道维度上批量提交 FFT 任务以提高硬件利用率、多流并行掩盖数据搬移延迟。
5.3 端到端加速收益
在典型的 Vision Transformer 或大核 CNN 中，将 11×11 卷积层替换为 FFT 版本，在昇腾 910 系列芯片上可获得1.5× ~ 3×的端到端 throughput 提升，具体取决于输入分辨率和 batch size。
6. 总结与资源
FFT 卷积的本质是利用卷积定理，将时域的 O(N²) 卷积操作转换为频域的 O(N log N) FFT 操作。对于大卷积核（K > 7）、stride 卷积和 depthwise 场景，这种转换带来的加速收益显著，是高性能深度学习模型优化的重要手段。
ops-fft 仓作为昇腾 NPU 上的 FFT 算子实现，封装了混合基数 FFT、原地补零优化、half-format 实数加速等关键特性，为 CANN 生态下的高性能卷积计算提供了开箱即用的方案。
仓库地址：https://atomgit.com/cann/ops-fft
该仓持续更新，涵盖 1D FFT、2D FFT 及多维 FFT 的完整实现，欢迎通过 Issue 或 Pull Request 参与共建。