Skip to content
Snippets Groups Projects
Unverified Commit c68a9a60 authored by Thibaut Lunet's avatar Thibaut Lunet Committed by GitHub
Browse files

FieldsIO performance fix (#550)

* TL: performance fix for fieldsIO

* TL: final update before PR

* TL: missing detail

* TL: removing warning

* TL: trying something

* TL: removing gusto stuff
parent 72fa09e2
No related branches found
No related tags found
No related merge requests found
Pipeline #272997 passed
...@@ -44,14 +44,13 @@ See :class:`pySDC.helpers.fieldsIO.writeFields_MPI` for an illustrative example. ...@@ -44,14 +44,13 @@ See :class:`pySDC.helpers.fieldsIO.writeFields_MPI` for an illustrative example.
Warning Warning
------- -------
To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.initMPI` (cf their docstring). To use MPI collective writing, you need to call first the class methods :class:`Rectilinear.setupMPI` (cf their docstring).
Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, whether the code is run in parallel or not. Also, `Rectilinear.setHeader` **must be given the global grids coordinates**, whether the code is run in parallel or not.
""" """
import os import os
import numpy as np import numpy as np
from typing import Type, TypeVar from typing import Type, TypeVar
import logging import logging
import itertools
T = TypeVar("T") T = TypeVar("T")
...@@ -61,11 +60,17 @@ try: ...@@ -61,11 +60,17 @@ try:
except ImportError: except ImportError:
pass pass
from mpi4py import MPI from mpi4py import MPI
from mpi4py.util.dtlib import from_numpy_dtype as MPI_DTYPE
except ImportError: except ImportError:
class MPI: class MPI:
COMM_WORLD = None COMM_WORLD = None
Intracomm = T Intracomm = T
File = T
Datatype = T
def MPI_DTYPE():
pass
# Supported data types # Supported data types
...@@ -412,6 +417,8 @@ class Rectilinear(Scalar): ...@@ -412,6 +417,8 @@ class Rectilinear(Scalar):
coords = self.setupCoords(*coords) coords = self.setupCoords(*coords)
self.header = {"nVar": int(nVar), "coords": coords} self.header = {"nVar": int(nVar), "coords": coords}
self.nItems = nVar * self.nDoF self.nItems = nVar * self.nDoF
if self.MPI_ON:
self.MPI_SETUP()
@property @property
def hInfos(self): def hInfos(self):
...@@ -433,6 +440,8 @@ class Rectilinear(Scalar): ...@@ -433,6 +440,8 @@ class Rectilinear(Scalar):
gridSizes = np.fromfile(f, dtype=np.int32, count=dim) gridSizes = np.fromfile(f, dtype=np.int32, count=dim)
coords = [np.fromfile(f, dtype=np.float64, count=n) for n in gridSizes] coords = [np.fromfile(f, dtype=np.float64, count=n) for n in gridSizes]
self.setHeader(nVar, coords) self.setHeader(nVar, coords)
if self.MPI_ON:
self.MPI_SETUP()
def reshape(self, fields: np.ndarray): def reshape(self, fields: np.ndarray):
"""Reshape the fields to a N-d array (inplace operation)""" """Reshape the fields to a N-d array (inplace operation)"""
...@@ -493,7 +502,6 @@ class Rectilinear(Scalar): ...@@ -493,7 +502,6 @@ class Rectilinear(Scalar):
# MPI-parallel implementation # MPI-parallel implementation
# ------------------------------------------------------------------------- # -------------------------------------------------------------------------
comm: MPI.Intracomm = None comm: MPI.Intracomm = None
_nCollectiveIO = None
@classmethod @classmethod
def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc): def setupMPI(cls, comm: MPI.Intracomm, iLoc, nLoc):
...@@ -513,21 +521,9 @@ class Rectilinear(Scalar): ...@@ -513,21 +521,9 @@ class Rectilinear(Scalar):
cls.comm = comm cls.comm = comm
cls.iLoc = iLoc cls.iLoc = iLoc
cls.nLoc = nLoc cls.nLoc = nLoc
cls.mpiFile = None cls.mpiFile: MPI.File = None
cls._nCollectiveIO = None cls.mpiType: MPI.Datatype = None
cls.mpiFileType: MPI.Datatype = None
@property
def nCollectiveIO(self):
"""
Number of collective IO operations over all processes, when reading or writing a field.
Returns:
--------
int: Number of collective IO accesses
"""
if self._nCollectiveIO is None:
self._nCollectiveIO = self.comm.allreduce(self.nVar * np.prod(self.nLoc[:-1]), op=MPI.MAX)
return self._nCollectiveIO
@property @property
def MPI_ON(self): def MPI_ON(self):
...@@ -543,6 +539,16 @@ class Rectilinear(Scalar): ...@@ -543,6 +539,16 @@ class Rectilinear(Scalar):
return True return True
return self.comm.Get_rank() == 0 return self.comm.Get_rank() == 0
def MPI_SETUP(self):
"""Setup subarray masks for each processes"""
self.mpiType = MPI_DTYPE(self.dtype)
self.mpiFileType = self.mpiType.Create_subarray(
[self.nVar, *self.gridSizes], # Global array sizes
[self.nVar, *self.nLoc], # Local array sizes
[0, *self.iLoc], # Global starting indices of local blocks
)
self.mpiFileType.Commit()
def MPI_FILE_OPEN(self, mode): def MPI_FILE_OPEN(self, mode):
"""Open the binary file in MPI mode""" """Open the binary file in MPI mode"""
amode = { amode = {
...@@ -567,7 +573,8 @@ class Rectilinear(Scalar): ...@@ -567,7 +573,8 @@ class Rectilinear(Scalar):
data : np.ndarray data : np.ndarray
Data to be written in the binary file. Data to be written in the binary file.
""" """
self.mpiFile.Write_at_all(offset, data) self.mpiFile.Set_view(disp=offset, etype=self.mpiType, filetype=self.mpiFileType)
self.mpiFile.Write_all(data)
def MPI_READ_AT_ALL(self, offset, data: np.ndarray): def MPI_READ_AT_ALL(self, offset, data: np.ndarray):
""" """
...@@ -581,7 +588,8 @@ class Rectilinear(Scalar): ...@@ -581,7 +588,8 @@ class Rectilinear(Scalar):
data : np.ndarray data : np.ndarray
Array on which to read the data from the binary file. Array on which to read the data from the binary file.
""" """
self.mpiFile.Read_at_all(offset, data) self.mpiFile.Set_view(disp=offset, etype=self.mpiType, filetype=self.mpiFileType)
self.mpiFile.Read_all(data)
def MPI_FILE_CLOSE(self): def MPI_FILE_CLOSE(self):
"""Close the binary file in MPI mode""" """Close the binary file in MPI mode"""
...@@ -632,33 +640,15 @@ class Rectilinear(Scalar): ...@@ -632,33 +640,15 @@ class Rectilinear(Scalar):
*self.nLoc, *self.nLoc,
), f"expected {(self.nVar, *self.nLoc)} shape, got {field.shape}" ), f"expected {(self.nVar, *self.nLoc)} shape, got {field.shape}"
offset0 = self.fileSize offset = self.fileSize
self.MPI_FILE_OPEN(mode="a") self.MPI_FILE_OPEN(mode="a")
nWrites = 0
nCollectiveIO = self.nCollectiveIO
if self.MPI_ROOT: if self.MPI_ROOT:
self.MPI_WRITE(np.array(time, dtype=T_DTYPE)) self.MPI_WRITE(np.array(time, dtype=T_DTYPE))
offset0 += self.tSize offset += self.tSize
self.MPI_WRITE_AT_ALL(offset, field)
for (iVar, *iBeg) in itertools.product(range(self.nVar), *[range(n) for n in self.nLoc[:-1]]):
offset = offset0 + self.iPos(iVar, iBeg) * self.itemSize
self.MPI_WRITE_AT_ALL(offset, field[(iVar, *iBeg)])
nWrites += 1
for _ in range(nCollectiveIO - nWrites):
# Additional collective write to catch up with other processes
self.MPI_WRITE_AT_ALL(offset0, field[:0])
self.MPI_FILE_CLOSE() self.MPI_FILE_CLOSE()
def iPos(self, iVar, iX):
iPos = iVar * self.nDoF
for axis in range(self.dim - 1):
iPos += (self.iLoc[axis] + iX[axis]) * np.prod(self.gridSizes[axis + 1 :])
iPos += self.iLoc[-1]
return iPos
def readField(self, idx): def readField(self, idx):
""" """
Read one field stored in the binary file, corresponding to the given Read one field stored in the binary file, corresponding to the given
...@@ -684,26 +674,15 @@ class Rectilinear(Scalar): ...@@ -684,26 +674,15 @@ class Rectilinear(Scalar):
return super().readField(idx) return super().readField(idx)
idx = self.formatIndex(idx) idx = self.formatIndex(idx)
offset0 = self.hSize + idx * (self.tSize + self.fSize) offset = self.hSize + idx * (self.tSize + self.fSize)
with open(self.fileName, "rb") as f: with open(self.fileName, "rb") as f:
t = float(np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset0)[0]) t = float(np.fromfile(f, dtype=T_DTYPE, count=1, offset=offset)[0])
offset0 += self.tSize offset += self.tSize
field = np.empty((self.nVar, *self.nLoc), dtype=self.dtype) field = np.empty((self.nVar, *self.nLoc), dtype=self.dtype)
self.MPI_FILE_OPEN(mode="r") self.MPI_FILE_OPEN(mode="r")
nReads = 0 self.MPI_READ_AT_ALL(offset, field)
nCollectiveIO = self.nCollectiveIO
for (iVar, *iBeg) in itertools.product(range(self.nVar), *[range(n) for n in self.nLoc[:-1]]):
offset = offset0 + self.iPos(iVar, iBeg) * self.itemSize
self.MPI_READ_AT_ALL(offset, field[(iVar, *iBeg)])
nReads += 1
for _ in range(nCollectiveIO - nReads):
# Additional collective read to catch up with other processes
self.MPI_READ_AT_ALL(offset0, field[:0])
self.MPI_FILE_CLOSE() self.MPI_FILE_CLOSE()
return t, field return t, field
......
...@@ -3,8 +3,8 @@ import numpy as np ...@@ -3,8 +3,8 @@ import numpy as np
from pySDC.core.collocation import CollBase from pySDC.core.collocation import CollBase
t_start = float(np.random.rand(1) * 0.2) t_start = float(np.random.rand(1)[0] * 0.2)
t_end = float(0.8 + np.random.rand(1) * 0.2) t_end = float(0.8 + np.random.rand(1)[0] * 0.2)
tolQuad = 1e-13 tolQuad = 1e-13
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment