TL: update to write parallel decomposition

Author: tlunet

pySDC/playgrounds/dedalus/problems/rbc.py

@@ -225,6 +225,7 @@ def initField(self, initPath, seed):
     def runSimulation(cls, runDir, tEnd, baseDt, tBeg=0, logEvery=100,
                       dtWrite=None, writeVort=False, writeTau=False,
                       timeScheme="RK443", timeParallel=False, groupTimeProcs=False,
+                      writeDecomposition=False,
                       **pParams):
 
         cls.log(f"RBC simulation in {runDir}")
@@ -243,7 +244,7 @@ def runSimulation(cls, runDir, tEnd, baseDt, tBeg=0, logEvery=100,
 
         if timeParallel:
             assert timeScheme == "SDC", "need timeScheme=SDC for timeParallel"
-            _, sComm, _ = SDCIMEX_MPI.initSpaceTimeComms(groupTime=groupTimeProcs)
+            tComm, sComm, gComm = SDCIMEX_MPI.initSpaceTimeComms(groupTime=groupTimeProcs)
             pParams.update(sComm=sComm)
             if timeParallel == "MPI":
                 TimeStepper = SDCIMEX_MPI
@@ -255,6 +256,30 @@ def runSimulation(cls, runDir, tEnd, baseDt, tBeg=0, logEvery=100,
 
         p = cls(**pParams)
 
+        if writeDecomposition:
+            if MPI_RANK == 0:
+                with open("{runDir}/distrib.txt", "r") as f:
+                    f.write("Parallel distribution on compute cores\n")
+                    f.write(f" -- space parallelization on {p.sComm.Get_size()} procs\n")
+                    if timeParallel:
+                        f.write(f" -- time parallelization on {tComm.Get_size()} procs\n")
+                        f.write(f" -- global parallelization on {gComm.Get_size()} procs\n")
+            coords = [p.grids[axis] for axis in p.axes]
+            labels = ["x", "y", "z"]
+            COMM_WORLD.Barrier()
+            sleep(0.0001*MPI_RANK)
+            with open("{runDir}/distrib.txt", "a") as f:
+                if timeParallel:
+                    out = f"P{gComm.Get_rank()}-S{sComm.Get_rank()}-T{tComm.Get_rank()} :\n"
+                else:
+                    out = f"P{MPI_RANK} :\n"
+                out += f" -- cpu {os.sched_getaffinity(0)} on {socket.gethostname()}\n"
+                out += "\n".join(
+                    f" -- {d} {c.shape} : [{c.min(initial=np.inf)}, {c.max(initial=-np.inf)}]"
+                    for d, c in zip(labels, coords)
+                    )
+            COMM_WORLD.Barrier()
+
         dt = baseDt/p.resFactor
         nSteps = round(float(tEnd-tBeg)/dt, ndigits=3)
         if float(tEnd-tBeg) != round(nSteps*dt, ndigits=3):

pySDC/playgrounds/dedalus/scripts/plotScaling.py

@@ -11,19 +11,20 @@
 
 folder = "_benchJusuf"
 
-methods = ["RK443", "SDC", "SDC-MPI"]
+schemes = ["RK443", "SDC", "SDC-MPI", "SDC-MPI2"]
 R = 2
 
 useNSpS = False
 nSpS = {
     "RK443": 23,
     "SDC": 17,
     "SDC-MPI": 17,
+    "SDC-MPI2": 17,
     }
 
 results = {}
 
-for scheme in methods:
+for scheme in schemes:
 
     files = glob.glob(f"{folder}/R{R}_{scheme}_*.json")
 
@@ -44,7 +45,9 @@
 
     results[scheme].sort(key=lambda p: p[0])
 
-symbols = ["o", "^", "s"]
+symbols = ["o", "^", "s", "p"]
+
+
 plt.figure("scaling"+"-nSpS"*useNSpS)
 for scheme, sym in zip(results.keys(), symbols):
     res = np.array(results[scheme]).T
@@ -58,3 +61,16 @@
 else:
     plt.ylabel("$t_{wall}/N_{DoF}/N_{steps}$")
 plt.tight_layout()
+
+
+plt.figure("PinT-speedup")
+schemes = ["SDC-MPI", "SDC-MPI2"]
+nProcSpace, tSDC = np.array(results["SDC"]).T
+for scheme, sym in zip(schemes, symbols):
+    _, tSDCPinT = np.array(results[scheme]).T
+    speedup = tSDC[:-2]/tSDCPinT
+    plt.semilogx(nProcSpace[:-2], speedup, label=scheme)
+plt.legend()
+plt.grid(True)
+plt.xlabel("$N_{p,S}$"), plt.ylabel("PinT Speedup")
+plt.tight_layout()

pySDC/playgrounds/dedalus/scripts/run_rbc.py

@@ -55,12 +55,15 @@
 parser.add_argument(
     "--groupTimeProcs", help="wether or not grouping the time processes",
     action="store_true")
+parser.add_argument(
+    "--writeDecomposition", help="write the parallel space-time decomposition in a file",
+    action="store_true")
 
 parser.add_argument(
     "--nNodesSDC", help="number of time nodes per step for SDC",
     default=4, type=int)
 parser.add_argument(
-    "--nSweepSDC", help="number of sweep per step for SDC",
+    "--nSweepsSDC", help="number of sweeps per time step for SDC",
     default=4, type=int)
 parser.add_argument(
     "--implSweepSDC", help="implicit sweep type for SDC",
@@ -80,9 +83,9 @@
 dim = params.pop("dim")
 ProblemClass = RBCProblem2D if dim == 2 else RBCProblem3D
 SDCIMEX.setParameters(
-    nNodes=params.pop("nNodesSDC"), 
+    nNodes=params.pop("nNodesSDC"),
     nodeType="LEGENDRE", quadType="RADAU-RIGHT",
-    nSweeps=params.pop("nSweepSDC"), 
+    nSweeps=params.pop("nSweepsSDC"),
     initSweep="COPY",
     implSweep=params.pop("implSweepSDC"), explSweep=params.pop("explSweepSDC"),
 )