Skip to content

Commit 3a87fad

Browse files
authored
Merge branch 'main' into JDBetteridge/amd_compression
2 parents 24f3775 + 4109b58 commit 3a87fad

9 files changed

Lines changed: 161 additions & 28 deletions

File tree

‎devito/ir/clusters/cluster.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,14 +118,12 @@ def exprs_dimensions(self):
118118
dims_implicit = {d for e in self.exprs for d in e.implicit_dims}
119119
return dims_explicit | dims_implicit
120120

121-
@cached_property
121+
@property
122122
def guards_dimensions(self):
123123
"""
124124
The Dimensions that appear explicitly in the guards.
125125
"""
126-
syms_guards = {d for e in self.guards.values() for d in e.free_symbols}
127-
dims_guards = {i for i in syms_guards if i.is_Dimension}
128-
return dims_guards
126+
return self.guards.dimensions
129127

130128
@cached_property
131129
def used_dimensions(self):

‎devito/ir/iet/nodes.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
AFFINE, INBOUND, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, SEQUENTIAL,
1919
VECTORIZED, Forward, PrefetchUpdate, Property, WithLock
2020
)
21-
from devito.symbolics import CallFromPointer, ListInitializer
21+
from devito.symbolics import CallFromPointer, ListInitializer, Macro
2222
from devito.tools import (
2323
Signer, as_tuple, ctypes_to_cstr, filter_ordered, filter_sorted, flatten
2424
)
@@ -61,6 +61,7 @@
6161
'Section',
6262
'Switch',
6363
'SyncSpot',
64+
'ThreadFence',
6465
'TimedList',
6566
'Transfer',
6667
'Using',
@@ -1302,6 +1303,31 @@ def size(self):
13021303
return self.ispace.size
13031304

13041305

1306+
class ThreadFence(Call):
1307+
1308+
"""
1309+
A memory fence, ordering a thread's accesses either side of it.
1310+
1311+
Threads exchanging work through a shared flag need memory ordering on both
1312+
sides: a release before publishing the flag, and an acquire after observing it.
1313+
This ensures the state associated with the flag is visible before it is used.
1314+
Without this ordering, updates may be observed in the wrong order and lost.
1315+
1316+
`__atomic_thread_fence` is a compiler builtin, so this renders the same in C
1317+
and in C++ and needs no header.
1318+
"""
1319+
1320+
def __init__(self, order):
1321+
assert order in ('acquire', 'release')
1322+
super().__init__('__atomic_thread_fence',
1323+
[Macro(f'__ATOMIC_{order.upper()}')])
1324+
self._order = order
1325+
1326+
@property
1327+
def order(self):
1328+
return self._order
1329+
1330+
13051331
class Prodder(Call):
13061332

13071333
"""

‎devito/ir/support/guards.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,15 @@
55
"""
66

77
from collections import Counter, defaultdict
8-
from functools import singledispatch
8+
from functools import cached_property, singledispatch
99
from operator import ge, gt, le, lt
1010

1111
import numpy as np
1212
from sympy import And, Expr, Ge, Gt, Le, Lt, Mul, true
1313
from sympy.logic.boolalg import BooleanFunction
1414

1515
from devito.ir.support.space import Forward, IterationDirection
16+
from devito.ir.support.utils import pull_dims
1617
from devito.symbolics import CondEq, CondNe, IntDiv, search
1718
from devito.symbolics.manipulation import _uxreplace_handle, _uxreplace_registry
1819
from devito.tools import Pickable, as_tuple, frozendict, split
@@ -280,6 +281,15 @@ class Guards(frozendict):
280281
def get(self, d, v=true):
281282
return super().get(d, v)
282283

