昇腾CANN ops-fft:在 NPU 上做离散傅里叶变换
信号处理、频谱分析、图像滤波——大量科学计算和工程应用都绕不开 FFT。ops-fft 把离散傅里叶变换DFT的计算搬到昇腾 NPU 上用 Vector 单元的并行能力甩开 CPU 几个量级。但 FFT 在 NPU 上的坑不比数学算子少。基数选择不当、维度顺序混了、频域到时域的归一化遗漏——每一个都能让结果面目全非。FFT 的硬件前提昇腾 NPU 的 Vector 单元一次能处理 256 个 float/fp16 的并行运算。FFT 的核心运算是蝶形运算Butterfly Operation每一级需要对 256 个数据点做交叉乘加。如果数据长度刚好是 256 的倍数Vector 单元能跑满。如果不是需要 padding 或者分段。ops-fft 内部用两种基数策略Radix-2每级蝶形 2 点输入 → 适合 2^n 长度的序列如 256、512、1024、2048Mixed-Radix混合基数2/4/8适合任意长度如 312、625、1250Mixed-Radix 的计算效率比 Radix-2 低 10-15%但覆盖了所有输入长度。四种算子接口ops-fft 面向四种高频场景// 一维 FFTops_fft_1d(input,length,directionFFT_FORWARD)// 例频谱分析// 二维 FFT逐列做一维ops_fft_2d(input,height,width,directionFFT_FORWARD)// 例2D 图像频域滤波// 任意维度 FFT通过参数指定维度ops_fftn(input,shape,axes,directionFFT_FORWARD)// 例多维信号联合分析// 卷积定理乘积 → FFT → 逐点乘 → IFFTops_fft_convolution(signal_x,signal_y)// 例信号卷积、图像滤波踩坑一频率分辨率和采样率的关系FFT 的输出频率分辨率 Δf 由采样率 Fs 和点数 N 决定Δf Fs / N。如果这两个参数配错了频谱的横轴直接歪了。错误写法importtorch_npuimportnumpyasnp# 采样率 44100 Hz采集 1024 个点Fs44100N1024# 生成一个 440 Hz 的正弦波tnp.arange(N)/Fs signalnp.sin(2*np.pi*440*t).astype(np.float32)# 直接做 FFT忽略了频率分辨率spectrumtorch_npu.ops_fft_1d(signal.npu(),N,directionforward)# 错频谱横轴的单位是“ bins”不是 Hz# 峰值应该在 bin 440 / (44100/1024) ≈ bin 10 处# 但代码里直接 print(spectrum)频率轴对不上实际物理意义正确写法显示计算并对齐频率轴。importtorch_npuimportnumpyasnp# 采样率和点数Fs44100N1024# 生成信号tnp.arange(N)/Fs signalnp.sin(2*np.pi*440*t).astype(np.float32)# 做 FFTspectrumtorch_npu.ops_fft_1d(signal.npu(),N,directionforward)# 计算频率轴Nyquist 频率之前的正值部分freq_axisnp.fft.fftfreq(N,d1.0/Fs)[:N//2]# 频谱只取正值部分负频率对称不需要重复spectrum_positivespectrum[:,:N//2]# 在 freq_axis 上定位峰值peak_idxnp.argmax(np.abs(spectrum_positive))peak_freqfreq_axis[peak_idx]print(f实际峰值频率:{peak_freq:.1f}Hz)# 输出 ≈ 440 HzC 侧原理ops-fft 输出的频率轴信息以元数据形式存储不是张量的一部分。需要用np.fft.fftfreq(N, d1.0/Fs)显式计算横轴。踩坑二2D FFT 的维度顺序二维 FFT 有两种理解方式先对行做一维 FFT再对列做一维 FFTRow-First或者反过来Col-First。 ops-fft 默认是 Row-First——但很多图像处理的代码习惯 Col-First导致结果差了转置。错误写importtorch_npu# 图像height × width 512 × 512# 想做 2D FFT 做频域滤波imagetorch.randn(512,512).npu().half()# 按 ops-fft 默认的 Row-First 做 2D FFT# 先对每一行做 FFT再对每一列做 FFTfreq_imagetorch_npu.ops_fft_2d(image,512,512,directionforward)# 错误如果代码后面对 ImageJ 的滤波器做卷积# 预期滤波是在 frequency domain 做的# 实际这个 frequency domain 的排列顺序取决于 filter 的生成方式# 不一致时filter 和 image 的频率轴对不上过滤波段就错了正确写法显式控制维度。importtorch_npuimporttorch.fft# 方法一用 PyTorch 的 fft2d底层调用 ops-fftfreq_imagetorch.fft.fft2(image)# 默认Row-First# 方法二显式指定 axes确保一致freq_imagetorch.fft.fft2(image,dim(0,1))# 先 dim0 再 dim1# 方法三如果习惯 Col-First手动 transposefreq_imagetorch.transpose(torch.fft.fft2(image,dim0),# 先对列 dim0 做 fft0,1# 再对行 fft实际变成了 Column-First)# 生成 filter 时确保维度一致filtertorch.ones_like(freq_image)filter[256-10:25610,256-10:25610]0# 去掉低频中心filtered_freqfreq_image*filter# 逐点乘filtered_imagetorch.fft.ifft2(filtered_freq)# 逆变换回去根因ops-fft 的 Row-First 是在每个ops_fft_2d调用内部固定的。改变维度顺序只能在 Python 侧通过 transpose 控制或者调用ops_fftn并显式指定 axes 参数。踩坑三频域卷积的归一化遗漏频域卷积的核心是时域的卷积等于频域的乘积。用 FFT 做卷积的标准流程是FFTx(A) × FFTx(B) → 逐点乘 → IFFTx → 除以 N归一化最后一步「除以 N」经常漏掉——因为数学公式里 IFFT 本身带 1/N但 FFT 的正向和逆向的 scaling 分配有多种标准如果不统一结果会差 N 倍。错误写法importtorch_npu# 两个有限脉冲响应各 128 点h1torch.randn(128).npu().half()h2torch.randn(128).npu().half()N128# FFT 后逐点乘H1torch_npu.ops_fft_1d(h1,N,directionforward)H2torch_npu.ops_fft_1d(h2,N,directionforward)productH1*H2# 错误直接 IFFT没有归一化# 结果的值比理论值大了 N128 倍result_no_normtorch_npu.ops_fft_1d(product,N,directioninverse)正确写法显式归一化。importtorch_npuimporttorch# 两个信号h1torch.randn(128).npu().half()h2torch.randn(128).npu().half()N128# FFT 后逐点乘H1torch_npu.ops_fft_1d(h1,N,directionforward)H2torch_npu.ops_fft_1d(h2,N,directionforward)productH1*H2# IFFT 后除以 N 归一化resulttorch_npu.ops_fft_1d(product,N,directioninverse)result_normalizedresult/float(N)# 验证和不使用 FFT 的直接卷积结果对比参考实现expectedtorch.nn.functional.conv1d(h1.unsqueeze(0).unsqueeze(0),# [1, 1, 128]h2.unsqueeze(0).unsqueeze(0),paddingsame).squeeze(0)# 误差应该 1e-3asserttorch.allclose(result_normalized,expected,atol1e-3)ops-fft 的内部 scaling默认策略是「Forward 1/sqrt(N)Inverse sqrt(N)」Parseval 归一化但这个默认值在某些信号处理教材里不常见。更通用的约定是「Forward 1Inverse 1/N」——ops-fft 的 scaling 策略可以通过参数选择。踩坑四复数输出的视图混淆FFT 的输出是复数。但在 Python 侧复数的实部和虚部访问方式有两种C 风格的 (real, imag交替存储和 Fortran 风格的连续实部、连续虚部。ops-fft 默认用 C style即 interleaved但 NumPy 的 fft 输出 Fortran styleseparate real array then imaginary array。错误写法importtorch_npuimportnumpyasnp signaltorch.randn(256).npu().half()spectrumtorch_npu.ops_fft_1d(signal,256)# 错误用访问 NumPy 复数数组的方式访问 ops-fft 的输出# NumPyspectrum.real 和 spectrum.imag 是分开的 viewnp_spectrumnp.fft.fft(signal.cpu().numpy())real_part_npnp_spectrum.real# ← view直接可用# ops-fftspectrum 是一个 interleave 存储的 tensor# 直接取 spectrum.real 会出错没有这个属性# 应该用 view_as_complex 或分量的 split正确写法确保前后端的复数表示一致。importtorch_npuimporttorch signaltorch.randn(256).npu().half()spectrumtorch_npu.ops_fft_1d(signal,256)# 方法一转为 PyTorch 的复数格式推荐spectrum_complextorch.view_as_complex(spectrum.reshape(256,2)# [N, 2] → [N])# 现在 spectrum_complex 是 torch.complex64 类型# 幅度谱magnitudetorch.abs(spectrum_complex)# 方法二拆分成实部和虚部real_partspectrum[...,0]# 实部交错存储的第一个imag_partspectrum[...,1]# 虚部# 如果要和 NumPy 结果对比也需要转 viewnp_spectrumnp.fft.fft(signal.cpu().numpy())torch_np_complextorch.from_numpy(np_spectrum)# 比对 magnitudeasserttorch.allclose(torch.abs(spectrum_complex.float()),torch.abs(torch_np_complex.float()),atol1e-2)根本原因ops-fft 为了和 Ascend C 的 Vector 单元数据排布兼容默认用 interleaved 格式复数的实部和虚部交错存储。这个设计对 Vector 单元的 align 最友好但和 Python 生态的 NumPy/PyTorch 不兼容需要一次 view 转换。使用性能在 Ascend 910 上跑 FFT 的吞吐量长度CPU (NumPy)NPU (ops-fft)加速比10240.12 ms0.008 ms15×40960.85 ms0.031 ms27×163847.2 ms0.18 ms40×6553658 ms0.92 ms63×随着点数增加Vector 单元的并行度优势越大加速比从 15× 提升到 63×。瓶颈不在计算在 Vector 单元上蝶形计算本身很快而在 HBM 的读写带宽——ops-fft 内部尽可能复用 L1 缓存数据只过一次 HBM。FFT 的坑集中在「计算外的辅助信息」上——频率轴的物理单位、2D 变换的维度顺序、IFFT 的归一化因子、复数的存储格式。这四个点每一个都是「算子本身没错但参数配置和环境适配导致结果偏离」的问题。ops-fft 提供的是算子的底层实现上层的调用姿势需要调用方自己理清楚——这和其他 CANN 算子库的坑的模式是一致的。