Constant Propagation Transform

loki/analyse/__init__.py

@@ -8,4 +8,4 @@
 Advanced analysis utilities, such as dataflow analysis functionalities.
 """
 
-from loki.analyse.analyse_dataflow import *  # noqa
+from loki.analyse.data_flow_analysis import *  # noqa

loki/analyse/abstract_dfa.py

@@ -0,0 +1,50 @@
+# (C) Copyright 2024- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+
+from loki import Transformer
+
+__all__ = ['AbstractDataflowAnalysis']
+
+class AbstractDataflowAnalysis(Transformer, ABC):
+    class _Attacher(Transformer):
+        pass
+
+    class _Detacher(Transformer):
+        pass
+
+    def get_attacher(self):
+        return self._Attacher()
+
+    def get_detacher(self):
+        return self._Detacher()
+
+    @abstractmethod
+    def attach_dataflow_analysis(self, module_or_routine):
+        pass
+
+    def detach_dataflow_analysis(self, module_or_routine):
+        """
+        Remove from each IR node the stored dataflow analysis metadata.
+
+        Accessing the relevant attributes afterwards raises :py:class:`RuntimeError`.
+        """
+
+        if hasattr(module_or_routine, 'spec'):
+            self.get_detacher().visit(module_or_routine.spec)
+        if hasattr(module_or_routine, 'body'):
+            self.get_detacher().visit(module_or_routine.body)
+
+    @contextmanager
+    def dataflow_analysis_attached(self, module_or_routine):
+        self.attach_dataflow_analysis(module_or_routine)
+        try:
+            yield module_or_routine
+        finally:
+            self.detach_dataflow_analysis(module_or_routine)

loki/analyse/tests/test_analyse_dataflow.py → loki/analyse/tests/test_data_flow_analyse.py

@@ -13,7 +13,7 @@
     Associate, Module
 )
 from loki.analyse import (
-    dataflow_analysis_attached, read_after_write_vars, loop_carried_dependencies
+    DataFlowAnalysis, read_after_write_vars, loop_carried_dependencies
 )
 from loki.frontend import available_frontends
 
@@ -55,7 +55,7 @@ def test_analyse_live_symbols(frontend):
         'v2': {'tmp', 'a', 'n', 'v1', 'v2', 'v3'}
     }
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert routine.body
 
         for assignment in assignments:
@@ -103,7 +103,7 @@ def test_analyse_defines_uses_symbols(frontend):
         for cond in conditionals:
             _ = cond.uses_symbols
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert fgen(routine) == ref_fgen
         assert len(FindNodes(Conditional).visit(routine.body)) == 2
         assert len(FindNodes(Loop).visit(routine.body)) == 1
@@ -164,7 +164,7 @@ def test_read_after_write_vars(frontend):
     pragmas = FindNodes(Pragma).visit(routine.body)
     assert len(pragmas) == 5
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         for pragma in pragmas:
             assert read_after_write_vars(routine.body, pragma) == vars_at_inspection_node[pragma.content]
 
@@ -210,7 +210,7 @@ def test_read_after_write_vars_conditionals(frontend):
     pragmas = FindNodes(Pragma).visit(routine.body)
     assert len(pragmas) == len(vars_at_inspection_node)
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         for pragma in pragmas:
             assert read_after_write_vars(routine.body, pragma) == vars_at_inspection_node[pragma.content]
 
@@ -237,7 +237,7 @@ def test_loop_carried_dependencies(frontend):
     loops = FindNodes(Loop).visit(routine.body)
     assert len(loops) == 1
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert loop_carried_dependencies(loops[0]) == {variable_map['b'], variable_map['c']}
 
 @pytest.mark.parametrize('frontend', available_frontends())
@@ -273,7 +273,7 @@ def test_analyse_interface(frontend):
     source = Sourcefile.from_source(fcode, frontend=frontend)
     routine = source['test']
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert len(routine.body.defines_symbols) == 0
         assert len(routine.body.uses_symbols) == 0
         assert len(routine.spec.uses_symbols) == 0
@@ -311,7 +311,7 @@ def test_analyse_imports(frontend, tmp_path):
     module = Module.from_source(fcode_module, frontend=frontend, xmods=[tmp_path])
     routine = Subroutine.from_source(fcode, frontend=frontend, definitions=module, xmods=[tmp_path])
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert len(routine.spec.defines_symbols) == 1
         assert 'random_call' in routine.spec.defines_symbols
 
@@ -346,7 +346,7 @@ def test_analyse_enriched_call(frontend):
     routine.enrich(source.all_subroutines)
     call = FindNodes(CallStatement).visit(routine.body)[0]
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert all(i in call.defines_symbols for i in ('v_out', 'v_inout'))
         assert all(i in call.uses_symbols for i in ('v_in', 'v_inout'))
 
@@ -370,7 +370,7 @@ def test_analyse_unenriched_call(frontend):
     routine = source['test']
     call = FindNodes(CallStatement).visit(routine.body)[0]
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert all(i in call.defines_symbols for i in ('v_out', 'v_inout', 'v_in'))
         assert all(i in call.uses_symbols for i in ('v_in', 'v_inout', 'v_in'))
 
@@ -394,7 +394,7 @@ def test_analyse_allocate_statement(frontend):
     """.strip()
 
     routine = Subroutine.from_source(fcode, frontend=frontend)
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert all(i not in routine.body.defines_symbols for i in ['m', 'n'])
         assert all(i in routine.body.uses_symbols for i in ['m', 'n'])
         assert 'a' in routine.body.defines_symbols
