python - 难以让 OpenMP 与 f2py 一起工作

Question

我正在为我的研究进行一些模拟工作，并且遇到了将 fortran 导入我的 python 脚本的障碍。作为背景，我已经使用 Python 多年，并且只在需要时才在 Fortran 内部闲逛。

我过去用 Fortran 完成了一些工作，实现了一些简单的 OpenMP 功能。我不是这方面的专家，但我之前已经掌握了基础知识。

我现在正在使用 f2py 创建一个库，我可以从我的 python 脚本中调用它。当我尝试编译 openmp 时，它编译正确并运行完成，但速度提高为零，查看顶部我看到 CPU 使用率表明只有一个线程在运行。

我已经搜索了 f2py 的文档(没有很好的记录)，并完成了正常的网络调查以寻找答案。我已经包含了我正在编译的 Fortran 代码以及调用它的简单 python 脚本。我还输入了我正在使用的编译命令。

目前我将模拟减少到 10^4 作为一个很好的基准。在我的系统上运行需要 3 秒。不过，最终我需要运行多个 10^6 粒子模拟，所以我需要稍微缩短时间。

如果有人能指出如何让我的代码正常工作的方向，我将不胜感激。我还可以根据需要尝试包含有关系统的任何详细信息。

干杯，雷尔坎

---

1)编译

f2py -c --f90flags='-fopenmp' -lgomp -m calc_accel_jerk calc_accel_jerk.f90

---

2) 调用Python脚本

import numpy as N
import calc_accel_jerk

# a is a (1e5,7) array with M,r,v information
a        = N.load('../test.npy')
a        = a[:1e4]

out = calc_accel_jerk.calc(a,a.shape[0])
print out[:10]

---

3) Fortran代码

subroutine calc (input_array, nrow, output_array)
implicit none
!f2py threadsafe
include "omp_lib.h"

integer, intent(in) :: nrow
double precision, dimension(nrow,7), intent(in) :: input_array
double precision, dimension(nrow,2), intent(out) :: output_array

! Calculation parameters with set values
double precision,parameter :: psr_M=1.55*1.3267297e20
double precision,parameter :: G_Msun=1.3267297e20
double precision,parameter :: pc_to_m=3.08e16

! Vector declarations
integer :: irow
double precision :: vfac
double precision, dimension(nrow) :: drx,dry,drz,dvx,dvy,dvz,rmag,jfac,az,jz

! Break up the input array for faster access
double precision,dimension(nrow) :: input_M
double precision,dimension(nrow) :: input_rx
double precision,dimension(nrow) :: input_ry
double precision,dimension(nrow) :: input_rz
double precision,dimension(nrow) :: input_vx
double precision,dimension(nrow) :: input_vy
double precision,dimension(nrow) :: input_vz

input_M(:)  = input_array(:,1)*G_Msun
input_rx(:) = input_array(:,2)*pc_to_m
input_ry(:) = input_array(:,3)*pc_to_m
input_rz(:) = input_array(:,4)*pc_to_m
input_vx(:) = input_array(:,5)*1000
input_vy(:) = input_array(:,6)*1000
input_vz(:) = input_array(:,7)*1000

!$OMP PARALLEL DO private(vfac,drx,dry,drz,dvx,dvy,dvz,rmag,jfac,az,jz) shared(output_array) NUM_THREADS(2)
DO irow = 1,nrow
    ! Get the i-th iteration
    vfac  = sqrt(input_M(irow)/psr_M)
    drx   = (input_rx-input_rx(irow))
    dry   = (input_ry-input_ry(irow))
    drz   = (input_rz-input_rz(irow))
    dvx   = (input_vx-input_vx(irow)*vfac)
    dvy   = (input_vy-input_vy(irow)*vfac)
    dvz   = (input_vz-input_vz(irow)*vfac)
    rmag  = sqrt(drx**2+dry**2+drz**2)
    jfac  = -3*drz/(drx**2+dry**2+drz**2)

    ! Calculate the acceleration and jerk
    az = input_M*(drz/rmag**3)
    jz = (input_M/rmag**3)*((dvx*drx*jfac)+(dvy*dry*jfac)+(dvz+dvz*drz*jfac))

    ! Remove bad index
    az(irow) = 0
    jz(irow) = 0

    output_array(irow,1) = sum(az)
    output_array(irow,2) = sum(jz)
END DO
!$OMP END PARALLEL DO

END subroutine calc

Answer 1

这是一个简单的检查，看看 OpenMP 线程是否确实在 Fortran 代码中可见:

module OTmod
  !$ use omp_lib
  implicit none

  public :: get_threads

contains

  function get_threads() result(nt)
    integer :: nt

    nt = 0
    !$ nt = omp_get_max_threads()

  end function get_threads

end module OTmod

编译:

> f2py -m OTfor --fcompiler=gfortran --f90flags='-fopenmp' -lgomp -c OTmod.f90

执行:

> python
>>> from OTfor import otmod
>>> otmod.get_threads()
12