Use hptt for transpose

.github/actions/pip_installation/action.yml

@@ -14,8 +14,18 @@ runs:
     - name: Install dependencies
       shell: bash
       run: |
-              export CC="mpicc" HDF5_MPI="ON"
-              python -m pip install -r requirements.txt
-              python -m pip install numba
-              python -m pip install pythran
-
+        export CC="mpicc" HDF5_MPI="ON"
+        python -m pip install -r requirements.txt
+        python -m pip install numba
+        python -m pip install pythran
+    - uses: actions/checkout@v5
+      with:
+        repository: EmilyBourne/hptt
+        path: hptt
+    - name: Install hptt
+      shell: bash
+      run: |
+        cd hptt
+        make -j
+        cd pythonAPI
+        python -m pip install .

fullSimulation.py

@@ -218,6 +218,12 @@ def my_print(rank, *args, **kwargs):
     average_output = 0
     startPrint = max(0, ti % saveStep)
     timeForLoop = True
+
+    timeQN = 0
+    timeFluxAdv = 0
+    timevParAdv = 0
+    timePolAdv = 0
+
     while (ti < tN and timeForLoop):
 
         full_loop_start = time.time()
@@ -228,49 +234,77 @@ def my_print(rank, *args, **kwargs):
         # Compute f^n+1/2 using Lie splitting
         distribFunc.setLayout('flux_surface')
         distribFunc.saveGridValues()
+        time_start = time.time()
         fluxAdv.gridStep(distribFunc)
+        timeFluxAdv += (time.time()-time_start)
         distribFunc.setLayout('v_parallel')
         phi.setLayout('v_parallel_1d')
+        time_start = time.time()
         vParAdv.gridStep(distribFunc, phi, parGrad, parGradVals, halfStep)
+        timevParAdv += (time.time()-time_start)
         distribFunc.setLayout('poloidal')
         phi.setLayout('poloidal')
+        time_start = time.time()
         polAdv.gridStep(distribFunc, phi, halfStep)
+        timePolAdv += (time.time()-time_start)
 
         # Find phi from f^n+1/2 by solving QN eq again
         distribFunc.setLayout('v_parallel')
+        time_start = time.time()
         density.getPerturbedRho(distribFunc, rho)
+        timeQN += (time.time()-time_start)
         QNSolver.getModes(rho)
         rho.setLayout('mode_solve')
         phi.setLayout('mode_solve')
+        time_start = time.time()
         QNSolver.solveEquation(phi, rho)
+        timeQN += (time.time()-time_start)
         phi.setLayout('v_parallel_2d')
         rho.setLayout('v_parallel_2d')
+        time_start = time.time()
         QNSolver.findPotential(phi)
+        timeQN += (time.time()-time_start)
 
         # Compute f^n+1 using strang splitting
         distribFunc.restoreGridValues()  # restored from flux_surface layout
+        time_start = time.time()
         fluxAdv.gridStep(distribFunc)
+        timeFluxAdv += (time.time()-time_start)
         distribFunc.setLayout('v_parallel')
         phi.setLayout('v_parallel_1d')
+        time_start = time.time()
         vParAdv.gridStep(distribFunc, phi, parGrad, parGradVals, halfStep)
+        timevParAdv += (time.time()-time_start)
         distribFunc.setLayout('poloidal')
         phi.setLayout('poloidal')
+        time_start = time.time()
         polAdv.gridStep(distribFunc, phi, fullStep)
+        timePolAdv += (time.time()-time_start)
         distribFunc.setLayout('v_parallel')
+        time_start = time.time()
         vParAdv.gridStepKeepGradient(distribFunc, parGradVals, halfStep)
+        timevParAdv += (time.time()-time_start)
         distribFunc.setLayout('flux_surface')
+        time_start = time.time()
         fluxAdv.gridStep(distribFunc)
+        timeFluxAdv += (time.time()-time_start)
 
         # Find phi from f^n by solving QN eq
         distribFunc.setLayout('v_parallel')
