前言

你有没有遇到过这种情况：做信号处理，用NumPy的FFT算个1024点变换，要算半天。后来发现昇腾CANN有个ops-fft库，专门加速FFT计算，同样的计算在NPU上只要几毫秒。这篇文章就来讲讲FFT是啥、为啥要优化、怎么用ops-fft库。

一、先搞懂：什么是FFT？

说白了，FFT就是快速算傅里叶变换的算法。

傅里叶变换是啥？它是把信号从"时域"变到"频域"的数学工具。比如你有一段音频（时域信号），想看看它包含哪些频率成分（频域信息），就用傅里叶变换。

举个例子：

• 时域信号：[sin(2π*10*t), sin(2π*20*t)]（两个正弦波叠加，频率10Hz和20Hz）
• 傅里叶变换后：[(10Hz, 幅值1.0), (20Hz, 幅值1.0)]（能看到两个频率成分）

FFT（Fast Fourier Transform）是算傅里叶变换的快速算法，不然按照原始定义算，复杂度O(N²)，慢得要死。FFT把复杂度降到O(N log N)，才算能用。

二、为什么FFT需要专门优化？

你可能会问：“NumPy不是有FFT吗？干嘛还要ops-fft？”

两个原因：

第一，CPU算FFT太慢了。
FFT是计算密集型任务，特别是2D/3D FFT（比如图像处理、科学计算），CPU要算半天。NPU有专门的FFT加速单元，能快很多。

第二，NumPy的FFT不支持GPU/NPU。
NumPy的FFT只能在CPU上算，你要算完再搬数据到NPU，浪费带宽。ops-fft直接在NPU上算，数据不用搬来搬去。

所以，ops-fft的核心价值就是：让FFT计算在NPU上跑到硬件极限性能，同时提供和NumPy一样的易用接口。

三、ops-fft的核心架构

ops-fft是昇腾CANN开源社区的FFT算子库，它在CANN五层架构中位于第二层——昇腾计算服务层，是AOL算子库的重要组成部分。

这个库的核心工作流程分三步：

第一步：参数解析

它先解析你给的FFT参数（比如变换维度、变换长度、变换类型等），然后生成对应的执行计划。

打个比方：你告诉厨师"我要做蛋炒饭"，厨师先搞清楚：用几个鸡蛋、几碗饭、要不要加葱、用啥锅…然后想好步骤（先炒蛋、再加饭、最后调味）。

第二步：算子选择

根据解析好的参数，选择合适的底层FFT算子（比如1D-FFT、2D-FFT、3D-FFT、实数FFT、复数FFT等）。

打个比方：厨师根据你的要求，选择合适的厨具（炒锅、平底锅、蒸锅…）。

第三步：执行计算

调用选好的底层FFT算子，在NPU上执行计算，然后把结果返回给你。

打个比方：厨师用选好的厨具，按照想好的步骤，把蛋炒饭做出来。

四、核心算子详解

1. 1D-FFT（一维FFT）

一维FFT是最基础的FFT算子，计算一维信号的傅里叶变换。

代码示例：

import torch
import ops_fft  # 导入ops-fft的Python接口
import numpy as np

# 1. 创建测试信号（时域）
# 信号包含10Hz和20Hz两个频率成分
fs = 1000  # 采样率1000Hz
t = np.linspace(0, 1, fs)
signal = np.sin(2 * np.pi * 10 * t) + np.sin(2 * np.pi * 20 * t)
signal_tensor = torch.from_numpy(signal).npu()

# 2. 执行1D-FFT
# 使用ops-fft的fft算子
fft_result = ops_fft.fft(signal_tensor)

# 3. 计算幅度谱
magnitude = fft_result.abs()

# 4. 打印结果（应该能看到10Hz和20Hz的峰值）
freqs = np.fft.fftfreq(len(signal), 1/fs)
print("频率成分（前20个）:")
for i in range(20):
    print("  频率: {:.1f} Hz, 幅度: {:.2f}".format(
        freqs[i], magnitude[i].item()
    ))