@@ -417,7 +417,7 @@ def test_analyse_import_kind(frontend):
     """.strip()
 
     routine = Subroutine.from_source(fcode, frontend=frontend)
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert 'real64' in routine.body.uses_symbols
         assert 'real64' not in routine.body.defines_symbols
         assert 'a' in routine.body.defines_symbols
@@ -446,7 +446,7 @@ def test_analyse_query_memory_attributes(frontend):
     """.strip()
 
     routine = Subroutine.from_source(fcode, frontend=frontend)
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert not 'a' in routine.body.uses_symbols
         assert 'a' in routine.body.defines_symbols
         assert not 'b' in routine.body.uses_symbols
@@ -483,7 +483,7 @@ def test_analyse_call_args_array_slicing(frontend):
     calls = FindNodes(CallStatement).visit(routine.body)
     routine.enrich(source.all_subroutines)
 
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert 'n' in calls[0].uses_symbols
         assert not 'n' in calls[0].defines_symbols
         assert 'b' in calls[1].uses_symbols
@@ -510,7 +510,7 @@ def test_analyse_multiconditional(frontend):
 
     routine = Subroutine.from_source(fcode, frontend=frontend)
     mcond = FindNodes(MultiConditional).visit(routine.body)[0]
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert len(mcond.bodies) == 2
         assert len(mcond.else_body) == 1
         for b in mcond.bodies:
@@ -549,7 +549,7 @@ def test_analyse_maskedstatement(frontend):
     routine = Subroutine.from_source(fcode, frontend=frontend)
     mask = FindNodes(MaskedStatement).visit(routine.body)[0]
     num_bodies = len(mask.bodies)
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert len(mask.uses_symbols) == 1
         assert len(mask.defines_symbols) == 1
         assert 'mask' in mask.uses_symbols
@@ -582,15 +582,15 @@ def test_analyse_whileloop(frontend):
     routine = Subroutine.from_source(fcode, frontend=frontend)
     loop = FindNodes(WhileLoop).visit(routine.body)[0]
     cond = FindNodes(Conditional).visit(routine.body)[0]
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         assert len(cond.uses_symbols) == 1
         assert 'flag' in cond.uses_symbols
         assert len(loop.uses_symbols) == 1
         assert len(loop.defines_symbols) == 2
         assert 'ij' in loop.uses_symbols
         assert all(v in loop.defines_symbols for v in ('ij', 'a'))
 
