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

Pad such an entry with its innermost size instead, 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),
+                                stretch=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,9 +478,16 @@ def __new__(cls, items, rule=None, tag=None):
 
 class ParTile(UnboundedMultiTuple, OptOption):
 
-    def __new__(cls, items, default=None, sparse=None, reduce=None):
+    # A short entry is padded to the nest it lands on rather than losing its
+    # outermost BlockDimensions. Only meaningful on a device
+    stretch = False
+
+    def __new__(cls, items, default=None, sparse=None, reduce=None,
+                stretch=False):
         if not items:
-            return UnboundedMultiTuple()
+            obj = UnboundedMultiTuple()
+            obj.stretch = False
+            return obj
         elif isinstance(items, bool):
             if not default:
                 raise ValueError("Expected `default` value, got None")
@@ -536,6 +543,7 @@ 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.stretch = stretch
 
         return obj
 

devito/passes/clusters/blocking.py

@@ -495,6 +495,7 @@ class BlockSizeGenerator:
 
     def __init__(self, par_tile):
         self.umt = par_tile
+        self.stretch = par_tile.stretch
 
         if par_tile.is_multi:
             # The user has supplied one specific par-tile per blocked nest
@@ -522,7 +523,25 @@ def __init__(self, par_tile):
             else:
                 self.umt_reduce = UnboundTuple(*par_tile.default, 1)
 
+    def stretched(self, umt, dims):
+        """
+        Pad `umt` with its innermost size until it covers `dims`.
+
+        A user-supplied par-tile entry may be shorter than the nest it lands on.
+        Consuming it as is would leave the outermost Dimensions without a block
+        size, and dropping their BlockDimensions would cost them their
+        parallelism -- on a device, that means iterating them on the host.
+        """
+        if not umt or len(umt) >= len(dims):
+            return umt
+
+        sizes = tuple(umt) + (umt[-1],) * (len(dims) - len(umt))
+
+        return UnboundTuple(*sizes)
+
     def schedule(self, dims, clusters):
+        stretch = False
+
         if any(c.properties.is_parallel_atomic(dims) for c in clusters):
             # Correctness -- enforce blocking where necessary.
             # See also issue #276:PRO
@@ -537,9 +556,14 @@ def schedule(self, dims, clusters):
 
         else:
             umt = self.umt
+            stretch = self.stretch and umt.is_multi
 
         umt.iter()
 
+        if stretch:
+            # `umt` walks this nest's own entry, which may be too short
+            return self.stretched(umt.curitem(), dims).reset()
+
         return umt
 
 
@@ -622,10 +646,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)
 

tests/test_gpu_openacc.py

@@ -144,6 +144,39 @@ 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'})
+
+        # The short entry is stretched by repeating its innermost size, so the
+        # nest stays fully blocked rather than losing its outermost Dimension
+        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))