# 5. 对比NumPy的结果（应该一致）
numpy_fft = np.fft.fft(signal)
numpy_magnitude = np.abs(numpy_fft)
error = (magnitude.cpu().numpy() - numpy_magnitude).max()
print("最大误差:", error)  # 应该很小（<1e-5）

这段代码展示了ops-fft的1D-FFT功能：你给它时域信号，它给你频域表示。

2. 2D-FFT（二维FFT）

二维FFT用于图像处理（比如滤波、压缩等）。

代码示例：

import torch
import ops_fft
import numpy as np
from PIL import Image

# 1. 读取图像（灰度图）
image = Image.open("test_image.jpg").convert("L")
image_array = np.array(image)
image_tensor = torch.from_numpy(image_array).npu().float()

# 2. 执行2D-FFT
fft2_result = ops_fft.fft2(image_tensor)

# 3. 计算幅度谱（并做log变换，方便显示）
magnitude = fft2_result.abs()
magnitude_log = torch.log(magnitude + 1)

# 4. 打印结果形状
print("图像形状:", image_tensor.shape)
print("FFT结果形状:", fft2_result.shape)

# 5. 频域滤波（低通滤波）
# 把高频成分设为0
h, w = fft2_result.shape
center_h, center_w = h // 2, w // 2
radius = 30  # 低通滤波半径

mask = torch.ones_like(fft2_result)
for i in range(h):
    for j in range(w):
        if (i - center_h)**2 + (j - center_w)**2 > radius**2:
            mask[i, j] = 0.0

filtered_fft = fft2_result * mask

# 6. 逆FFT（回到时域）
filtered_image = ops_fft.ifft2(filtered_fft).abs()

# 7. 保存结果
filtered_image_np = filtered_image.cpu().numpy().astype(np.uint8)
Image.fromarray(filtered_image_np).save("filtered_image.jpg")

print("滤波完成，结果保存为 filtered_image.jpg")

这段代码展示了ops-fft的2D-FFT功能：你可以做频域滤波（比如低通滤波、高通滤波等）。

3. 实数FFT（RFFT）

实数FFT是专门针对实数信号的FFT，能节省一半的计算量和显存。

代码示例：

import torch
import ops_fft
import numpy as np

# 1. 创建实数信号
fs = 1000
t = np.linspace(0, 1, fs)
signal = np.sin(2 * np.pi * 10 * t) + np.sin(2 * np.pi * 20 * t)
signal_tensor = torch.from_numpy(signal).npu().float()

# 2. 执行实数FFT（RFFT）
rfft_result = ops_fft.rfft(signal_tensor)

# 3. 对比普通FFT的结果
fft_result = ops_fft.fft(signal_tensor)

print("RFFT结果形状:", rfft_result.shape)  # 应该是 (fs//2 + 1,)
print("FFT结果形状:", fft_result.shape)     # 应该是 (fs,)

# 4. 验证等价性（RFFT的结果和FFT的前半部分一致）
error = (rfft_result - fft_result[:len(rfft_result)]).abs().max().item()
print("最大误差:", error)  # 应该很小（<1e-5）

# 5. 逆RFFT（回到时域）
reconstructed_signal = ops_fft.irfft(rfft_result)

# 6. 验证重建误差
reconstruction_error = (reconstructed_signal - signal_tensor).abs().max().item()
print("重建误差:", reconstruction_error)  # 应该很小（<1e-5）

这段代码展示了ops-fft的实数FFT功能：对于实数信号，用RFFT能节省一半的计算量和显存。

五、性能优化技巧

1. 使用混合精度

import torch
import ops_fft

# 方案1：FP32计算（高精度）
signal_fp32 = torch.randn(1024).npu().float()
fft_fp32 = ops_fft.fft(signal_fp32, dtype=torch.complex64)

# 方案2：FP16计算（高性能）
signal_fp16 = signal_fp32.half()
fft_fp16 = ops_fft.fft(signal_fp16, dtype=torch.complex32)

