Skip to content

Commit

Permalink
compiler: Added pass to abridge SubDimension names where possible in …
Browse files Browse the repository at this point in the history
…the generated C
  • Loading branch information
EdCaunt authored and Edward Caunt committed Dec 4, 2023
1 parent 25fa68b commit 10b7e8c
Show file tree
Hide file tree
Showing 12 changed files with 175 additions and 90 deletions.
10 changes: 5 additions & 5 deletions devito/passes/clusters/aliases.py
Original file line number Diff line number Diff line change
Expand Up @@ -781,10 +781,10 @@ def optimize_schedule_rotations(schedule, sregistry):
iis = candidate.lower
iib = candidate.upper

ii = ModuloDimension('%sii' % d, ds, iis, incr=iib)
cd = CustomDimension(name='%s%s' % (d, d), symbolic_min=ii, symbolic_max=iib,
symbolic_size=n)
dsi = ModuloDimension('%si' % ds, cd, cd + ds - iis, n)
ii = ModuloDimension('%sii' % d.root.name, ds, iis, incr=iib)
cd = CustomDimension(name='%s%s' % (d.root.name, d.root.name), symbolic_min=ii,
symbolic_max=iib, symbolic_size=n)
dsi = ModuloDimension('%si' % ds.root.name, cd, cd + ds - iis, n)

mapper = OrderedDict()
for i in g:
Expand All @@ -795,7 +795,7 @@ def optimize_schedule_rotations(schedule, sregistry):
try:
md = mapper[v]
except KeyError:
name = sregistry.make_name(prefix='%sr' % d.name)
name = sregistry.make_name(prefix='%sr' % d.root.name)
md = mapper.setdefault(v, ModuloDimension(name, ds, v, n))
mds.append(md)
indicess = [indices[:ridx] + [md] + indices[ridx + 1:]
Expand Down
41 changes: 40 additions & 1 deletion devito/passes/iet/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from devito.finite_differences import Max, Min
from devito.ir import (Any, Forward, Iteration, List, Prodder, FindApplications,
FindNodes, FindSymbols, Transformer, Uxreplace,
filter_iterations, retrieve_iteration_tree)
filter_iterations, retrieve_iteration_tree, pull_dims)
from devito.passes.iet.engine import iet_pass
from devito.symbolics import evalrel, has_integer_args
from devito.tools import as_mapper, filter_ordered, split
Expand Down Expand Up @@ -171,8 +171,11 @@ def minimize_symbols(iet):
* Remove redundant ModuloDimensions (e.g., due to using the
`save=Buffer(2)` API)
* Abridge SubDimension names where possible to declutter generated
loop nests and shrink indices
"""
iet = remove_redundant_moddims(iet)
iet = abridge_dim_names(iet)

return iet, {}

Expand Down Expand Up @@ -204,3 +207,39 @@ def remove_redundant_moddims(iet):
iet = Transformer(subs, nested=True).visit(iet)

return iet


def abridge_dim_names(iet):
subs = {}
mapper = {}
tree = retrieve_iteration_tree(iet)
for it in tree:
indexeds = FindSymbols('indexeds').visit(it.inner)
inds = set().union(*[pull_dims(i, flag=False) for i in indexeds])
inds = [i for i in inds if any([d.is_Sub for d in i._defines])]
inds = [i for i in inds if not i.is_SubIterator]
names = [i.root.name for i in inds]

mapper.update({i: i._rebuild(i.root.name) for i in inds
if i.root not in it.dimensions and names.count(i.root.name) < 2})

if not mapper:
return iet

for it in tree:
uinds = set().union(*[i.uindices for i in it])
uinds = [i for i in uinds if i.is_Incr and i.parent in mapper]
mapper.update({i: i._rebuild(parent=mapper[i.parent]) for i in uinds})

for it in tree:
dims = FindSymbols('dimensions').visit(it)
dims = [d for d in dims if d.is_Modulo and d.parent in mapper]
mapper.update({d: d._rebuild(parent=mapper[d.parent]) for d in dims})

for it in tree:
subs.update({i: Uxreplace(mapper).visit(i) for i in it})

if subs:
iet = Transformer(subs, nested=True).visit(iet)

return iet
6 changes: 6 additions & 0 deletions devito/tools/abc.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,12 @@ def __init__(self, a, b, c=4):
args += tuple(getattr(self, i[1:]))
else:
args += (getattr(self, i),)

args = list(args)
for k in list(kwargs):
if k in self.__rargs__:
args[self.__rargs__.index(k)] = kwargs.pop(k)

kwargs.update({i: getattr(self, i) for i in self.__rkwargs__ if i not in kwargs})

# Should we use a constum reconstructor?
Expand Down
15 changes: 13 additions & 2 deletions devito/types/grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -549,7 +549,10 @@ class MultiSubDimension(SubDimension):
A special SubDimension for graceful lowering of MultiSubDomains.
"""

