如何通过 cython 将 MPI 通信器从 python 传递到 C?

How to pass an MPI communicator from python to C via cython?

我正在尝试通过 cython 包装一个以 MPI_Comm 通信器句柄作为参数的 C 函数。因此,我希望能够从 python 调用该函数,并向其传递一个 mpi4py.MPI.Comm 对象。我想知道的是,如何将 mpi4py.MPI.Comm 转换为 MPI_Comm.

为了演示,我使用了一个简单的 "Hello World!" 类型的函数:

helloworld.h:

#ifndef HELLOWORLD
#define HELLOWORLD
#include <mpi.h>

void sayhello(MPI_Comm comm);

#endif

helloworld.c:

#include <stdio.h>
#include "helloworld.h"

void sayhello(MPI_Comm comm){
    int size, rank;
    MPI_Comm_size(comm, &size);
    MPI_Comm_rank(comm, &rank);
    printf("Hello, World! "
           "I am process %d of %d.\n",
           rank, size);
}

我现在想从 python 调用此函数,如下所示:

from_python.py:

import mpi4py
import helloworld_wrap

helloworld_wrap.py_sayhello(mpi4py.MPI.COMM_WORLD)

意思是 mpirun -np 4 python2 from_python.py 应该给出类似的东西:

Hello, World! I am process 0 of 4.
Hello, World! I am process 1 of 4.
Hello, World! I am process 2 of 4.
Hello, World! I am process 3 of 4.

但是如果我像这样通过 cython 实现这个:

helloworld_wrap.pyx:

cimport mpi4py.MPI as MPI
cimport mpi4py.libmpi as libmpi

cdef extern from "helloworld.h":
   void sayhello(libmpi.MPI_Comm comm)

def py_sayhello(MPI.Comm comm):
    sayhello(comm)

和:

setup.py:

import os
from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext

mpi_compile_args = os.popen("mpicc --showme:compile").read().strip().split(' ')
mpi_link_args    = os.popen("mpicc --showme:link").read().strip().split(' ')

ext_modules=[
    Extension("helloworld_wrap",
              sources            = ["helloworld_wrap.pyx", "helloworld.c"],
              language           = 'c',
              extra_compile_args = mpi_compile_args,
              extra_link_args    = mpi_link_args,
          )
]

setup(
  name = "helloworld_wrap",
  cmdclass = {"build_ext": build_ext},
  ext_modules = ext_modules
)

我收到以下错误消息:

helloworld_wrap.pyx:8:13: Cannot convert Python object to 'MPI_Comm'

表示mpi4py.MPI.Comm无法转换为MPI_Comm。那么如何将 mpi4py.MPI.Comm 转换为 MPI_Comm 以使包装器正常工作?

转换相当简单,因为 mpi4py.MPI.Comm-对象在内部将 MPI_Comm 句柄存储为成员 ob_mpi1。因此,如果将 helloworld_wrap.pyx 的最后一行更改为传递 comm.ob_mpi 而不是 comm,模块将按预期编译和工作:

helloworld_wrap.pyx:

cimport mpi4py.MPI as MPI
cimport mpi4py.libmpi as libmpi

cdef extern from "helloworld.h":
   void sayhello(libmpi.MPI_Comm comm)

def py_sayhello(MPI.Comm comm):
    sayhello(comm.ob_mpi)

令人惊讶的是,我没有找到这方面的任何文档,而是在研究 sources of mpi4py.MPI.Comm 时才意识到这一点。我不确定这是否是处理此问题的预期方式,但我无法使其正常工作。


1 实际上,mpi4py.MPI 中的大多数(如果不是全部)对象在 C 中模拟相应的 MPI 句柄,并持有相应的句柄作为成员 ob_mpi.