FortranからMPIを使ってPythonにnumpy配列の並列処理をさせたい！: forpyの利用その4

この記事はFortranからPythonを扱う方法であるforpyを解説する記事の第四弾です。
FortranからPythonを使いたい！:forpy
Fortranで機械学習がしたい：FortranからのPyTorchの利用
Fortranからnumpy配列を読み書きしたい！: forpyの利用その3
が以前の記事です。

今回は、MPI並列計算の中からPythonを実行してみましょう。特にnumpy配列をPythonに渡して処理してもらって、それをFortranが受け取ってまた加工し、そしてまたnumpyにしてPythonに送る、という処理をしたいと思います。
これの利点の一つは、Pythonをmpi4pyなどを使わなくてもFortranを経由する事でPythonの処理を並列化できる点です。
もう一つの利点は、FortranのMPI並列計算コードが既に手元にあるときに、PythonにFortranだと面倒なこと（機械学習とか）をやらせることができる、という点です。
やってみると全然難しくなくて、numpy配列をメモリコピーなしでFortranの配列に変換する方法だけ気をつけておけば大丈夫でした。

バージョン

• gcc version 10.2.0 (Homebrew GCC 10.2.0_4)
• macOS 10.15.17

コード

subroutine test()
    use forpy_mod
    use mpi
    use,intrinsic::iso_fortran_env
    implicit none
    integer(int32)::ierror
    type(tuple) :: args
    integer(int32)::ierr,myrank
    type(module_py) :: testmodule
    type(list) :: paths
    integer(int32)::N,i
    real(real64),allocatable::x(:)!,y(:)
    real(real64),dimension(:), pointer ::y
    type(ndarray) ::x_py,y_py
    type(object)::ret

    call mpi_init(ierr)
    call mpi_comm_rank(mpi_comm_world,myrank,ierr)

    ierror = forpy_initialize()

    ierror = get_sys_path(paths)
    ierror = paths%append(".")
    ierror = import_py(testmodule, "numpympitest")

    N = 10
    allocate(x(1:N))
    do i=1,N
        x(i) = i
    end do
    write(*,*) "x = ",x

    ierror = ndarray_create_nocopy(x_py, x)

    ierror = tuple_create(args, 2)
    ierror = args%setitem(0, myrank)
    ierror = args%setitem(1, x_py)

    write(*,*) "x_py = "
    ierror =  print_py(x_py)
    
    ierror = call_py(ret,testmodule,"each_sum",args)
    !ierror =  print_py(ret)

    ierror = ndarray_create_empty(y_py, [N], dtype="float64")
    ierror = cast(y_py,ret)
    ierror = y_py%get_data(y)

    write(*,*) "myrank, y = ",myrank, y

    ierror =  print_py(ret)

    call MPI_Allreduce(y, x, 10, MPI_DOUBLE_PRECISION,MPI_SUM, mpi_comm_world, ierr)
    write(*,*) "sum = ",x
    ierror = call_py(ret,testmodule,"each_sum",args)
    ierror =  print_py(ret)

    call forpy_finalize
    call mpi_finalize(ierr)
end subroutine

program main
    integer :: ierror 
    call test() 
end program

コードはこんな感じになりました。Pythonのコードは

numpympitest.py

import numpy as np

def each_sum(myrank,y0):
    print("myrank ",myrank)
    yout = y0*(myrank+1)
    #print(yout)
    return yout

です。

解説

1. call mpi_comm_rank(mpi_comm_world,myrank,ierr)でそれぞれのプロセスが異なるmyrankを持つことになります
2. 長さ10のFortranの配列を作って、1から10までを入れます
3. ierror = ndarray_create_nocopy(x_py, x)でFortranの配列xをPythonのnumpy配列x_pyにします（コピーなし）
4. 引数の用意

    ierror = tuple_create(args, 2)
    ierror = args%setitem(0, myrank)
    ierror = args%setitem(1, x_py)

でPythonに投げ渡すための引数を作っています。Python側の引数の数が二つなので、タプルを作成しています。
4. ierror = call_py(ret,testmodule,"each_sum",args)でeach_sumというPythonの関数を呼び出します。myrankを引数にとっているために、プロセスごとにPython関数は異なるアウトプットを返します
5. 帰ってくる変数retはPythonのオブジェクトなので、これをnumpyに直します。

ierror = ndarray_create_empty(y_py, [N], dtype="float64")
ierror = cast(y_py,ret)

で空のnumpy配列y_pyを用意して、cast(y_py,ret)でretからy_pyにデータを変換します。
6. ierror = y_py%get_data(y)でFortranの配列yにnumpy配列y_pyの内容を入れます。これで、プロセスごとに異なったyが手に入りました。
7. call MPI_Allreduce(y, x, 10, MPI_DOUBLE_PRECISION,MPI_SUM, mpi_comm_world, ierr)で、プロセスごとに異なったyの全ての和をとってxに入れる、という通信を行います。
8. さらにierror = call_py(ret,testmodule,"each_sum",args)ではx_pyがxとメモリアドレスを共有していることを利用して、先ほど和の計算をした結果が入った引数xをまたPythonの関数に投げました。またプロセスごとに異なった結果が得られます。
9. 最後に ierror = print_py(ret)でプロセスごとにどんな値が入っているのかをみました。

もし2並列でやっていたら、最終結果は

[ 3.  6.  9. 12. 15. 18. 21. 24. 27. 30.]
[ 6. 12. 18. 24. 30. 36. 42. 48. 54. 60.]

となります。