def test_multiple_subnests(self):
grid = Grid(shape=(3, 3, 3))
x, y, z = grid.dimensions
t = grid.stepping_dim
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid)
eqn = Eq(
u.forward,
(u[t, x, y, z] * u[t, x + 1, y + 1, z + 1] * 3 * f +
u[t, x + 2, y + 2, z + 2] * u[t, x + 3, y + 3, z + 3] * 3 * f +
1))
op = Operator(eqn,
dse='aggressive',
dle=('advanced', {
'openmp': True
}))
trees = retrieve_iteration_tree(op._func_table['bf0'].root)
assert len(trees) == 2
assert trees[0][0] is trees[1][0]
assert trees[0][0].pragmas[0].value ==\
'omp for collapse(1) schedule(static,1)'
assert trees[0][2].pragmas[0].value ==\
('omp parallel for collapse(1) schedule(static,1) num_threads(%d)'
% nhyperthreads())
assert trees[1][2].pragmas[0].value ==\
('omp parallel for collapse(1) schedule(static,1) num_threads(%d)'
% nhyperthreads())
def test_multiple_subnests(self):
grid = Grid(shape=(3, 3, 3))
x, y, z = grid.dimensions
t = grid.stepping_dim
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid)
eqn = Eq(u.forward, (u[t, x, y, z]*u[t, x+1, y+1, z+1]*3*f +
u[t, x+2, y+2, z+2]*u[t, x+3, y+3, z+3]*3*f + 1))
op = Operator(eqn, dse='aggressive', dle=('advanced', {'openmp': True}))
trees = retrieve_iteration_tree(op._func_table['bf0'].root)
assert len(trees) == 2
assert trees[0][0] is trees[1][0]
assert trees[0][0].pragmas[0].value ==\
'omp for collapse(1) schedule(static,1)'
assert trees[0][2].pragmas[0].value ==\
('omp parallel for collapse(1) schedule(static,1) num_threads(%d)'
% nhyperthreads())
assert trees[1][2].pragmas[0].value ==\
('omp parallel for collapse(1) schedule(static,1) num_threads(%d)'
% nhyperthreads())
def test_collapsing(self):
grid = Grid(shape=(3, 3, 3))
u = TimeFunction(name='u', grid=grid)
op = Operator(Eq(u.forward, u + 1), dle=('blocking', 'openmp'))
# Does it compile? Honoring the OpenMP specification isn't trivial
assert op.cfunction
# Does it produce the right result
op.apply(t_M=9)
assert np.all(u.data[0] == 10)
iterations = FindNodes(Iteration).visit(op._func_table['bf0'])
assert iterations[0].pragmas[0].value == 'omp for collapse(2) schedule(static,1)'
assert iterations[2].pragmas[0].value ==\
('omp parallel for collapse(2) schedule(static,1) num_threads(%d)'
% nhyperthreads())
def test_collapsing(self):
grid = Grid(shape=(3, 3, 3))
u = TimeFunction(name='u', grid=grid)
op = Operator(Eq(u.forward, u + 1), dle=('blocking', 'openmp'))
# Does it compile? Honoring the OpenMP specification isn't trivial
assert op.cfunction
# Does it produce the right result
op.apply(t_M=9)
assert np.all(u.data[0] == 10)
iterations = FindNodes(Iteration).visit(op._func_table['bf0'])
assert iterations[0].pragmas[
0].value == 'omp for collapse(2) schedule(static,1)'
assert iterations[2].pragmas[0].value ==\
('omp parallel for collapse(2) schedule(static,1) num_threads(%d)'
% nhyperthreads())