Skip to content

Commit ad24b75

Browse files
committed
compiler: Don't treat a subsampled Inc as a reduction
`Inc` into a factor-subsampled Function was classified PARALLEL_IF_ATOMIC, so an OpenMP array reduction was emitted over the whole Function. With nine such Functions the thread-private copies came to ~13 MiB and the kernel segfaulted; four arrays worked, six did not, which is the size of the default stack. The dependence is reported with an all-zero distance vector, so the write is plainly independent across iterations. What pushed it into the reduction branch was `is_storage_related`, which treats any NonlinearDerived Dimension as capable of mapping several iterations onto one location. That is true of a SteppingDimension, where `time % 2` deliberately revisits a slot, but not of a factor-based ConditionalDimension: the guard admits only `parent % factor == 0`, and `i -> i//factor` is one-to-one on those. Add `Dimension.is_Injective`, true exactly for that case, and exempt it. An arbitrary `condition`, or a factor combined with a condition, carries no such promise and stays non-injective.
1 parent 7330b23 commit ad24b75

3 files changed

Lines changed: 153 additions & 3 deletions

File tree

‎devito/ir/support/basic.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -869,8 +869,15 @@ def is_storage_related(self, dims=None):
869869
"""
870870
for d in self.findices:
871871
if d._defines & set(as_tuple(dims)): # noqa: SIM102
872-
if any(i.is_NonlinearDerived for i in d._defines) or \
873-
self.is_const(d):
872+
# An injective Dimension maps distinct iterations to distinct
873+
# locations by construction -- e.g. a factor-based
874+
# ConditionalDimension, whose guard admits only
875+
# `parent % factor == 0` and whose `i -> i//factor` is
876+
# one-to-one on those. Unlike, say, a SteppingDimension, where
877+
# `time % 2` deliberately revisits the same slot
878+
nonlinear = [i for i in d._defines
879+
if i.is_NonlinearDerived and not i.is_Injective]
880+
if nonlinear or self.is_const(d):
874881
return True
875882
return False
876883

‎devito/types/dimension.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ class Dimension(ArgProvider):
116116
is_Custom = False
117117
is_Derived = False
118118
is_NonlinearDerived = False
119+
is_Injective = False
119120
is_AbstractSub = False
120121
is_Sub = False
121122
is_MultiSub = False
@@ -1090,6 +1091,19 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
10901091
self._indirect = indirect
10911092
self._relation = relation
10921093

1094+
@property
1095+
def is_Injective(self):
1096+
"""
1097+
True if distinct iterations of the parent that pass the guard map to
1098+
distinct indices of `self`, so that no two iterations can access the
1099+
same memory location through it.
1100+
1101+
A pure factor subsampling is injective: the guard admits only
1102+
`parent % factor == 0`, and `i -> i//factor` is one-to-one on those.
1103+
An arbitrary `condition` carries no such promise.
1104+
"""
1105+
return self._factor is not None and self._condition is None
1106+
10931107
@property
10941108
def uses_symbolic_factor(self):
10951109
return self._factor is not None

‎tests/test_subsampling.py‎

Lines changed: 130 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from sympy import Or
66

77
from devito import (
8-
CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Operator, SubGrid,
8+
CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Inc, Operator, SubGrid,
99
TimeFunction, switchconfig
1010
)
1111
from devito.ir.iet import Conditional, Expression, FindNodes, Iteration
@@ -915,3 +915,132 @@ def test_time_subsampling(self, mode):
915915

916916
for i in range((nt - 2)//factor + 1):
917917
assert np.allclose(vsave.data[i], i*factor)
918+
919+
920+
class TestIncReduction:
921+
922+
"""
923+
`Inc` into a factor-subsampled Function must not be treated as a
924+
reduction. The write index `i -> i//factor` is one-to-one on the guarded
925+
points, so no two iterations touch the same location and the loop is
926+
plainly parallel.
927+
928+
Treating it as one emitted an OpenMP array reduction over the whole
929+
Function; with several such Functions the private copies overflowed the
930+
stack and the kernel segfaulted.
931+
"""
932+
933+
@pytest.mark.parametrize('factor', [1, 2, 3])
934+
def test_no_reduction_clause(self, factor):
935+
n = 12
936+
grid = Grid(shape=(n, n))
937+
x, y = grid.dimensions
938+
cd = ConditionalDimension(name='xs', parent=x, factor=factor)
939+
940+
f = Function(name='f', shape=((n + factor - 1)//factor, n),
941+
dimensions=(cd, y), grid=grid, space_order=0)
942+
943+
with switchconfig(language='openmp'):
944+
op = Operator(Inc(f, 1))
945+
946+
assert 'reduction' not in str(op.ccode)
947+
948+
@pytest.mark.parametrize('factor', [1, 2, 3])
949+
def test_dimensions_are_parallel(self, factor):
950+
"""
951+
The properties must match what the same `Inc` gets on a plain
952+
Function, i.e. PARALLEL rather than PARALLEL_IF_ATOMIC.
953+
"""
954+
n = 12
955+
grid = Grid(shape=(n, n))
956+
x, y = grid.dimensions
957+
cd = ConditionalDimension(name='xs', parent=x, factor=factor)
958+
959+
f = Function(name='f', shape=((n + factor - 1)//factor, n),
960+
dimensions=(cd, y), grid=grid, space_order=0)
961+
962+
op = Operator(Inc(f, 1))
963+
964+
for i in FindNodes(Iteration).visit(op):
965+
assert i.is_Parallel
966+
967+
@pytest.mark.parametrize('factor', [1, 2, 3])
968+
def test_accumulates(self, factor):
969+
"""Dropping the reduction must not drop the accumulation."""
970+
n = 12
971+
grid = Grid(shape=(n, n))
972+
x, y = grid.dimensions
973+
cd = ConditionalDimension(name='xs', parent=x, factor=factor)
974+
975+
f = Function(name='f', shape=((n + factor - 1)//factor, n),
976+
dimensions=(cd, y), grid=grid, space_order=0)
977+
f.data[:] = 0.
978+
979+
op = Operator(Inc(f, 1))
980+
op.apply()
981+
op.apply()
982+
983+
assert np.all(f.data == 2.)
984+
985+
def test_many_functions(self):
986+
"""
987+
The original failure: nine subsampled Functions incremented together.
988+
The spurious array reductions needed ~13 MiB of thread-private copies
989+
and overflowed the stack.
990+
"""
991+
n = 40
992+
grid = Grid(shape=(n, n, n))
993+
dims = tuple(ConditionalDimension(name=f'{d.name}s', parent=d, factor=1)
994+
for d in grid.dimensions)
995+
996+
fs = [Function(name=f'og_{i}', grid=grid, shape=(n, n, n),
997+
dimensions=dims) for i in range(9)]
998+
for f in fs:
999+
f.data[:] = 0.
1000+
1001+
with switchconfig(language='openmp'):
1002+
op = Operator([Inc(f, 1) for f in fs])
1003+
assert 'reduction' not in str(op.ccode)
1004+
op.apply()
1005+
1006+
for f in fs:
1007+
assert np.all(f.data == 1.)
1008+
1009+
1010+
class TestInjectivity:
1011+
1012+
"""`is_Injective` must be claimed only where it holds."""
1013+
1014+
def test_factor_is_injective(self):
1015+
grid = Grid(shape=(8, 8))
1016+
x, _ = grid.dimensions
1017+
1018+
assert ConditionalDimension(name='a', parent=x, factor=2).is_Injective
1019+
1020+
def test_stepping_is_not(self):
1021+
"""`time % 2` deliberately revisits the same slot."""
1022+
grid = Grid(shape=(8, 8))
1023+
1024+
assert not grid.stepping_dim.is_Injective
1025+
1026+
def test_condition_is_not(self):
1027+
"""An arbitrary condition carries no one-to-one promise."""
1028+
grid = Grid(shape=(8, 8))
1029+
x, _ = grid.dimensions
1030+
cd = ConditionalDimension(name='b', parent=x, condition=CondEq(x, 0))
1031+
1032+
assert not cd.is_Injective
1033+
1034+
def test_factor_with_condition_is_not(self):
1035+
grid = Grid(shape=(8, 8))
1036+
x, _ = grid.dimensions
1037+
cd = ConditionalDimension(name='c', parent=x, factor=2,
1038+
condition=CondEq(x % 4, 0))
1039+
1040+
assert not cd.is_Injective
1041+
1042+
def test_plain_dimension_is_not(self):
1043+
grid = Grid(shape=(8, 8))
1044+
x, _ = grid.dimensions
1045+
1046+
assert not x.is_Injective

0 commit comments

Comments
 (0)