FFT速度比較 #

import numpy as np
import scipy as sp
import cupy as cp
import torch
import pyfftw

# - データ生成
N = 2**10  # FFT点数
fx = np.random.random(N) + 1.0j * np.random.random(N)

# - Numpy fft
gk = np.fft.fft(fx)   # FFTを実行
gx = np.fft.ifft(gk)  # IFFTを実行

# - Scipy fft
gk = sp.fft.fft(fx)   # FFTを実行
gx = sp.fft.ifft(gk)  # IFFTを実行

# - Self-made fft
gk = fft_my(fx)       # FFTを実行
gx = ifft_my(gk) / N  # IFFTを実行 (規格化も実施)

# - Cupy fft
fx_cp = cp.array(fx)               # NumPy配列をCuPy配列に変換 (非同期処理)
gk_cp = cp.fft.fft(fx_cp)             # FFTを実行 (非同期処理)
gx_cp = cp.fft.ifft(gk_cp)            # IFFTを実行 (非同期処理)
cp.cuda.Stream.null.synchronize()  # GPU計算が終了するのを待つ (ここで同期)
gk = cp.asnumpy(gk_cp)             # 結果をCPU (NumPy) に戻す
gx = cp.asnumpy(gx_cp)             # 結果をCPU (NumPy) に戻す

# - CuPy (Plan)
fx_cp = cp.array(fx)               # NumPy配列をCuPy配列に変換（非同期処理）
cp.cuda.Stream.null.synchronize()  # GPUの転送が完了するのを待つ
plan = cp.cuda.cufft.Plan1d(N, cp.cuda.cufft.CUFFT_Z2Z, 1) # Plan1dを使用して計画を作成
gk_cp = cp.empty_like(fx_cp)       # デバイスメモリ上に入力と出力の配列を準備
gx_cp = cp.empty_like(fx_cp)       # デバイスメモリ上に入力と出力の配列を準備
plan.fft(fx_cp, gk_cp, cp.cuda.cufft.CUFFT_FORWARD)  # FFTを実行  (gk = F[fx])
plan.fft(gk_cp, gx_cp, cp.cuda.cufft.CUFFT_INVERSE)  # IFFTを実行 (gx = F^{-1}[gk])
gx_cp /= N  # IFFTの結果を規格化
cp.cuda.Stream.null.synchronize()  # GPU計算が終了するのを待つ
gk = cp.asnumpy(gk_cp)  # 結果をCPU（NumPy）に戻す
gx = cp.asnumpy(gx_cp)  # 結果をCPU（NumPy）に戻す

# - PyTorch (CPU)
fx_tc = torch.tensor(fx, dtype=torch.complex128)  # NumPy配列をPyTorchの複素数テンソルに変換
gk_tc = torch.fft.fft(fx_tc)   # FFTを実行
gx_tc = torch.fft.ifft(gk_tc)  # IFFTを実行
gk = gk_tc.numpy()  # PyTorchの複素数テンソルをNumPy配列に変換(.numpy())
gx = gx_tc.numpy()  # PyTorchの複素数テンソルをNumPy配列に変換(.numpy())

# - PyTorch (GPU)
fx_tc = torch.tensor(fx, dtype=torch.complex128, device="cuda") # NumPy配列をPyTorchの複素数テンソルに変換（complex128）し、GPUへ転送
gk_tc = torch.fft.fft(fx_tc)   # FFTを実行
gx_tc = torch.fft.ifft(gk_tc)  # IFFTを実行
gk = gk_tc.cpu().numpy()  # 結果をCPUに戻し(.cpu()), NumPy配列に変換(.numpy())
gx = gx_tc.cpu().numpy()  # 結果をCPUに戻し(.cpu()), NumPy配列に変換(.numpy())

