compiler: Add a node for standalone macro statements

devito/ir/iet/nodes.py

@@ -49,6 +49,7 @@
     'Lambda',
     'List',
     'ListMajor',
+    'MacroStatement',
     'MetaCall',
     'MultiTraversable',
     'Node',
@@ -1012,20 +1013,20 @@ def __init__(self, timer, lname, body):
         self._name = lname
         self._timer = timer
 
-        super().__init__(header=c.Line(f'START({lname})'),
+        super().__init__(header=MacroStatement(f'START({lname})'),
                          body=body,
-                         footer=c.Line(f'STOP({lname},{timer.name})'))
+                         footer=MacroStatement(f'STOP({lname},{timer.name})'))
 
     @classmethod
     def _start_timer_header(cls):
         return ('START(S)', ('struct timeval start_ ## S , end_ ## S ; '
-                             'gettimeofday(&start_ ## S , NULL);'))
+                             'gettimeofday(&start_ ## S , NULL)'))
 
     @classmethod
     def _stop_timer_header(cls):
         return ('STOP(S,T)', ('gettimeofday(&end_ ## S, NULL); T->S += (double)'
                               '(end_ ## S .tv_sec-start_ ## S.tv_sec)+(double)'
-                              '(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000;'))
+                              '(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000'))
 
     @property
     def name(self):
@@ -1331,6 +1332,25 @@ def periodic(self):
         return self._periodic
 
 
