Merge pull request #2974 from devitocodes/fix-cire-subdomain-parallelism

compiler: Fix regressed OpenMP parallelism with CIRE and subdomains

Author: FabioLuporini

devito/passes/clusters/aliases.py

@@ -939,7 +939,7 @@ def optimize_schedule_maxpar_compact(schedule):
         guards = sa.guards
 
         # Do we need to introduce guards to preserve the original semantics?
-        for i in sa.indices:
+        for i in sa.active_dims:
             for d in i._defines:
                 try:
                     int0, int1 = sa.ispace[d], ispace_union[d]
@@ -973,6 +973,7 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype, opt_min
     subs = {}
     for sa in schedule:
         pivot, writeto, ispace, aliaseds, indicess, guards = sa
+        active_dims = sa.active_dims
 
         name = sregistry.make_name()
 
@@ -1003,7 +1004,7 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype, opt_min
 
             # The indices used to populate the Array
             indices = []
-            for interval, i, s in zip(writeto, sa.indices, shift, strict=True):
+            for interval, i, s in zip(writeto, active_dims, shift, strict=True):
                 try:
                     # E.g., `xs`
                     di, = writeto.sub_iterators[i]
@@ -1043,15 +1044,15 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype, opt_min
             try:
                 if any(i.is_Modulo for i in ispace.sub_iterators[d]):
                     properties[d] = normalize_properties(v, {SEQUENTIAL})
-                elif d not in writeto.itdims:
+                elif d not in active_dims:
                     properties[d] = normalize_properties(v, {PARALLEL_IF_PVT})
             except KeyError:
                 # Non-dimension key such as (x, y) for diagonal stencil u(x+i hx, y+i hy)
                 pass
 
         # Track star-shaped stencils for potential future optimization
         if len(writeto) > 1 and schedule.is_frame:
-            properties[Hyperplane(writeto.itdims)] = {SEPARABLE}
+            properties[Hyperplane(active_dims)] = {SEPARABLE}
 
         # Finally, build the alias Cluster
         clusters.append(Cluster(expression, ispace, guards, properties))
@@ -1516,9 +1517,9 @@ def dimensions(self):
         return self.writeto.itdims
 
     @property
-    def indices(self):
+    def active_dims(self):
         """
-        The indices used to populate the temporary.
+        The active Dimensions used to populate the temporary.
 
         The write-to space may be relaxed for storage sizing, while the
         temporary still has to be indexed with the active iteration Dimension.

tests/test_dle.py

@@ -765,6 +765,37 @@ def test_collapsing_v2(self):
         assert iterations[2].is_Sequential
         assert iterations[3].is_Sequential
 
+    def test_collapsing_cire_subdomain(self):
+        """
+        Test that CIRE temporaries over a SubDomain retain SIMD parallelism.
+        """
+        grid = Grid(shape=(16, 16, 16))
+
+        u = Function(name='u', grid=grid, space_order=8)
+        out = TimeFunction(name='out', grid=grid, space_order=8)
+        a = Function(name='a', grid=grid)
+        b = Function(name='b', grid=grid)
+        c = Function(name='c', grid=grid)
+
+        eq = Eq(out.forward,
+                a*u.dx.dx + b*u.dy.dy + c*u.dz.dz,
+                subdomain=grid.interior)
+        op = Operator(eq, opt=('advanced', {'blockeager': False,
+                                            'blocklazy': False,
+                                            'cire-mingain': 0,
+                                            'openmp': True,
+                                            'par-collapse-ncores': 1,
+                                            'par-collapse-work': 0}))
+
+        iterations = FindNodes(Iteration).visit(op)
+        assert len(iterations) == 7
+
+        # The CIRE producer should collapse x and y, then vectorize z
+        assert iterations[1].ncollapsed == 2
+        assert iterations[3].is_Vectorized
+
+        assert op.cfunction
+
     def test_scheduling(self):
         """
         Affine iterations -> #pragma omp ... schedule(dynamic,1) ...