def __init_finalize__(self, name, parent, msd):
__rargs__ = ('name', 'parent', 'msd')
__rkwargs__ = ('thickness',)

def __init_finalize__(self, name, parent, msd, thickness=None):
# NOTE: a MultiSubDimension stashes a reference to the originating MultiSubDomain.
# This creates a circular pattern as the `msd` itself carries references to
# its MultiSubDimensions. This is currently necessary because during compilation
Expand All @@ -558,7 +561,15 @@ def __init_finalize__(self, name, parent, msd):
# definitely possible, but not straightforward
self.msd = msd

lst, rst = self._symbolic_thickness(name)
if not thickness:
lst, rst = self._symbolic_thickness(name)
else: # Used for rebuilding. Reuse thickness symbols rather than making new ones
try:
((lst, _), (rst, _)) = thickness
except ValueError:
raise ValueError("Invalid thickness specification: %s does not match"
"expected format ((left_symbol, left_thickness),"
" (right_symbol, right_thickness))" % thickness)
left = parent.symbolic_min + lst
right = parent.symbolic_max - rst

Expand Down
89 changes: 43 additions & 46 deletions examples/cfd/01_convection_revisited.ipynb

Large diffs are not rendered by default.

6 changes: 3 additions & 3 deletions examples/performance/00_overview.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -1360,12 +1360,12 @@
" {\n",
" for (int y = y0_blk0, ys = 0, yr0 = (ys)%(5), yr1 = (ys + 3)%(5), yr2 = (ys + 4)%(5), yr3 = (ys + 1)%(5), yii = -2; y <= MIN(y_M, y0_blk0 + y0_blk0_size - 1); y += 1, ys += 1, yr0 = (ys)%(5), yr1 = (ys + 3)%(5), yr2 = (ys + 4)%(5), yr3 = (ys + 1)%(5), yii = 2)\n",
" {\n",
" for (int yy = yii, ysi = (yy + ys + 2)%(5); yy <= 2; yy += 1, ysi = (yy + ys + 2)%(5))\n",
" for (int yy = yii, yi = (yy + ys + 2)%(5); yy <= 2; yy += 1, yi = (yy + ys + 2)%(5))\n",
" {\n",
" #pragma omp simd aligned(u:32)\n",
" for (int z = z_m; z <= z_M; z += 1)\n",
" {\n",
" r2[ysi][z] = r1*(8.33333333e-2F*(u[t0][x + 4][y + yy + 2][z + 4] - u[t0][x + 4][y + yy + 6][z + 4]) + 6.66666667e-1F*(-u[t0][x + 4][y + yy + 3][z + 4] + u[t0][x + 4][y + yy + 5][z + 4]));\n",
" r2[yi][z] = r1*(8.33333333e-2F*(u[t0][x + 4][y + yy + 2][z + 4] - u[t0][x + 4][y + yy + 6][z + 4]) + 6.66666667e-1F*(-u[t0][x + 4][y + yy + 3][z + 4] + u[t0][x + 4][y + yy + 5][z + 4]));\n",
" }\n",
" }\n",
" #pragma omp simd aligned(f,u:32)\n",
Expand Down Expand Up @@ -1744,7 +1744,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.16"
"version": "3.11.5"
},
"varInspector": {
"cols": {
Expand Down
19 changes: 11 additions & 8 deletions tests/test_dimension.py
Original file line number Diff line number Diff line change
Expand Up @@ -403,14 +403,14 @@ def test_flow_detection_interior(self):
@pytest.mark.parametrize('exprs,expected,', [
# Carried dependence in both /t/ and /x/
(['Eq(u[t+1, x, y], u[t+1, x-1, y] + u[t, x, y])'], 'y'),
(['Eq(u[t+1, x, y], u[t+1, x-1, y] + u[t, x, y], subdomain=interior)'], 'i0y'),
(['Eq(u[t+1, x, y], u[t+1, x-1, y] + u[t, x, y], subdomain=interior)'], 'y'),
# Carried dependence in both /t/ and /y/
(['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y])'], 'x'),
(['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y], subdomain=interior)'], 'i0x'),
(['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y], subdomain=interior)'], 'x'),
# Carried dependence in /y/, leading to separate /y/ loops, one
# going forward, the other backward
(['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y], subdomain=interior)',
'Eq(u[t+1, x, y], u[t+1, x, y+1] + u[t, x, y], subdomain=interior)'], 'i0x'),
'Eq(u[t+1, x, y], u[t+1, x, y+1] + u[t, x, y], subdomain=interior)'], 'x'),
])
def test_iteration_property_parallel(self, exprs, expected):
"""Tests detection of sequental and parallel Iterations when applying
Expand All @@ -435,7 +435,7 @@ def test_iteration_property_parallel(self, exprs, expected):
@skipif(['device'])
@pytest.mark.parametrize('exprs,expected,', [
# All parallel, the innermost Iteration gets vectorized
(['Eq(u[time, x, yleft], u[time, x, yleft] + 1.)'], ['yleft']),
(['Eq(u[time, x, yleft], u[time, x, yleft] + 1.)'], ['y']),
# All outers are parallel, carried dependence in `yleft`, so the middle
# Iteration over `x` gets vectorized
(['Eq(u[time, x, yleft], u[time, x, yleft+1] + 1.)'], ['x']),
Expand Down Expand Up @@ -637,9 +637,11 @@ def test_subdimmiddle_subdimleft_blocked(self):

# Check generated code -- expected loop blocking over x and y, with the
# two z loops, zi and zl, within y
# Note that the zi and zl iterators are renamed to remove clutter given the loop
# structure
assert_structure(op,
['t,x0_blk0,y0_blk0,x,y,zi', 't,x0_blk0,y0_blk0,x,y,zl'],
't,x0_blk0,y0_blk0,x,y,zi,zl')
['t,x0_blk0,y0_blk0,x,y,z', 't,x0_blk0,y0_blk0,x,y,z'],
't,x0_blk0,y0_blk0,x,y,z,z')

op.apply(time_M=0)
assert np.all(u.data[0] == 0)
Expand All @@ -660,9 +662,10 @@ def test_subdim_blocked_w_partial_reuse(self):
op = Operator(eqns)

# Check generated code -- expected loop blocking
# Note that loop structure means zr is renamed to z to remove clutter
assert_structure(op,
['t,x0_blk0,y0_blk0,x,y,z', 't,x0_blk0,y0_blk0,x,y,zr'],
't,x0_blk0,y0_blk0,x,y,z,zr')
['t,x0_blk0,y0_blk0,x,y,z', 't,x0_blk0,y0_blk0,x,y,z'],
't,x0_blk0,y0_blk0,x,y,z,z')

def test_subdim_fd(self):
"""
Expand Down
16 changes: 11 additions & 5 deletions tests/test_dle.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,11 @@ def test_cache_blocking_structure_subdims():
tree[2].dim.root is x
assert tree[3].dim.is_Block and tree[3].dim.parent is tree[1].dim and\
tree[3].dim.root is y
assert not tree[4].dim.is_Block and tree[4].dim is zi and tree[4].dim.parent is z
# zi is rebuilt with name z, so check symbolic max and min are preserved
# Also check the zi was rebuilt
assert not tree[4].dim.is_Block and tree[4].dim is not zi and\
tree[4].dim.symbolic_min is zi.symbolic_min and\
tree[4].dim.symbolic_max is zi.symbolic_max and tree[4].dim.parent is z


