Merge pull request #3021 from devitocodes/blocking-keep-outer-dim-parallel

compiler: Keep the outer Dimension parallel on a short par-tile

Author: mloubout

devito/core/gpu.py

@@ -90,7 +90,8 @@ def _normalize_kwargs(cls, **kwargs):
         # GPU parallelism
         o['par-tile'] = ParTile(oo.pop('par-tile', False), default=(32, 4, 4),
                                 sparse=oo.pop('par-tile-sparse', None),
-                                reduce=oo.pop('par-tile-reduce', None))
+                                reduce=oo.pop('par-tile-reduce', None),
+                                unbound=True)
         o['par-collapse-ncores'] = 1  # Always collapse (meaningful if `par-tile=False`)
         o['par-collapse-work'] = 1  # Always collapse (meaningful if `par-tile=False`)
         o['par-chunk-nonaffine'] = oo.pop('par-chunk-nonaffine', cls.PAR_CHUNK_NONAFFINE)

devito/core/operator.py

@@ -478,7 +478,8 @@ def __new__(cls, items, rule=None, tag=None):
 
 class ParTile(UnboundedMultiTuple, OptOption):
 
-    def __new__(cls, items, default=None, sparse=None, reduce=None):
+    def __new__(cls, items, default=None, sparse=None, reduce=None,
+                unbound=False):
         if not items:
             return UnboundedMultiTuple()
         elif isinstance(items, bool):
@@ -536,9 +537,6 @@ def __new__(cls, items, default=None, sparse=None, reduce=None):
         obj.default = as_tuple(default)
         obj.sparse = as_tuple(sparse)
         obj.reduce = as_tuple(reduce)
+        obj.unbound = unbound
 
         return obj
-
-    @property
-    def is_multi(self):
-        return len(self) > 1

devito/passes/clusters/blocking.py

@@ -495,6 +495,7 @@ class BlockSizeGenerator:
 
     def __init__(self, par_tile):
         self.umt = par_tile
+        self.unbound = par_tile.unbound
 
         if par_tile.is_multi:
             # The user has supplied one specific par-tile per blocked nest
@@ -523,6 +524,8 @@ def __init__(self, par_tile):
                 self.umt_reduce = UnboundTuple(*par_tile.default, 1)
 
     def schedule(self, dims, clusters):
+        unbound = False
+
         if any(c.properties.is_parallel_atomic(dims) for c in clusters):
             # Correctness -- enforce blocking where necessary.
             # See also issue #276:PRO
@@ -534,12 +537,17 @@ def schedule(self, dims, clusters):
         elif all(c.properties.avoid_tuning(dims) for c in clusters):
             # Performance heuristics -- use a smaller par-tile
             umt = self.umt_small
+            unbound = self.unbound and umt.is_multi
 
         else:
             umt = self.umt
+            unbound = self.unbound and umt.is_multi
 
         umt.iter()
 
+        if unbound:
+            return umt.curitem()
+
         return umt
 
 
@@ -622,10 +630,11 @@ def apply_par_tiles(clusters, options, **kwargs):
     Use the par-tile parameter to replace the symbolic BlockDimension sizes
     with actual integer numbers representing the block shape.
     """
-    if not options['par-tile']:
+    par_tile = options['par-tile']
+    if not par_tile:
         return clusters
 
-    blk_size_gen = BlockSizeGenerator(options['par-tile'])
+    blk_size_gen = BlockSizeGenerator(par_tile)
 
     key = lambda c: c.ispace.project(lambda d: d.is_Block)
 

devito/tools/data_structures.py

@@ -880,6 +880,10 @@ def __new__(cls, *items, **kwargs):
         obj.current = None
         return obj
 
+    @property
+    def is_multi(self):
+        return len(self) > 1
+
     def reset(self):
         self.current = None
         return self

tests/test_gpu_openacc.py

@@ -144,6 +144,37 @@ def test_multiple_tile_sizes(self, par_tile):
         assert trees[3][1].pragmas[0].ccode.value ==\
             f'acc parallel loop {sclause} present(src,src_gp,src_wx,src_wy,src_wz,u)'
 
+    def test_short_multi_tile_keeps_outer_dim_blocked(self):
+        """
+        A multi `par-tile` entry shorter than the nest it lands on must not cost
+        the outermost Dimension its BlockDimension: on a device, dropping it
+        would leave `x` iterated outside the offloaded nest.
+        """
+        grid = Grid(shape=(8, 8, 8))
+
+        u = TimeFunction(name="u", grid=grid, space_order=4)
+        v = TimeFunction(name="v", grid=grid, space_order=4)
+
+        eqns = [Eq(u.forward, u.dx),
+                Eq(v.forward, u.forward.dx)]
+
+        # The second entry is 2D, while the nest it lands on is 3D
+        par_tile = ((32, 4, 4), (16, 4))
+
+        op = Operator(eqns, platform='nvidiaX', language='openacc',
+                      opt=(
+                          'advanced',
+                          {'par-tile': par_tile, 'blocklevels': 1, 'blockinner': True}))
+
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0'})
+
+        expected = ((4, 4, 32), (4, 4, 16))
+        for root, v in zip(bns.values(), expected, strict=True):
+            iters = FindNodes(Iteration).visit(root)
+            iters = [i for i in iters if i.dim.is_Block and i.dim._depth == 1]
+            assert len(iters) == len(v)
+            assert all(i.step == j for i, j in zip(iters, v, strict=True))
+
     def test_multi_tile_blocking_structure(self):
         grid = Grid(shape=(8, 8, 8))