mpi: Simplify custom approach using 'factorint' and 'array_split'

devito/mpi/distributed.py

@@ -2,7 +2,7 @@
 from ctypes import c_int, c_void_p, sizeof
 from itertools import groupby, product
 from math import ceil, pow
-from sympy import primefactors
+from sympy import factorint
 
 import atexit
 
@@ -634,16 +634,15 @@ def __new__(cls, items, input_comm):
             alloc_procs = np.prod([i for i in items if i != '*'])
             rem_procs = int(input_comm.size // alloc_procs)
 
-            # Start by using the max prime factor at the first starred position,
-            # then iteratively decompose as evenly as possible until decomposing
-            # to the number of `rem_procs`
-            star_vals = [1] * len(items)
-            star_i = 0
-            while rem_procs > 1:
-                prime_factors = primefactors(rem_procs)
-                rem_procs //= max(prime_factors)
-                star_vals[star_i] *= max(prime_factors)
-                star_i = (star_i + 1) % nstars
+            # List of all factors of rem_procs in decreasing order
+            factors = factorint(rem_procs)
+            vals = [k for (k, v) in factors.items() for _ in range(v)][::-1]
+
+            # Split in number of stars
+            split = np.array_split(vals, nstars)
+
+            # Reduce
+            star_vals = [int(np.prod(s)) for s in split]
 
             # Apply computed star values to the processed
             for index, value in zip(star_pos, star_vals):

tests/test_mpi.py

@@ -219,7 +219,7 @@ def test_custom_topology(self):
         (256, ('*', '*', 2), (16, 8, 2)),
         (256, ('*', 32, 2), (4, 32, 2)),
     ])
-    def test_custom_topology_3d_dummy(self, comm_size, topology, dist_topology):
+    def test_custom_topology_v2(self, comm_size, topology, dist_topology):
         dummy_comm = Bunch(size=comm_size)
         custom_topology = CustomTopology(topology, dummy_comm)
         assert custom_topology == dist_topology