test_visitors.py

import cgen as c
from sympy import Mod
import pytest

from devito import Grid, Eq, Function, TimeFunction, Operator, sin
from devito.ir.equations import DummyEq
from devito.ir.iet import (Block, Expression, Callable, FindNodes, FindSections,
                           FindSymbols, IsPerfectIteration, Transformer,
                           Conditional, printAST, Iteration, MapNodes, Call)
from devito.types import SpaceDimension, Array


@pytest.fixture(scope="module")
def dims():
    x, y, z = Grid((3, 3, 3)).dimensions
    return {'i': SpaceDimension(name='i'),
            'j': SpaceDimension(name='j'),
            'k': SpaceDimension(name='k'),
            'l': SpaceDimension(name='l'),
            's': SpaceDimension(name='s'),
            'q': SpaceDimension(name='q'),
            'x': x, 'y': y, 'z': z}


@pytest.fixture(scope="module")
def iters(dims):
    return [lambda ex: Iteration(ex, dims['i'], (0, 3, 1)),
            lambda ex: Iteration(ex, dims['j'], (0, 5, 1)),
            lambda ex: Iteration(ex, dims['k'], (0, 7, 1)),
            lambda ex: Iteration(ex, dims['s'], (0, 4, 1)),
            lambda ex: Iteration(ex, dims['q'], (0, 4, 1)),
            lambda ex: Iteration(ex, dims['l'], (0, 6, 1)),
            lambda ex: Iteration(ex, dims['x'], (0, 5, 1)),
            lambda ex: Iteration(ex, dims['y'], (0, 7, 1)),
            lambda ex: Iteration(ex, dims['z'], (0, 7, 1))]


@pytest.fixture(scope="module")
def exprs(dims):
    a = Array(name='a', shape=(3,), dimensions=(dims["i"],)).indexify()
    b = Array(name='b', shape=(3,), dimensions=(dims["i"],)).indexify()
    return [Expression(DummyEq(a, a + b + 5.)),
            Expression(DummyEq(a, b - a)),
            Expression(DummyEq(a, 4 * (b * a))),
            Expression(DummyEq(a, (6. / b) + (8. * a)))]


@pytest.fixture(scope="module")
def block1(exprs, iters):
    # Perfect loop nest:
    # for i
    #   for j
    #     for k
    #       expr0
    return iters[0](iters[1](iters[2](exprs[0])))


@pytest.fixture(scope="module")
def block2(exprs, iters):
    # Non-perfect simple loop nest:
    # for i
    #   expr0
    #   for j
    #     for k
    #       expr1
    return iters[0]([exprs[0], iters[1](iters[2](exprs[1]))])


@pytest.fixture(scope="module")
def block3(exprs, iters):
    # Non-perfect non-trivial loop nest:
    # for i
    #   for s
    #     expr0
    #   for j
    #     for k
    #       expr1
    #       expr2
    #   for p
    #     expr3
    return iters[0]([iters[3](exprs[0]),
                     iters[1](iters[2]([exprs[1], exprs[2]])),
                     iters[4](exprs[3])])


@pytest.fixture(scope="module")
def block4(exprs, iters, dims):
    # Non-perfect loop nest due to conditional
    # for i
    #   if i % 2 == 0
    #   for j
    return iters[0](Conditional(Eq(Mod(dims['i'], 2), 0), iters[1](exprs[0])))


def test_printAST(block1, block2, block3, block4):
    str1 = printAST(block1)
    assert str1 in """
<Iteration i::i::(0, 3, 1)>
  <Iteration j::j::(0, 5, 1)>
    <Iteration k::k::(0, 7, 1)>
      <Expression a[i] = a[i] + b[i] + 5.0>
"""

    str2 = printAST(block2)
    assert str2 in """
<Iteration i::i::(0, 3, 1)>
  <Expression a[i] = a[i] + b[i] + 5.0>
  <Iteration j::j::(0, 5, 1)>
    <Iteration k::k::(0, 7, 1)>
      <Expression a[i] = -a[i] + b[i]>
"""

    str3 = printAST(block3)
    assert str3 in """
<Iteration i::i::(0, 3, 1)>
  <Iteration s::s::(0, 4, 1)>
    <Expression a[i] = a[i] + b[i] + 5.0>
  <Iteration j::j::(0, 5, 1)>
    <Iteration k::k::(0, 7, 1)>
      <Expression a[i] = -a[i] + b[i]>
      <Expression a[i] = 4*a[i]*b[i]>
  <Iteration q::q::(0, 4, 1)>
    <Expression a[i] = 8.0*a[i] + 6.0/b[i]>
"""

    str4 = printAST(block4)
    assert str4 in """
<Iteration i::i::(0, 3, 1)>
  <If Eq(Mod(i, 2), 0)>
    <Iteration j::j::(0, 5, 1)>
      <Expression a[i] = a[i] + b[i] + 5.0>
"""


