From 29ee29e2c4f88b253bb91b190992233ea222a54b Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 8 Sep 2026 14:11:23 -0400 Subject: [PATCH] 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. --- devito/core/gpu.py | 3 ++- devito/core/operator.py | 8 +++----- devito/passes/clusters/blocking.py | 13 +++++++++++-- devito/tools/data_structures.py | 4 ++++ tests/test_gpu_openacc.py | 31 ++++++++++++++++++++++++++++++ 5 files changed, 51 insertions(+), 8 deletions(-) diff --git a/devito/core/gpu.py b/devito/core/gpu.py index 8f186df0f7a..536a06422e2 100644 --- a/devito/core/gpu.py +++ b/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) diff --git a/devito/core/operator.py b/devito/core/operator.py index 6cf20767b29..ff9f42987b0 100644 --- a/devito/core/operator.py +++ b/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 diff --git a/devito/passes/clusters/blocking.py b/devito/passes/clusters/blocking.py index e4d68444996..fea502fbfbe 100644 --- a/devito/passes/clusters/blocking.py +++ b/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) diff --git a/devito/tools/data_structures.py b/devito/tools/data_structures.py index 2bf901c3f48..ced3fd2d9b0 100644 --- a/devito/tools/data_structures.py +++ b/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 diff --git a/tests/test_gpu_openacc.py b/tests/test_gpu_openacc.py index a335d55289b..2dba0976841 100644 --- a/tests/test_gpu_openacc.py +++ b/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))