284+
@cached_property
285+
def dimensions(self):
286+
"""
287+
The Dimensions the guards read, that is those a guarded object must
288+
be evaluated within.
289+
"""
290+
return frozenset({d for v in self.values()
291+
for d in pull_dims(v, flag=False)})
292+
283293
def has(self, d, cls):
284294
"""
285295
True if the guard registered for `d` contains an instance of `cls`.

‎devito/passes/clusters/aliases.py‎

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@
99
from devito.finite_differences import EvalDerivative, IndexDerivative, Weights
1010
from devito.ir import (
1111
PARALLEL_IF_PVT, SEPARABLE, SEQUENTIAL, Cluster, ClusterGroup, ExprGeometry, Forward,
12-
Interval, IntervalGroup, IterationSpace, LabeledVector, Queue, Vector, extrema,
13-
maximum, minimum, normalize_properties, relax_properties, unbounded, vmax, vmin
12+
Interval, IntervalGroup, IterationSpace, LabeledVector, Properties, Queue, Vector,
13+
extrema, maximum, minimum, normalize_properties, relax_properties, unbounded, vmax,
14+
vmin
1415
)
1516
from devito.passes.clusters.cse import _cse
1617
from devito.passes.clusters.utils import expose_tuning_knobs
@@ -19,8 +20,8 @@
1920
uxreplace
2021
)
2122
from devito.tools import (
22-
Reconstructable, Stamp, as_mapper, as_tuple, flatten, frozendict, generator,
23-
is_integer, split, timed_pass
23+
Reconstructable, Stamp, as_mapper, as_tuple, flatten, generator, is_integer, split,
24+
timed_pass
2425
)
2526
from devito.types import (
2627
CustomDimension, Eq, Hyperplane, IncrDimension, Indexed, ModuloDimension, Size,
@@ -288,6 +289,7 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None):
288289
free_symbols = i.free_symbols
289290
if {a.function for a in free_symbols} & exclude:
290291
continue
292+
291293
mapper.add(i, make, terms)
292294

293295
return mapper
@@ -304,6 +306,31 @@ def __init__(self, sregistry, options, platform):
304306
def process(self, clusters):
305307
return self._process_fatd(clusters, 1, xtracted=[])
306308

309+
@classmethod
310+
def _make_exclude(cls, clusters, d, p):
311+
"""
312+
The symbols an extraction must not touch.
313+
"""
314+
# Rule out extractions that would break data dependencies
315+
exclude = set().union(*[c.scope.writes for c in clusters])
316+
317+
# Rule out extractions that depend on the Dimension currently investigated,
318+
# as they clearly wouldn't be invariants
319+
exclude.update({d, *p.sub_iterators})
320+
321+
# An extraction is hoisted out of `d`, but it inherits its guard as it
322+
# is, so it must not be hoisted past any Dimension the guard reads, or
323+
# the guard would be evaluated where it is not defined. Excluding those
324+
# Dimensions keeps the extraction inside the loops defining them. A
325+
# SEQUENTIAL Dimension is exempt, its loop enclosing the extraction's own
326+
for c in clusters:
327+
exclude.update(
328+
i for i in c.guards.dimensions
329+
if not c.properties.is_sequential(i._defines)
330+
)
331+
332+
return exclude
333+
307334
def callback(self, clusters, prefix, xtracted=None):
308335
if not prefix:
309336
return clusters
@@ -314,12 +341,7 @@ def callback(self, clusters, prefix, xtracted=None):
314341
if d.is_Virtual:
315342
return clusters
316343

317-
# Rule out extractions that would break data dependencies
318-
exclude = set().union(*[c.scope.writes for c in clusters])
319-
320-
# Rule out extractions that depend on the Dimension currently investigated,
321-
# as they clearly wouldn't be invariants
322-
exclude.update({d, *p.sub_iterators})
344+
exclude = self._make_exclude(clusters, d, p)
323345

324346
key = lambda c: self._lookup_key(c, d)
325347
processed = list(clusters)
@@ -343,7 +365,7 @@ def callback(self, clusters, prefix, xtracted=None):
343365
def _lookup_key(self, c, d):
344366
ispace = c.ispace.reset()
345367
intervals = c.ispace.intervals.drop(d).reset()
346-
properties = frozendict({d: relax_properties(v) for d, v in c.properties.items()})
368+
properties = Properties({d: relax_properties(v) for d, v in c.properties.items()})
347369

348370
return AliasKey(ispace, intervals, c.dtype, c.guards, properties)
349371

‎devito/passes/iet/asynchrony.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from devito.ir import (
88
AsyncCall, AsyncCallable, BlankLine, Call, Callable, Conditional, DummyEq, DummyExpr,
99
EntryFunction, FindNodes, FindSymbols, Increment, Iteration, List, PointerCast,
10-
Return, ThreadCallable, Transformer, While, make_callable, maybe_alias
10+
Return, ThreadCallable, ThreadFence, Transformer, While, make_callable, maybe_alias
1111
)
1212
from devito.passes.iet.definitions import DataManager
1313
from devito.passes.iet.engine import iet_pass
@@ -86,6 +86,11 @@ def _lower_async_objs(iet, tracker=None, sregistry=None, **kwargs):
8686
arguments.append(i)
8787
activation.extend([DummyExpr(FieldFromComposite(i.base, sdata[d]), i)
8888
for i in arguments])
89+
90+
# Publishing the request must happen last. Everything the thread depends on
91+
# must be visible before this flag is set. `volatile` does not provide that
92+
# ordering, so weakly ordered targets could otherwise observe stale state.
93+
activation.append(ThreadFence('release'))
8994
activation.append(
9095
DummyExpr(FieldFromComposite(sdata.symbolic_flag, sdata[d]), 2)
9196
)
@@ -138,12 +143,17 @@ def _(iet, key=None, tracker=None, sregistry=None, **kwargs):
138143
tbase = threads.indexed
139144

140145
# Prepend the SharedData fields available upon thread activation
141-
preactions = [DummyExpr(i, FieldFromPointer(i.base, sbase)) for i in ncfields]
146+
preactions = [ThreadFence('acquire')]
147+
preactions.extend(DummyExpr(i, FieldFromPointer(i.base, sbase))
148+
for i in ncfields)
142149
preactions.append(BlankLine)
143150

144151
# Append the flag reset
145152
postactions = [List(body=[
146153
BlankLine,
154+
# Whatever the task produced -- data in a buffer, a lock handed back --
155+
# has to be visible before the flag says it is done.
156+
ThreadFence('release'),
147157
DummyExpr(FieldFromPointer(sdata.symbolic_flag, sbase), 1)
148158
])]
149159

‎devito/passes/iet/orchestration.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from devito.exceptions import CompilationError
88
from devito.ir.iet import (
99
AsyncCall, AsyncCallable, BlankLine, Block, BusyWait, Call, Callable, Conditional,
10-
DummyExpr, List, SyncSpot, Transformer, derive_parameters, make_callable
10+
DummyExpr, List, SyncSpot, ThreadFence, Transformer, derive_parameters, make_callable
1111
)
1212
from devito.ir.iet.visitors import Visitor
1313
from devito.ir.support import (
@@ -53,7 +53,11 @@ def _make_waitlock(self, iet, sync_ops, *args):
5353
def _make_releaselock(self, iet, sync_ops, *args):
5454
pre = []
5555
pre.append(BusyWait(Or(*[CondNe(s.handle, 2) for s in sync_ops])))
56+
pre.append(ThreadFence('acquire'))
5657
pre.extend(DummyExpr(s.handle, 0) for s in sync_ops)
58+
# Hand the lock back before anything that follows can be observed --
59+
# in particular before the request that asks the thread to refill it.
60+
pre.append(ThreadFence('release'))
5761

5862
name = self.sregistry.make_name(prefix="release_lock")
5963
parameters = derive_parameters(pre, ordering='canonical')

‎tests/test_dse.py‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
)
1111
from devito import ( # noqa
1212
NODE, Abs, ConditionalDimension, Constant, DefaultDimension, Derivative, Dimension,
13-
Eq, Function, Ge, Grid, Inc, Lt, Operator, SparseTimeFunction, SubDimension,
13+
Eq, Function, Ge, Grid, Inc, Lt, Max, Operator, SparseTimeFunction, SubDimension,
1414
TimeFunction, configuration, cos, dimensions, div, exp, first_derivative, floor, grad,
1515
norm, sin, solve, sqrt, switchconfig, transpose
1616
)
@@ -2512,6 +2512,43 @@ def test_contraction_with_conditional(self):
25122512
assert len(FindNodes(Conditional).visit(op)) == 1
25132513
assert np.all(u.data[6:] == 1.42)
25142514

2515+
def test_no_extraction_guarded_by_unspanned_dimension(self):
2516+
"""
2517+
An alias is scheduled over the Dimensions it spans, but it inherits its
2518+
guard as it is. A guard reading a Dimension the alias does not span
2519+
would then be evaluated in a loop nest that does not define it, giving
2520+
code that does not compile.
2521+
2522+
Here `sin(...)` depends on `y` alone while the condition reads both `x`
2523+
and `y`, which is what an immersed boundary condition produces.
2524+
"""
2525+
grid = Grid(shape=(16, 16))
2526+
x, y = grid.dimensions
2527+
2528+
sdf = Function(name='sdf', grid=grid)
2529+
sdf.data[:] = 1.
2530+
2531+
cond = ConditionalDimension(name='inside', parent=y, condition=Ge(sdf, 0))
2532+
2533+
u = TimeFunction(name='u', grid=grid, space_order=4)
2534+
u.data[:] = 1.
2535+
2536+
# Expensive enough to be extracted, and `y`-only
2537+
prof = sin(Max(0., 1. - y)) + sin(Max(0., 1. + y))
2538+
2539+
eqn = Eq(u.forward, u.laplace + prof*u, implicit_dims=[cond])
2540+
2541+
op = Operator(eqn, opt=('advanced', {'cire-mingain': 0, 'openmp': False}))
2542+
2543+
# No temporary may be created over fewer Dimensions than the guard reads
2544+
for i in FindSymbols().visit(op):
2545+
if i.is_Array:
2546+
assert {x, y}.issubset(set(i.dimensions))
2547+
2548+
# Used to fail to compile with "'x' undeclared"
2549+
op.apply(time_M=2)
2550+
assert np.all(np.isfinite(u.data[:]))
2551+
25152552
def test_collection_from_conditional(self):
25162553
nt = 10
25172554
grid = Grid(shape=(10, 10))

‎tests/test_gpu_common.py‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -413,11 +413,16 @@ def test_tasking_in_isolation(self, opt):
413413
assert str(sections[0].body[0].body[0].body[0].body[0]) == 'while(lock0[0] == 0);'
414414
body = op._func_table['release_lock0'].root.body
415415
assert str(body.body[0].condition) == 'Ne(lock0[0], 2)'
416-
assert str(body.body[1]) == 'lock0[0] = 0;'
416+
# An acquire fence pairs with the thread's release before the lock is
417+
# read back, and a release fence publishes it before the next request
418+
assert 'atomic_thread_fence' in str(body.body[1])
419+
assert str(body.body[2]) == 'lock0[0] = 0;'
420+
assert 'atomic_thread_fence' in str(body.body[3])
417421
body = op._func_table['activate0'].root.body
418422
assert str(body.body[0].condition) == 'Ne(sdata0[0].flag, 1)'
419423
assert str(body.body[1]) == 'sdata0[0].time = time;'
420-
assert str(body.body[2]) == 'sdata0[0].flag = 2;'
424+
assert 'atomic_thread_fence' in str(body.body[2])
425+
assert str(body.body[3]) == 'sdata0[0].flag = 2;'
421426

422427
op.apply(time_M=nt-2)
423428

@@ -529,12 +534,15 @@ def test_tasking_forcefuse(self):
529534
'while(lock0[0] == 0 || lock1[0] == 0);') # Wait-lock
530535
body = op._func_table['release_lock0'].root.body
531536
assert str(body.body[0].condition) == 'Ne(lock0[0], 2) | Ne(lock1[0], 2)'
532-
assert str(body.body[1]) == 'lock0[0] = 0;' # Set-lock
533-
assert str(body.body[2]) == 'lock1[0] = 0;' # Set-lock
537+
assert 'atomic_thread_fence' in str(body.body[1])
538+
assert str(body.body[2]) == 'lock0[0] = 0;' # Set-lock
539+
assert str(body.body[3]) == 'lock1[0] = 0;' # Set-lock
540+
assert 'atomic_thread_fence' in str(body.body[4])
534541
body = op._func_table['activate0'].root.body
535542
assert str(body.body[0].condition) == 'Ne(sdata0[0].flag, 1)' # Wait-thread
536543
assert str(body.body[1]) == 'sdata0[0].time = time;'
537-
assert str(body.body[2]) == 'sdata0[0].flag = 2;'
544+
assert 'atomic_thread_fence' in str(body.body[2])
545+
assert str(body.body[3]) == 'sdata0[0].flag = 2;'
538546
assert len(op._func_table) == 5
539547
exprs = FindNodes(Expression).visit(op._func_table['copy_to_host0'].root)
540548
b = 21 if configuration['language'] == 'openacc' else 20 # No `qid` w/ OMP
@@ -597,8 +605,9 @@ def test_tasking_multi_output(self):
597605
assert str(sections[0].body[0].body[0].body[0].body[0]) ==\
598606
'while(lock0[t2] == 0);'
599607
body = op1._func_table['release_lock0'].root.body
608+
assert 'atomic_thread_fence' in str(body.body[1])
600609
for i in range(3):
601-
assert 'lock0[t' in str(body.body[1 + i]) # Set-lock
610+
assert 'lock0[t' in str(body.body[2 + i]) # Set-lock
602611
body = op1._func_table['activate0'].root.body
603612
assert str(body.body[-1]) == 'sdata0[wi0].flag = 2;'
604613
assert len(op1._func_table) == 5

‎tests/test_iet.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from devito.ir.iet import (
1313
Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr,
1414
ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List,
15-
Switch, Transformer, filter_iterations, make_callable, make_efunc,
15+
Switch, ThreadFence, Transformer, filter_iterations, make_callable, make_efunc,
1616
retrieve_iteration_tree
1717
)
1818
from devito.ir.iet.visitors import sorted_efuncs
@@ -118,6 +118,23 @@ def test_nested_calls_cgen():
118118
assert str(code) == 'foo(bar());'
119119

120120

121+
def test_thread_fence_cgen():
122+
"""
123+
A fence renders as the builtin, which is valid in C and C++ alike.
124+
125+
Threads that hand work to each other through a shared flag need one either
126+
side of the handshake; without them the stores can be observed out of order
127+
and a hand-off is lost.
128+
"""
129+
assert str(CGen().visit(ThreadFence('acquire'))) == \
130+
'__atomic_thread_fence(__ATOMIC_ACQUIRE);'
131+
assert str(CGen().visit(ThreadFence('release'))) == \
132+
'__atomic_thread_fence(__ATOMIC_RELEASE);'
133+
134+
with pytest.raises(AssertionError):
135+
ThreadFence('seq_cst')
136+
137+
121138
@pytest.mark.parametrize('mode,expected', [
122139
('basics', '["x"]'),
123140
('symbolics', '["f"]')

0 commit comments

Comments
 (0)