C++ 扩展python(五）mpi传递numpy（使用SWIG）

现在我们可以尝试着用C/C++来实现mpi传递一个python的numpy数组了。

 

首先先实现我们的get_states程序：

get_states.h

#include <mpi.h>

void get_states(MPI_Comm comm, double* message, int size);

get_states.c

#define MPICH_SKIP_MPICXX 1
#define OMPI_SKIP_MPICXX  1

#include <mpi.h>
#include <stdio.h>
#include "get_states.h"

void get_states(MPI_Comm comm, double* message, int size) {
  int comm_size, comm_rank;
  int value;
  MPI_Status status;
  char pname[MPI_MAX_PROCESSOR_NAME]; int len;
  if (comm == MPI_COMM_NULL) {
    printf("You passed MPI_COMM_NULL !!!\n");
    return ;
  }
   if (message == NULL) {
    printf("You passed NULL array !!!\n");
    return ;
  }
  MPI_Comm_size(comm, &comm_size);
  MPI_Comm_rank(comm, &comm_rank);
  MPI_Get_processor_name(pname, &len);
  pname[len] = 0;
  int x;
  if(comm_rank == 0){
        value = 2;
        MPI_Send(message, size, MPI_DOUBLE, 1, 0, comm);
  }
  else{
          MPI_Recv(message, size, MPI_DOUBLE, 0, 0, comm, &status);
  }

}

接口文件get_states.i：

/*  Example of wrapping a C function that takes a C double array as input using
 *  numpy typemaps for SWIG. */

%module get_states
%{
    /* the resulting C file should be built as a python extension */
    #define SWIG_FILE_WITH_INIT
    /*  Includes the header in the wrapper code */
    #include "get_states.h"
%}

/*  include the numpy typemaps */
%include "numpy.i"

%include "mpi4py/mpi4py.i"

/*  need this for correct module initialization */
%init %{
    import_array();
%}
%mpi4py_typemap(Comm, MPI_Comm);

/*  typemaps for the two arrays, the second will be modified in-place */
%apply (double* IN_ARRAY1, int DIM1) {(double * a, int size_a)}

/*  Wrapper for cos_doubles that massages the types */
%inline %{
    /*  takes as input two numpy arrays */
    void get_states_func(MPI_Comm comm, double * a, int size_a) {
        /*  calls the original funcion, providing only the size of the first */
        get_states(comm, a, size_a);
    }
%}

接下来让我们编译：

写个setup.py出了问题。

问题来了，就是用setup.py写的时候，include_dirs就是个无用的东西，一直提示我找不到mpi4py.i和numpy.i文件。唉，也不知道咋解决，于是选择了原始编译法。

输入命令：

swig -python -I/public/home/---/software/anaconda3/include/python3.7m -I/public/home/---/software/anaconda3/lib/python3.7/site-packages/mpi4py/include  get_states.i

(我的numpy.i直接放在了跟get_states.*一起，所以我没有链接上numpy的问题。）

将get_states.c和get_states_wrap.c一起编译生成.so文件

mpicc -I/public/home/---/software/anaconda3/include/python3.7m -I/public/home/---/software/anaconda3/lib/python3.7/site-packages/mpi4py/include  -o _get_states.so get_states.c get_states_wrap.c -fPIC -shared -lpthread -ldl -lutil

接下来写个test.py文件做测试：

from mpi4py import MPI
import get_states as gs
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

x = np.arange(0, 2*np.pi, 0.1)
print(x)

if rank == 0:
   print(" rank 0")
   gs.get_states_func(comm, x)
else:
   y = np.empty_like(x)
   gs.get_states_func(comm, y)
   print("ddd")
   print(y)

看看测试结果，恩，总算是可以了。

 

然后我就有些懵逼，这么简单的一个事我为啥现在才做完，在家的效率果然是极低极低。惆怅！