-
Notifications
You must be signed in to change notification settings - Fork 222
/
test_fission.py
96 lines (65 loc) · 2.74 KB
/
test_fission.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
from conftest import assert_structure
from devito import (Eq, Inc, Grid, Function, TimeFunction, SubDimension, SubDomain,
Operator, solve)
def test_issue_1725():
class ToyPMLLeft(SubDomain):
name = 'toypmlleft'
def define(self, dimensions):
x, y = dimensions
return {x: x, y: ('left', 2)}
class ToyPMLRight(SubDomain):
name = 'toypmlright'
def define(self, dimensions):
x, y = dimensions
return {x: x, y: ('right', 2)}
subdomains = [ToyPMLLeft(), ToyPMLRight()]
grid = Grid(shape=(20, 20), subdomains=subdomains)
u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
eqns = [Eq(u.forward, solve(u.dt2 - u.laplace, u.forward), subdomain=sd)
for sd in subdomains]
op = Operator(eqns, opt='fission')
# 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')
def test_nofission_as_unprofitable():
"""
Test there's no fission if not gonna increase number of collapsable loops.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
t = grid.stepping_dim
yl = SubDimension.left(name='yl', parent=y, thickness=4)
yr = SubDimension.right(name='yr', parent=y, thickness=4)
u = TimeFunction(name='u', grid=grid)
eqns = [Eq(u.forward, u[t + 1, x, y + 1] + 1.).subs(y, yl),
Eq(u.forward, u[t + 1, x, y - 1] + 1.).subs(y, yr)]
op = Operator(eqns, opt='fission')
assert_structure(op, ['t,x,yl', 't,x,yr'], 't,x,yl,yr')
def test_nofission_as_illegal():
"""
Test there's no fission if dependencies would break.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
f = Function(name='f', grid=grid, dimensions=(y,), shape=(20,))
u = TimeFunction(name='u', grid=grid)
v = TimeFunction(name='v', grid=grid)
eqns = [Inc(f, v + 1.),
Eq(u.forward, f[y + 1] + 1.)]
op = Operator(eqns, opt='fission')
assert_structure(op, ['t,x,y', 't,x,y'], 't,x,y,y')
def test_fission_partial():
"""
Test there's no fission if not gonna increase number of collapsable loops.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
t = grid.stepping_dim
yl = SubDimension.left(name='yl', parent=y, thickness=4)
yr = SubDimension.right(name='yr', parent=y, thickness=4)
u = TimeFunction(name='u', grid=grid)
eqns = [Eq(u.forward, u[t + 1, x, y + 1] + 1.).subs(y, yl),
Eq(u.forward, u[t + 1, x, y - 1] + 1.).subs(y, yr),
Eq(u.forward, u[t + 1, x, y] + 1.)]
op = Operator(eqns, opt='fission')
assert_structure(op, ['t,x,yl', 't,x,yr', 't,x,y'], 't,x,yl,yr,x,y')