diff --git a/firedrake/preconditioners/asm.py b/firedrake/preconditioners/asm.py index 19738fe438..f3eaa037d9 100644 --- a/firedrake/preconditioners/asm.py +++ b/firedrake/preconditioners/asm.py @@ -6,6 +6,7 @@ from firedrake.dmhooks import get_function_space from firedrake.logging import warning from tinyasm import _tinyasm as tinyasm +from mpi4py import MPI import numpy @@ -33,7 +34,7 @@ def initialize(self, pc): # Extract function space and mesh to obtain plex and indexing functions V = get_function_space(dm) - # Obtain patches from user defined funtion + # Obtain patches from user defined function ises = self.get_patches(V) # PCASM expects at least one patch, so we define an empty one on idle processes if len(ises) == 0: @@ -90,6 +91,20 @@ def initialize(self, pc): asmpc.setFromOptions() self.asmpc = asmpc + self._patch_statistics = [] + if opts.getBool("view_patch_sizes", default=False): + # Compute and stash patch statistics + mpi_comm = pc.comm.tompi4py() + max_local_patch = max(is_.getSize() for is_ in ises) + min_local_patch = min(is_.getSize() for is_ in ises) + sum_local_patch = sum(is_.getSize() for is_ in ises) + max_global_patch = mpi_comm.allreduce(max_local_patch, op=MPI.MAX) + min_global_patch = mpi_comm.allreduce(min_local_patch, op=MPI.MIN) + sum_global_patch = mpi_comm.allreduce(sum_local_patch, op=MPI.SUM) + avg_global_patch = sum_global_patch / mpi_comm.allreduce(len(ises) if sum_local_patch > 0 else 0, op=MPI.SUM) + msg = f"Minimum / average / maximum patch sizes : {min_global_patch} / {avg_global_patch} / {max_global_patch}\n" + self._patch_statistics.append(msg) + @abc.abstractmethod def get_patches(self, V): ''' Get the patches used for PETSc PCASM @@ -103,6 +118,9 @@ def get_patches(self, V): def view(self, pc, viewer=None): self.asmpc.view(viewer=viewer) + if viewer is not None: + for msg in self._patch_statistics: + viewer.printfASCII(msg) def update(self, pc): # This is required to update an inplace ILU factorization diff --git a/tests/firedrake/regression/test_star_pc.py b/tests/firedrake/regression/test_star_pc.py index ba1f065eca..eb02c665e8 100644 --- a/tests/firedrake/regression/test_star_pc.py +++ b/tests/firedrake/regression/test_star_pc.py @@ -41,6 +41,7 @@ def test_star_equivalence(problem_type, backend): star_params = {"mat_type": "aij", "snes_type": "ksponly", + "snes_view": None, "ksp_type": "richardson", "pc_type": "mg", "pc_mg_type": "multiplicative", @@ -50,7 +51,8 @@ def test_star_equivalence(problem_type, backend): "mg_levels_ksp_max_it": 1, "mg_levels_pc_type": "python", "mg_levels_pc_python_type": "firedrake.ASMStarPC", - "mg_levels_pc_star_construct_dim": 0} + "mg_levels_pc_star_construct_dim": 0, + "mg_levels_pc_star_view_patch_sizes": None} comp_params = {"mat_type": "aij", "snes_type": "ksponly",