+class MacroStatement(ExprStmt, Node):
+
+    """
+    A standalone macro invocation, such as ``MY_MACRO;``.
+
+    Parameters
+    ----------
+    name : str
+        The macro name or complete invocation (including any arguments),
+        without a trailing semicolon.
+    """
+
+    def __init__(self, name):
+        self.name = name
+
+    def __repr__(self):
+        return f"<MacroStatement({self.name})>"
+
+
 class Using(Node):
 
     """

devito/ir/iet/visitors.py

@@ -563,7 +563,9 @@ def visit_Block(self, o):
     def visit_List(self, o):
         body = flatten(self._visit(i) for i in self._blankline_logic(o.children))
         body = c.Line(' '.join(str(i) for i in body)) if o.inline else c.Collection(body)
-        return c.Module(o.header + (body,) + o.footer)
+        header = tuple(self._visit(i) if isinstance(i, Node) else i for i in o.header)
+        footer = tuple(self._visit(i) if isinstance(i, Node) else i for i in o.footer)
+        return c.Module(header + (body,) + footer)
 
     def visit_Section(self, o):
         body = flatten(self._visit(i) for i in o.children)
@@ -572,6 +574,9 @@ def visit_Section(self, o):
     def visit_Break(self, o):
         return c.Statement('break')
 
+    def visit_MacroStatement(self, o):
+        return c.Statement(o.name)
+
     def visit_Return(self, o):
         v = 'return'
         if o.value is not None:

devito/passes/iet/errors.py

@@ -6,7 +6,7 @@
 
 from devito.ir.iet import (
     Break, Call, Conditional, DummyExpr, EntryFunction, FindNodes, FindSymbols, Iteration,
-    KernelLaunch, List, Return, Transformer, make_callable
+    KernelLaunch, List, MacroStatement, Return, Transformer, make_callable
 )
 from devito.mpi.distributed import MPICommObject
 from devito.passes.iet.engine import iet_pass
@@ -134,9 +134,9 @@ def _check_launch(iet, definition=None, options=None, **kwargs):
     if options['mpi'] and \
        isinstance(iet, EntryFunction) and \
        any(isinstance(i, MPICommObject) for i in iet.parameters):
-        check = [List(body=c.Statement('CHECK_LAUNCH_RETURN'))]
+        check = [MacroStatement('CHECK_LAUNCH_RETURN')]
     else:
-        check = [List(body=c.Statement('CHECK_LAUNCH'))]
+        check = [MacroStatement('CHECK_LAUNCH')]
 
     mapper = {}
     for i in iterations:

devito/passes/iet/misc.py

@@ -7,8 +7,8 @@
 from devito.finite_differences import Max, Min
 from devito.finite_differences.differentiable import SafeInv
 from devito.ir import (
-    Any, DummyExpr, EmptyList, FindApplications, FindNodes, FindSymbols, Forward,
-    Iteration, Prodder, Transformer, Uxreplace, filter_iterations, pull_dims,
+    Any, DummyExpr, FindApplications, FindNodes, FindSymbols, Forward, Iteration, List,
+    MacroStatement, Prodder, Transformer, Uxreplace, filter_iterations, pull_dims,
     retrieve_iteration_tree
 )
 from devito.ir.iet.efunc import DeviceFunction, EntryFunction
@@ -49,12 +49,14 @@ def avoid_denormals(iet, platform=None, **kwargs):
     if iet.is_ElementalFunction:
         return iet, {}
 
-    header = (cgen.Comment('Flush denormal numbers to zero in hardware'),
-              cgen.Statement('_MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON)'),
-              cgen.Statement('_MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON)'),
-              cgen.Line())
+    header = List(
+        header=cgen.Comment('Flush denormal numbers to zero in hardware'),
+        body=(MacroStatement('_MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON)'),
+              MacroStatement('_MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON)')),
+        footer=cgen.Line()
+    )
 
-    body = iet.body._rebuild(body=(EmptyList(header=header),) + iet.body.body)
+    body = iet.body._rebuild(body=(header,) + iet.body.body)
     iet = iet._rebuild(body=body)
 
     return iet, {'includes': ('xmmintrin.h', 'pmmintrin.h')}

tests/test_iet.py

@@ -12,8 +12,8 @@
 from devito.ir.iet import (
     Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr,
     ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List,
-    Switch, Transformer, Uxreplace, filter_iterations, make_callable, make_efunc,
-    retrieve_iteration_tree
+    MacroStatement, MapExprStmts, Switch, Transformer, Uxreplace, filter_iterations,
+    make_callable, make_efunc, retrieve_iteration_tree
 )
 from devito.ir.iet.visitors import sorted_efuncs
 from devito.passes.iet.engine import Graph, reuse_efuncs
@@ -120,6 +120,35 @@ def test_nested_calls_cgen():
     assert str(code) == 'foo(bar());'
 
 
+def test_macro_statement_codegen():
+    macro = MacroStatement('MY_MACRO')
+
+    assert str(CGen().visit(macro)) == 'MY_MACRO;'
+    assert macro.view_cir == 'MY_MACRO;'
+
+    func = Callable('foo', [macro, Call('bar')], 'void', (), prefix=())
+    assert str(func) == """\
+void foo()
+{
+  MY_MACRO;
+  bar();
+}"""
+
+
+def test_macro_statement_visitors():
+    macro = MacroStatement('MY_MACRO')
+    body = List(body=[macro, Call('bar')])
+
+    assert FindNodes(MacroStatement).visit(body) == [macro]
+    assert macro in MapExprStmts().visit(body)
+    assert FindSymbols('basics').visit(body) == []
+
+    replacement = macro._rebuild(name='OTHER_MACRO')
+    rebuilt = Transformer({macro: replacement}).visit(body)
+    assert str(rebuilt) == 'OTHER_MACRO;\nbar();'
+    assert str(body) == 'MY_MACRO;\nbar();'
+
+
 def test_uxreplace_dereference_offset():
     d = CustomDimension(name='d', symbolic_size=4)
     array = Array(name='a0', dimensions=(d,), shape=(4,), scope='shared')

tests/test_macro_statements.py

@@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+
+from devito import Eq, Grid, Operator, TimeFunction
+from devito.ir.iet import (
+    Call, EntryFunction, FindNodes, Iteration, KernelLaunch, MacroStatement, TimedList
+)
+from devito.mpi.distributed import MPICommObject
+from devito.passes.iet.engine import Graph
+from devito.passes.iet.errors import check_launch
+from devito.types import Dimension, Symbol
+
+
+@pytest.mark.parametrize('language', ['C', 'openmp'])
+def test_operator_timer_macros(language):
+    grid = Grid(shape=(4, 4))
+    u = TimeFunction(name='u', grid=grid)
+    op = Operator(Eq(u.forward, u + 1), language=language)
+
+    timers = FindNodes(TimedList).visit(op)
+    assert timers
+    for timer in timers:
+        assert isinstance(timer.header[0], MacroStatement)
+        assert isinstance(timer.footer[0], MacroStatement)
+        assert f'START({timer.name});' in str(op)
+        assert f'STOP({timer.name},{timer.timer.name});' in str(op)
+
+    summary = op.apply(time_M=1)
+    assert np.all(u.data[0] == 2)
+    assert np.all(u.data[1] == 1)
+    assert summary
+
+
+def test_operator_denormal_macros():
+    grid = Grid(shape=(4, 4))
+    u = TimeFunction(name='u', grid=grid)
+    op = Operator(Eq(u.forward, u + 1), platform='intel64', language='C')
+
+    macros = FindNodes(MacroStatement).visit(op)
+    names = [macro.name for macro in macros]
+    assert '_MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON)' in names
+    assert '_MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON)' in names
+    code = str(op)
+    assert '_MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON);' in code
+    assert '_MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON);' in code
+
+
+@pytest.mark.parametrize('mpi, with_comm', [(False, False), (True, False), (True, True)])
+@pytest.mark.parametrize('errctl', [False, 'basic'])
+@pytest.mark.parametrize('inside_loop', [False, True])
+def test_kernel_launch_macros(errctl, inside_loop, mpi, with_comm):
+    launch = KernelLaunch('kernel', Symbol(name='blocks'), Symbol(name='threads'))
+    body = Iteration(launch, Dimension(name='t'), (0, 1, 1)) if inside_loop else launch
+    if with_comm:
+        pytest.importorskip('mpi4py')
+        comm = MPICommObject()
+        body = [Call('MPI_Barrier', comm), body]
+        parameters = (comm,)
+    else:
+        parameters = ()
+    graph = Graph(EntryFunction('entry', body, 'int', parameters))
+    langbb = {'peek-error': lambda: Call('peek_error'), 'error-none': 'SUCCESS'}
+
+    check_launch(graph, options={'errctl': errctl, 'mpi': mpi}, langbb=langbb)
+
+    macros = FindNodes(MacroStatement).visit(graph.root)
+    if errctl and inside_loop:
+        name = 'CHECK_LAUNCH_RETURN' if with_comm else 'CHECK_LAUNCH'
+        assert [macro.name for macro in macros] == [name]
+        code = str(graph.root)
+        assert f'{name};' in code
+        assert code.index('kernel<<<') < code.index(f'{name};')
+        assert ('CHECK_LAUNCH', 'if (peek_error() != SUCCESS) {break;}') in graph.headers
+    else:
+        assert not macros
+        assert not graph.headers