在 mpi4py 中包装和调用 C MPI 程序

本文从本人简书博客同步过来

在上一篇中我们介绍了使用 scalapy 调用 ScaLAPACK 进行分布式内存的线性代数运算。Python 作为一种胶水语言，可以非常容易地包装和调用其它计算机语言已有的程序代码和工具库，如果我们有用 C，C++，Fortran 或其它计算机语言编写的 MPI 计算程序，也能很容易地将其包装后在 mpi4py 中进行调用。另外我们也可以用这些计算机语言编写一些运算速度更快的 MPI 扩展模块供 mpi4py 程序调用。在下面的章节中我们将介绍怎么在 mpi4py 程序中包装和调用 C，C++，Fortran 的 MPI 程序。下面我们首先介绍直接使用 Python/C API 进行 C 语言 MPI 程序包装的方法。

Python 文档中有使用 C/C++ 写扩展模块的详细文档和完整的 Python/C API 参考，此处不会专门讲解怎么使用 Python/C API 编写扩展模块，只会介绍如何在 mpi4py 程序中包装和调用已经存在的 C 语言 MPI 程序代码。

假设我们有如下使用 Python/C API 编写的扩展 C 语言 MPI 程序，其中定义了 sayhello 函数，这个函数接受一个 MPI 通信子作为参数，在其中调用 MPI 的相关函数完成计算工作：

/* helloworld.c */

#define MPICH_SKIP_MPICXX 1
#define OMPI_SKIP_MPICXX  1
#include <mpi4py/mpi4py.h>

/* -------------------------------------------------------------------------- */

static void
sayhello(MPI_Comm comm) {
  int size, rank;
  char pname[MPI_MAX_PROCESSOR_NAME]; int len;
  if (comm == MPI_COMM_NULL) {
    printf("You passed MPI_COMM_NULL !!!\n");
    return;
  }
  MPI_Comm_size(comm, &size);
  MPI_Comm_rank(comm, &rank);
  MPI_Get_processor_name(pname, &len);
  pname[len] = 0;
  printf("Hello, World! I am process %d of %d on %s.\n",
         rank, size, pname);
}

/* -------------------------------------------------------------------------- */

static PyObject *
hw_sayhello(PyObject *self, PyObject *args)
{
  PyObject *py_comm = NULL;
  MPI_Comm *comm_p  = NULL;

  if (!PyArg_ParseTuple(args, "O:sayhello", &py_comm))
    return NULL;

  comm_p = PyMPIComm_Get(py_comm);
  if (comm_p == NULL)
    return NULL;

  sayhello(*comm_p);

  Py_INCREF(Py_None);
  return Py_None;
}

static struct PyMethodDef hw_methods[] = {
  {"sayhello", (PyCFunction)hw_sayhello, METH_VARARGS, NULL},
  {NULL,       NULL,                     0,            NULL} /* sentinel */
};

#if PY_MAJOR_VERSION < 3
/* --- Python 2 --- */

PyMODINIT_FUNC inithelloworld(void)
{
  PyObject *m = NULL;

  /* Initialize mpi4py C-API */
  if (import_mpi4py() < 0) goto bad;

  /* Module initialization  */
  m = Py_InitModule("helloworld", hw_methods);
  if (m == NULL) goto bad;

  return;

 bad:
  return;
}

#else
/* --- Python 3 --- */

static struct PyModuleDef hw_module = {
  PyModuleDef_HEAD_INIT,
  "helloworld", /* m_name */
  NULL,         /* m_doc */
  -1,           /* m_size */
  hw_methods    /* m_methods */,
  NULL,         /* m_reload */
  NULL,         /* m_traverse */
  NULL,         /* m_clear */
  NULL          /* m_free */
};

PyMODINIT_FUNC
PyInit_helloworld(void)
{
  PyObject *m = NULL;

  /* Initialize mpi4py's C-API */
  if (import_mpi4py() < 0) goto bad;

  /* Module initialization  */
  m = PyModule_Create(&hw_module);
  if (m == NULL) goto bad;

  return m;

 bad:
  return NULL;
}

#endif

/* -------------------------------------------------------------------------- */

/*
  Local variables:
  c-basic-offset: 2
  indent-tabs-mode: nil
  End:
*/

我们希望在 mpi4py 程序中调用以上定义的 C 语言函数 sayhello：

# test.py

from mpi4py import MPI
from mpi4py import MPI
import helloworld as hw

null = MPI.COMM_NULL
hw.sayhello(null)

comm = MPI.COMM_WORLD
hw.sayhello(comm)

try:
    hw.sayhello(list())
except:
    pass
else:
    assert 0, "exception not raised"

为此我们需要将以上的 C 语言程序编译成一个扩展模块 helloworld.so (此处以 linux 为例，其它系统会有不同的后缀)。因为是 MPI 程序，我们需要使用一个支持 MPI 的 C 编译器，如 mpicc，例外还需要指定编译过程中需要用到的头文件的路径和需要链接的库文件，此处用到了 Python/C API 和 mpi4py.h 头文件，因此需要指定 Python.h 头文件所在的路径和 mpi4py.h 头文件所在的路径，链接 python2.7 等库文件，例外需加上指令 -fPIC -shared 以指导其编译成一个动态链接库。编译使用类似下面的命令（注意将其中的头文件路径改成你的系统中实际的路径）：

$ mpicc -I/path/to/python/include/python2.7 -I/path/to/python/lib/python2.7/site-packages/mpi4py/include -o helloworld.so helloworld.c -fPIC -shared -lpthread -ldl -lutil -lm -lpython2.7