@pytest.mark.parallel(mode=[(1, 'full')]) # Shortcut to put loops in nested efuncs
Expand Down Expand Up @@ -1280,17 +1284,19 @@ def test_nested_cache_blocking_structure_subdims(self, blocklevels):
tree[3].dim.root is y

if blocklevels == 1:
assert not tree[4].dim.is_Block and tree[4].dim is zi and\
tree[4].dim.parent is z
assert not tree[4].dim.is_Block and tree[4].dim is not zi and\
tree[4].dim.symbolic_min is zi.symbolic_min and\
tree[4].dim.symbolic_max is zi.symbolic_max and tree[4].dim.parent is z
elif blocklevels == 2:
assert tree[3].dim.is_Block and tree[3].dim.parent is tree[1].dim and\
tree[3].dim.root is y
assert tree[4].dim.is_Block and tree[4].dim.parent is tree[2].dim and\
tree[4].dim.root is x
assert tree[5].dim.is_Block and tree[5].dim.parent is tree[3].dim and\
tree[5].dim.root is y
assert not tree[6].dim.is_Block and tree[6].dim is zi and\
tree[6].dim.parent is z
assert not tree[6].dim.is_Block and tree[6].dim is not zi and\
tree[6].dim.symbolic_min is zi.symbolic_min and\
tree[6].dim.symbolic_max is zi.symbolic_max and tree[6].dim.parent is z

