diff --git a/devito/passes/clusters/blocking.py b/devito/passes/clusters/blocking.py index e4d6844499..29ae52b1e2 100644 --- a/devito/passes/clusters/blocking.py +++ b/devito/passes/clusters/blocking.py @@ -1,3 +1,4 @@ +import contextlib from itertools import groupby from sympy import sympify @@ -67,7 +68,11 @@ def blocking(clusters, sregistry, options): clusters = AnalyzeSkewing().process(clusters) if options['blocklevels'] > 0: - clusters = SynthesizeBlocking(sregistry, options).process(clusters) + if options['blockrelax'] == 'device-aware': + synthesizer = SynthesizeBlockingDeviceAware(sregistry, options) + else: + synthesizer = SynthesizeBlocking(sregistry, options) + clusters = synthesizer.process(clusters) if options['skewing']: clusters = SynthesizeSkewing(options).process(clusters) @@ -162,6 +167,14 @@ def __init__(self, options): self.gpu_fit = options.get('gpu-fit', ()) + def _process_fatd(self, clusters, level, prefix=None): + processed = [] + for _, group in groupby(clusters, key=lambda c: c.ispace): + g = list(group) + processed.extend(Queue._process_fatd(self, g, level, prefix)) + + return processed + def _make_key_hook(self, cluster, level): return (is_on_device(cluster.functions, self.gpu_fit),) @@ -315,18 +328,21 @@ def callback(self, clusters, prefix): return processed -class SynthesizeBlocking(Queue): +class SynthesizeBlockingBase(Queue): + + mapper = None + """ + A mapping from a tuple of (Dimension, number of stencil points) to a tuple + of BlockDimensions, so that we can reuse existing BlockDimensions to avoid + unnecessary `steps`. Disabled by default, to be enabled in subclasses that + need it (e.g., SynthesizeBlocking). + """ def __init__(self, sregistry, options): self.sregistry = sregistry self.levels = options['blocklevels'] - # Track the BlockDimensions created so far so that we can reuse them - # in case of Clusters that are different but share the same number of - # stencil points - self.mapper = {} - super().__init__() def process(self, clusters): @@ -340,11 +356,11 @@ def _make_key_guards(self, cluster, ispace): if not cluster.properties.is_blockable(i.dim)) def _derive_block_dims(self, clusters, prefix, d): - # Can I reuse existing BlockDimensions to avoid a proliferation of steps? + # Can I reuse existing BlockDimensions to avoid unnecessary `steps`? k = stencil_footprint(clusters, d) try: return self.mapper[k] - except KeyError: + except (KeyError, TypeError): pass base = self.sregistry.make_name(prefix=d.root.name) @@ -362,7 +378,10 @@ def _derive_block_dims(self, clusters, prefix, d): bd = BlockDimension(d.name, bd, bd, bd + bd.step - 1, 1, size=step) block_dims.append(bd) - retval = self.mapper[k] = tuple(block_dims), bd + retval = tuple(block_dims), bd + + with contextlib.suppress(TypeError): + self.mapper[k] = retval return retval @@ -404,6 +423,28 @@ def callback(self, clusters, prefix): return processed +class SynthesizeBlocking(SynthesizeBlockingBase): + + def __init__(self, sregistry, options): + # Track the BlockDimensions created so far so that we can reuse them + # in case of Clusters that are different but share the same number of + # stencil points + self.mapper = {} + + super().__init__(sregistry, options) + + +class SynthesizeBlockingDeviceAware(SynthesizeBlockingBase): + + def _process_fdta(self, clusters, level, prefix=None): + processed = [] + for _, group in groupby(clusters, key=lambda c: c.ispace): + g = list(group) + processed.extend(Queue._process_fdta(self, g, level, prefix)) + + return processed + + def stencil_footprint(clusters, d): """ Compute the number of stencil points in the given Dimension `d` across the diff --git a/tests/test_dle.py b/tests/test_dle.py index dd3ebeb947..cd26e9eafd 100644 --- a/tests/test_dle.py +++ b/tests/test_dle.py @@ -264,7 +264,8 @@ def test_leftright_subdims(self): op = Operator(eqns, opt=('fission', 'blocking', {'blockrelax': 'device-aware'})) - bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'x2_blk0'}) + bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'x2_blk0', + 'x3_blk0', 'x4_blk0'}) assert all(IsPerfectIteration().visit(i) for i in bns.values()) assert all(len(FindNodes(Iteration).visit(i)) == 4 for i in bns.values())