def test_create_cgen_tree(block1, block2, block3):
    assert str(Callable('foo', block1, 'void', ()).ccode) == """\
void foo()
{
  for (int i = 0; i <= 3; i += 1)
  {
    for (int j = 0; j <= 5; j += 1)
    {
      for (int k = 0; k <= 7; k += 1)
      {
        a[i] = a[i] + b[i] + 5.0F;
      }
    }
  }
}"""
    assert str(Callable('foo', block2, 'void', ()).ccode) == """\
void foo()
{
  for (int i = 0; i <= 3; i += 1)
  {
    a[i] = a[i] + b[i] + 5.0F;

    for (int j = 0; j <= 5; j += 1)
    {
      for (int k = 0; k <= 7; k += 1)
      {
        a[i] = -a[i] + b[i];
      }
    }
  }
}"""
    assert str(Callable('foo', block3, 'void', ()).ccode) == """\
void foo()
{
  for (int i = 0; i <= 3; i += 1)
  {
    for (int s = 0; s <= 4; s += 1)
    {
      a[i] = a[i] + b[i] + 5.0F;
    }
    for (int j = 0; j <= 5; j += 1)
    {
      for (int k = 0; k <= 7; k += 1)
      {
        a[i] = -a[i] + b[i];
        a[i] = 4*a[i]*b[i];
      }
    }
    for (int q = 0; q <= 4; q += 1)
    {
      a[i] = 8.0F*a[i] + 6.0F/b[i];
    }
  }
}"""


def test_find_sections(exprs, block1, block2, block3):
    finder = FindSections()

    sections = finder.visit(block1)
    assert len(sections) == 1

    sections = finder.visit(block2)
    assert len(sections) == 2
    found = list(sections.values())
    assert len(found[0]) == 1
    assert len(found[1]) == 1

    sections = finder.visit(block3)
    assert len(sections) == 3
    found = list(sections.values())
    assert len(found[0]) == 1
    assert len(found[1]) == 2
    assert len(found[2]) == 1


def test_is_perfect_iteration(block1, block2, block3, block4):
    checker = IsPerfectIteration()

    assert checker.visit(block1) is True
    assert checker.visit(block1.nodes[0]) is True
    assert checker.visit(block1.nodes[0].nodes[0]) is True

    assert checker.visit(block2) is False
    assert checker.visit(block2.nodes[1]) is True
    assert checker.visit(block2.nodes[1].nodes[0]) is True

    assert checker.visit(block3) is False
    assert checker.visit(block3.nodes[0]) is True
    assert checker.visit(block3.nodes[1]) is True
    assert checker.visit(block3.nodes[2]) is True

    assert checker.visit(block4) is False
    assert checker.visit(block4.nodes[0]) is False
    assert checker.visit(block4.nodes[0].then_body) is True


def test_transformer_wrap(exprs, block1, block2, block3):
    """Basic transformer test that wraps an expression in comments"""
    line1 = '// This is the opening comment'
    line2 = '// This is the closing comment'
    wrapper = lambda n: Block(c.Line(line1), n, c.Line(line2))
    transformer = Transformer({exprs[0]: wrapper(exprs[0])})

    for block in [block1, block2, block3]:
        newblock = transformer.visit(block)
        newcode = str(newblock.ccode)
        oldnumlines = len(str(block.ccode).split('\n'))
        newnumlines = len(newcode.split('\n'))
        assert newnumlines >= oldnumlines + 2
        assert line1 in newcode
        assert line2 in newcode
        assert "a[i] = a[i] + b[i] + 5.0F;" in newcode


def test_transformer_replace(exprs, block1, block2, block3):
    """Basic transformer test that replaces an expression"""
    line1 = '// Replaced expression'
    replacer = Block(c.Line(line1))
    transformer = Transformer({exprs[0]: replacer})

    for block in [block1, block2, block3]:
        newblock = transformer.visit(block)
        newcode = str(newblock.ccode)
        oldnumlines = len(str(block.ccode).split('\n'))
        newnumlines = len(newcode.split('\n'))
        assert newnumlines >= oldnumlines
        assert line1 in newcode
        assert "a[i0] = a[i0] + b[i0] + 5.0F;" not in newcode