# 方案3：混合精度（FP16计算，FP32输出）
fft_mixed = ops_fft.fft(signal_fp16, dtype=torch.complex64)

2. 批量处理

import torch
import ops_fft

# 方案1：逐个处理（慢）
signals = torch.randn(100, 1024).npu()
results = []
for i in range(100):
    result = ops_fft.fft(signals[i])
    results.append(result)

# 方案2：批量处理（快）
# ops-fft支持批量FFT（通过指定dim参数）
batch_fft = ops_fft.fft(signals, dim=1)

3. 合理利用显存

import torch
import ops_fft

# 对于大尺寸FFT，显存可能不够
# 解决方案1：拆分计算（分块FFT）
def block_fft(signal, block_size):
    num_blocks = (len(signal) + block_size - 1) // block_size
    results = []
    for i in range(num_blocks):
        block = signal[i*block_size:(i+1)*block_size]
        result = ops_fft.fft(block)
        results.append(result)
    return results

# 解决方案2：使用实数FFT（节省一半显存）
signal = torch.randn(1048576).npu()  # 1M点FFT
rfft_result = ops_fft.rfft(signal)  # 只用一半显存

六、实际应用场景

场景1：音频处理（降噪）

import torch
import ops_fft
import numpy as np
from scipy.io import wavfile

# 1. 读取音频文件
fs, audio = wavfile.read("noisy_audio.wav")
audio_tensor = torch.from_numpy(audio).npu().float()

# 2. 执行FFT（转到频域）
fft_result = ops_fft.fft(audio_tensor)

# 3. 频域降噪（去除高频噪声）
# 假设噪声在频率>4000Hz的区域
freqs = torch.fft.fftfreq(len(audio), 1/fs).npu()
noise_mask = (freqs.abs() > 4000)
fft_result[noise_mask] = 0

# 4. 逆FFT（回到时域）
denoised_audio = ops_fft.ifft(fft_result).abs()

# 5. 保存降噪后的音频
wavfile.write(
    "denoised_audio.wav",
    fs,
    denoised_audio.cpu().numpy().astype(np.int16)
)

print("降噪完成，结果保存为 denoised_audio.wav")

场景2：图像处理（边缘检测）

import torch
import ops_fft
import numpy as np
from PIL import Image

# 1. 读取图像（灰度图）
image = Image.open("test_image.jpg").convert("L")
image_array = np.array(image)
image_tensor = torch.from_numpy(image_array).npu().float()

# 2. 执行2D-FFT
fft2_result = ops_fft.fft2(image_tensor)

# 3. 频域高通滤波（提取边缘）
h, w = fft2_result.shape
center_h, center_w = h // 2, w // 2
radius = 30  # 高通滤波半径

mask = torch.ones_like(fft2_result)
for i in range(h):
    for j in range(w):
        if (i - center_h)**2 + (j - center_w)**2 <= radius**2:
            mask[i, j] = 0.0

highpass_fft = fft2_result * mask

# 4. 逆FFT（回到时域）
edge_image = ops_fft.ifft2(highpass_fft).abs()

# 5. 保存结果
edge_image_np = edge_image.cpu().numpy().astype(np.uint8)
Image.fromarray(edge_image_np).save("edge_image.jpg")

print("边缘检测完成，结果保存为 edge_image.jpg")

七、总结

ops-fft是昇腾CANN生态中非常重要的FFT算子库，核心价值在于：

1. 高性能：针对昇腾NPU的FFT加速单元做了深度优化
2. 易用性：Python接口和NumPy/PyTorch无缝集成，改几行代码就能用上
3. 全面性：支持1D/2D/3D FFT、实数FFT、批量FFT等各种场景

实际用下来，在信号处理、图像处理、科学计算等领域，这个库能带来显著的性能提升。特别是大尺寸FFT，性能提升非常明显。

当然，这个库也不是万能的。有些特别新的FFT变种可能还没实现，需要你自己参考现有算子开发。但这种参考的过程，也是深入理解FFT优化的好机会。

更多技术细节和最新进展，可以去仓库看看：https://atomgit.com/cann/ops-fft