diff --git a/devito/core/cpu.py b/devito/core/cpu.py index 79d1b409e7..3435a472e6 100644 --- a/devito/core/cpu.py +++ b/devito/core/cpu.py @@ -47,8 +47,8 @@ def _normalize_kwargs(cls, **kwargs): o['buf-async-degree'] = oo.pop('buf-async-degree', None) o['buf-reuse'] = oo.pop('buf-reuse', None) - # Fusion - o['fuse-tasks'] = oo.pop('fuse-tasks', False) + # Tasking + o['npthreads'] = oo.pop('npthreads', None) # Flops minimization o['cse-min-cost'] = oo.pop('cse-min-cost', cls.CSE_MIN_COST) diff --git a/devito/core/gpu.py b/devito/core/gpu.py index f907d4be87..0f0354f663 100644 --- a/devito/core/gpu.py +++ b/devito/core/gpu.py @@ -61,8 +61,8 @@ def _normalize_kwargs(cls, **kwargs): o['buf-async-degree'] = oo.pop('buf-async-degree', None) o['buf-reuse'] = oo.pop('buf-reuse', None) - # Fusion - o['fuse-tasks'] = oo.pop('fuse-tasks', False) + # Tasking + o['npthreads'] = oo.pop('npthreads', None) # Flops minimization o['cse-min-cost'] = oo.pop('cse-min-cost', cls.CSE_MIN_COST) diff --git a/devito/core/operator.py b/devito/core/operator.py index 2494dbd7cf..f8f34b0a24 100644 --- a/devito/core/operator.py +++ b/devito/core/operator.py @@ -247,6 +247,15 @@ def _check_kwargs(cls, **kwargs): if oo['mpi'] and oo['mpi'] not in cls.MPI_MODES: raise InvalidOperator(f"Unsupported MPI mode `{oo['mpi']}`") + npthreads = oo['npthreads'] + if npthreads is not None and ( + isinstance(npthreads, bool) or + not is_integer(npthreads) or npthreads <= 0 + ): + raise InvalidOperator( + "`npthreads` must be a positive integer" + ) + if oo['cire-maxpar'] not in (False, 'basic', 'compact'): raise InvalidOperator("Illegal `cire-maxpar` value") diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index bc4b5dc48c..4e63438dad 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1511,8 +1511,7 @@ class BusyWait(While): class SyncSpot(List): """ - A node representing one or more synchronization operations, e.g., WaitLock, - withLock, etc. + A node coupling synchronization operations with an IET body. """ is_SyncSpot = True diff --git a/devito/ir/support/syncs.py b/devito/ir/support/syncs.py index d7f4f86bfe..57280fa28f 100644 --- a/devito/ir/support/syncs.py +++ b/devito/ir/support/syncs.py @@ -26,11 +26,42 @@ class SyncOp(Pickable): + """ + Metadata for a synchronization operation attached to a `Cluster` or `SyncSpot`. + + Parameters + ---------- + handle : object + The symbolic object identifying or controlling the operation, such as + an entry in a `Lock`. May be None when no handle is required. + target : AbstractFunction + The `Function` whose access is synchronized. For buffered data movements, + this is the compiler-generated buffer. + tindex : Expr, optional + The index into `target` involved in the operation. + function : AbstractFunction, optional + The original `Function` represented by a compiler-generated `target`. It + is the source of a `SyncCopyIn` and the destination of a `SyncCopyOut`. + findex : Expr, optional + The index into `function` corresponding to `tindex`. + dim : Dimension, optional + The `Dimension` along which `tindex` and `findex` are defined. + size : int, optional + The extent associated with the operation along `dim`. Defaults to 1. + origin : Indexed, optional + The original `Indexed` access from which the operation was derived. + gid : Stamp, optional + The `Stamp` identifying the asynchronous task group to which the operation + belongs. + """ + __rargs__ = ('handle', 'target') - __rkwargs__ = ('tindex', 'function', 'findex', 'dim', 'size', 'origin') + __rkwargs__ = ( + 'tindex', 'function', 'findex', 'dim', 'size', 'origin', 'gid' + ) def __init__(self, handle, target, tindex=None, function=None, findex=None, - dim=None, size=1, origin=None): + dim=None, size=1, origin=None, gid=None): self.handle = handle self.target = target @@ -40,6 +71,7 @@ def __init__(self, handle, target, tindex=None, function=None, findex=None, self.dim = dim self.size = size self.origin = origin + self.gid = gid def __eq__(self, other): return (type(self) is type(other) and @@ -50,11 +82,13 @@ def __eq__(self, other): self.findex == other.findex and self.dim is other.dim and self.size == other.size and - self.origin == other.origin) + self.origin == other.origin and + self.gid == other.gid) def __hash__(self): return hash((self.__class__, self.handle, self.target, self.tindex, - self.function, self.findex, self.dim, self.size, self.origin)) + self.function, self.findex, self.dim, self.size, self.origin, + self.gid)) def __repr__(self): return f"{self.__class__.__name__}<{self.handle.name}>" diff --git a/devito/passes/clusters/asynchrony.py b/devito/passes/clusters/asynchrony.py index e32190ddef..e37e334927 100644 --- a/devito/passes/clusters/asynchrony.py +++ b/devito/passes/clusters/asynchrony.py @@ -45,6 +45,17 @@ def memcpy_key(c): return task_key, memcpy_key +def make_gid(ispace, d): + """ + Make a group ID from the stamps carried by the non-trigger Intervals. + """ + gids = {i.stamp for i in ispace if not i.dim._defines & d._defines} + if len(gids) != 1: + return None + else: + return gids.pop() + + @timed_pass(name='tasking') def tasking(clusters, key0, sregistry): """ @@ -147,6 +158,8 @@ def _schedule_waitlocks(self, c0, d, clusters, locks, syncs): return protected def _schedule_withlocks(self, c0, d, protected, locks, syncs): + gid = make_gid(c0.ispace, d) + for target in protected: lock = locks[target] @@ -169,7 +182,7 @@ def _schedule_withlocks(self, c0, d, protected, locks, syncs): for i in indices: syncs[c0][d].update([ ReleaseLock(lock[i], target), - WithLock(lock[i], target, i, function, findex, d) + WithLock(lock[i], target, i, function, findex, d, gid=gid) ]) @@ -274,9 +287,12 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry): guard1 = GuardBoundNext(function.indices[d], direction) guards = c.guards.impose(d, guard0 & guard1) + gid = make_gid(ispace, d) + syncs = {d: [ ReleaseLock(handle, target), - PrefetchUpdate(handle, target, tindex, function, findex, d, 1, e.rhs) + PrefetchUpdate(handle, target, tindex, function, findex, d, + origin=e.rhs, gid=gid) ]} syncs = {**c.syncs, **syncs} diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py index 3bf69ef1d3..5f72b8d746 100644 --- a/devito/passes/clusters/buffering.py +++ b/devito/passes/clusters/buffering.py @@ -18,6 +18,7 @@ timed_pass ) from devito.types import Array, CustomDimension, Eq, ModuloDimension +from devito.warnings import warn __all__ = ['buffering'] @@ -39,7 +40,7 @@ def buffering(clusters, key, sregistry, options, **kwargs): The symbol registry, to create unique names for buffers and Dimensions. options : dict The optimization options. - Accepted: ['buf-async-degree']. + Accepted: ['buf-async-degree', 'buf-reuse', 'npthreads']. * 'buf-async-degree': Specify the size of the buffer. By default, the buffer size is the minimal one, inferred from the memory accesses in the ``clusters`` themselves. An asynchronous degree equals to `k` @@ -49,6 +50,9 @@ def buffering(clusters, key, sregistry, options, **kwargs): implemented by other passes). * 'buf-reuse': If True, the pass will try to reuse existing Buffers for different buffered Functions. By default, False. + * 'npthreads': Number of pthreads for asynchronous tasks. The tasks are + divided into this many balanced groups. By default, None, which uses + one pthread per task. **kwargs Additional compilation options. Accepted: ['opt_init_onwrite', 'opt_buffer']. @@ -252,15 +256,18 @@ def callback(self, clusters, prefix): processed.append(Cluster(expr, ispace, guards, properties, syncs)) # Lift {write,read}-only buffers into separate IterationSpaces - if not self.options['fuse-tasks']: - processed = self._optimize(processed, descriptors) + processed = self._optimize(processed, descriptors) - if self.options['buf-reuse']: - init, processed = self._reuse(init, processed, descriptors) + # Reuse existing Buffers for buffering candidates, if requested + init, processed = self._reuse(init, processed, descriptors) return init + processed def _optimize(self, clusters, descriptors): + npthreads = self.options['npthreads'] + + stamps = self._make_task_groups(descriptors) + for b, v in descriptors.items(): if v.is_writeonly: # `b` might be written by multiple, potentially mutually @@ -268,20 +275,22 @@ def _optimize(self, clusters, descriptors): # will have complementary guards, hence only one will be # executed. In such a case, we can split the equations over # separate IterationSpaces - key0 = lambda: Stamp() + key0 = lambda: stamps[b] if npthreads else Stamp() # noqa: B023 elif v.is_readonly: # `b` is read multiple times -- this could just be the case of # coupled equations, so we more cautiously perform a # "buffer-wise" splitting of the IterationSpaces (i.e., only # relevant if there are at least two read-only buffers) - stamp = Stamp() - key0 = lambda: stamp # noqa: B023 + stamp_fixed = Stamp() + key0 = lambda: stamps[b] if npthreads else stamp_fixed # noqa: B023 else: continue processed = [] for c in clusters: - if b not in c.functions: + f = v.f if npthreads else b + + if f not in c.functions: processed.append(c) continue @@ -294,12 +303,57 @@ def _optimize(self, clusters, descriptors): return clusters + def _make_task_groups(self, descriptors): + """ + Assign task buffers to at most `npthreads` balanced groups. + """ + npthreads = self.options['npthreads'] + + stamps = {} + if not npthreads: + return stamps + + task_sets = ( + [b for b, v in descriptors.items() if v.is_writeonly], + [b for b, v in descriptors.items() if v.is_readonly], + ) + + for tasks in task_sets: + if not tasks: + continue + + # Calculate the number of groups and their sizes, so that the tasks + # are divided into balanced groups + ntasks = len(tasks) + ngroups = min(npthreads, ntasks) + base, remainder = divmod(ntasks, ngroups) + sizes = (base + 1,) * remainder + (base,) * (ngroups - remainder) + + if npthreads > ntasks: + warn( + f"`npthreads={npthreads}` exceeds the {ntasks} available " + f"tasks; using `npthreads={ntasks}` instead" + ) + + # Create a unique Stamp for each group and assign it to the tasks + # in that group + start = 0 + for n in sizes: + stamp = Stamp() + stamps.update({b: stamp for b in tasks[start:start + n]}) + start += n + + return stamps + def _reuse(self, init, clusters, descriptors): """ Reuse existing Buffers for buffering candidates. """ buf_reuse = self.options['buf-reuse'] + if not buf_reuse: + return init, clusters + if callable(buf_reuse): cbk = lambda v: [i for i in v if buf_reuse(descriptors[i])] else: diff --git a/devito/passes/clusters/fusion.py b/devito/passes/clusters/fusion.py index fcc3d9affa..bbd4694614 100644 --- a/devito/passes/clusters/fusion.py +++ b/devito/passes/clusters/fusion.py @@ -152,7 +152,7 @@ def __init__(self, toposort, options=None): options = options or {} self.toposort = toposort - self.fuse_tasks = options.get('fuse-tasks', False) + self.fuse_tasks = options.get('npthreads') is not None super().__init__() diff --git a/devito/passes/iet/orchestration.py b/devito/passes/iet/orchestration.py index 3bea70fb0f..52c0821992 100644 --- a/devito/passes/iet/orchestration.py +++ b/devito/passes/iet/orchestration.py @@ -1,30 +1,59 @@ -from collections import OrderedDict +from collections import OrderedDict, defaultdict, namedtuple from contextlib import suppress -from functools import singledispatch +from functools import cached_property, singledispatch from sympy import Or from devito.exceptions import CompilationError from devito.ir.iet import ( - AsyncCall, AsyncCallable, BlankLine, BusyWait, Call, Callable, DummyExpr, FindNodes, - List, SyncSpot, Transformer, derive_parameters, make_callable + AsyncCall, AsyncCallable, BlankLine, Block, BusyWait, Call, Callable, Conditional, + DummyExpr, FindNodes, List, SyncSpot, Transformer, derive_parameters, make_callable ) +from devito.ir.iet.visitors import LazyVisitor from devito.ir.support import ( InitArray, PrefetchUpdate, ReleaseLock, SnapIn, SnapOut, SyncArray, WaitLock, WithLock ) from devito.passes.iet.engine import iet_pass from devito.passes.iet.langbase import LangBB from devito.symbolics import CondEq, CondNe -from devito.tools import DAG, as_mapper, as_tuple +from devito.tools import DAG, as_mapper from devito.types import HostLayer -__init__ = ['Orchestrator'] +__all__ = ['Orchestrator'] + + +Task = namedtuple('Task', 'spot guard sync_ops releases') + + +class SyncSpotRegion(List): + + """A sequence of SyncSpots treated as one unit during orchestration.""" + + def __init__(self, body): + super().__init__(body=body) + assert self.body and all(isinstance(i, SyncSpot) for i in self.body) + + @property + def sync_spots(self): + return self.body + + @cached_property + def optype(self): + """ + The type of the synchronization operation in this region. + """ + optypes = {type(op) for spot in self.sync_spots for op in spot.sync_ops} + assert len(optypes) == 1, ( + "Expected a SyncSpotRegion to contain exactly one type of " + "synchronization operation" + ) + return optypes.pop() class Orchestrator: """ - Lower the SyncSpot in IET for efficient host-device asynchronous computation. + Lower synchronization nodes for efficient host-device asynchronous computation. """ langbb = LangBB @@ -32,8 +61,45 @@ class Orchestrator: The language used to implement host-device data movements. """ - def __init__(self, sregistry=None, **kwargs): + def __init__(self, sregistry=None, options=None, **kwargs): self.sregistry = sregistry + self.npthreads = (options or {}).get('npthreads') + + def _fuse_tasks(self, iet): + """ + Group compatible task SyncSpots into SyncSpotRegions. + """ + if self.npthreads is None: + return iet + + insertions = defaultdict(list) + subs1 = {} + + for tasks, anchor in CollectTasks().visit(iet): + spots = [] + for task in tasks: + if task.guard is None: + # Preserve a local scope when there is no Conditional + scope = Block(body=task.spot.body) + else: + scope = Conditional(task.guard, task.spot.body) + spots.append(SyncSpot(task.sync_ops, body=scope)) + + insertions[anchor].append(SyncSpotRegion(spots)) + subs1.update({task.spot: SyncSpot(task.releases) for task in tasks}) + + if not subs1: + return iet + + # These substitutions cannot be merged because a Transformer does not + # revisit a replacement, while task spots may be nested below an anchor + subs0 = {anchor: List(body=(anchor, *regions)) + for anchor, regions in insertions.items()} + + iet = Transformer(subs0).visit(iet) + iet = Transformer(subs1).visit(iet) + + return iet def _make_waitlock(self, iet, sync_ops, *args): waitloop = List( @@ -53,6 +119,9 @@ def _make_releaselock(self, iet, sync_ops, *args): parameters = derive_parameters(pre, ordering='canonical') efunc = Callable(name, pre, 'void', parameters, 'static') + if isinstance(iet, SyncSpot) and not iet.body: + iet = List() + iet = List(body=[Call(name, efunc.parameters)] + [iet]) return iet, [efunc] @@ -60,17 +129,15 @@ def _make_releaselock(self, iet, sync_ops, *args): def _make_withlock(self, iet, sync_ops, layer): body, prefix = withlock(layer, iet, sync_ops, self.langbb, self.sregistry) - # Turn `iet` into an AsyncCallable so that subsequent passes know - # that we're happy for this Callable to be executed asynchronously + return self._make_async_callable(body, prefix) + + def _make_async_callable(self, body, prefix): name = self.sregistry.make_name(prefix=prefix) body = List(body=body) parameters = derive_parameters(body, ordering='canonical') efunc = AsyncCallable(name, body, parameters=parameters) - # The corresponding AsyncCall - iet = AsyncCall(name, efunc.parameters) - - return iet, [efunc] + return AsyncCall(name, efunc.parameters), [efunc] def _make_callable(self, name, iet, *args): name = self.sregistry.make_name(prefix=name) @@ -108,23 +175,52 @@ def _make_syncarray(self, iet, sync_ops, layer): def _make_prefetchupdate(self, iet, sync_ops, layer): body, prefix = prefetchupdate(layer, iet, sync_ops, self.langbb, self.sregistry) - # Turn `iet` into an AsyncCallable so that subsequent passes know - # that we're happy for this Callable to be executed asynchronously - name = self.sregistry.make_name(prefix=prefix) - body = List(body=body) - parameters = derive_parameters(body, ordering='canonical') - efunc = AsyncCallable(name, body, parameters=parameters) + return self._make_async_callable(body, prefix) - # The corresponding AsyncCall - iet = AsyncCall(name, efunc.parameters) + def _make_region(self, iet): + """Lower a SyncSpotRegion into one asynchronous callable.""" + sync_ops = tuple(op for spot in iet.sync_spots + for op in spot.sync_ops) + callback = { + PrefetchUpdate: prefetchupdate, + WithLock: withlock + }[iet.optype] - return iet, [efunc] + layer = infer_sync_layer(sync_ops) + + body = [] + for spot in iet.sync_spots: + scope, = spot.body + if isinstance(scope, Conditional): + task_body = scope.then_body + else: + assert isinstance(scope, Block) + task_body = scope.body + + task_body, prefix = callback( + layer, List(body=task_body), spot.sync_ops, self.langbb, + self.sregistry + ) + body.append(scope._rebuild(task_body)) + + return self._make_async_callable(body, prefix) @iet_pass def process(self, iet): - sync_spots = FindNodes(SyncSpot).visit(iet) - if not sync_spots: - return iet, {} + # Group compatible task SyncSpots into SyncSpotRegions, if requested + # by the user + iet = self._fuse_tasks(iet) + + # Lower regions first so their member SyncSpots are not processed + # independently by the generic SyncSpot lowering below + efuncs = [] + subs = {} + for region in FindNodes(SyncSpotRegion).visit(iet): + call, efuncs1 = self._make_region(region) + subs[region] = call + efuncs.extend(efuncs1) + + iet = Transformer(subs).visit(iet) # The SyncOps are to be processed in a given order callbacks = OrderedDict([ @@ -137,30 +233,26 @@ def process(self, iet): (PrefetchUpdate, self._make_prefetchupdate), (ReleaseLock, self._make_releaselock), ]) - key = lambda s: list(callbacks).index(s) + key = tuple(callbacks).index # The SyncSpots may be nested, so we compute a topological ordering # so that they are processed in a bottom-up fashion. This is necessary # because e.g. an inner SyncSpot may generate new objects (e.g., a new # Queue), which in turn must be visible to the outer SyncSpot to # generate the correct parameters list - efuncs = [] while True: sync_spots = FindNodes(SyncSpot).visit(iet) if not sync_spots: break n0 = ordered(sync_spots).pop(0) - mapper = as_mapper(n0.sync_ops, lambda i: type(i)) + mapper = as_mapper(n0.sync_ops, type) subs = {} for t in sorted(mapper, key=key): sync_ops = mapper[t] - layers = {infer_layer(s.function) for s in sync_ops} - if len(layers) != 1: - raise CompilationError("Unsupported streaming case") - layer = layers.pop() + layer = infer_sync_layer(sync_ops) n1, v = callbacks[t](subs.get(n0, n0), sync_ops, layer) @@ -175,12 +267,70 @@ def process(self, iet): def ordered(sync_spots): dag = DAG(nodes=sync_spots) for n0 in sync_spots: - for n1 in as_tuple(FindNodes(SyncSpot).visit(n0.body)): + for n1 in FindNodes(SyncSpot).visit(n0.body): dag.add_edge(n1, n0) return dag.topological_sort() +class CollectTasks(LazyVisitor): + + """Collect and group compatible asynchronous SyncSpots.""" + + def _post_visit(self, ret): + groups = defaultdict(list) + anchors = {} + for key, task, anchor in ret: + groups[key].append(task) + # Activate the group after its last task in program order + anchors[key] = anchor + return tuple((tasks, anchors[key]) for key, tasks in groups.items()) + + def visit_Iteration(self, o, **kwargs): + kwargs['iteration'] = o + yield from self._visit(o.children, **kwargs) + + def visit_Conditional(self, o, **kwargs): + kwargs['condition'] = o + yield from self._visit(o.children, **kwargs) + + def visit_SyncSpot(self, o, iteration=None, condition=None, + in_snapshot=False): + if iteration is not None: + syncs = as_mapper(o.sync_ops, type) + + optypes = set(syncs) + for optype in (PrefetchUpdate, WithLock): + if optypes != {optype, ReleaseLock}: + continue + + guard = None + if condition is not None: + # Task SyncSpots inherit a guard without an `else` branch + # from the originating Cluster + assert not condition.else_body + guard = condition.condition + + sync_ops = syncs[optype] + + gid, = {i.gid for i in sync_ops} + if gid is not None: + task = Task(o, guard, sync_ops, syncs[ReleaseLock]) + key = (iteration, gid, optype, in_snapshot) + yield key, task, condition or o + + break + + if any(isinstance(i, SnapOut) for i in o.sync_ops): + # Do not mix composite tasks with other compatible groups + in_snapshot = True + + yield from self._visit( + o.children, iteration=iteration, condition=condition, + in_snapshot=in_snapshot + ) + + # Task handlers layer_host = HostLayer('host') @@ -194,6 +344,20 @@ def infer_layer(f): return layer_host +def infer_sync_layer(sync_ops): + """ + Infer the unique storage layer used by a sequence of SyncOps. + """ + layers = {infer_layer(i.function) for i in sync_ops} + if len(layers) != 1: + found = ', '.join(sorted(str(i) for i in layers)) or 'none' + raise CompilationError( + "Expected synchronization operations to use exactly one storage " + f"layer, but found: {found}" + ) + return layers.pop() + + @singledispatch def withlock(layer, iet, sync_ops, lang, sregistry): raise NotImplementedError diff --git a/tests/test_gpu_common.py b/tests/test_gpu_common.py index 079b5bdf1b..e3297f221d 100644 --- a/tests/test_gpu_common.py +++ b/tests/test_gpu_common.py @@ -496,7 +496,7 @@ def test_tasking_forcefuse(self): Eq(v.forward, tmp1, subdomain=bundle0)] op = Operator(eqns, opt=('tasking', 'fuse', 'orchestrate', - {'fuse-tasks': True, 'linearize': False})) + {'npthreads': 1, 'linearize': False})) # Check generated code assert len(retrieve_iteration_tree(op)) == 3 @@ -646,7 +646,7 @@ def test_streaming_basic(self, opt, ntmps): @pytest.mark.parametrize('opt,ntmps', [ (('buffering', 'streaming', 'orchestrate'), 14), - (('buffering', 'streaming', 'fuse', 'orchestrate', {'fuse-tasks': True}), 8), + (('buffering', 'streaming', 'fuse', 'orchestrate', {'npthreads': 1}), 8), ]) def test_streaming_two_buffers(self, opt, ntmps): nt = 10 @@ -678,6 +678,40 @@ def test_streaming_two_buffers(self, opt, ntmps): assert np.all(u.data[0] == 56) assert np.all(u.data[1] == 72) + @pytest.mark.parametrize('npthreads', [1, 3, 4]) + def test_streaming_fuse_groups(self, npthreads): + nt = 4 + grid = Grid(shape=(4, 4)) + + u = TimeFunction(name='u', grid=grid) + fsaves = [TimeFunction(name=f'fsave{i}', grid=grid, save=nt) + for i in range(9)] + + for i, f in enumerate(fsaves): + for t in range(nt): + f.data[t, :] = i + t + + eqn = Eq(u.forward, u + sum(fsaves)) + + op0 = Operator(eqn, opt=('noop', {'gpu-fit': tuple(fsaves)})) + op1 = Operator( + eqn, + opt=('buffering', 'streaming', 'fuse', 'orchestrate', + {'npthreads': npthreads}) + ) + + symbols = FindSymbols().visit(op1) + threads = [i for i in symbols if isinstance(i, PThreadArray)] + assert all(i.size == 1 for i in threads) + assert op1.npthreads == npthreads + + op0.apply(time_M=nt-2) + + u1 = TimeFunction(name='u', grid=grid) + op1.apply(time_M=nt-2, u=u1) + + assert np.all(u.data == u1.data) + def test_streaming_fused(self): nt = 10 grid = Grid(shape=(4, 4)) @@ -961,7 +995,7 @@ def test_composite_buffering_tasking_multi_output(self): {'buf-async-degree': async_degree})) op2 = Operator(eqns, opt=('buffering', 'tasking', 'topofuse', 'orchestrate', {'buf-async-degree': async_degree, - 'fuse-tasks': True})) + 'npthreads': 1})) # Check generated code -- thanks to buffering only expect 1 lock! assert len(retrieve_iteration_tree(op0)) == 2 @@ -1002,6 +1036,53 @@ def test_composite_buffering_tasking_multi_output(self): assert np.all(usave.data == usave1.data) assert np.all(vsave.data == vsave1.data) + @pytest.mark.parametrize('npthreads', [1, 3, 4]) + def test_composite_buffering_tasking_fuse_groups(self, npthreads): + nt = 4 + bundle0 = Bundle() + grid = Grid(shape=(4, 4, 4), subdomains=bundle0) + + u = TimeFunction(name='u', grid=grid, time_order=2) + fsaves = [TimeFunction(name=f'fsave{i}', grid=grid, save=nt) + for i in range(9)] + + eqns = [Eq(u.forward, u + 1)] + eqns.extend(Eq(f, u + i, subdomain=bundle0) + for i, f in enumerate(fsaves)) + + op0 = Operator(eqns, opt=('noop', {'gpu-fit': tuple(fsaves)})) + + opt = ('buffering', 'tasking', 'topofuse', 'orchestrate', + {'npthreads': npthreads}) + op1 = Operator(eqns, opt=opt) + + symbols = FindSymbols().visit(op1) + threads = [i for i in symbols if isinstance(i, PThreadArray)] + assert all(i.size == 1 for i in threads) + assert op1.npthreads == npthreads + + tasks = [v.root for k, v in op1._func_table.items() + if k.startswith('copy_to_host')] + # Pthreads handling the same number of tasks reuse a Callable + widths = [] + for task in tasks: + writes = {i.write for i in FindNodes(Expression).visit(task)} + widths.append(len([i for i in writes if i.is_TimeFunction])) + expected = {1: [9], 3: [3], 4: [2, 3]} + assert sorted(widths) == expected[npthreads] + + op0.apply(time_M=nt-1) + + u1 = TimeFunction(name='u', grid=grid, time_order=2) + fsaves1 = [TimeFunction(name=f'fsave{i}', grid=grid, save=nt) + for i in range(9)] + kwargs = {f.name: f1 for f, f1 in zip(fsaves, fsaves1, strict=True)} + op1.apply(time_M=nt-1, u=u1, **kwargs) + + assert np.all(u.data == u1.data) + assert all(np.all(f.data == f1.data) + for f, f1 in zip(fsaves, fsaves1, strict=True)) + def test_composite_full_0(self): nt = 10 grid = Grid(shape=(10, 10, 10)) @@ -1327,7 +1408,7 @@ def test_streaming_complete(self): op1 = Operator(eqns, opt=('buffering', 'streaming', 'orchestrate')) op2 = Operator(eqns, opt=('buffering', 'streaming', 'fuse', 'orchestrate')) op3 = Operator(eqns, opt=('buffering', 'streaming', 'fuse', 'orchestrate', - {'fuse-tasks': True})) + {'npthreads': 1})) # Check generated code diff = int(configuration['language'] == 'openmp') diff --git a/tests/test_operator.py b/tests/test_operator.py index 806b091e24..b3f7ea295e 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -135,6 +135,12 @@ def test_opt_options(self): except InvalidOperator: assert True + # `npthreads` is a direct count, not a boolean switch + for npthreads in (False, True, 0, -1, 1.5): + with pytest.raises(InvalidOperator, match='positive integer'): + Operator(Eq(u, u + 1), + opt=('advanced', {'npthreads': npthreads})) + def test_compiler_uniqueness(self): grid = Grid(shape=(3, 3, 3))