compiler: Call finalize_args once at the end of the lowering

Author: FabioLuporini

devito/operator/operator.py

@@ -32,8 +32,8 @@
 from devito.operator.registry import operator_selector
 from devito.parameters import configuration
 from devito.passes import (
-    Graph, error_mapper, generate_implicit, generate_macros, is_on_device, lower_dtypes,
-    lower_index_derivatives, minimize_symbols, optimize_pows, unevaluate
+    Graph, error_mapper, finalize_args, generate_implicit, generate_macros, is_on_device,
+    lower_dtypes, lower_index_derivatives, minimize_symbols, optimize_pows, unevaluate
 )
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (
@@ -522,6 +522,9 @@ def _lower_iet(cls, uiet, **kwargs):
         # Target-independent optimizations
         minimize_symbols(graph)
 
+        # Finalize helper signatures after all IET transformations have settled.
+        finalize_args(graph)
+
         return graph.root, graph
 
     # Read-only properties exposed to the outside world

devito/passes/iet/definitions.py

@@ -444,7 +444,7 @@ def _inject_definitions(self, iet, storage):
 
         return processed, flatten(efuncs)
 
-    @iet_pass
+    @iet_pass(updates_args=True)
     def place_definitions(self, iet, globs=None, **kwargs):
         """
         Create a new IET where all symbols have been declared, allocated, and
@@ -518,7 +518,7 @@ def place_definitions(self, iet, globs=None, **kwargs):
                      'globals': as_tuple(globs),
                      'includes': as_tuple(sorted(storage.includes))}
 
-    @iet_pass(updates_args=False)
+    @iet_pass
     def place_casts(self, iet, **kwargs):
         """
         Create a new IET with the necessary type casts.
@@ -669,7 +669,7 @@ def place_transfers(self, iet, data_movs=None, ctx=None, **kwargs):
 
         return iet, {'efuncs': efuncs}
 
-    @iet_pass(updates_args=False)
+    @iet_pass
     def place_devptr(self, iet, **kwargs):
         """
         Transform `iet` such that device pointers are used in DeviceCalls.

devito/passes/iet/engine.py

@@ -28,7 +28,7 @@
 from devito.types.dense import DiscreteFunction
 from devito.types.dimension import AbstractIncrDimension, BlockDimension
 
-__all__ = ['Graph', 'iet_pass', 'iet_visit']
+__all__ = ['Graph', 'finalize_args', 'iet_pass', 'iet_visit']
 
 
 class Byproduct:
@@ -128,10 +128,22 @@ def sync_mapper(self):
 
         return found
 
-    def apply(self, func, *, updates_args=True, **kwargs):
+    def apply(self, func, *, updates_args=False, **kwargs):
         """
-        Apply `func` to all nodes in the Graph. This changes the state of the Graph.
+        Apply ``func`` to all nodes in the Graph.
+
+        Parameters
+        ----------
+        updates_args : bool, optional
+            If True, reconcile Callable parameters and Call arguments before
+            the graph walk and after each changed node. This is only needed by
+            passes whose transformation logic depends on already-updated
+            signatures while the pass is still running. Otherwise, argument
+            reconciliation is intentionally deferred to ``finalize_args``.
         """
+        if updates_args:
+            _update_args(self)
+
         dag = create_call_graph(self.root.name, as_hashable(self.efuncs))
 
         # Apply `func`
@@ -161,10 +173,8 @@ def apply(self, func, *, updates_args=True, **kwargs):
             efuncs[i] = efunc
             efuncs.update(dict([(i.name, i) for i in new_efuncs]))
 
-            # Update the parameters / arguments lists if the pass may have
-            # introduced or removed objects.
             if updates_args:
-                efuncs = update_args(efunc, efuncs, dag)
+                efuncs = _update_args_efunc(efunc, efuncs, dag)
 
         # Minimize code size
         if len(efuncs) > len(self.efuncs):
@@ -209,7 +219,37 @@ def filter(self, key):
         )
 
 
-def iet_pass(func=None, *, updates_args=True):
+@timed_pass(name='finalize_args')
+def finalize_args(graph):
+    """
+    Finalize Callable parameter lists and Call argument lists across ``graph``.
+
+    IET passes may temporarily leave helper signatures stale while introducing
+    or eliminating symbols. This pass reconciles the whole call graph once,
+    after lowering has settled.
+    """
+    _update_args(graph)
+
+
+def _update_args(graph):
+    dag = create_call_graph(graph.root.name, as_hashable(graph.efuncs))
+
+    efuncs = graph.efuncs
+    for i in dag.topological_sort():
+        efuncs = _update_args_efunc(efuncs[i], efuncs, dag)
+
+    graph.efuncs = efuncs
+
+
+def iet_pass(func=None, *, updates_args=False):
+    """
+    Decorate an IET pass.
+
+    ``updates_args=True`` is an opt-in for passes that must observe up-to-date
+    Callable/Call signatures before and during their own graph walk. Most
+    passes should leave it False and rely on ``finalize_args`` at the end of
+    IET lowering.
+    """
     if func is None:
         return partial(iet_pass, updates_args=updates_args)
 
@@ -702,7 +742,7 @@ def _(i, mapper, sregistry):
     mapper[i] = i._rebuild(name=sregistry.make_name(prefix='nthreads'))
 
 
-def update_args(root, efuncs, dag):
+def _update_args_efunc(root, efuncs, dag):
     """
     Re-derive the parameters of `root` and apply the changes in cascade through
     the `efuncs`.
@@ -800,6 +840,6 @@ def _filter(v, efunc=None):
             continue
 
         efuncs[n] = efunc
-        efuncs = update_args(efunc, efuncs, dag)
+        efuncs = _update_args_efunc(efunc, efuncs, dag)
 
     return efuncs

devito/passes/iet/instrument.py

@@ -27,7 +27,7 @@ def instrument(graph, **kwargs):
     sync_sections(graph, **kwargs)
 
 
-@iet_pass(updates_args=False)
+@iet_pass
 def track_subsections(iet, **kwargs):
     """
     Add sub-Sections to the `profiler`. Sub-Sections include:
@@ -122,7 +122,7 @@ def instrument_sections(iet, **kwargs):
     return piet, {'headers': headers}
 
 
-@iet_pass(updates_args=False)
+@iet_pass
 def sync_sections(iet, langbb=None, profiler=None, **kwargs):
     """
     Wrap sections within global barriers if deemed necessary by the profiler.

devito/passes/iet/linearization.py

@@ -46,7 +46,7 @@ def linearize(graph, **kwargs):
     linearization(graph, key=key, tracker=tracker, **kwargs)
 
 
-@iet_pass
+@iet_pass(updates_args=True)
 def linearization(iet, key=None, tracker=None, **kwargs):
     """
     Carry out the actual work of `linearize`.

devito/passes/iet/misc.py

@@ -29,7 +29,7 @@
 ]
 
 
-@iet_pass(updates_args=False)
+@iet_pass
 def avoid_denormals(iet, platform=None, **kwargs):
     """
     Introduce nodes in the Iteration/Expression tree that will expand to C
@@ -60,7 +60,7 @@ def avoid_denormals(iet, platform=None, **kwargs):
     return iet, {'includes': ('xmmintrin.h', 'pmmintrin.h')}
 
 
-@iet_pass(updates_args=False)
+@iet_pass
 def hoist_prodders(iet):
     """
     Move Prodders within the outer levels of an Iteration tree.
@@ -151,7 +151,7 @@ def generate_macros(graph, **kwargs):
     _generate_macros(graph, tracker={}, **kwargs)
 
 
-@iet_pass(updates_args=False)
+@iet_pass
 def _generate_macros(iet, tracker=None, langbb=None, printer=CPrinter, **kwargs):
     # Derive the Macros necessary for the FIndexeds
     iet = _generate_macros_findexeds(iet, tracker=tracker, **kwargs)

tests/test_iet.py

@@ -540,20 +540,20 @@ def test_complex_array():
         "float _Complex **restrict a_vec __attribute__ ((aligned (64)));"
 
 
-def test_iet_pass_skip_update_args(monkeypatch):
+def test_iet_pass_does_not_update_args(monkeypatch):
     x = Symbol(name='x')
     y = Symbol(name='y')
 
     foo = Callable('foo', DummyExpr(x, y), 'void', parameters=(x, y))
     graph = Graph(foo)
 
-    @iet_pass(updates_args=False)
+    @iet_pass
     def inject_expr(iet):
         body = iet.body._rebuild(body=iet.body.body + (DummyExpr(x, x),))
         return iet._rebuild(body=body), {}
 
-    monkeypatch.setattr(iet_engine, 'update_args',
-                        lambda *args, **kwargs: pytest.fail("update_args called"))
+    monkeypatch.setattr(iet_engine, '_update_args',
+                        lambda *args, **kwargs: pytest.fail("_update_args called"))
 
     inject_expr(graph)