-    with dataflow_analysis_attached(cond):
+    with DataFlowAnalysis().dataflow_analysis_attached(cond):
         assert len(loop.uses_symbols) == 1
         assert len(loop.defines_symbols) == 2
         assert 'ij' in loop.uses_symbols
@@ -621,7 +621,7 @@ def test_analyse_associate(frontend):
     routine = Subroutine.from_source(fcode, frontend=frontend)
     associates = FindNodes(Associate).visit(routine.body)
     assigns = FindNodes(Assignment).visit(routine.body)
-    with dataflow_analysis_attached(routine):
+    with DataFlowAnalysis().dataflow_analysis_attached(routine):
         # check that associates use variables names in outer scope
         assert associates[0].uses_symbols == {'in_var'}
         assert associates[0].defines_symbols == {'a', 'b', 'c'}

loki/ir/nodes.py

@@ -187,6 +187,12 @@ def ir_graph(self, show_comments=False, show_expressions=False, linewidth=40, sy
 
         return ir_graph(self, show_comments, show_expressions,linewidth, symgen)
 
+    @property
+    def constants_map(self):
+        if self.__dict__['_constants_map'] is None:
+            raise RuntimeError('Need to run constant propagation analysis on the IR first.')
+        return self.__dict__['_constants_map']
+
     @property
     def live_symbols(self):
         """
@@ -195,9 +201,9 @@ def live_symbols(self):
         graph.
 
         This property is attached to the Node by
-        :py:func:`loki.analyse.analyse_dataflow.attach_dataflow_analysis` or
+        :py:func:`loki.analyse.LiveVariableAnalysis.attach_dataflow_analysis` or
         when using the
-        :py:func:`loki.analyse.analyse_dataflow.dataflow_analysis_attached`
+        :py:func:`loki.analyse.LiveVariableAnalysis.dataflow_analysis_attached`
         context manager.
         """
         if self.__dict__['_live_symbols'] is None:
@@ -210,9 +216,9 @@ def defines_symbols(self):
         Yield the list of symbols (potentially) defined by this node.
 
         This property is attached to the Node by
-        :py:func:`loki.analyse.analyse_dataflow.attach_dataflow_analysis` or
+        :py:func:`loki.analyse.LiveVariableAnalysis.attach_dataflow_analysis` or
         when using the
-        :py:func:`loki.analyse.analyse_dataflow.dataflow_analysis_attached`
+        :py:func:`loki.analyse.LiveVariableAnalysis.dataflow_analysis_attached`
         context manager.
         """
         if self.__dict__['_defines_symbols'] is None:
@@ -226,9 +232,9 @@ def uses_symbols(self):
         Yield the list of symbols used by this node before defining it.
 
         This property is attached to the Node by
-        :py:func:`loki.analyse.analyse_dataflow.attach_dataflow_analysis` or
+        :py:func:`loki.analyse.LiveVariableAnalysis.attach_dataflow_analysis` or
         when using the
-        :py:func:`loki.analyse.analyse_dataflow.dataflow_analysis_attached`
+        :py:func:`loki.analyse.LiveVariableAnalysis.dataflow_analysis_attached`
         context manager.
         """
         if self.__dict__['_uses_symbols'] is None:

loki/ir/tests/test_ir_graph.py

@@ -10,7 +10,7 @@
 import pytest
 
 from loki import Sourcefile, graphviz_present
-from loki.analyse import dataflow_analysis_attached
+from loki.analyse import DataFlowAnalysis
 from loki.ir import Node, FindNodes, ir_graph, GraphCollector
 
 
@@ -324,7 +324,7 @@ def test_ir_graph_writes_correct_graphs(testdir, test_file, tmp_path):
 
 
 @pytest.mark.parametrize("test_file", test_files)
-def test_ir_graph_dataflow_analysis_attached(testdir, test_file, tmp_path):
+def test_ir_graph_live_variable_analysis_attached(testdir, test_file, tmp_path):
     source = Sourcefile.from_file(testdir / test_file, xmods=[tmp_path])
 
     def find_lives_defines_uses(text):
@@ -349,7 +349,7 @@ def apply_conversion(text):
         )
 
     for routine in source.all_subroutines:
