diff options
author | 2021-07-09 20:13:56 -0700 | |
---|---|---|
committer | 2021-07-09 20:13:56 -0700 | |
commit | f462adc324cf466b60d62bc8c06c035221b75861 (patch) | |
tree | 3b100b97e40559ee9d809698e7a348b1c033220c /Python/pywarpx/_libwarpx.py | |
parent | b28004a648e5b729f8ef52c89d8df7a9423e368c (diff) | |
download | WarpX-f462adc324cf466b60d62bc8c06c035221b75861.tar.gz WarpX-f462adc324cf466b60d62bc8c06c035221b75861.tar.zst WarpX-f462adc324cf466b60d62bc8c06c035221b75861.zip |
Feature pass MPI comm from Python script to AMReX during initialization (#2034)
* Added functionality to pass mpi comm from python script to amrex during initialization
* Fixed missing _ in MPI._sizeof()
* Added functions to get the current processor's rank and total number of processors
* Renamed MPI_Comm to _MPI_Comm_type and defined _MPI_Comm_type in except statement
* Updated comment to explain why mpi4py needs to be imported before loading libwarpx
* Removed ifdef flags that prevent amrex_init_with_inited_mpi from being declared when MPI is off
* Changed amrex_init_with_inited_mpi to be declared even when not using mpi, but will be defined to be functionally the same as amrex_init
* Defined MPI = None to signify whether MPI is used, to add another check when initializing amrex
* Changed ifdef blocks in WarpXWrappers.cpp/h to fix compile errors.
Added ifdef block to conditionally declare amrex_init_with_inited_mpi in WarpXWrappers.h to prevent compile error when not using MPI. Removed ifdef block to declare/define same function in WarpXWrappers.cpp since function needs to be declared even when MPI is not used, but will never be called in that case.
* Changed BL_USE_MPI to AMREX_USE_MPI and removed incorrect MPI=None statement
* Changed BL_USE_MPI to AMREX_USE_MPI
* Added test to verify correct passing of MPI communicator to amrex
* Added ability to pass mpi_comm to sim.step
* Change test to check for differeing outputs when passed different inputs
* Removed obsolete comments refactored program to use more shared code
* Refactored comments
* Updated description to match test
* Removed unecessary imports and updated comments
Diffstat (limited to 'Python/pywarpx/_libwarpx.py')
-rwxr-xr-x | Python/pywarpx/_libwarpx.py | 26 |
1 files changed, 20 insertions, 6 deletions
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index d51bdfa9d..f1cc6564d 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -18,10 +18,16 @@ from .Geometry import geometry try: # --- If mpi4py is going to be used, this needs to be imported - # --- before libwarpx is loaded (though don't know why) + # --- before libwarpx is loaded, because mpi4py calls MPI_Init from mpi4py import MPI + # --- Change MPI Comm type depending on MPICH (int) or OpenMPI (void*) + if MPI._sizeof(MPI.Comm) == ctypes.sizeof(ctypes.c_int): + _MPI_Comm_type = ctypes.c_int + else: + _MPI_Comm_type = ctypes.c_void_p except ImportError: - pass + MPI = None + _MPI_Comm_type = ctypes.c_void_p # --- Is there a better way of handling constants? clight = 2.99792458e+8 # m/s @@ -136,6 +142,7 @@ def _array1d_from_pointer(pointer, dtype, size): # set the arg and return types of the wrapped functions libwarpx.amrex_init.argtypes = (ctypes.c_int, _LP_LP_c_char) +libwarpx.amrex_init_with_inited_mpi.argtypes = (ctypes.c_int, _LP_LP_c_char, _MPI_Comm_type) libwarpx.warpx_getParticleStructs.restype = _LP_particle_p libwarpx.warpx_getParticleArrays.restype = _LP_LP_c_particlereal libwarpx.warpx_getEfield.restype = _LP_LP_c_real @@ -208,6 +215,8 @@ libwarpx.warpx_getdt.restype = c_real libwarpx.warpx_maxStep.restype = ctypes.c_int libwarpx.warpx_stopTime.restype = c_real libwarpx.warpx_finestLevel.restype = ctypes.c_int +libwarpx.warpx_getMyProc.restype = ctypes.c_int +libwarpx.warpx_getNProcs.restype = ctypes.c_int libwarpx.warpx_EvolveE.argtypes = [c_real] libwarpx.warpx_EvolveB.argtypes = [c_real] @@ -234,7 +243,7 @@ def get_nattr(): # --- The -3 is because the comps include the velocites return libwarpx.warpx_nComps() - 3 -def amrex_init(argv): +def amrex_init(argv, mpi_comm=None): # --- Construct the ctype list of strings to pass in argc = len(argv) argvC = (_LP_c_char * (argc+1))() @@ -242,9 +251,14 @@ def amrex_init(argv): enc_arg = arg.encode('utf-8') argvC[i] = ctypes.create_string_buffer(enc_arg) - libwarpx.amrex_init(argc, argvC) + if mpi_comm is None or MPI is None: + libwarpx.amrex_init(argc, argvC) + else: + comm_ptr = MPI._addressof(mpi_comm) + comm_val = _MPI_Comm_type.from_address(comm_ptr) + libwarpx.amrex_init_with_inited_mpi(argc, argvC, comm_val) -def initialize(argv=None): +def initialize(argv=None, mpi_comm=None): ''' Initialize WarpX and AMReX. Must be called before @@ -253,7 +267,7 @@ def initialize(argv=None): ''' if argv is None: argv = sys.argv - amrex_init(argv) + amrex_init(argv, mpi_comm) libwarpx.warpx_ConvertLabParamsToBoost() libwarpx.warpx_ReadBCParams() if geometry_dim == 'rz': |