diff --git a/doc/generate_profiling_figures.sh b/doc/generate_profiling_figures.sh index cd0beb0b8..e80524794 100755 --- a/doc/generate_profiling_figures.sh +++ b/doc/generate_profiling_figures.sh @@ -30,8 +30,7 @@ if needle not in content: content = content.replace( needle, "env = EnvironmentOptions(\n" - " profiling_activated=True,\n" - " profiling_trace=True,\n" + " deactivate_profiling=False,\n" ")", ) open(path, "w").write(content) diff --git a/doc/sections/userguide.rst b/doc/sections/userguide.rst index 071b6914b..27d2b20bd 100644 --- a/doc/sections/userguide.rst +++ b/doc/sections/userguide.rst @@ -780,13 +780,52 @@ Struphy's simulation-wide profiler is configured in The relevant switches live in :class:`~struphy.EnvironmentOptions`: 1. ``profiling_activated=True`` enables profiling data collection. -2. ``profiling_trace=True`` additionally records a time trace of profiling - regions. The profiler is set up automatically in ``Simulation.__init__()`` and finalized when ``Simulation.run()`` finishes. The simulation code already wraps key work -inside regions such as ``model.integrate`` via -``ProfileManager.profile_region(...)``. +inside regions via ``ProfileManager.profile_region(...)``. Since the profiler is +active from the end of ``Simulation.__init__()``, the setup phase is covered as +well, and the following regions are recorded out of the box: + +1. Setup: ``setup: allocate`` (total allocation time), with the nested regions + ``setup: feec`` (``setup: derham``, ``setup: mass ops``, ``setup: basis ops``, + ``setup: projected equil``), ``setup: variables`` (one + ``setup var: .`` region per model variable, so that e.g. + marker drawing shows up per particle species), ``setup: propagators`` (one + ``setup prop: `` region per propagator) and + ``setup: helpers``. +2. Remaining run preparation: ``setup: run metadata``, ``setup: data storage``, + ``setup: geometry vtk``, ``setup: plasma params``, + ``setup: initial diagnostics``, ``setup: hdf5 datasets`` and, for restarted + runs, ``setup: restart``. +3. Time loop: ``model.integrate``, ``diagnostics``, ``save data`` and + ``sort particles``. + +Inside ``model.integrate`` the regions nest as follows: + +1. ``prop: ``, one per propagator call (twice per step for the + half steps of Strang splitting). +2. Particle pushing: ``pusher: `` for a full + :class:`~struphy.pic.pushing.pusher.Pusher` call, containing one + ``kernel: `` region per pusher, init and eval kernel call. +3. Accumulation: ``accum: `` for a full + :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` call, + containing the ``kernel: `` region of the accumulation kernel + and ``accum comm: `` for the assembly/ghost-region exchange and + the inter-clone ``Allreduce``. +4. Particle bookkeeping and communication, recorded wherever they are called + from: ``mpi_sort_markers``, ``apply_kinetic_bc``, ``put_particles_in_boxes`` + and ``do_sort``. +5. Linear solves: ``solve: SchurSolver``, ``solve: SchurSolverFull``, + ``solve: SchurSolverFull3``, ``solve: SaddlePointSolver``, + ``solve: ODEsolverFEEC`` for the shared solver classes, and + ``solve: `` for propagators that call a + ``feectools`` inverse operator directly. +6. ``update_feec_variables`` for writing back FEEC coefficients (includes the + ghost-region update). + +Since regions nest, the sum over all regions exceeds the wall-clock time; use +the flame graph (below) to read the containment. Example configuration: diff --git a/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py b/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py index cc10abe4d..c087731bc 100644 --- a/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py +++ b/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py @@ -76,7 +76,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_1",profiling_activated=True, profiling_trace=True, restart=False) +env = EnvironmentOptions(sim_folder="sim_1",profiling_activated=True, restart=False) # Time stepping time_opts = Time(dt=0.001, Tend=0.01, split_algo="LieTrotter") diff --git a/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py b/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py index c3ecb2368..0ad798b95 100644 --- a/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py +++ b/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py @@ -78,7 +78,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_1", profiling_activated=True, profiling_trace=True, restart=False) +env = EnvironmentOptions(sim_folder="sim_1", profiling_activated=True, restart=False) # Time stepping time_opts = Time(dt=5.0, Tend=500.0, split_algo="LieTrotter") diff --git a/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py b/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py index 39f9e4819..1ae5f7bc0 100644 --- a/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py +++ b/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py @@ -74,7 +74,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_1", profiling_activated=True, profiling_trace=True, restart=False) +env = EnvironmentOptions(sim_folder="sim_1", profiling_activated=True, restart=False) # Time stepping time_opts = Time(dt=0.01, Tend=51.0, split_algo="LieTrotter") diff --git a/params_LinearMHDDriftkineticCC.py b/params_LinearMHDDriftkineticCC.py deleted file mode 100644 index f9f5ecd72..000000000 --- a/params_LinearMHDDriftkineticCC.py +++ /dev/null @@ -1,116 +0,0 @@ -# import model, set verbosity -from struphy.models.hybrid import LinearMHDDriftkineticCC - -from struphy import main -from struphy.fields_background import equils -from struphy.geometry import domains -from struphy.initial import perturbations -from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time -from struphy.kinetic_background import maxwellians -from struphy.pic.utilities import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, -) -from struphy.topology import grids - -# environment options -env = EnvironmentOptions() - -# units -base_units = BaseUnits() - -# time stepping -time_opts = Time() - -# geometry -domain = domains.Cuboid() - -# fluid equilibrium (can be used as part of initial conditions) -equil = equils.HomogenSlab() - -# grid -grid = grids.TensorProductGrid(Nel=(16, 16, 16)) - -# derham options -derham_opts = DerhamOptions() - -# light-weight model instance -model = LinearMHDDriftkineticCC() - -# species parameters -model.mhd.set_phys_params() -model.energetic_ions.set_phys_params() - -loading_params = LoadingParameters(ppc=1000) -weights_params = WeightsParameters() -boundary_params = BoundaryParameters() -model.energetic_ions.set_markers( - loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, -) -model.energetic_ions.set_sorting_boxes() -model.energetic_ions.set_save_data() - -# propagator options -model.propagators.push_bxe.options = model.propagators.push_bxe.Options( - b_tilde=model.em_fields.b_field, -) -model.propagators.push_parallel.options = model.propagators.push_parallel.Options( - b_tilde=model.em_fields.b_field, -) -model.propagators.shearalfen_cc5d.options = model.propagators.shearalfen_cc5d.Options( - energetic_ions=model.energetic_ions.var, -) -model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( - b_field=model.em_fields.b_field, -) -model.propagators.cc5d_density.options = model.propagators.cc5d_density.Options( - energetic_ions=model.energetic_ions.var, - b_tilde=model.em_fields.b_field, -) -model.propagators.cc5d_gradb.options = model.propagators.cc5d_gradb.Options( - b_tilde=model.em_fields.b_field, -) -model.propagators.cc5d_curlb.options = model.propagators.cc5d_curlb.Options( - b_tilde=model.em_fields.b_field, -) - -# background, perturbations and initial conditions -model.mhd.velocity.add_background(FieldsBackground()) -model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=0)) -model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=1)) -model.mhd.velocity.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=2)) -maxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, None), equil=equil) -maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None), equil=equil) -background = maxwellian_1 + maxwellian_2 -model.energetic_ions.var.add_background(background) - -# if .add_initial_condition is not called, the background is the kinetic initial condition -perturbation = perturbations.TorusModesCos() -maxwellian_1pt = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil) -init = maxwellian_1pt + maxwellian_2 -model.energetic_ions.var.add_initial_condition(init) - -# optional: exclude variables from saving -# model.energetic_ions.var.save_data = False - -if __name__ == "__main__": - # start run - verbose = True - - main.run( - model, - params_path=__file__, - env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) diff --git a/profiling/examples/Poisson/cube_strong_scaling/params_poisson.py b/profiling/examples/Poisson/cube_strong_scaling/params_poisson.py index b479bcd67..6f609e3dd 100644 --- a/profiling/examples/Poisson/cube_strong_scaling/params_poisson.py +++ b/profiling/examples/Poisson/cube_strong_scaling/params_poisson.py @@ -65,7 +65,6 @@ env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", profiling_activated=True, - profiling_trace=True, restart=False, ) diff --git a/profiling/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py b/profiling/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py index f7624f1a4..1ab145ebc 100644 --- a/profiling/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py +++ b/profiling/examples/ToyGyrokinetic/diocotron_instability/params_diocotron.py @@ -85,7 +85,6 @@ env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", profiling_activated=True, - profiling_trace=True, restart=False ) diff --git a/pyproject.toml b/pyproject.toml index 11cab22a9..9a928bdf3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ dependencies = [ "pytest-testmon<=2.2.0", "ruff==0.15.0, <=0.16.0", "line_profiler<=5.0.2", - "scope-profiler==0.2.6, <=0.2.6", + "scope-profiler<=0.2.8", ] [project.license] diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index bb849b2bf..84a7b7515 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -315,6 +315,9 @@ class EnvironmentOptions(OptionsBase): Folder in ``out_folders/`` for the current simulation (default= ``sim_1/`` ). Will create the folder if it does not exist OR cleans the folder for new runs. + sim_label: str | None, optional + Label for the simulation (default=None) + restart : bool Whether to restart a run (default=False). @@ -332,20 +335,17 @@ class EnvironmentOptions(OptionsBase): profiling_activated: bool, optional Activate profiling with scope-profiler (default=False) - - profiling_trace: bool, optional - Save time-trace of each profiling region (default=False) """ out_folders: str = os.getcwd() sim_folder: str = "sim_1" + sim_label: str | None = None restart: bool = False max_runtime: int = 300 save_step: int = 1 sort_step: int = 0 num_clones: int = 1 profiling_activated: bool = False - profiling_trace: bool = False def __post_init__(self): self.path_out: str = os.path.join(self.out_folders, self.sim_folder) diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index 40430489e..2bdbc05bd 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -7,6 +7,7 @@ from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace from feectools.linalg.direct_solvers import SparseSolver from feectools.linalg.solvers import inverse +from scope_profiler import ProfileManager from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms @@ -249,6 +250,7 @@ def Apre(self, a): elif self._variant == "Inverse_Solver": self._Apre = a + @ProfileManager.profile("solve: SaddlePointSolver") def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): """ Solves the saddle-point problem using the Uzawa algorithm. diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index 8789b6f08..36c0c7956 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -2,6 +2,7 @@ from feectools.linalg.block import BlockLinearOperator, BlockVector from feectools.linalg.solvers import inverse from line_profiler import profile +from scope_profiler import ProfileManager from struphy.linear_algebra.solver import SolverParameters @@ -108,6 +109,7 @@ def BC(self, bc): self._BC = bc @profile + @ProfileManager.profile("solve: SchurSolver") def __call__(self, xn, Byn, dt, out=None): """Solves the 2x2 block matrix linear system. @@ -228,6 +230,7 @@ def __init__(self, M, solver_name, **solver_params): self._rhs = self._A.codomain.zeros() @profile + @ProfileManager.profile("solve: SchurSolverFull") def dot(self, v, out=None): """Solves the 2x2 block matrix linear system. @@ -346,6 +349,7 @@ def __init__(self, M, solver_name, **solver_params): self._rhs2 = self._A.codomain.zeros() @profile + @ProfileManager.profile("solve: SchurSolverFull3") def dot(self, v, out=None): """Solves the 3x3 block matrix linear system. diff --git a/src/struphy/models/tests/utils_testing.py b/src/struphy/models/tests/utils_testing.py index 4d205e650..f9ea3fa89 100644 --- a/src/struphy/models/tests/utils_testing.py +++ b/src/struphy/models/tests/utils_testing.py @@ -55,7 +55,6 @@ def call_test(model: StruphyModel, test_profiling: bool = False): out_folders=test_folder, sim_folder=f"{model_name}", profiling_activated=test_profiling, - profiling_trace=test_profiling, ) # read parameters diff --git a/src/struphy/ode/solvers.py b/src/struphy/ode/solvers.py index 8dce21d41..ed89bd098 100644 --- a/src/struphy/ode/solvers.py +++ b/src/struphy/ode/solvers.py @@ -3,6 +3,7 @@ import cunumpy as xp from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector +from scope_profiler import ProfileManager from struphy.ode.utils import ButcherTableau @@ -59,6 +60,7 @@ def __init__( self._yn = [v.copy() for v in self.y] self._ystar = [v.copy() for v in self.y] + @ProfileManager.profile("solve: ODEsolverFEEC") def __call__(self, tn, h): a = self.butcher.a b = self.butcher.b diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 062f6056e..0deedc6d4 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -111,6 +111,10 @@ def __init__( self._derham = mass_ops.derham self._args_domain = args_domain + # profiling region names (precomputed, they are looked up on every call) + self._region_name = "accum: " + kernel.name + self._comm_region_name = "accum comm: " + kernel.name + self._symmetry = symmetry self._form = self.derham.space_to_form[space_id] @@ -210,6 +214,11 @@ def __call__(self, *optional_args, **args_control): args_control : any Keyword arguments for an analytical control variate correction in the accumulation step. Possible keywords are 'control_vec' for a vector correction or 'control_mat' for a matrix correction. Values are a 1d (vector) or 2d (matrix) list with callables or xp.ndarrays used for the correction. """ + with ProfileManager.profile_region(self._region_name): + self._accumulate(*optional_args, **args_control) + + def _accumulate(self, *optional_args, **args_control): + """Body of :meth:`__call__`, see there.""" # flags for break vec_finished = False @@ -232,8 +241,9 @@ def __call__(self, *optional_args, **args_control): # apply filter if self.accfilter.params.use_filter is not None: for vec in self._vectors: - vec.exchange_assembly_data() - vec.update_ghost_regions() + with ProfileManager.profile_region(self._comm_region_name): + vec.exchange_assembly_data() + vec.update_ghost_regions() self.accfilter(vec) vec_finished = True @@ -244,12 +254,13 @@ def __call__(self, *optional_args, **args_control): num_clones = self.particles.clone_config.num_clones if num_clones > 1: - for data_array in self._args_data: - self.particles.clone_config.inter_comm.Allreduce( - MPI.IN_PLACE, - data_array, - op=MPI.SUM, - ) + with ProfileManager.profile_region(self._comm_region_name): + for data_array in self._args_data: + self.particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + data_array, + op=MPI.SUM, + ) # add analytical contribution (control variate) to vector if "control_vec" in args_control and len(self._vectors) > 0: @@ -270,15 +281,17 @@ def __call__(self, *optional_args, **args_control): # finish vector: accumulate ghost regions and update ghost regions if not vec_finished: - for vec in self._vectors: - vec.exchange_assembly_data() - vec.update_ghost_regions() + with ProfileManager.profile_region(self._comm_region_name): + for vec in self._vectors: + vec.exchange_assembly_data() + vec.update_ghost_regions() # finish matrix: accumulate ghost regions, update ghost regions and copy data for symmetric/antisymmetric block matrices if not mat_finished: - for op in self._operators: - op.matrix.exchange_assembly_data() - op.matrix.update_ghost_regions() + with ProfileManager.profile_region(self._comm_region_name): + for op in self._operators: + op.matrix.exchange_assembly_data() + op.matrix.update_ghost_regions() if self.symmetry == "symm": self._operators[0].matrix[0, 1].transpose( @@ -492,6 +505,10 @@ def __init__( self._derham = mass_ops.derham self._args_domain = args_domain + # profiling region names (precomputed, they are looked up on every call) + self._region_name = "accum: " + kernel.name + self._comm_region_name = "accum comm: " + kernel.name + self._form = self.derham.space_to_form[space_id] # initialize vectors @@ -558,6 +575,11 @@ def __call__(self, *optional_args, **args_control): Possible keywords are 'control_vec' for a vector correction or 'control_mat' for a matrix correction. Values are a 1d (vector) or 2d (matrix) list with callables or xp.ndarrays used for the correction. """ + with ProfileManager.profile_region(self._region_name): + self._accumulate(*optional_args, **args_control) + + def _accumulate(self, *optional_args, **args_control): + """Body of :meth:`__call__`, see there.""" # flags for break vec_finished = False @@ -579,8 +601,9 @@ def __call__(self, *optional_args, **args_control): # apply filter if self.accfilter.params.use_filter is not None: for vec in self._vectors: - vec.exchange_assembly_data() - vec.update_ghost_regions() + with ProfileManager.profile_region(self._comm_region_name): + vec.exchange_assembly_data() + vec.update_ghost_regions() self.accfilter(vec) vec_finished = True @@ -591,12 +614,13 @@ def __call__(self, *optional_args, **args_control): num_clones = self.particles.clone_config.num_clones if num_clones > 1: - for data_array in self._args_data: - self.particles.clone_config.inter_comm.Allreduce( - MPI.IN_PLACE, - data_array, - op=MPI.SUM, - ) + with ProfileManager.profile_region(self._comm_region_name): + for data_array in self._args_data: + self.particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + data_array, + op=MPI.SUM, + ) # add analytical contribution (control variate) to vector if "control_vec" in args_control and len(self._vectors) > 0: @@ -609,9 +633,10 @@ def __call__(self, *optional_args, **args_control): # finish vector: accumulate ghost regions and update ghost regions if not vec_finished: - for vec in self._vectors: - vec.exchange_assembly_data() - vec.update_ghost_regions() + with ProfileManager.profile_region(self._comm_region_name): + for vec in self._vectors: + vec.exchange_assembly_data() + vec.update_ghost_regions() @property def particles(self): diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index e5f8110d0..ebcf50fb0 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -20,6 +20,7 @@ class Intracomm: from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI from line_profiler import profile +from scope_profiler import ProfileManager from sympy.ntheory import factorint from struphy.bsplines.bsplines import quadrature_grid @@ -1622,6 +1623,7 @@ def show_distribution_function(self, components, bin_edges): plt.show() @profile + @ProfileManager.profile("mpi_sort_markers") def mpi_sort_markers( self, apply_bc: bool = True, @@ -1703,6 +1705,7 @@ def mpi_sort_markers( self._Barrier() @profile + @ProfileManager.profile("apply_kinetic_bc") def apply_kinetic_bc(self, newton=False): """ Apply boundary conditions to markers that are outside of the logical unit cube. @@ -1810,6 +1813,7 @@ def set_velocities_comp(self, velocity, comp): self._markers[self.valid_mks, slice(3 + c, 3 + c + 1)] = new @profile + @ProfileManager.profile("put_particles_in_boxes") def put_particles_in_boxes(self): """Assign the right box to the particles and the list of the particles to each box. If sorting_boxes was instantiated with an MPI comm, then the particles in the @@ -1840,6 +1844,7 @@ def put_particles_in_boxes(self): # logger.info(f"Number of markers in box {i} is {n_mks_box}") @profile + @ProfileManager.profile("do_sort") def do_sort(self, use_numpy_argsort=False): """Assign the particles to their sorting boxes and reorder the markers array accordingly, so that markers in the same box occupy contiguous rows. diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 2ce1ff007..a224a5dcf 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -14,6 +14,11 @@ logger = logging.getLogger("struphy") +def _kernel_name(kernel) -> str: + """Name of a pyccelized kernel, which can be a bare pyccel function or a PyccelKernel.""" + return getattr(kernel, "name", None) or getattr(kernel, "__name__", type(kernel).__name__) + + class Pusher: r""" Class for solving particle ODEs @@ -153,6 +158,10 @@ def __init__( self._init_kernels = init_kernels self._eval_kernels = eval_kernels + # profiling region names (cached, they are looked up on every call) + self._region_name = "pusher: " + self.kernel.name + self._kernel_region_names = {} + self._residuals = xp.zeros(self.particles.markers.shape[0]) self._converged_loc = self._residuals == 1.0 self._not_converged_loc = self._residuals == 0.0 @@ -168,6 +177,19 @@ def __call__(self, dt: float): Applies the chosen pusher kernel by a time step dt, applies kinetic boundary conditions and performs MPI sorting. """ + with ProfileManager.profile_region(self._region_name): + self._push(dt) + + def _kernel_region(self, kernel) -> str: + """Cached name of the profiling region of an init/eval kernel.""" + name = self._kernel_region_names.get(id(kernel)) + if name is None: + name = "kernel: " + _kernel_name(kernel) + self._kernel_region_names[id(kernel)] = name + return name + + def _push(self, dt: float): + """Body of :meth:`__call__`, see there.""" # some idx and slice markers = self.particles.markers @@ -203,14 +225,15 @@ def __call__(self, dt: float): comps = ker_args[2] add_args = ker_args[3] - ker( - xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), - column_nr, - comps, - self.particles.args_markers, - self._args_domain, - *add_args, - ) + with ProfileManager.profile_region(self._kernel_region(ker)): + ker( + xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + column_nr, + comps, + self.particles.args_markers, + self._args_domain, + *add_args, + ) # update boxes if self._box_comm: @@ -250,14 +273,15 @@ def __call__(self, dt: float): ) # evaluate - ker( - alpha, - column_nr, - comps, - self.particles.args_markers, - self._args_domain, - *add_args, - ) + with ProfileManager.profile_region(self._kernel_region(ker)): + ker( + alpha, + column_nr, + comps, + self.particles.args_markers, + self._args_domain, + *add_args, + ) # update boxes if self._box_comm: diff --git a/src/struphy/propagators/adiabatic_phi.py b/src/struphy/propagators/adiabatic_phi.py index c0971ceea..9dbb28380 100644 --- a/src/struphy/propagators/adiabatic_phi.py +++ b/src/struphy/propagators/adiabatic_phi.py @@ -2,6 +2,7 @@ from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.feec.mass import L2Projector, WeightedMassOperator, WeightedMassOperators @@ -183,7 +184,8 @@ def __call__(self, dt): self._rhs += self._rho # solve - out = self._solver.solve(self._rhs, out=self._tmp) + with ProfileManager.profile_region(self._solve_region): + out = self._solver.solve(self._rhs, out=self._tmp) info = self._solver._info if self._lin_solver["info"]: diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 00e135ce8..b7a81272e 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -8,6 +8,7 @@ import cunumpy as xp from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector +from scope_profiler import ProfileManager from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators @@ -96,12 +97,20 @@ def __call__(self, dt: float): Time step size. """ + @property + def _solve_region(self) -> str: + """Name of the profiling region for the linear solve(s) of this propagator.""" + if not hasattr(self, "_solve_region_name"): + self._solve_region_name = "solve: " + self.__class__.__name__ + return self._solve_region_name + def show_options(self): """Print the options of the propagator.""" logger.info(f"\nOptions for propagator '{self.__class__.__name__}':") for k, v in self.options.__dict__.items(): logger.info(f" {k + ':':<20}{v}") + @ProfileManager.profile("update_feec_variables") def update_feec_variables(self, **new_coeffs): r"""Return max_diff = max(abs(new - old)) for each new_coeffs, update feec coefficients and update ghost regions. diff --git a/src/struphy/propagators/curl_curl_solve.py b/src/struphy/propagators/curl_curl_solve.py index 48a7967fe..0e9773c61 100644 --- a/src/struphy/propagators/curl_curl_solve.py +++ b/src/struphy/propagators/curl_curl_solve.py @@ -8,6 +8,7 @@ from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec.mass import L2Projector, WeightedMassOperator from struphy.io.options import LiteralOptions @@ -324,7 +325,8 @@ def __call__(self, dt): self._solver.linop = self._diffusion_op - self._sigma * self._stab_mat # solve - out = self._solver.solve(self._rhs, out=self._tmp) + with ProfileManager.profile_region(self._solve_region): + out = self._solver.solve(self._rhs, out=self._tmp) info = self._solver._info if self._info: diff --git a/src/struphy/propagators/current_coupling_5d_density.py b/src/struphy/propagators/current_coupling_5d_density.py index 3f3c858d8..a17d0b452 100644 --- a/src/struphy/propagators/current_coupling_5d_density.py +++ b/src/struphy/propagators/current_coupling_5d_density.py @@ -5,6 +5,7 @@ from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.io.options import LiteralOptions, OptionsBase @@ -230,7 +231,8 @@ def __call__(self, dt): rhs = rhs.dot(un, out=self._rhs_v) self._A_inv.linop = lhs - _u = self._A_inv.solve(rhs, out=self._u_new) + with ProfileManager.profile_region(self._solve_region): + _u = self._A_inv.solve(rhs, out=self._u_new) info = self._A_inv._info diffs = self.update_feec_variables(u=_u) diff --git a/src/struphy/propagators/current_coupling_6d_density.py b/src/struphy/propagators/current_coupling_6d_density.py index 19ac34394..828625ff0 100644 --- a/src/struphy/propagators/current_coupling_6d_density.py +++ b/src/struphy/propagators/current_coupling_6d_density.py @@ -5,6 +5,7 @@ from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.io.options import LiteralOptions, OptionsBase @@ -286,7 +287,8 @@ def __call__(self, dt): rhs = rhs.dot(un, out=self._rhs_v) self._solver.linop = lhs - un1 = self._solver.solve(rhs, out=self._u_new) + with ProfileManager.profile_region(self._solve_region): + un1 = self._solver.solve(rhs, out=self._u_new) info = self._solver._info # write new coeffs into Propagator.variables diff --git a/src/struphy/propagators/hall.py b/src/struphy/propagators/hall.py index e3ef95fd3..03a59a719 100644 --- a/src/struphy/propagators/hall.py +++ b/src/struphy/propagators/hall.py @@ -4,6 +4,7 @@ from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.io.options import LiteralOptions, OptionsBase @@ -165,7 +166,8 @@ def __call__(self, dt): rhs = rhs.dot(bn, out=self._rhs_b) self._solver.linop = lhs - bn1 = self._solver.solve(rhs, out=self._b_new) + with ProfileManager.profile_region(self._solve_region): + bn1 = self._solver.solve(rhs, out=self._b_new) info = self._solver._info # write new coeffs into self.feec_vars diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index fb09cd349..f5a19b6ef 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -7,6 +7,7 @@ from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec.mass import L2Projector, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase @@ -461,7 +462,8 @@ def __call__(self, dt): self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op # solve - out = self._solver.solve(rhs, out=self._tmp) + with ProfileManager.profile_region(self._solve_region): + out = self._solver.solve(rhs, out=self._tmp) info = self._solver._info if self._info: diff --git a/src/struphy/propagators/jxb_cold.py b/src/struphy/propagators/jxb_cold.py index 9b8308dec..da8034939 100644 --- a/src/struphy/propagators/jxb_cold.py +++ b/src/struphy/propagators/jxb_cold.py @@ -3,6 +3,7 @@ from feectools.linalg.solvers import inverse from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.io.options import LiteralOptions, OptionsBase @@ -144,7 +145,8 @@ def __call__(self, dt): self._solver.linop = lhs # solve linear system for updated j coefficients (in-place) - jn1 = self._solver.solve(rhsv, out=self._j_new) + with ProfileManager.profile_region(self._solve_region): + jn1 = self._solver.solve(rhsv, out=self._j_new) info = self._solver._info # write new coeffs into Propagator.variables diff --git a/src/struphy/propagators/variational_momentum_advection.py b/src/struphy/propagators/variational_momentum_advection.py index 7d77ea0f0..162e02810 100644 --- a/src/struphy/propagators/variational_momentum_advection.py +++ b/src/struphy/propagators/variational_momentum_advection.py @@ -5,6 +5,7 @@ from feectools.ddm.mpi import mpi as MPI from feectools.linalg.solvers import inverse from line_profiler import profile +from scope_profiler import ProfileManager from struphy.feec import preconditioner from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner @@ -208,7 +209,8 @@ def __call_newton(self, dt): break # Newton step - pc_diff = self._Mrho_inv.dot(diff, out=self._tmp__pc_diff) + with ProfileManager.profile_region(self._solve_region): + pc_diff = self._Mrho_inv.dot(diff, out=self._tmp__pc_diff) update = self.inv_derivative.dot(pc_diff, out=self._tmp_update) if self._info: logger.info( @@ -262,7 +264,8 @@ def __call_picard(self, dt): mn1 -= advection # Inverse the mass matrix to get the velocity - un1 = self._Mrho_inv.dot(mn1, out=self._tmp_un1) + with ProfileManager.profile_region(self._solve_region): + un1 = self._Mrho_inv.dot(mn1, out=self._tmp_un1) if it == self.options.nonlin_solver.maxiter - 1 or xp.isnan(err): logger.info( diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index cf6215fed..ce0fc2f72 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -149,7 +149,7 @@ def __init__( self.name = name self.description = description self.params_path = params_path - self.env = env + self.env = env # Set name first since it's used as a label self.time_opts = time_opts self.domain = domain self.equil = equil @@ -174,17 +174,17 @@ def __init__( # Abstract methods # ---------------- - def _setup_profiling(self): + def _setup_profiling(self, label: str = ""): # setup profiling agent ProfileManager.setup( - profiling_activated=self.env.profiling_activated, - time_trace=self.env.profiling_trace, + deactivate_profiling=not self.env.profiling_activated, use_likwid=False, file_path=os.path.join( self.env.out_folders, self.env.sim_folder, "profiling_data.h5", ), + label=label, ) def show_parameters(self): @@ -239,17 +239,22 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") - # feec - self._allocate_feec(self.grid, self.derham_opts) + with ProfileManager.profile_region("setup: allocate"): + # feec + with ProfileManager.profile_region("setup: feec"): + self._allocate_feec(self.grid, self.derham_opts) - # allocate model variables - self._allocate_variables() + # allocate model variables + with ProfileManager.profile_region("setup: variables"): + self._allocate_variables() - # pass info to propagators - self._allocate_propagators() + # pass info to propagators + with ProfileManager.profile_region("setup: propagators"): + self._allocate_propagators() - # allocate helper fields and perform initial solves if needed - self.model.allocate_helpers() + # allocate helper fields and perform initial solves if needed + with ProfileManager.profile_region("setup: helpers"): + self.model.allocate_helpers() logger.debug("... Done.") @@ -626,16 +631,20 @@ def run(self, one_time_step: bool = False): # equation paramters self.allocate() - self._write_run_metadata(one_time_step=one_time_step) + with ProfileManager.profile_region("setup: run metadata"): + self._write_run_metadata(one_time_step=one_time_step) # output - self.initialize_data_storage() + with ProfileManager.profile_region("setup: data storage"): + self.initialize_data_storage() # peek view into geometry - self.save_geometry_and_equil_vtk() + with ProfileManager.profile_region("setup: geometry vtk"): + self.save_geometry_and_equil_vtk() # plasma parameters - self.compute_plasma_params() + with ProfileManager.profile_region("setup: plasma params"): + self.compute_plasma_params() # print info on mpi procs if self.comm_size < 32: @@ -665,7 +674,8 @@ def run(self, one_time_step: bool = False): # set initial conditions for all variables if self.env.restart: - self._initialize_from_restart(self.data) + with ProfileManager.profile_region("setup: restart"): + self._initialize_from_restart(self.data) with h5py.File(self.data.file_path, "a") as file: self.time_state["value"][0] = file["restart/time/value"][-1] @@ -688,13 +698,15 @@ def run(self, one_time_step: bool = False): total_steps_str = str(total_steps) # compute initial scalars and kinetic data, pass time state to all propagators - self.model.update_scalar_quantities() - self.model.update_markers_to_be_saved() - self.model.update_distr_functions() - self._add_time_state(self.time_state["value"]) + with ProfileManager.profile_region("setup: initial diagnostics"): + self.model.update_scalar_quantities() + self.model.update_markers_to_be_saved() + self.model.update_distr_functions() + self._add_time_state(self.time_state["value"]) # add all variables to be saved to data object - save_keys_all, save_keys_end = self._initialize_hdf5_datasets(self.data, self.comm_size) + with ProfileManager.profile_region("setup: hdf5 datasets"): + save_keys_all, save_keys_end = self._initialize_hdf5_datasets(self.data, self.comm_size) # ======================== main time loop ====================== self.model.update_scalar_quantities() @@ -722,7 +734,8 @@ def run(self, one_time_step: bool = False): if break_cond_1 or break_cond_2: # save restart data (other data already saved below) - self.data.save_data(keys=save_keys_end) + with ProfileManager.profile_region("save data"): + self.data.save_data(keys=save_keys_end) end_time = time.time() logger.info(f"\nTime steps done: {int(self.time_state['index'][0])}") logger.info(f"wall-clock time of simulation [sec]: {end_time - self.start_time}") @@ -731,9 +744,10 @@ def run(self, one_time_step: bool = False): if self.env.sort_step and int(self.time_state["index"][0]) % self.env.sort_step == 0: t0 = time.time() - for key, val in self.model.pointer.items(): - if isinstance(val, Particles): - val.do_sort() + with ProfileManager.profile_region("sort particles"): + for key, val in self.model.pointer.items(): + if isinstance(val, Particles): + val.do_sort() t1 = time.time() message = "Particles sorted | wall clock [s]: {0:8.4f} | sorting duration [s]: {1:8.4f}".format( run_time_now * 60, @@ -760,22 +774,24 @@ def run(self, one_time_step: bool = False): # update diagnostics data and save data if int(self.time_state["index"][0]) % self.env.save_step == 0: # compute scalars and kinetic data - self.model.update_scalar_quantities() - self.model.update_markers_to_be_saved() - self.model.update_distr_functions() - - # extract FEEC coefficients - feec_species = self.model.field_species | self.model.fluid_species | self.model.diagnostic_species - for species, val in feec_species.items(): - assert isinstance(val, Species) - for variable, subval in val.variables.items(): - assert isinstance(subval, FEECVariable) - spline = subval.spline - # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - spline.extract_coeffs(update_ghost_regions=False) + with ProfileManager.profile_region("diagnostics"): + self.model.update_scalar_quantities() + self.model.update_markers_to_be_saved() + self.model.update_distr_functions() + + # extract FEEC coefficients + feec_species = self.model.field_species | self.model.fluid_species | self.model.diagnostic_species + for species, val in feec_species.items(): + assert isinstance(val, Species) + for variable, subval in val.variables.items(): + assert isinstance(subval, FEECVariable) + spline = subval.spline + # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! + spline.extract_coeffs(update_ghost_regions=False) # save data (everything but restart data) - self.data.save_data(keys=save_keys_all) + with ProfileManager.profile_region("save data"): + self.data.save_data(keys=save_keys_all) # print current time and scalar quantities to screen step = str(int(self.time_state["index"][0])).zfill(len(total_steps_str)) @@ -832,7 +848,29 @@ def run(self, one_time_step: bool = False): if self.clone_config is not None: self.clone_config.free() - ProfileManager.finalize() + # ProfileManager.finalize(verbose=True) + results = ProfileManager.finalize(return_results=True, verbose=False) + + # one table per region family; the last group catches everything not matched above, + # so that no recorded region is silently missing from the printed summary + groups = ( + ("Setup", [r"^setup:", r"^setup prop:", r"^setup var:"]), + ("Model propagation", [r"^model\.integrate", r"^prop:"]), + ("Pusher", [r"^pusher:"]), + ("Kernel", [r"^kernel:"]), + ("Accumulation", [r"^accum:", r"^accum comm:"]), + ("Linear solves", [r"^solve:"]), + ( + "Particle sorting and communication", + [r"^mpi_sort_markers$", r"^apply_kinetic_bc$", r"^put_particles_in_boxes$", r"^do_sort$"], + ), + ) + all_patterns = [pattern for _, include in groups for pattern in include] + for title, include in groups + (("Other", None),): + kwargs = {"include": include} if include is not None else {"exclude": all_patterns} + if not results.get_regions(**kwargs): + continue + results.print_summary(title=title, suppress_notes=True, **kwargs) def pproc( self, @@ -1151,47 +1189,51 @@ def _allocate_feec(self, grid: grids.TensorProductGrid, derham_opts: DerhamOptio logger.debug(f"\n{grid=}, {derham_opts=}: no Derham object set up.") self._derham = None else: - self._derham = Derham( - grid, - derham_opts, - comm=derham_comm, - domain=self.domain, - ) + with ProfileManager.profile_region("setup: derham"): + self._derham = Derham( + grid, + derham_opts, + comm=derham_comm, + domain=self.domain, + ) # create weighted mass and basis operators if self.derham is None: self._mass_ops = None self._basis_ops = None else: - self._mass_ops = WeightedMassOperators(self.derham, self.domain, eq_mhd=self.equil) + with ProfileManager.profile_region("setup: mass ops"): + self._mass_ops = WeightedMassOperators(self.derham, self.domain, eq_mhd=self.equil) - self._basis_ops = BasisProjectionOperators( - self.derham, - self.domain, - eq_mhd=self.equil, - ) + with ProfileManager.profile_region("setup: basis ops"): + self._basis_ops = BasisProjectionOperators( + self.derham, + self.domain, + eq_mhd=self.equil, + ) # create projected equilibrium if self.derham is None: self._projected_equil = None else: - if isinstance(self.equil, MHDequilibrium): - self._projected_equil = ProjectedMHDequilibrium( - self.equil, - self.derham, - ) - elif isinstance(self.equil, FluidEquilibriumWithB): - self._projected_equil = ProjectedFluidEquilibriumWithB( - self.equil, - self.derham, - ) - elif isinstance(self.equil, FluidEquilibrium): - self._projected_equil = ProjectedFluidEquilibrium( - self.equil, - self.derham, - ) - else: - self._projected_equil = None + with ProfileManager.profile_region("setup: projected equil"): + if isinstance(self.equil, MHDequilibrium): + self._projected_equil = ProjectedMHDequilibrium( + self.equil, + self.derham, + ) + elif isinstance(self.equil, FluidEquilibriumWithB): + self._projected_equil = ProjectedFluidEquilibriumWithB( + self.equil, + self.derham, + ) + elif isinstance(self.equil, FluidEquilibrium): + self._projected_equil = ProjectedFluidEquilibrium( + self.equil, + self.derham, + ) + else: + self._projected_equil = None @profile def _allocate_variables(self): @@ -1204,11 +1246,12 @@ def _allocate_variables(self): assert isinstance(spec, FieldSpecies) for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) - v.allocate( - derham=self.derham, - domain=self.domain, - equil=self.equil, - ) + with ProfileManager.profile_region(f"setup var: {species}.{k}"): + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + ) # allocate memory for FE coeffs of fluid variables if self.model.fluid_species: @@ -1216,11 +1259,12 @@ def _allocate_variables(self): assert isinstance(spec, FluidSpecies) for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) - v.allocate( - derham=self.derham, - domain=self.domain, - equil=self.equil, - ) + with ProfileManager.profile_region(f"setup var: {species}.{k}"): + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + ) # allocate memory for marker arrays of kinetic variables if self.model.particle_species: @@ -1228,20 +1272,22 @@ def _allocate_variables(self): assert isinstance(spec, ParticleSpecies) for k, v in spec.variables.items(): if isinstance(v, PICVariable): - v.allocate( - clone_config=self.clone_config, - derham=self.derham, - domain=self.domain, - equil=self.equil, - projected_equil=self.projected_equil, - ) + with ProfileManager.profile_region(f"setup var: {species}.{k}"): + v.allocate( + clone_config=self.clone_config, + derham=self.derham, + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + ) if isinstance(v, SPHVariable): - v.allocate( - derham=self.derham, - domain=self.domain, - equil=self.equil, - projected_equil=self.projected_equil, - ) + with ProfileManager.profile_region(f"setup var: {species}.{k}"): + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + ) # allocate memory for FE coeffs of fluid variables if self.model.diagnostic_species: @@ -1249,11 +1295,12 @@ def _allocate_variables(self): assert isinstance(spec, DiagnosticSpecies) for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) - v.allocate( - derham=self.derham, - domain=self.domain, - equil=self.equil, - ) + with ProfileManager.profile_region(f"setup var: {species}.{k}"): + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + ) # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: @@ -1291,7 +1338,8 @@ def _allocate_propagators(self): assert len(self.model.prop_list) > 0, "No propagators in this model, check the model class." for prop in self.model.prop_list: assert isinstance(prop, Propagator) - prop.allocate() + with ProfileManager.profile_region("setup prop: " + prop.__class__.__name__): + prop.allocate() logger.debug(f"\nAllocated propagator '{prop.__class__.__name__}'.") @profile @@ -1782,7 +1830,7 @@ def env(self, value: EnvironmentOptions): # create output folders self._setup_folders() - self._setup_profiling() + self._setup_profiling(label=self.name) @property def time_opts(self):