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)

看看測試結果，恩，總算是可以了。

 

然後我就有些懵逼，這麼簡單的一個事我爲啥現在才做完，在家的效率果然是極低極低。惆悵！