Add functionality to print patch statistics in ASMStarPC and friends (#…

…3875) * Add functionality to print patch statistics in ASMStarPC and friends --------- Co-authored-by: Pablo Brubeck <[email protected]>
firedrakeproject · Jan 15, 2025 · cd68f9d · cd68f9d
1 parent cbfc2b2
commit cd68f9d
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git 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
@@ -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",