def test_transformer_replace_function_body(block1, block2):
    """Create a Function and replace its body with another."""
    args = FindSymbols().visit(block1)
    f = Callable('foo', block1, 'void', args)
    assert str(f.ccode) == """void foo(float *restrict a_vec, float *restrict b_vec)
{
  for (int i = 0; i <= 3; i += 1)
  {
    for (int j = 0; j <= 5; j += 1)
    {
      for (int k = 0; k <= 7; k += 1)
      {
        a[i] = a[i] + b[i] + 5.0F;
      }
    }
  }
}"""

    f = Transformer({block1: block2}).visit(f)
    assert str(f.ccode) == """void foo(float *restrict a_vec, float *restrict b_vec)
{
  for (int i = 0; i <= 3; i += 1)
  {
    a[i] = a[i] + b[i] + 5.0F;

    for (int j = 0; j <= 5; j += 1)
    {
      for (int k = 0; k <= 7; k += 1)
      {
        a[i] = -a[i] + b[i];
      }
    }
  }
}"""


def test_transformer_add_replace(exprs, block2, block3):
    """Basic transformer test that adds one expression and replaces another"""
    line1 = '// Replaced expression'
    line2 = '// Adding a simple line'
    replacer = Block(c.Line(line1))
    adder = lambda n: Block(c.Line(line2), n)
    transformer = Transformer({exprs[0]: replacer,
                               exprs[1]: adder(exprs[1])})

    for block in [block2, block3]:
        newblock = transformer.visit(block)
        newcode = str(newblock.ccode)
        oldnumlines = len(str(block.ccode).split('\n'))
        newnumlines = len(newcode.split('\n'))
        assert newnumlines >= oldnumlines + 1
        assert line1 in newcode
        assert line2 in newcode
        assert "a[i0] = a[i0] + b[i0] + 5.0F;" not in newcode


def test_nested_transformer(exprs, iters, block2):
    """When created with the kwarg ``nested=True``, a Transformer performs
    nested replacements. This test simultaneously replace an inner expression
    and an Iteration sorrounding it."""
    target_loop = block2.nodes[1]
    target_expr = target_loop.nodes[0].nodes[0]
    mapper = {target_loop: iters[3](target_loop.nodes[0]),
              target_expr: exprs[3]}
    processed = Transformer(mapper, nested=True).visit(block2)
    assert printAST(processed) == """<Iteration i::i::(0, 3, 1)>
  <Expression a[i] = a[i] + b[i] + 5.0>
  <Iteration s::s::(0, 4, 1)>
    <Iteration k::k::(0, 7, 1)>
      <Expression a[i] = 8.0*a[i] + 6.0/b[i]>"""


def test_find_symbols():
    grid = Grid(shape=(4, 4))
    x, y = grid.dimensions

    f = Function(name='f', grid=grid)

    op = Operator(Eq(f, f[x-1, y] + f[x+1, y] + 1))

    symbols = FindSymbols('indexeds').visit(op)

    assert len(symbols) == 3


def test_find_symbols_with_duplicates():
    grid = Grid(shape=(4, 4))
    x, y = grid.dimensions

    f = TimeFunction(name='f', grid=grid)
    g = Function(name='g', grid=grid)

    eq = Eq(f.forward, sin(g)*f + g)

    op = Operator(eq)

    exprs = FindNodes(Expression).visit(op)

    assert len(exprs) == 2

    # Two syntactically identical indexeds r0[x, y], but they're different
    # objects because they are constructed in two different places during
    # the CIRE pass
    r0a = exprs[0].output
    r0b = exprs[1].expr.rhs.args[0].args[0]
    assert r0a.function is r0b.function
    assert r0a.name == r0b.name == "r0"
    assert hash(r0a) == hash(r0b)
    assert r0a is not r0b

    # So we expect FindSymbols to catch five Indexeds in total
    symbols = FindSymbols('indexeds').visit(op)
    assert len(symbols) == 5


def test_map_nodes(block1):
    """
    Tests MapNodes visitor. When MapNodes is created with mode='groupby',
    matching ancestors are grouped together under a single key.
    This can be useful, for example, when applying transformations to the
    outermost Iteration containing a specific node.
    """
    map_nodes = MapNodes(Iteration, Expression, mode='groupby').visit(block1)

    assert len(map_nodes.keys()) == 1

    for iters, (expr,) in map_nodes.items():
        # Replace the outermost `Iteration` with a `Call`
        callback = Callable('solver', iters[0], 'void', ())
        processed = Transformer({iters[0]: Call(callback.name)}).visit(block1)

    assert str(processed) == 'solver();'