# - fftw (planner_effort = "FFTW_MEASURE" の場合)
fft_objects = pyfftw.builders.fft(fx, planner_effort="FFTW_MEASURE")    # fxの型、サイズでFFT実行用オブジェクト生成
ifft_objects = pyfftw.builders.ifft(fx, planner_effort="FFTW_MEASURE")  # fxの型、サイズでIFFT実行用オブジェクト生成
gk = fft_object(fx)   # FFTを実行
gx = ifft_object(gk)  # IFFTを実行

  # Transfer from CPU to GPU
  start_tmp = time.perf_counter()
  cp_array = cp.array(fx)  # Convert NumPy array to CuPy array (asynchronous operation)
  cp.cuda.Stream.null.synchronize()
  end_tmp = time.perf_counter()
  benchmarks[fft_lib]["cpu2gpu_transfer_time_ms"].append((end_tmp - start_tmp) * 1000)

  start = time.perf_counter()
  for fft_repeat in range(fft_repeats):
     gk = cp.fft.fft(cp_array)  # Compute FFT (asynchronous operation)
     gx = cp.fft.ifft(gk)  # Compute IFFT (asynchronous operation)
     cp.cuda.Stream.null.synchronize()  # Wait for GPU computations to finish (synchronization occurs here)
  end = time.perf_counter()
  benchmarks[fft_lib]["elapsed_times_ms"].append(((end - start) * 1000) / fft_repeats) # Average

  # Transfer from GPU to CPU
  start_tmp = time.perf_counter()
  gx = cp.asnumpy(gx)  # Transfer results back to CPU (NumPy)
  cp.cuda.Stream.null.synchronize()
  end_tmp = time.perf_counter()
  benchmarks[fft_lib]["gpu2cpu_transfer_time_ms"].append((end_tmp - start_tmp) * 1000)

  # Plan making using Plan1d
  plan = cp.cuda.cufft.Plan1d(N, cp.cuda.cufft.CUFFT_Z2Z, 1)

  # Prepare input and output arrays on GPU memory
  gk = cp.empty_like(cp_array)
  gx = cp.empty_like(cp_array)

  start = time.perf_counter()
  for fft_repeat in range(fft_repeats):
      plan.fft(cp_array, gk, cp.cuda.cufft.CUFFT_FORWARD)  # Compute FFT (asynchronous operation)
      plan.fft(gk, gx, cp.cuda.cufft.CUFFT_INVERSE)  # Compute IFFT (asynchronous operation)
      gx /= N  # Normalize
      cp.cuda.Stream.null.synchronize()  # Wait for GPU computations to finish (synchronization occurs here)
  end = time.perf_counter()

> nvidia-smi
Tue Jan 14 00:11:30 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 566.03                 Driver Version: 566.03         CUDA Version: 12.7     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3060 Ti   WDDM  |   00000000:0D:00.0  On |                  N/A |
| 33%   56C    P0             47W /  200W |    1285MiB /   8192MiB |      2%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
...

pip install cupy-cuda12x

import cupy as cp

# CUDAデバイスの数を取得
print("CUDAデバイス数:", cp.cuda.runtime.getDeviceCount())

# CUDAバージョンを取得
print("CUDAバージョン:", cp.cuda.runtime.runtimeGetVersion() / 1000)

# GPUの名前を取得
device_id = 0  # 使用するGPUのID
device_properties = cp.cuda.runtime.getDeviceProperties(device_id)
print("GPU名:", device_properties["name"])

pip install pyfftw

> nvidia-smi
Tue Jan 14 00:11:30 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 566.03                 Driver Version: 566.03         CUDA Version: 12.7     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3060 Ti   WDDM  |   00000000:0D:00.0  On |                  N/A |
...

def fft_my(x):
    # https://pythonnumericalmethods.studentorg.berkeley.edu/notebooks/chapter24.03-Fast-Fourier-Transform.html
    """
    A recursive implementation of
    the 1D Cooley-Tukey FFT, the
    input should have a length of
    power of 2.
    """
    N = len(x)

    if N == 1:
        return x
    else:
        X_even = fft_my(x[::2])
        X_odd = fft_my(x[1::2])
        factor = np.exp(-2j * np.pi * np.arange(N) / N)

        X = np.concatenate([X_even + factor[: int(N / 2)] * X_odd, X_even + factor[int(N / 2) :] * X_odd])
        return X


def ifft_my(x):
    # https://pythonnumericalmethods.studentorg.berkeley.edu/notebooks/chapter24.03-Fast-Fourier-Transform.html
    """
    A recursive implementation of
    the 1D Cooley-Tukey FFT, the
    input should have a length of
    power of 2.
    """
    N = len(x)

    if N == 1:
        return x
    else:
        X_even = ifft_my(x[::2])
        X_odd = ifft_my(x[1::2])
        factor = np.exp(2j * np.pi * np.arange(N) / N)

        X = np.concatenate([X_even + factor[: int(N / 2)] * X_odd, X_even + factor[int(N / 2) :] * X_odd])
        return X