Repository: stfc/PSyclone
Pull request #2970 (2025-05-01 11:03:51)

(towards #2966) Add `--enble-cache` psyclone command flag


Patch:
diff --git a/doc/user_guide/psyclone_command.rst b/doc/user_guide/psyclone_command.rst
index f9680a1c3..593eb67e1 100644
--- a/doc/user_guide/psyclone_command.rst
+++ b/doc/user_guide/psyclone_command.rst
@@ -73,6 +73,9 @@ by the command:
                            filename of a PSyclone optimisation recipe
      -I INCLUDE, --include INCLUDE
                            path to Fortran INCLUDE or module files
+     --enable-cache        whether to enable caching of imported module dependencies (if
+                           enabled, it will generate a .psycache file of each imported
+                           module in the same location as the imported source file).
      -l {off,all,output}, --limit {off,all,output}
                            limit the Fortran line length to 132 characters (default 'off').
                            Use 'all' to apply limit to both input and output Fortran. Use
@@ -141,12 +144,11 @@ section. With a transformation recipe the command looks like:
 Fortran INCLUDE Files and Modules
 ---------------------------------
 
-If the source code to be processed by PSyclone
-contains INCLUDE statements then the location of any INCLUDE'd files
-*must* be supplied to PSyclone via the ``-I`` or ``--include``
-option. (This is necessary because INCLUDE lines are a part of the
-Fortran language and must therefore be parsed - they are not handled
-by any pre-processing step.) Multiple locations may be specified by
+If the source code to be processed by PSyclone contains INCLUDE statements
+then the location of any INCLUDE'd files *must* be supplied to PSyclone via
+the ``-I`` or ``--include`` option. (This is necessary because INCLUDE lines
+are a part of the Fortran language and must therefore be parsed - they are not
+handled by any pre-processing step.) Multiple locations may be specified by
 using multiple ``-I`` flags, e.g.:
 
 .. code-block:: console
@@ -164,6 +166,15 @@ an appropriate error. For example:
     psyclone -I nonexisting test.f90
     PSyclone configuration error: Include path 'nonexisting' does not exist
 
+The `-I` locations will also be used when a script requests to follow module
+dependencies in order to obtain more information about the code symbols (see
+:ref:`sec_script_globals`). But note that if the whole program has many
+dependencies and the imports happen from multiple files, it can increase the
+psyclone processing time considerably. In this case it is recommended to use
+the `--enable-cache` flag. This will creates a `filename.psycache` file in the
+same location as the original file for every import followed. The next time
+the same import is requested, if the hashes match, the cached file will be used.
+
 Currently, the PSyKAl-based APIs (LFRic and GOcean - see below) will ignore
 (but preserve) INCLUDE statements in algorithm-layer code. However, INCLUDE
 statements in kernels will, in general, cause the kernel parsing to fail
diff --git a/doc/user_guide/user_scripts.rst b/doc/user_guide/user_scripts.rst
index 45d754498..f2caab5e0 100644
--- a/doc/user_guide/user_scripts.rst
+++ b/doc/user_guide/user_scripts.rst
@@ -90,6 +90,8 @@ The :ref:`examples section<examples>` provides a list of psyclone user scripts
 and associated usage instructions for multiple applications.
 
 
+.. _sec_script_globals:
+
 Script Global Variables
 -----------------------
 
@@ -102,9 +104,11 @@ is applied). These are:
     # List of all files that psyclone will skip processing
     FILES_TO_SKIP = ["broken_file1.f90", "broken_file2.f90"]
 
-    # Boolean to decide whether PSyclone should chase external modules while
-    # creating a PSyIR tree in order to obtain better external symbol information.
-    # It can also be a list of module names for more precise control
+    # Whether to chase the imported modules to improve symbol information (it can
+    # also be a list of module filenames to limit the chasing to only specific
+    # modules). This has to be used in combination with '-I' command flag in order
+    # to point to the module location directory. We also strongly recommend using
+    # the '--enable-cache' flag to reduce the performance overhead.
     RESOLVE_IMPORTS = ["relevant_module1.f90", "relevant_module2.f90"]
 
     def trans(psyir):
