tests: simplify by ret ParallelBlocks

tests/conftest.py

@@ -12,7 +12,7 @@
 from devito.finite_differences.differentiable import EvalDerivative
 from devito.arch import Device, sniff_mpi_distro
 from devito.arch.compiler import compiler_registry
-from devito.ir.iet import retrieve_iteration_tree, FindNodes, Iteration
+from devito.ir.iet import retrieve_iteration_tree, FindNodes, Iteration, ParallelBlock
 from devito.tools import as_tuple
 
 try:
@@ -213,12 +213,12 @@ def _R(expr):
 def assert_structure(operator, exp_trees=None, exp_iters=None):
     """
     Utility function that helps to check loop structure of IETs. Retrieves trees from an
-    Operator and check that blocking structure is as expected. Trees and Iterations are
-    returned for further use in tests.
+    Operator and check that the blocking structure is as expected. Trees and Iterations
+    are returned for further use in tests.
 
-    Example:
-
-    To check the following structure:
+    Examples
+    --------
+    To check that an Iteration tree has the following structure:
 
     .. code-block:: python
 
@@ -230,9 +230,10 @@ def assert_structure(operator, exp_trees=None, exp_iters=None):
 
     we call:
 
-    `trees, iters = assert_structure(op, ['t,x,y', 't,f,y'], 't,x,y,f,y')`
-    """
+    .. code-block:: python
 
