From 02b8c59beabd667c4eda06eda5dfa04d26968aaa Mon Sep 17 00:00:00 2001 From: ada-wh525 Date: Tue, 27 Oct 2026 21:47:00 +0800 Subject: [PATCH 1/2] compiler: Add a node for standalone macro statements Introduce MacroStatement as a leaf IET statement and generate standalone macro invocations through CGen. Cover code generation, statement mapping, node discovery and transformations. Fixes devitocodes/devito#2853. --- devito/ir/iet/nodes.py | 19 +++++++++++++++++++ devito/ir/iet/visitors.py | 3 +++ tests/test_iet.py | 33 +++++++++++++++++++++++++++++++-- 3 files changed, 53 insertions(+), 2 deletions(-) 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"]') From be6c434c489929fd55da59039d24df8d67bf53af Mon Sep 17 00:00:00 2001 From: ada-wh525 Date: Tue, 3 Nov 2026 22:25:00 +0800 Subject: [PATCH 2/2] compiler: Use macro statements in compiler-generated code --- devito/ir/iet/nodes.py | 11 ++--- devito/ir/iet/visitors.py | 4 +- devito/passes/iet/errors.py | 6 +-- devito/passes/iet/misc.py | 16 +++---- tests/test_macro_statements.py | 76 ++++++++++++++++++++++++++++++++++ 5 files changed, 97 insertions(+), 16 deletions(-) create mode 100644 tests/test_macro_statements.py diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index f8b76761a9..3fcfbd113a 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1013,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): @@ -1340,7 +1340,8 @@ class MacroStatement(ExprStmt, Node): Parameters ---------- name : str - The macro name, without a trailing semicolon. + The macro name or complete invocation (including any arguments), + without a trailing semicolon. """ def __init__(self, name): diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index b411d61067..ebe3981095 100644 --- a/devito/ir/iet/visitors.py +++ b/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) diff --git a/devito/passes/iet/errors.py b/devito/passes/iet/errors.py index d9f8be1eef..b4433a6409 100644 --- a/devito/passes/iet/errors.py +++ b/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: diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index 4facf6646d..fe53b5ec65 100644 --- a/devito/passes/iet/misc.py +++ b/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')} diff --git a/tests/test_macro_statements.py b/tests/test_macro_statements.py new file mode 100644 index 0000000000..eddad2d567 --- /dev/null +++ b/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