diff --git a/examples/nemo/scripts/acc_kernels_trans.py b/examples/nemo/scripts/acc_kernels_trans.py
index 905a04eb3..2466733e2 100755
--- a/examples/nemo/scripts/acc_kernels_trans.py
+++ b/examples/nemo/scripts/acc_kernels_trans.py
@@ -77,8 +77,11 @@ Loop.set_loop_type_inference_rules({
         "tracers": {"variable": "jt"}
 })
 
-# List of all module names that PSyclone will chase during the creation of the
-# PSyIR tree in order to use the symbol information from those modules
+# Whether to chase the imported modules to improve symbol information (it can
+# also be a list of module filenames to limit the chasing to only specific
+# modules). This has to be used in combination with '-I' command flag in order
+# to point to the module location directory. We also strongly recommend using
+# the '--enable-cache' flag to reduce the performance overhead.
 RESOLVE_IMPORTS = NEMO_MODULES_TO_IMPORT
 
 # Get the PSyclone transformations we will use
diff --git a/examples/nemo/scripts/acc_loops_trans.py b/examples/nemo/scripts/acc_loops_trans.py
index c66747979..4c2addd5c 100755
--- a/examples/nemo/scripts/acc_loops_trans.py
+++ b/examples/nemo/scripts/acc_loops_trans.py
@@ -47,8 +47,11 @@ from psyclone.transformations import (
 # Enable the insertion of profiling hooks during the transformation script
 PROFILING_ENABLED = True
 
-# List of all module names that PSyclone will chase during the creation of the
-# PSyIR tree in order to use the symbol information from those modules
+# Whether to chase the imported modules to improve symbol information (it can
+# also be a list of module filenames to limit the chasing to only specific
+# modules). This has to be used in combination with '-I' command flag in order
+# to point to the module location directory. We also strongly recommend using
+# the '--enable-cache' flag to reduce the performance overhead.
 RESOLVE_IMPORTS = NEMO_MODULES_TO_IMPORT
 
 # List of all files that psyclone will skip processing
diff --git a/examples/nemo/scripts/omp_cpu_trans.py b/examples/nemo/scripts/omp_cpu_trans.py
index 403f023e1..6e3f824be 100755
--- a/examples/nemo/scripts/omp_cpu_trans.py
+++ b/examples/nemo/scripts/omp_cpu_trans.py
@@ -40,16 +40,19 @@ directives into Nemo code. Tested with ECMWF Nemo 4.0 code. '''
 import os
 from utils import (
     insert_explicit_loop_parallelism, normalise_loops, add_profiling,
-    enhance_tree_information, PARALLELISATION_ISSUES, NEMO_MODULES_TO_IMPORT,
-    PRIVATISATION_ISSUES)
+    enhance_tree_information, PARALLELISATION_ISSUES, PRIVATISATION_ISSUES,
+    NEMO_MODULES_TO_IMPORT)
 from psyclone.psyir.nodes import Routine
 from psyclone.transformations import OMPLoopTrans
 
 # Enable the insertion of profiling hooks during the transformation script
 PROFILING_ENABLED = False
 
-# List of all module names that PSyclone will chase during the creation of the
-# PSyIR tree in order to use the symbol information from those modules
+# Whether to chase the imported modules to improve symbol information (it can
+# also be a list of module filenames to limit the chasing to only specific
+# modules). This has to be used in combination with '-I' command flag in order
+# to point to the module location directory. We also strongly recommend using
+# the '--enable-cache' flag to reduce the performance overhead.
 RESOLVE_IMPORTS = NEMO_MODULES_TO_IMPORT
 
 # A environment variable can inform if this is targeting NEMOv4, in which case
diff --git a/examples/nemo/scripts/omp_gpu_trans.py b/examples/nemo/scripts/omp_gpu_trans.py
index f965e63b1..aa8adfe0a 100755
--- a/examples/nemo/scripts/omp_gpu_trans.py
+++ b/examples/nemo/scripts/omp_gpu_trans.py
@@ -41,7 +41,7 @@ import os
 from utils import (
     add_profiling, inline_calls, insert_explicit_loop_parallelism,
     normalise_loops, enhance_tree_information, PARALLELISATION_ISSUES,
-    NEMO_MODULES_TO_IMPORT, PRIVATISATION_ISSUES)
+    PRIVATISATION_ISSUES, NEMO_MODULES_TO_IMPORT)
 from psyclone.psyir.nodes import Loop, Routine
 from psyclone.psyir.transformations import OMPTargetTrans
 from psyclone.transformations import (
@@ -58,8 +58,11 @@ INLINING_ENABLED = os.environ.get('ENABLE_INLINING', False)
 # array privatisation is disabled and some more files excluded
 NEMOV4 = os.environ.get('NEMOV4', False)
 
-# List of all module names that PSyclone will chase during the creation of the
-# PSyIR tree in order to use the symbol information from those modules
+# Whether to chase the imported modules to improve symbol information (it can
+# also be a list of module filenames to limit the chasing to only specific
+# modules). This has to be used in combination with '-I' command flag in order
+# to point to the module location directory. We also strongly recommend using
+# the '--enable-cache' flag to reduce the performance overhead.
 RESOLVE_IMPORTS = NEMO_MODULES_TO_IMPORT
 
 # List of all files that psyclone will skip processing
diff --git a/examples/nemo/scripts/utils.py b/examples/nemo/scripts/utils.py
index 5a669e593..4249355b9 100755
--- a/examples/nemo/scripts/utils.py
+++ b/examples/nemo/scripts/utils.py
@@ -47,7 +47,6 @@ from psyclone.psyir.transformations import (
     Reference2ArrayRangeTrans, ScalarisationTrans)
 from psyclone.transformations import TransformationError
 
-
 # USE statements to chase to gather additional symbol information.
 NEMO_MODULES_TO_IMPORT = [
     "oce", "par_oce", "par_kind", "dom_oce", "phycst", "ice",
@@ -180,8 +179,8 @@ def _it_should_be(symbol, of_type, instance):
 
 def enhance_tree_information(schedule):
     ''' Manually fix some PSyIR issues produced by not having enough symbol
-    information from external modules. Setting NEMO_MODULES_TO_IMPORT above
-    improve the situation but its not complete (not all symbols are imported)
+    information from external modules. Using RESOLVE_IMPORTS improves the
+    situation but it's not complete (not all symbols are imported)
     and it is not transitive (imports that inside import other symbols).
 
     :param schedule: the PSyIR Schedule to transform.
diff --git a/src/psyclone/generator.py b/src/psyclone/generator.py
index f4e3583fc..f4f635f58 100644
--- a/src/psyclone/generator.py
+++ b/src/psyclone/generator.py
@@ -437,6 +437,12 @@ def main(arguments):
     parser.add_argument(
         '-I', '--include', default=[], action="append",
         help='path to Fortran INCLUDE or module files')
+    parser.add_argument(
+        '--enable-cache', action="store_true", default=False,
+        help='whether to enable caching of imported module dependencies (if '
+             'enabled, it will generate a .psycache file of each imported '
+             'module in the same location as the imported source file).'
+    )
     parser.add_argument(
         '-l', '--limit', dest='limit', default='off',
         choices=['off', 'all', 'output'],
@@ -505,6 +511,10 @@ def main(arguments):
                   "specify the output destination of each psykal layer.")
             sys.exit(1)
 
+    # This has be be before the Config.get, because otherwise that creates a
+    # ModuleManager Singleton without caching
+    _ = ModuleManager.get(cache_active=args.enable_cache)
+
     # If no config file name is specified, args.config is none
     # and config will load the default config file.
     Config.get().load(args.config)
diff --git a/src/psyclone/psyir/frontend/fparser2.py b/src/psyclone/psyir/frontend/fparser2.py
index d215dbde8..af4257669 100644
--- a/src/psyclone/psyir/frontend/fparser2.py
+++ b/src/psyclone/psyir/frontend/fparser2.py
@@ -1605,6 +1605,13 @@ class Fparser2Reader():
                     container.name.lower() in self._modules_to_resolve):
                 parent.symbol_table.resolve_imports([container])
 