+        time_start = time.time()
         density.getPerturbedRho(distribFunc, rho)
+        timeQN += (time.time()-time_start)
         QNSolver.getModes(rho)
         rho.setLayout('mode_solve')
         phi.setLayout('mode_solve')
+        time_start = time.time()
         QNSolver.solveEquation(phi, rho)
+        timeQN += (time.time()-time_start)
         phi.setLayout('v_parallel_2d')
         rho.setLayout('v_parallel_2d')
+        time_start = time.time()
         QNSolver.findPotential(phi)
+        timeQN += (time.time()-time_start)
 
         diagnostic_start = time.time()
         # Calculate diagnostic quantities
@@ -301,8 +335,6 @@ def my_print(rank, *args, **kwargs):
         timeForLoop = comm.allreduce((time.time(
         ) - setup_time_start + 2*average_loop + 2*average_output) < stopTime, op=MPI.LAND)
 
-    full_loop_time += (time.time()-full_loop_start)
-
     output_start = time.time()
 
     if (ti % saveStep != 0):
@@ -327,15 +359,31 @@ def my_print(rank, *args, **kwargs):
         print(s.getvalue(), file=open("profile/l2Test{}.txt".format(rank), "w"))
 
     if (rank == 0):
-        if (not os.path.isdir("timing")):
-            os.mkdir("timing")
+        if (not os.path.isdir(f"{foldername}/timing")):
+            os.mkdir(f"{foldername}/timing")
 
     MPI.COMM_WORLD.Barrier()
 
-    print("{loop:16.10e}   {output:16.10e}   {setup:16.10e}   {diagnostic:16.10e}".
-          format(loop=full_loop_time, output=output_time, setup=setup_time,
-                 diagnostic=diagnostic_time),
-          file=open("timing/{}_l2Test{}.txt".format(MPI.COMM_WORLD.Get_size(), rank), "w"))
+    mpi_time = distribFunc._layout_manager._mpi_time
+    transp_time = distribFunc._layout_manager._transpose_time - mpi_time
+
+    timeOther = full_loop_time - sum((timeQN, timeFluxAdv, timevParAdv, timePolAdv, transp_time, mpi_time))
+    with open(f"{foldername}/timing/{MPI.COMM_WORLD.Get_size()}_l2Test{rank}.txt", "w") as timing_file:
+        for t in (timeQN, timeFluxAdv, timevParAdv, timePolAdv, transp_time, mpi_time, timeOther, full_loop_time, output_time, setup_time, diagnostic_time):
+            print(f"{t:16.10e}   ", end="", file=timing_file)
+        print(file=timing_file)
+
+    if rank == 0:
+        print("Timings")
+        print("-------")
+        print("QN  | Flux | vPar | Pol | MPI | transpOther | Total")
+        for t in (timeQN, timeFluxAdv, timevParAdv, timePolAdv, mpi_time, transp_time, timeOther):
+            print(round(t / full_loop_time, 3), end=' | ')
+        print(1)
+        print()
+        for t in (timeQN, timeFluxAdv, timevParAdv, timePolAdv, timeOther):
+            print(t, end=' | ')
+        print(full_loop_time)
 
 
 if __name__ == "__main__":

pygyro/model/grid.py

@@ -94,6 +94,16 @@ def getGlobalIndices(self, *indices: 'ints'):
             result[self._layout.dims_order[i]] = indices[i]+toAdd
         return result
 
+    def get3DSlice(self, *slices: 'ints'):
+        """ get the 3D slice at the provided list of coordinates
+        """
+        assert len(slices) == self._nDims-3
+        slices = slices + (slice(self._nGlobalCoords[self._layout.dims_order[-3]]),
+                           slice(
+                               self._nGlobalCoords[self._layout.dims_order[-2]]),
+                           slice(self._nGlobalCoords[self._layout.dims_order[-1]]))
+        return self._f[tuple(slices)]
+
     def get2DSlice(self, *slices: 'ints'):
         """ get the 2D slice at the provided list of coordinates
         """