Add functionality to print patch statistics in ASMStarPC and friends

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,16 @@ def initialize(self, pc):
         asmpc.setFromOptions()
         self.asmpc = asmpc
 
+        self._patch_statistics = []
+        if opts.getBool("print_patch_statistics", default=False):
+            # Compute and stash patch statistics
+            max_local_patch = max(is_.getSize() for is_ in ises)
+            min_local_patch = min(is_.getSize() for is_ in ises)
+            max_global_patch = pc.comm.tompi4py().allreduce(max_local_patch, op=MPI.MAX)
+            min_global_patch = pc.comm.tompi4py().allreduce(min_local_patch, op=MPI.MIN)
+            msg = f"Minimum / maximum patch sizes: {min_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 +114,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

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,6 +51,7 @@ 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_print_patch_statistics": None,
                        "mg_levels_pc_star_construct_dim": 0}
 
         comp_params = {"mat_type": "aij",