+        trees, iters = assert_structure(op, ['t,x,y', 't,f,y'], 't,x,y,f,y')`
+    """
     mapper = {'time': 't'}
     trees = retrieve_iteration_tree(operator)
     iters = FindNodes(Iteration).visit(operator)
@@ -251,13 +252,13 @@ def assert_structure(operator, exp_trees=None, exp_iters=None):
     return trees, iters
 
 
-def get_blocked_nests(operator, exp_nests):
+def assert_blocking(operator, exp_nests):
     """
     Utility function that helps to check existence of blocked nests. The structure of the
-    operator is not used. Trees are returned for further use in tests.
-
-    Example:
+    operator is not used.
 
+    Examples
+    --------
     For the following structure:
 
     .. code-block:: python
@@ -270,30 +271,42 @@ def get_blocked_nests(operator, exp_nests):
 
     we call:
 
-    `trees, bns = get_blocked_nests(op, {'x0_blk0', 'x1_blk0'})`
+    .. code-block:: python
+
+       bns, pbs = assert_blocking(op, {'x0_blk0', 'x1_blk0'})
+
+    to assert the existence of 'x0_blk0', 'x1_blk0' and then the function returns a
+    dictionary with the blocking Iterations that start a blocking subtree:
 
-    to assert the existence of 'x0_blk0', 'x1_blk0' and then the function returns the
-    operator trees and a dictionary with the first blocking Iterations that start a
-    blocking subtree
+    ['x0_blk0': Iteration x0_blk0..., 'x1_blk0': Iteration x1_blk0...]
 
+    and the ParallelBlock that encapsules the above blocking subtree
+
+    ['x0_blk0': ParallelBlock encapsulating Iteration x0_blk0,
+     'x1_blk0': ParallelBlock encapsulating Iteration x1_blk0]
     """
     bns = {}
-    iterations = []
+    pbs = {}
     trees = retrieve_iteration_tree(operator)
     for tree in trees:
         iterations = [i for i in tree if i.dim.is_Incr]  # Collect Incr dimensions
-        if len(iterations):
+        parallel_blocks = FindNodes(ParallelBlock).visit(tree)
+        if iterations:
             # If Incr dimensions exist map the first one to its name in the dict
             bns[iterations[0].dim.name] = iterations[0]
-            iterations = []
+            try:
+                pbs[iterations[0].dim.name] = parallel_blocks[0]
+            except IndexError:
+                pbs[iterations[0].dim.name] = tree[0]
 
     # Return if no Incr dimensions, ensuring that no Incr expected
     if not bns and not exp_nests:
-        return trees, bns
+        return {}, {}
 
     # Assert Incr dimensions found as expected
     assert bns.keys() == exp_nests
-    return trees, bns
+
+    return bns, pbs
 
 
 # A list of optimization options/pipelines to be used in testing

tests/test_autotuner.py

@@ -1,7 +1,7 @@
 import pytest
 import numpy as np
 
-from conftest import get_blocked_nests, skipif
+from conftest import assert_blocking, skipif
 from devito import Grid, TimeFunction, Eq, Operator, configuration, switchconfig
 from devito.data import LEFT
 from devito.core.autotuning import options  # noqa
@@ -238,7 +238,7 @@ def test_multiple_blocking():
                   opt=('blocking', {'openmp': False}))
 
     # First of all, make sure there are indeed two different loop nests
-    _ = get_blocked_nests(op, {'x0_blk0', 'x1_blk0'})
+    _, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0'})
 
     # 'basic' mode
     op.apply(time_M=0, autotune='basic')

tests/test_dimension.py

@@ -4,7 +4,7 @@
 from sympy import And
 import pytest
 
-from conftest import get_blocked_nests, skipif, opts_tiling
+from conftest import assert_blocking, skipif, opts_tiling
 from devito import (ConditionalDimension, Grid, Function, TimeFunction, SparseFunction,  # noqa
                     Eq, Operator, Constant, Dimension, SubDimension, switchconfig,
                     SubDomain, Lt, Le, Gt, Ge, Ne, Buffer)
@@ -1057,7 +1057,8 @@ def test_no_fusion_simple(self):
 
         op = Operator(eqns)
 
-        _, bns = get_blocked_nests(op, {'x0_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0'})
+
         exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 4
         assert exprs[1].expr.rhs is exprs[0].output
@@ -1072,7 +1073,9 @@ def test_no_fusion_simple(self):
                 Eq(g, f + 1, implicit_dims=[ctime])]
 
         op = Operator(eqns)
-        _, bns = get_blocked_nests(op, {'x0_blk0', 'x1_blk0'})
+
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0'})
+
         exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 3
         assert exprs[1].expr.rhs is exprs[0].output
@@ -1104,7 +1107,8 @@ def test_no_fusion_convoluted(self):
 
         op = Operator(eqns)
 
-        _, bns = get_blocked_nests(op, {'x0_blk0', 'x1_blk0', 'x2_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'x2_blk0'})
+
         exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 3
         assert exprs[1].expr.rhs is exprs[0].output
@@ -1175,7 +1179,8 @@ def test_topofusion_w_subdims_conddims(self):
 
         # Check generated code -- expect the gsave equation to be scheduled together
         # in the same loop nest with the fsave equation
-        _, bns = get_blocked_nests(op, {'x0_blk0', 'x1_blk0', 'i0x0_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0', 'x1_blk0', 'i0x0_blk0'})
+
         exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is f
@@ -1216,7 +1221,8 @@ def test_topofusion_w_subdims_conddims_v2(self):
 
         # Check generated code -- expect the gsave equation to be scheduled together
         # in the same loop nest with the fsave equation
-        _, bns = get_blocked_nests(op, {'i0x0_blk0', 'x0_blk0'})
+        bns, _ = assert_blocking(op, {'i0x0_blk0', 'x0_blk0'})
+
         assert len(FindNodes(Expression).visit(bns['i0x0_blk0'])) == 3
         exprs = FindNodes(Expression).visit(bns['x0_blk0'])
         assert len(exprs) == 2
@@ -1249,7 +1255,8 @@ def test_topofusion_w_subdims_conddims_v3(self):
 
         # Check generated code -- expect the gsave equation to be scheduled together
         # in the same loop nest with the fsave equation
-        _, bns = get_blocked_nests(op, {'i0x0_blk0', 'x0_blk0', 'i0x1_blk0'})
+        bns, _ = assert_blocking(op, {'i0x0_blk0', 'x0_blk0', 'i0x1_blk0'})
+
         exprs = FindNodes(Expression).visit(bns['i0x0_blk0'])
         assert len(exprs) == 2
         assert exprs[0].write is f

tests/test_dle.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pytest
 
-from conftest import assert_structure, get_blocked_nests, _R
+from conftest import assert_structure, assert_blocking, _R
 from devito import (Grid, Function, TimeFunction, SparseTimeFunction, SpaceDimension,
                     Dimension, SubDimension, Eq, Inc, Operator, info)
 from devito.exceptions import InvalidArgument
@@ -96,9 +96,9 @@ def test_cache_blocking_structure(blockinner):
                                               'par-collapse-ncores': 1}))
     for opi in [op, op2]:
         if blockinner:
-            trees, _ = assert_structure(opi, ['t,x0_blk0,y0_blk0,z0_blk0,x,y,z'])
+            _, _ = assert_structure(opi, ['t,x0_blk0,y0_blk0,z0_blk0,x,y,z'])
         else:
-            trees, _ = assert_structure(opi, ['t,x0_blk0,y0_blk0,x,y,z'])
+            _, _ = assert_structure(opi, ['t,x0_blk0,y0_blk0,x,y,z'])
 
     # Check presence of openmp pragmas at the right place
     trees = retrieve_iteration_tree(op2)
@@ -125,11 +125,24 @@ def test_cache_blocking_structure_subdims():
 
     # Local SubDimension -> no blocking expected
     op = Operator(Eq(f[t+1, xl, y, z], f[t, xl, y, z] + 1))
-    _, _ = get_blocked_nests(op, {})
+
+    _, _ = assert_blocking(op, {})
 
     # Non-local SubDimension -> blocking expected
     op = Operator(Eq(f.forward, f + 1, subdomain=grid.interior))
-    _, _ = assert_structure(op, ['t,i0x0_blk0,i0y0_blk0,i0x,i0y,i0z'])
+
+    bns, _ = assert_blocking(op, {'i0x0_blk0'})
+
+    trees = retrieve_iteration_tree(bns['i0x0_blk0'])
+    tree = trees[0]
+    assert len(tree) == 5
+    assert tree[0].dim.is_Incr and tree[0].dim.parent is xi and tree[0].dim.root is x
+    assert tree[1].dim.is_Incr and tree[1].dim.parent is yi and tree[1].dim.root is y
+    assert tree[2].dim.is_Incr and tree[2].dim.parent is tree[0].dim and\
+        tree[2].dim.root is x
+    assert tree[3].dim.is_Incr and tree[3].dim.parent is tree[1].dim and\
+        tree[3].dim.root is y
+    assert not tree[4].dim.is_Incr and tree[4].dim is zi and tree[4].dim.parent is z
 
 
 @pytest.mark.parallel(mode=[(1, 'full')])  # Shortcut to put loops in nested efuncs
@@ -151,8 +164,8 @@ def test_cache_blocking_structure_distributed():
 
     op = Operator(eqns)
 
-    _, bns0 = get_blocked_nests(op._func_table['compute0'].root, {'x0_blk0'})
-    _, bns1 = get_blocked_nests(op, {'x1_blk0'})
+    bns0, _ = assert_blocking(op._func_table['compute0'].root, {'x0_blk0'})
+    bns1, _ = assert_blocking(op, {'x1_blk0'})
 
     for i in [bns0['x0_blk0'], bns1['x1_blk0']]:
         iters = FindNodes(Iteration).visit(i)
@@ -264,7 +277,8 @@ def test_cache_blocking_imperfect_nest(blockinner):
     op1 = Operator(eqns, opt=('advanced', {'blockinner': blockinner}))
 
     # First, check the generated code
-    _, bns = get_blocked_nests(op1, {'x0_blk0'})
+    bns, _ = assert_blocking(op1, {'x0_blk0'})
+
     trees = retrieve_iteration_tree(bns['x0_blk0'])
     assert len(trees) == 2
     assert len(trees[0]) == len(trees[1])
@@ -310,7 +324,8 @@ def test_cache_blocking_imperfect_nest_v2(blockinner):
     op2 = Operator(eq, opt=('advanced-fsg', {'blockinner': blockinner}))
 
     # First, check the generated code
-    _, bns = get_blocked_nests(op2, {'x0_blk0'})
+    bns, _ = assert_blocking(op2, {'x0_blk0'})
+
     trees = retrieve_iteration_tree(bns['x0_blk0'])
     assert len(trees) == 2
     assert len(trees[0]) == len(trees[1])
@@ -619,7 +634,8 @@ def test_basic(self):
         assert np.all(u.data[0] == 10)
         assert op.arguments(t_M=9, nthreads_nested=2)['nthreads_nested'] == 2
 
-        _, bns = get_blocked_nests(op, ({'x0_blk0'}))
+        bns, _ = assert_blocking(op, {'x0_blk0'})
+
         iterations = FindNodes(Iteration).visit(bns['x0_blk0'])
         assert iterations[0].pragmas[0].value == 'omp for collapse(1) schedule(dynamic,1)'
         assert iterations[2].pragmas[0].value == ('omp parallel for collapse(1) '
@@ -644,7 +660,8 @@ def test_collapsing(self):
         op.apply(t_M=9)
         assert np.all(u.data[0] == 10)
 
-        _, bns = get_blocked_nests(op, ({'x0_blk0'}))
+        bns, _ = assert_blocking(op, {'x0_blk0'})
+
         iterations = FindNodes(Iteration).visit(bns['x0_blk0'])
         assert iterations[0].pragmas[0].value == 'omp for collapse(2) schedule(dynamic,1)'
         assert iterations[2].pragmas[0].value == ('omp parallel for collapse(2) '
@@ -667,7 +684,8 @@ def test_multiple_subnests_v0(self):
                                              'par-collapse-ncores': 1,
                                              'par-dynamic-work': 0}))
 
-        _, bns = get_blocked_nests(op, ({'x0_blk0'}))
+        bns, _ = assert_blocking(op, {'x0_blk0'})
+
         trees = retrieve_iteration_tree(bns['x0_blk0'])
         assert len(trees) == 2
 
@@ -702,7 +720,8 @@ def test_multiple_subnests_v1(self):
                                              'par-collapse-ncores': 1,
                                              'par-dynamic-work': 0}))
 
-        _, bns = get_blocked_nests(op, ({'x0_blk0'}))
+        bns, _ = assert_blocking(op, {'x0_blk0'})
+
         trees = retrieve_iteration_tree(bns['x0_blk0'])
         assert len(trees) == 2
 
@@ -743,7 +762,9 @@ def test_nested_cache_blocking_structure_subdims(self, blocklevels):
                            {'par-nested': 0, 'blocklevels': blocklevels,
                             'par-collapse-ncores': 2,
                             'par-dynamic-work': 0}))
-        _, bns = get_blocked_nests(op, ({'i0x0_blk0'}))
+
+        bns, _ = assert_blocking(op, {'i0x0_blk0'})
+
         trees = retrieve_iteration_tree(bns['i0x0_blk0'])
         assert len(trees) == 1
         tree = trees[0]