f2pyでPython forループの高速化

はじめに

Pythonでforループでの処理はなかなか計算時間がかかります。
量子アニーリングのQUBO行列だと4重ループくらいの計算をした際にかなり時間がかかりました。
CやFortranでは数秒で終わる計算でもPythonでは数十分かかる場合があります。

そこでループ計算部分のみ、C、Frotranで書いて高速化してみました。
numpyパッケージに含まれているf2pyを使うとPythonに対応したFortranのバイナリを生成して高速計算できるので記事にまとめてみました。

Python版のコード

今回のベースとなるコード。
4重ループ計算でgoogle colabで実行すると20秒くらいかかります。

import numpy as np
import time

def build_Q(M, N, V, lam1, lam2):
    # Qは (M*N, M*N) の2次元配列
    Q = np.zeros((M * N, M * N), dtype=float)

    def delta(i, j):
        return 1.0 if i == j else 0.0

    for k in range(M):        # 0..M-1
        for l in range(N):    # 0..N-1
            for i in range(M):
                for j in range(N):
                    ij = N * i + j          # Pythonは0始まり
                    kl = N * k + l
                    Q[ij, kl] = delta(i, k) * lam1 * V[i, j] * V[k, l]
                    if i == k and j == l:
                        Q[ij, kl] += lam2
    return Q

N = 64 # N行に並べるqbit数
M = 64 # M行に並べるqbit数

lam1 = 1.0
lam2 = 2.0

V = np.random.rand(N, M)  # 時間や重量など、何かの値。今回は乱数で作っておく。
Q = np.zeros((N*M,N*M),dtype=np.float64, order="F")  # QUBO行列 dtype以降の部分は無くても動くと思うが念の為書いています

start_calcQUBO_time = time.time() # 時間計測 start

Q = Q = build_Q(M, N, V, lam1, lam2)

calcQUBO_time = time.time() - start_calcQUBO_time # 時間計測
print("calc QUBO time: {0}".format(calcQUBO_time) + " [sec]")

Pythonでの実行時間

google colabで実行して見た時間は下記でした。

calc QUBO time: 15.123576164245605 [sec]

f2pyを使った高速化

ループ部分をFortran90で作りPython側から呼び出します。
まずは、ループ部分をサブルーチンとして作りました。OpenMPの並列化も含めて高速してみています。

Fortran90部分

! create QUBO using fortran with OpenMP accelerate
! build command:FC=gfortran f2py -m calc_qubo --f90flags="-O3 -fopenmp" -lgomp -c calc_qubo.f90
subroutine qubo(V, Q, lam1, lam2, N, M)
  use omp_lib
  implicit none
  integer,intent(in) :: N, M
  real,intent(in) :: lam1, lam2
  real(8),dimension(N,M),intent(in) :: V
  real(8),dimension(N*M,N*M),intent(out) :: Q

  integer :: i, j, k, l, ij, kl

  Q = 0.0

  !$omp parallel do default (shared) private(i,j,l,ij,kl) collapse(2)
  do k=1,M
    do l=1,N
      do i=1,M
        do j=1,N
          ij = N*(i-1) + j
          kl = N*(k-1) + l
          Q(ij,kl) =  delta(i,k) * lam1 * V(i,j) * V(k,l)
          if ( i==k .and. j==l ) then
            Q(ij,kl) = Q(ij,kl) + lam2
          end if
        end do
      end do
    end do
  end do
  !$omp end parallel do

contains
  real function delta(i,j)
    implicit none
    integer i, j

    if(i==j) then
      delta = 1.0
    else
      delta = 0.0
    end if
    return
  end function delta

end subroutine qubo

f2pyでFortran90部分のコンパイル

下記のコマンドでビルドします。

$ FC=gfortran f2py -m calc_qubo --f90flags="-O3 -fopenmp" -lgomp -c calc_qubo.f90

エラーがでず、バイナリが生成できればコンパイル成功です。

$ ls
calc_qubo.cpython-312-x86_64-linux-gnu.so  calc_qubo.f90

google colabで実行する場合

2025年9月末時点のgoogle colabのPython 3.12.11の環境では、下記のパッケージが追加で必要でした。
下記をインストール後、上述のf2pyコマンドでビルドできました。

!pip install meson ninja

Python側のコード

import numpy as np
import time
import calc_qubo

N = 64 # N行に並べるqbit数
M = 64 # M行に並べるqbit数

lam1 = 1.0
lam2 = 2.0

V = np.random.rand(N, M)  # 時間や重量など、何かの値。今回はデータがないので乱数で作っておく。
Q = np.zeros((N*M,N*M),dtype=np.float64, order="F") 

start_calcQUBO_time = time.time() # 時間計測 start

# Fortranのqubo関数呼び出し。intent(out)にしている部分は自明なので書かない仕様みたい
Q = calc_qubo.qubo(V,lam1,lam2,N,M) 

calcQUBO_time = time.time() - start_calcQUBO_time # 時間計測
print("calc QUBO time: {0}".format(calcQUBO_time) + " [sec]")

f2pyでの実行時間

google colab環境でexport OMP_NUM_THREADS=2を設定して下記の実行時間でした。
ほぼ待たずに結果が得られるようになりました。

calc QUBO time: 0.06850457191467285 [sec]

先のPythonでの実行時間と比較すると約220倍の高速化です。

calc_time_ratio = 15.123576164245605 / 0.06850457191467285 
print(calc_time_ratio)
-> 220.76739904430113

まとめ

ループをFortran90で記述してf2pyで作ったバイナリをPythonから呼ぶとかなり高速化します。
Fortran90のデバッグは嵌まると時間がかかるので最小限の開発はPythonで動作確認して、規模の大きい計算を行う場合はf2pyを使った方法に移植するのが良いと思います。