Create shared MPI fractal functions

这个提交包含在:
Nathan Mannall
2025-05-07 17:26:03 +01:00
父节点 548cf6599b
当前提交 5ee013d9fe
共有 3 个文件被更改,包括 84 次插入116 次删除

查看文件

@@ -10,6 +10,7 @@ from scipy import fftpack
from gprMax import config
from gprMax.cython.fractals_generate import generate_fractal2D
from gprMax.fractals.mpi_utilities import calculate_starts_and_subshape, create_mpi_type
from gprMax.utilities.mpi import Dim, Dir, get_relative_neighbour
logger = logging.getLogger(__name__)
@@ -346,13 +347,13 @@ class MPIFractalSurface(FractalSurface):
dirs = np.full(2, Dir.NONE)
starts, subshape = self.calculate_starts_and_subshape(
starts, subshape = calculate_starts_and_subshape(
A_shape, -negative_offset, -positive_offset, dirs, sending=True
)
ends = starts + subshape
A_local = A[starts[0] : ends[0], starts[1] : ends[1]]
starts, subshape = self.calculate_starts_and_subshape(
starts, subshape = calculate_starts_and_subshape(
local_shape, negative_offset, positive_offset, dirs
)
ends = starts + subshape
@@ -383,12 +384,13 @@ class MPIFractalSurface(FractalSurface):
# Check if any data to send
if all(
np.logical_or(
np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset <= 0, positive_offset <= 0],
dirs == Dir.NONE,
np.where(dirs == Dir.NEG, negative_offset <= 0, positive_offset <= 0),
)
):
mpi_type = self.create_mpi_type(
mpi_type = create_mpi_type(
A_shape, -negative_offset, -positive_offset, dirs, sending=True
)
@@ -399,12 +401,13 @@ class MPIFractalSurface(FractalSurface):
# Check if any data to receive
if all(
np.logical_or(
np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset > 0, positive_offset > 0],
dirs == Dir.NONE,
np.where(dirs == Dir.NEG, negative_offset > 0, positive_offset > 0),
)
):
mpi_type = self.create_mpi_type(local_shape, negative_offset, positive_offset, dirs)
mpi_type = create_mpi_type(local_shape, negative_offset, positive_offset, dirs)
logger.debug(
f"Receiving fractal surface from rank {rank}, MPI type={mpi_type.decode()}"
@@ -424,53 +427,3 @@ class MPIFractalSurface(FractalSurface):
logger.debug(
f"Generated fractal surface: start={self.start}, stop={self.stop}, size={self.size}, fractalrange={self.fractalrange}"
)
def calculate_starts_and_subshape(
self,
shape: npt.NDArray[np.int32],
negative_offset: npt.NDArray[np.int32],
positive_offset: npt.NDArray[np.int32],
dirs: npt.NDArray[np.int32],
sending: bool = False,
) -> Tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
negative_offset = np.where(
dirs == Dir.NONE,
np.maximum(negative_offset, 0),
np.abs(negative_offset),
)
positive_offset = np.where(
dirs == Dir.NONE,
np.maximum(positive_offset, 0),
np.abs(positive_offset),
)
starts = np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[0, shape - positive_offset - sending],
default=negative_offset,
)
subshape = np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset + sending, positive_offset + sending],
default=shape - negative_offset - positive_offset,
)
return starts, subshape
def create_mpi_type(
self,
shape: npt.NDArray[np.int32],
negative_offset: npt.NDArray[np.int32],
positive_offset: npt.NDArray[np.int32],
dirs: npt.NDArray[np.int32],
sending: bool = False,
) -> MPI.Datatype:
starts, subshape = self.calculate_starts_and_subshape(
shape, negative_offset, positive_offset, dirs, sending
)
mpi_type = MPI.FLOAT.Create_subarray(shape.tolist(), subshape.tolist(), starts.tolist())
mpi_type.Commit()
return mpi_type

查看文件