+            if visibility_map:
+                # Some of the imported symbols could have explicit visibility
+                # statements, so set the visibilities of all existing symbols
+                for symbol in parent.symbol_table.symbols_dict.values():
+                    if symbol.name.lower() in visibility_map:
+                        symbol.visibility = visibility_map[symbol.name.lower()]
+
     def _process_type_spec(self, parent, type_spec):
         '''
         Processes the fparser2 parse tree of a type specification in order to


Test patch:
diff --git a/src/psyclone/tests/generator_test.py b/src/psyclone/tests/generator_test.py
index 95567b2ff..4274f62cc 100644
--- a/src/psyclone/tests/generator_test.py
+++ b/src/psyclone/tests/generator_test.py
@@ -62,6 +62,7 @@ from psyclone.domain.lfric.transformations import LFRicLoopFuseTrans
 from psyclone.errors import GenerationError
 from psyclone.generator import (
     generate, main, check_psyir, add_builtins_use)
+from psyclone.parse import ModuleManager
 from psyclone.parse.algorithm import parse
 from psyclone.parse.utils import ParseError
 from psyclone.profiler import Profiler
@@ -1314,6 +1315,65 @@ def test_invalid_kern_naming():
     assert "but got 'not-a-scheme'" in str(err.value)
 
 
+def test_enable_cache_flag(capsys, tmpdir, monkeypatch):
+    ''' Check that if the --enable-cache flag is provided, resolve imports will
+    create .psycache files for each imported module.
+
+    '''
+    module1 = '''
+        module module1
+            integer :: a
+        end module module1
+    '''
+    module2 = '''
+        module module2
+            integer :: b
+        end module module2
+    '''
+    code = '''
+        module test
+            use module1
+            use module2
+            real :: result
+        contains
+            subroutine mytest()
+                result = a + b
+            end subroutine mytest
+        end module test
+    '''
+    recipe = '''
+
+RESOLVE_IMPORTS = True
+
+def trans(psyir):
+    pass
+    '''
+    recipe_name = "test_cache.py"
+    for filename, content in [("module1.f90", module1),
+                              ("module2.f90", module2),
+                              ("code.f90", code),
+                              (recipe_name, recipe)]:
+        with open(tmpdir.join(filename), "w", encoding='utf-8') as my_file:
+            my_file.write(content)
+
+    # If enable-cache not used, no .psycache files exist
+    monkeypatch.chdir(tmpdir)
+    ModuleManager._instance = None
+    main(["code.f90", "-s", recipe_name])
+    assert not os.path.isfile("module1.psycache")
+    assert not os.path.isfile("module2.psycache")
+    assert not ModuleManager.get()._cache_active
+
+    # If enable-cache is used, it will generate .psycache files for each module
+    ModuleManager._instance = None
+    main(["code.f90", "-s", recipe_name, "--enable-cache"])
+    assert os.path.isfile(tmpdir.join("module1.psycache"))
+    assert os.path.isfile(tmpdir.join("module2.psycache"))
+    assert ModuleManager.get()._cache_active
+
+    ModuleManager._instance = None
+
+
 def test_main_include_invalid(capsys, tmpdir):
     '''Check that the main function complains if a non-existent location
     is specified as a search path for INCLUDE files.