assert trees[0][0].pragmas[0].value ==\
'omp for collapse(2) schedule(dynamic,1)'
Expand Down
2 changes: 1 addition & 1 deletion tests/test_dse.py
Original file line number Diff line number Diff line change
Expand Up @@ -2652,7 +2652,7 @@ def test_issue_2163(self):
u = TimeFunction(name="u", grid=grid)
op = Operator(Eq(u.forward, u.dy.dy.subs(mapper),
subdomain=grid.interior))
assert_structure(op, ['t,i0x,i0y'], 'ti0xi0y')
assert_structure(op, ['t,x,y'], 'txy')

def test_dtype_aliases(self):
a = np.arange(64).reshape((8, 8))
Expand Down
10 changes: 5 additions & 5 deletions tests/test_fission.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,12 @@ def define(self, dimensions):

# Note the `x` loop is fissioned, so now both loop nests can be collapsed
# for maximum parallelism
assert_structure(op, ['t,x,i1y', 't,x,i2y'], 't,x,i1y,x,i2y')
assert_structure(op, ['t,x,y', 't,x,y'], 't,x,y,x,y')


def test_nofission_as_unprofitable():
"""
Test there's no fission if not gonna increase number of collapsable loops.
Test there's no fission if no increase in number of collapsable loops.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
Expand All @@ -54,7 +54,7 @@ def test_nofission_as_unprofitable():

op = Operator(eqns, opt='fission')

assert_structure(op, ['t,x,yl', 't,x,yr'], 't,x,yl,yr')
assert_structure(op, ['t,x,y', 't,x,y'], 't,x,y,y')


def test_nofission_as_illegal():
Expand All @@ -78,7 +78,7 @@ def test_nofission_as_illegal():

def test_fission_partial():
"""
Test there's no fission if not gonna increase number of collapsable loops.
Test there's no fission if no increase in number of collapsable loops.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
Expand All @@ -95,7 +95,7 @@ def test_fission_partial():

op = Operator(eqns, opt='fission')

assert_structure(op, ['t,x,yl', 't,x,yr', 't,x,y'], 't,x,yl,yr,x,y')
assert_structure(op, ['t,x,y', 't,x,y', 't,x,y'], 't,x,y,y,x,y')


def test_issue_1921():
Expand Down
4 changes: 2 additions & 2 deletions tests/test_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -1518,7 +1518,7 @@ def test_no_fission_as_illegal(self, exprs):
# Here the difference is that we're using SubDimensions
(('Eq(tv[t,xi,yi,zi], tu[t,xi-1,yi,zi] + tu[t,xi+1,yi,zi])',
'Eq(tu[t+1,xi,yi,zi], tu[t,xi,yi,zi] + tv[t,xi-1,yi,zi] + tv[t,xi+1,yi,zi])'),
'+++++++', ['ti0xi0yi0z', 'ti0xi0yi0z'], 'ti0xi0yi0zi0xi0yi0z'),
'+++++++', ['txyz', 'txyz'], 'txyzxyz'),
# 16) RAW 3->1; expected=2
# Time goes backward, but the third equation should get fused with
# the first one, as the time dependence is loop-carried
Expand All @@ -1545,7 +1545,7 @@ def test_no_fission_as_illegal(self, exprs):
(('Eq(tv[t+1,x,y,z], tu[t,x,y,z] + tu[t,x+1,y,z])',
'Eq(tu[t+1,xi,yi,zi], tv[t+1,xi,yi,zi] + tv[t+1,xi+1,yi,zi])',
'Eq(tw[t+1,x,y,z], tv[t+1,x,y,z] + tv[t+1,x+1,y,z])'),
'++++++++++', ['txyz', 'ti0xi0yi0z', 'txyz'], 'txyzi0xi0yi0zxyz'),
'++++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzxyz'),
])
def test_consistency_anti_dependences(self, exprs, directions, expected, visit):
"""
Expand Down
Loading

0 comments on commit 10b7e8c

Please sign in to comment.