-        with dataflow_analysis_attached(routine):
+        with DataFlowAnalysis().dataflow_analysis_attached(routine):
             for node in FindNodes(Node).visit(routine.body):
                 node_info, _ = GraphCollector(show_comments=True).visit(node)[0]
                 lives, defines, uses = find_lives_defines_uses(node_info["label"])

loki/transformations/__init__.py

@@ -16,6 +16,7 @@
 from loki.transformations.build_system import * # noqa
 from loki.transformations.argument_shape import * # noqa
 from loki.transformations.data_offload import * # noqa
+from loki.transformations.constant_propagation import * # noqa
 from loki.transformations.drhook import * # noqa
 from loki.transformations.extract import * # noqa
 from loki.transformations.field_api import * # noqa

loki/transformations/array_indexing.py

@@ -15,7 +15,7 @@
 
 from loki.batch import Transformation, ProcedureItem
 from loki.logging import info
-from loki.analyse import dataflow_analysis_attached
+from loki.analyse import DataFlowAnalysis
 from loki.expression import symbols as sym, simplify, symbolic_op, is_constant
 from loki.ir import (
     nodes as ir, Assignment, Loop, VariableDeclaration, FindNodes,
@@ -305,7 +305,7 @@ def promote_variables(routine, variable_names, pos, index=None, size=None):
         # Create a copy of the tree and apply promotion in-place
         routine.body = Transformer().visit(routine.body)
 
-        with dataflow_analysis_attached(routine):
+        with DataFlowAnalysis().dataflow_analysis_attached(routine):
             for node, var_list in FindVariables(unique=False, with_ir_node=True).visit(routine.body):
                 # All the variables marked for promotion that appear in this IR node
                 var_list = [v for v in var_list if v.name.lower() in variable_names]

loki/transformations/constant_propagation.py

@@ -0,0 +1,41 @@
+# (C) Copyright 2024- ECMWF.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from loki.analyse.constant_propagation_analysis import ConstantPropagationAnalysis
+from loki import Transformer, Subroutine
+
+__all__ = ['ConstantPropagationTransformer']
+
+class ConstantPropagationTransformer(Transformer):
+
+    def __init__(self, fold_floats=True, unroll_loops=True):
+        self.fold_floats = fold_floats
+        self.unroll_loops = unroll_loops
+        super().__init__()
+
+    def visit(self, expr, *args, **kwargs):
+        const_prop = ConstantPropagationAnalysis(self.fold_floats, self.unroll_loops, True)
+        constants_map = kwargs.get('constants_map', dict())
+        try:
+            declarations_map = const_prop.generate_declarations_map(expr)
+            # If a user specifies their own map, they probably want it to override these
+            declarations_map.update(constants_map)
+            constants_map = declarations_map
+        except AttributeError:
+            pass
+
+        is_routine = isinstance(expr, Subroutine)
+        target = expr.body if is_routine else expr
+
+        target = const_prop.get_attacher().visit(target, constants_map=constants_map)
+        target = const_prop.get_detacher().visit(target)
+
+        if is_routine:
+            expr.body = target
+            return expr
+
+        return target

loki/transformations/data_offload/field_offload.py

@@ -5,7 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-from loki.analyse import dataflow_analysis_attached
+from loki.analyse import DataFlowAnalysis
 from loki.batch import Transformation
 from loki.expression import Array, symbols as sym
 from loki.ir import (
@@ -72,7 +72,7 @@ def process_driver(self, driver):
         remove_field_api_view_updates(driver, self.field_group_types)
 
         with pragma_regions_attached(driver):
-            with dataflow_analysis_attached(driver):
+            with DataFlowAnalysis().dataflow_analysis_attached(driver):
                 for region in FindNodes(ir.PragmaRegion).visit(driver.body):
                     # Only work on active `!$loki data` regions
                     if not region.pragma or not is_loki_pragma(region.pragma, starts_with='data'):