diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 7c12b39647..f8b76761a9 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -49,6 +49,7 @@ 'Lambda', 'List', 'ListMajor', + 'MacroStatement', 'MetaCall', 'MultiTraversable', 'Node', @@ -1331,6 +1332,24 @@ def periodic(self): return self._periodic +class MacroStatement(ExprStmt, Node): + + """ + A standalone macro invocation, such as ``MY_MACRO;``. + + Parameters + ---------- + name : str + The macro name, without a trailing semicolon. + """ + + def __init__(self, name): + self.name = name + + def __repr__(self): + return f"" + + class Using(Node): """ diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 6186ff4e76..b411d61067 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -572,6 +572,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: diff --git a/tests/test_iet.py b/tests/test_iet.py index 99a94ed24c..5b45a57f0b 100644 --- a/tests/test_iet.py +++ b/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, filter_iterations, make_callable, make_efunc, - retrieve_iteration_tree + MacroStatement, MapExprStmts, Switch, Transformer, 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();' + + @pytest.mark.parametrize('mode,expected', [ ('basics', '["x"]'), ('symbolics', '["f"]')