diff --git a/src/psyclone/tests/psyir/frontend/fparser2_test.py b/src/psyclone/tests/psyir/frontend/fparser2_test.py
index f51493ced..1bfa31967 100644
--- a/src/psyclone/tests/psyir/frontend/fparser2_test.py
+++ b/src/psyclone/tests/psyir/frontend/fparser2_test.py
@@ -1698,6 +1698,8 @@ def test_process_use_stmts_resolving_external_imports(
     reader = FortranStringReader('''
     module test
         use other1
+        private
+        public a_func
     contains
         subroutine test_function()
             use other2, only : an_array
@@ -1722,10 +1724,15 @@ def test_process_use_stmts_resolving_external_imports(
     # the right symbol kind and datatype
     assert isinstance(symtab.lookup("other1"), ContainerSymbol)
     assert isinstance(symtab.lookup("a_func"), RoutineSymbol)
+    assert isinstance(symtab.lookup("N"), DataSymbol)
     assert isinstance(symtab.lookup("unused_array"), DataSymbol)
     assert symtab.lookup("n").datatype == INTEGER_TYPE
     # But not the private symbols
     assert "private_array" not in symtab
+    # The local symbols respect the local visibility statements
+    assert symtab.lookup("other1").visibility == Symbol.Visibility.PRIVATE
+    assert symtab.lookup("N").visibility == Symbol.Visibility.PRIVATE
+    assert symtab.lookup("a_func").visibility == Symbol.Visibility.PUBLIC
 
     routine = psyir.children[0]
     innersymtab = routine.symbol_table