compiler: Tweak CIRE max-par lowering

Author: FabioLuporini

devito/passes/clusters/aliases.py

@@ -139,6 +139,10 @@ def _aliases_from_clusters(self, cgroup, exclude, meta):
         # [Schedule]_m -> Schedule (s.t. best memory/flops trade-off)
         schedule, exprs = self._select(variants)
 
+        # Schedule -> Schedule (optimization)
+        if self.opt_maxpar:
+            schedule = optimize_schedule_maxpar(schedule)
+
         # Schedule -> Schedule (optimization)
         if self.opt_rotate:
             schedule = optimize_schedule_rotations(schedule, self.sregistry)
@@ -664,7 +668,6 @@ def lower_aliases(aliases, meta, maxpar):
     """
     Create a Schedule from an AliasList.
     """
-    stampcache = {}
     dmapper = {}
     processed = []
     for a in aliases:
@@ -704,12 +707,6 @@ def lower_aliases(aliases, meta, maxpar):
             # use `<1>` as stamp, which is what appears in `ispace`
             interval = interval.lift(i.stamp)
 
-            # We further bump the interval stamp if we were requested to trade
-            # fusion for more collapse-parallelism
-            if maxpar:
-                stamp = stampcache.setdefault(interval.dim, Stamp())
-                interval = interval.lift(stamp)
-
             writeto.append(interval)
             intervals.append(interval)
 
@@ -853,6 +850,30 @@ def optimize_schedule_rotations(schedule, sregistry):
     return schedule.rebuild(*processed, rmapper=rmapper)
 
 
+def optimize_schedule_maxpar(schedule):
+    """
+    Bump the IterationSpace' stamp trading fusion for more collapse-parallelism.
+    """
+    key = lambda i: (i.writeto, i.ispace)
+
+    processed = []
+    for (writeto0, ispace0), group in groupby(schedule, key=key):
+        g = list(group)
+
+        stamp = Stamp()
+        dims = writeto0.itdims
+
+        writeto = writeto0.lift(dims, stamp)
+        ispace = ispace0.lift(dims, stamp)
+
+        processed.extend([
+            ScheduledAlias(pivot, writeto, ispace, aliaseds, indicess)
+            for pivot, _, _, aliaseds, indicess in g
+        ])
+
+    return schedule.rebuild(*processed)
+
+
 def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
                    opt_minmem):
     """