@@ -28,6 +28,7 @@ from scipy import fftpack
from gprMax import config
from gprMax.cython.fractals_generate import generate_fractal3D
from gprMax.fractals.mpi_utilities import calculate_starts_and_subshape, create_mpi_type
from gprMax.utilities.mpi import Dim, Dir, get_relative_neighbour
from gprMax.utilities.utilities import round_value
@@ -439,13 +440,13 @@ class MPIFractalVolume(FractalVolume):
dirs = np.full(3, Dir.NONE)
starts, subshape = self.calculate_starts_and_subshape(
starts, subshape = calculate_starts_and_subshape(
A_shape, -negative_offset, -positive_offset, dirs, sending=True
)
ends = starts + subshape
A_local = A[starts[0] : ends[0], starts[1] : ends[1], starts[2] : ends[2]]
starts, subshape = self.calculate_starts_and_subshape(
starts, subshape = calculate_starts_and_subshape(
local_shape, negative_offset, positive_offset, dirs
)
ends = starts + subshape
@@ -479,12 +480,13 @@ class MPIFractalVolume(FractalVolume):
# Check if any data to send
if all(
np.logical_or(
np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset <= 0, positive_offset <= 0],
dirs == Dir.NONE,
np.where(dirs == Dir.NEG, negative_offset <= 0, positive_offset <= 0),
)
):
mpi_type = self.create_mpi_type(
mpi_type = create_mpi_type(
A_shape, -negative_offset, -positive_offset, dirs, sending=True
)
@@ -493,12 +495,13 @@ class MPIFractalVolume(FractalVolume):
# Check if any data to receive
if all(
np.logical_or(
np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset > 0, positive_offset > 0],
dirs == Dir.NONE,
np.where(dirs == Dir.NEG, negative_offset > 0, positive_offset > 0),
)
):
mpi_type = self.create_mpi_type(local_shape, negative_offset, positive_offset, dirs)
mpi_type = create_mpi_type(local_shape, negative_offset, positive_offset, dirs)
logger.debug(
f"Receiving fractal volume from rank {rank}, MPI type={mpi_type.decode()}"
@@ -518,53 +521,3 @@ class MPIFractalVolume(FractalVolume):
)
return True
def calculate_starts_and_subshape(
self,
shape: npt.NDArray[np.int32],
negative_offset: npt.NDArray[np.int32],
positive_offset: npt.NDArray[np.int32],
dirs: npt.NDArray[np.int32],
sending: bool = False,
) -> Tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
negative_offset = np.where(
dirs == Dir.NONE,
np.maximum(negative_offset, 0),
np.abs(negative_offset),
)
positive_offset = np.where(
dirs == Dir.NONE,
np.maximum(positive_offset, 0),
np.abs(positive_offset),
)
starts = np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[0, shape - positive_offset - sending],
default=negative_offset,
)
subshape = np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset + sending, positive_offset + sending],
default=shape - negative_offset - positive_offset,
)
return starts, subshape
def create_mpi_type(
self,
shape: npt.NDArray[np.int32],
negative_offset: npt.NDArray[np.int32],
positive_offset: npt.NDArray[np.int32],
dirs: npt.NDArray[np.int32],
sending: bool = False,
) -> MPI.Datatype:
starts, subshape = self.calculate_starts_and_subshape(
shape, negative_offset, positive_offset, dirs, sending
)
mpi_type = MPI.FLOAT.Create_subarray(shape.tolist(), subshape.tolist(), starts.tolist())
mpi_type.Commit()
return mpi_type

查看文件

@@ -0,0 +1,62 @@
from typing import Tuple
import numpy as np
import numpy.typing as npt
from mpi4py import MPI
from gprMax.utilities.mpi import Dir
def calculate_starts_and_subshape(
shape: npt.NDArray[np.int32],
negative_offset: npt.NDArray[np.int32],
positive_offset: npt.NDArray[np.int32],
dirs: npt.NDArray[np.int32],
sending: bool = False,
) -> Tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]:
negative_offset = np.where(
dirs == Dir.NONE,
np.maximum(negative_offset, 0),
np.abs(negative_offset),
)
positive_offset = np.where(
dirs == Dir.NONE,
np.maximum(positive_offset, 0),
np.abs(positive_offset),
)
starts = np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[0, shape - positive_offset - sending],
default=negative_offset,
)
subshape = np.select(
[dirs == Dir.NEG, dirs == Dir.POS],
[negative_offset + sending, positive_offset + sending],
default=shape - negative_offset - positive_offset,
)
return starts, subshape
def create_mpi_type(
shape: npt.NDArray[np.int32],
negative_offset: npt.NDArray[np.int32],
positive_offset: npt.NDArray[np.int32],
dirs: npt.NDArray[np.int32],
sending: bool = False,
) -> MPI.Datatype:
starts, subshape = calculate_starts_and_subshape(
shape, negative_offset, positive_offset, dirs, sending
)
print(
f"create_mpi_type: shape={shape}, negative_offset={negative_offset}, positive_offset={positive_offset}, dirs={dirs}, sending={sending}"
)
print(f"create_mpi_type: starts={starts}, subshape={subshape}")
mpi_type = MPI.FLOAT.Create_subarray(shape.tolist(), subshape.tolist(), starts.tolist())
mpi_type.Commit()
return mpi_type