如果不知道自己系统中的 Python.h 头文件路径和 mpi4py.h 头文件路径，可以用下面的命令获得：

$ python -c "import sysconfig; print( sysconfig.get_path('include') )"
$ python -c "import mpi4py; print( mpi4py.get_include() )"

编译成功后会生成扩展模块 helloworld.so，然后就可以在我们的 mpi4py 程序中像使用其它 Python 模块一样导入该模块并调用该模块中定义的 sayhello 函数，可以向此函数传递一个 mpi4py 中定义的通信子，如 MPI.COMM_WORLD 或者其它通信子对象。

执行以上 test.py 的结果如下：

$ mpiexec -n 4 python test.py
You passed MPI_COMM_NULL !!!
Hello, World! I am process 2 of 4 on node4.
You passed MPI_COMM_NULL !!!
Hello, World! I am process 3 of 4 on node4.
You passed MPI_COMM_NULL !!!
Hello, World! I am process 0 of 4 on node4.
You passed MPI_COMM_NULL !!!
Hello, World! I am process 1 of 4 on node4.

以上设置头文件路径和指定编译库的编译方法比较麻烦且容易出错，也不适用于比较大型的程序项目，一种更好的方法是写一个 Makefile，将编译的命令及编译的依赖关系写在此 Makefile 中，使用 GNU make 工具来完成项目的编译工作。另外，头文件的路径也可以使用程序自动查找和设定。我们可以使用使用以下的 python-config 文件和 Makefile 文件。

#!/usr/bin/env python
# -*- python -*-

# python-config

import sys, os
import getopt
try:
    import sysconfig
except ImportError:
    from distutils import sysconfig

valid_opts = ['help', 'prefix', 'exec-prefix', 'includes', 'libs', 'cflags',
              'ldflags', 'extension-suffix', 'abiflags', 'configdir']

def exit_with_usage(code=1):
    sys.stderr.write("Usage: %s [%s]\n" % (
        sys.argv[0], '|'.join('--'+opt for opt in valid_opts)))
    sys.exit(code)

try:
    opts, args = getopt.getopt(sys.argv[1:], '', valid_opts)
except getopt.error:
    exit_with_usage()

if not opts:
    exit_with_usage()

getvar = sysconfig.get_config_var
pyver = getvar('VERSION')
try:
    abiflags = sys.abiflags
except AttributeError:
    abiflags = ''

opt_flags = [flag for (flag, val) in opts]

if '--help' in opt_flags:
    exit_with_usage(code=0)

for opt in opt_flags:
    if opt == '--prefix':
        print(getvar('prefix'))

    elif opt == '--exec-prefix':
        print(getvar('exec_prefix'))

    elif opt in ('--includes', '--cflags'):
        try:
            include = sysconfig.get_path('include')
            platinclude = sysconfig.get_path('platinclude')
        except AttributeError:
            include = sysconfig.get_python_inc()
            platinclude = sysconfig.get_python_inc(plat_specific=True)
        flags = ['-I' + include]
        if include != platinclude:
            flags.append('-I' + platinclude)
        if opt == '--cflags':
            flags.extend(getvar('CFLAGS').split())
        print(' '.join(flags))

    elif opt in ('--libs', '--ldflags'):
        libs = getvar('LIBS').split() + getvar('SYSLIBS').split()
        libs.append('-lpython' + pyver + abiflags)
        if opt == '--ldflags':
            if not getvar('Py_ENABLE_SHARED'):
                libs.insert(0, '-L' + getvar('LIBPL'))
            if not getvar('PYTHONFRAMEWORK'):
                libs.extend(getvar('LINKFORSHARED').split())
        print(' '.join(libs))

    elif opt == '--extension-suffix':
        ext_suffix = getvar('EXT_SUFFIX')
        if ext_suffix is None:
            ext_suffix = getvar('SO')
        print(ext_suffix)

    elif opt == '--abiflags':
        print(abiflags)

    elif opt == '--configdir':
        print(getvar('LIBPL'))

# Makefile

.PHONY: default
default: build test clean

PYTHON = python
PYTHON_CONFIG = ${PYTHON} ./python-config
MPI4PY_INCLUDE = ${shell ${PYTHON} -c 'import mpi4py; print( mpi4py.get_include() )'}


MPICC = mpicc
CFLAGS = -fPIC ${shell ${PYTHON_CONFIG} --includes}
LDFLAGS = -shared ${shell ${PYTHON_CONFIG} --libs}
SO = ${shell ${PYTHON_CONFIG} --extension-suffix}
.PHONY: build
build: helloworld${SO}
helloworld${SO}: helloworld.c
    ${MPICC} ${CFLAGS} -I${MPI4PY_INCLUDE} -o $@ $< ${LDFLAGS}


MPIEXEC = mpiexec
NP_FLAG = -n
NP = 5
.PHONY: test
test: build
    ${MPIEXEC} ${NP_FLAG} ${NP} ${PYTHON} test.py


.PHONY: clean
clean:
    ${RM} helloworld${SO}

有了以上的 python-config 和 Makefile，编译扩展库，执行程序及清理可以分别使用如下命令，非常方便：

$ make build
$ make test
$ make clean

以上我们介绍了在 mpi4py 中包装和调用 C 语言 MPI 程序方法。在实际应用中直接使用 Python/C API 编写 Python 扩展模块是比较麻烦的，要求对 Python/C API 非常熟悉才能很好地运用，更常用的做法是使用像 SWIG 这样的工具来包装 C/C++ 程序文件，在下一篇中我们将介绍用 SWIG 包装 C 语言 MPI 程序以供 mpi4py 调用的方法。