C++ 扩展python(五）mpi传递numpy（使用SWIG）
现在我们可以尝试着用C/C++来实现mpi传递一个python的numpy数组了。
首先先实现我们的get_states程序：
get_states.h
#include
void get_states(MPI_Comm comm, double* message, int size);
get_states.c
#define MPICH_SKIP_MPICXX 1
#define OMPI_SKIP_MPICXX 1
#include
#include
#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)
看看测试结果，恩，总算是可以了。
然后我就有些懵逼，这么简单的一个事我为啥现在才做完，在家的效率果然是极低极低。惆怅！