compiler: Fix CustomTopology

Author: FabioLuporini

devito/mpi/distributed.py

@@ -261,7 +261,7 @@ def nprocs_local(self):
 
     @property
     def topology(self):
-        return DimensionTuple(*self._topology, getters=self.dimensions)
+        return self._topology
 
     @property
     def topology_logical(self):
@@ -353,7 +353,9 @@ def __init__(self, shape, dimensions, input_comm=None, topology=None):
                 self._topology = compute_dims(self._input_comm.size, len(shape))
             else:
                 # A custom topology may contain integers or the wildcard '*'
-                self._topology = CustomTopology(topology, self._input_comm)
+                self._topology = CustomTopology(
+                    topology, self._input_comm, getters=dimensions
+                )
 
             if self._input_comm is not input_comm:
                 # By default, Devito arranges processes into a cartesian topology.
@@ -896,7 +898,7 @@ def _arg_values(self, *args, **kwargs):
             return self._arg_defaults()
 
 
-class CustomTopology(tuple):
+class CustomTopology(DimensionTuple):
 
     """
     The CustomTopology class provides a mechanism to describe parametric domain
@@ -954,7 +956,7 @@ class CustomTopology(tuple):
         'xy': ('*', '*', 1),
     }
 
-    def __new__(cls, items, input_comm):
+    def __new__(cls, items, input_comm, **kwargs):
         # Keep track of nstars and already defined decompositions
         nstars = items.count('*')
 
@@ -992,11 +994,15 @@ def __new__(cls, items, input_comm):
         # Final check that topology matches the communicator size
         assert np.prod(processed) == input_comm.size
 
-        obj = super().__new__(cls, processed)
+        obj = super().__new__(cls, *processed, **kwargs)
         obj.logical = items
 
         return obj
 
+    def __repr__(self):
+        return (f"CustomTopology(logical={self.logical}, "
+                f"physical={super().__repr__()})")
+
 
 def compute_dims(nprocs, ndim):
     # We don't do anything clever here. In fact, we do something very basic --