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

When a multi `par-tile` entry is shorter than the blocked nest, the
innermost Dimensions consume the available block sizes and the outermost
one runs into a StopIteration, raised by the top-level tuple once the
entry is exhausted. It was then dropped into `compact`, which promotes it
back to its root Dimension, discarding its BlockDimension.

On a device that is a serialization: the Dimension ends up outside the
blocked nest, `filter_iterations` rejects it as non-parallel, and the
kernel is launched over a 2D grid from a host loop iterating the outer
Dimension one slice at a time.

Walk this nest's own entry instead, which repeats its innermost size for
as long as the nest needs it, so the nest stays fully blocked. Whether to
do so is decided when the par-tile is built, where the target is already
known, rather than by inspecting a cluster back in the pass. It applies
to multi par-tiles on a device only: a single user-supplied par-tile
defines the block rank on purpose, and on a host a short entry is the
documented way to ask for 2.5D blocking.

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))