diff --git a/doc/user_guide/user_scripts.rst b/doc/user_guide/user_scripts.rst index 36c677c044..6f2a757515 100644 --- a/doc/user_guide/user_scripts.rst +++ b/doc/user_guide/user_scripts.rst @@ -89,6 +89,26 @@ and/or kernels) contained within the provided tree. The :ref:`examples section` provides a list of psyclone user scripts and associated usage instructions for multiple applications. +.. _script_kwargs: + +Arguments for Scripts +--------------------- +Scripts can take optional keyword arguments specified on the command line +using the option `--script-kwargs`. The keyword arguments are specified +as separate `keyword:value` pairs, separate by `,`. For example: + +.. code-block:: shell + + psyclone -s ./optimise.py input_source.f90 \ + --scripts-kwargs "omp: True, tiling: [4,4]" + +This will result in the additional keyword arguments for any transformation call: + +.. code-block:: python + + def trans(psyir, omp: bool, tiling: list[int]): + # Modify psyir tree + .. _sec_script_globals: @@ -142,3 +162,8 @@ associated with the merged invoke. An example of the use of a script making use of the ``trans_alg`` function can be found in ``examples/gocean/eg7``. + +Note that the ``trans_alg`` function will receive the same keyword arguments +as the ``trans`` function if the PSyclone command line option +``--script-kwargs`` is used (see :ref:`script_kwargs`). It is therefore +important that both functions accept the same keyword arguments. diff --git a/examples/nemo/eg8/Makefile b/examples/nemo/eg8/Makefile index 3a393495bb..f5858fe64d 100644 --- a/examples/nemo/eg8/Makefile +++ b/examples/nemo/eg8/Makefile @@ -60,7 +60,8 @@ PSYCLONE_PROFILING_LIB ?= ${PSYCLONE_PROFILING_DIR}/libsimple_timing.a PSYCLONE_PROFILING_LIBS ?= -L${PSYCLONE_PROFILING_DIR} -lsimple_timing transform: - ENABLE_PROFILING=1 ${PSYCLONE} -s ./omp_gpu_profile_trans.py ../code/tra_adv.F90 -o traadv_instrumented.F90 + ${PSYCLONE} -s ./omp_gpu_profile_trans.py --script-kwargs "profiling: True" \ + ../code/tra_adv.F90 -o traadv_instrumented.F90 compile: transform traadv.exe diff --git a/examples/nemo/eg8/README.md b/examples/nemo/eg8/README.md index b0e976e271..1ad6577c04 100644 --- a/examples/nemo/eg8/README.md +++ b/examples/nemo/eg8/README.md @@ -18,7 +18,8 @@ make transform or explicitly: ```sh -ENABLE_PROFILING=1 ${PSYCLONE} -s ./omp_gpu_profile_trans.py ../code/tra_adv.F90 -o traadv_instrumented.F90 +${PSYCLONE} -s ./omp_gpu_profile_trans.py --script-kwargs "profiling: True" \ + ../code/tra_adv.F90 -o traadv_instrumented.F90 ``` This emits transformed Fortran code with PSyData profiling around OpenMP target diff --git a/examples/nemo/eg8/omp_gpu_profile_trans.py b/examples/nemo/eg8/omp_gpu_profile_trans.py index 50792c097e..7ab26415a9 100644 --- a/examples/nemo/eg8/omp_gpu_profile_trans.py +++ b/examples/nemo/eg8/omp_gpu_profile_trans.py @@ -33,13 +33,12 @@ # POSSIBILITY OF SUCH DAMAGE. # ----------------------------------------------------------------------------- -import os import pathlib import sys from typing import List, Union from psyclone.psyir.nodes import ( - Assignment, IfBlock, Node, OMPDirective, OMPTargetDirective, ProfileNode, - Routine, Schedule) + Assignment, FileContainer, IfBlock, Node, OMPDirective, + OMPTargetDirective, ProfileNode, Routine, Schedule) from psyclone.psyir.transformations import OMPTargetTrans, ProfileTrans from psyclone.transformations import OMPLoopTrans, TransformationError @@ -50,9 +49,6 @@ sys.path.insert(0, str(NEMO_SCRIPTS_DIR)) -PROFILING_ENABLED = os.environ.get("ENABLE_PROFILING", False) - - def add_omp_region_profiling_markers(children: Union[List[Node], Schedule]): """Insert profiling markers around all top-level OpenMP directives. @@ -94,8 +90,16 @@ def add_omp_region_profiling_markers(children: Union[List[Node], Schedule]): add_omp_region_profiling_markers(child.children) -def trans(psyir): - """Apply OpenMP offloading and insert profiling around target regions.""" +def trans(psyir: FileContainer, profiling=False): + """ + Apply OpenMP offloading and insert profiling around target regions. + + :param psyir: the PSyIR of the file container to modify. + :param profiling: if set to True (using the PSyclone command line option + --scripts-kwargs "profiling: True"), also adds profiling + instrumentation to the generated code. + """ + from utils import normalise_loops, insert_explicit_loop_parallelism omp_target_trans = OMPTargetTrans() @@ -119,5 +123,5 @@ def trans(psyir): collapse=True, enable_reductions=True ) - if PROFILING_ENABLED: + if profiling: add_omp_region_profiling_markers(subroutine.children) diff --git a/src/psyclone/generator.py b/src/psyclone/generator.py index f0005709a2..79b888b85d 100644 --- a/src/psyclone/generator.py +++ b/src/psyclone/generator.py @@ -47,13 +47,13 @@ ''' import argparse +import importlib +import logging import os +import shutil import sys import traceback -import importlib -import shutil from typing import Callable, Iterable, List, Optional, Tuple, Union -import logging from fparser.api import get_reader from fparser.two import Fortran2003 @@ -77,13 +77,14 @@ from psyclone.parse.kernel import get_kernel_filepath from psyclone.parse.utils import ParseError, parse_fp2 from psyclone.profiler import Profiler -from psyclone.psyGen import PSyFactory +from psyclone.psyGen import PSyFactory, Transformation from psyclone.psyir.backend.fortran import FortranWriter from psyclone.psyir.frontend.fortran import FortranReader from psyclone.psyir.frontend.fparser2 import Fparser2Reader -from psyclone.psyir.nodes import Loop, Container, Routine +from psyclone.psyir.nodes import Container, FileContainer, Loop, Node, Routine from psyclone.psyir.symbols import UnresolvedInterface from psyclone.psyir.transformations import TransformationError +from psyclone.utils import parse_kwargs from psyclone.version import __VERSION__ # TODO issue #1618 remove temporary LFRIC_TESTING flag, associated @@ -113,20 +114,26 @@ def load_script( - script_name: str, function_name: str = "trans", + script_name: str, + kwargs_str: Optional[str] = None, + function_name: str = "trans", is_optional: bool = False -) -> Tuple[Callable, List[str], Union[bool, List[str]]]: +) -> Tuple[Optional[Callable], + List[str], + Union[bool, List[str]], + dict[str, str]]: ''' Loads the specified script containing a psyclone recipe. We also prepend the script path to the sys.path, so that the script itself and any imports that it has from the same directory can be found. :param script_name: name of the script to load. + :param kwargs_str: the kwargs argument from the command line. :param function_name: the name of the function to call in the script. :param is_optional: whether the function is optional or not. Defaults to False. :returns: callable recipe, list of files to skip, whether to resolve - modules (or which ones). + modules (or which ones), the kwargs dictionary. :raises IOError: if the file is not found. :raises GenerationError: if the file does not have .py extension. @@ -134,8 +141,15 @@ def load_script( be called. ''' + filepath, filename = os.path.split(script_name) + if kwargs_str is not None: + kwargs = parse_kwargs(kwargs_str) + else: + kwargs = {} + module_name, fileext = os.path.splitext(filename) + # the file must either be: # a) at the given path or, given no path, in the current directory; or # b) given no path, in the system path @@ -156,11 +170,13 @@ def load_script( sys.path.insert(0, filepath) recipe_module = importlib.import_module(module_name) + files_to_skip: list[str] if hasattr(recipe_module, "FILES_TO_SKIP"): files_to_skip = recipe_module.FILES_TO_SKIP else: files_to_skip = [] + imports_to_resolve: list[str] if hasattr(recipe_module, "RESOLVE_IMPORTS"): imports_to_resolve = recipe_module.RESOLVE_IMPORTS # If the imports_to_resolve has the list of explicit filenames, respect @@ -173,12 +189,13 @@ def load_script( imports_to_resolve = [] if hasattr(recipe_module, function_name): - transformation_recipe = getattr(recipe_module, function_name) + transformation_recipe: Callable = getattr(recipe_module, function_name) if callable(transformation_recipe): # Everything is good, return recipe and files_to_skip - return transformation_recipe, files_to_skip, imports_to_resolve + return (transformation_recipe, files_to_skip, + imports_to_resolve, kwargs) elif is_optional: - return None, files_to_skip, imports_to_resolve + return None, files_to_skip, imports_to_resolve, {} raise GenerationError( f"generator: attempted to use specified PSyclone " f"transformation module '{module_name}' but it does not " @@ -189,14 +206,15 @@ def generate(filename: str, api: str = "", kernel_paths: Optional[list[str]] = None, script_name: Optional[str] = None, + kwargs_str: Optional[str] = None, line_length: bool = False, - distributed_memory: bool = None, + distributed_memory: Optional[bool] = None, kern_out_path: str = "", keep_comments: bool = False, keep_directives: bool = False, keep_conditional_openmp_statements: bool = False, free_form: bool = True - ) -> Tuple[str, str]: + ) -> Tuple[Optional[str], str]: # pylint: disable=too-many-arguments, too-many-statements # pylint: disable=too-many-branches, too-many-locals '''Takes a PSyclone algorithm specification as input and outputs the @@ -285,8 +303,10 @@ def generate(filename: str, .create(invoke_info) if script_name is not None: # Apply provided recipe to PSyIR - recipe, _, _ = load_script(script_name) - recipe(psy.container.root) + trans_func, _, _, kwargs = load_script(script_name, kwargs_str) + # trans_func is always defined, otherwise an exception is raised + assert trans_func + trans_func(psy.container.root, **kwargs) alg_gen = None elif api in GOCEAN_API_NAMES or (api in LFRIC_API_NAMES and LFRIC_TESTING): @@ -336,13 +356,13 @@ def generate(filename: str, if script_name is not None: # Call the optimisation script for algorithm optimisations - recipe, _, _ = load_script(script_name, "trans_alg", - is_optional=True) + recipe, _, _, kwargs = load_script(script_name, kwargs_str, + "trans_alg", is_optional=True) if recipe: - recipe(psyir) + recipe(psyir, **kwargs) # For each kernel called from the algorithm layer - kernels = {} + kernels: dict[int, dict[int, Node]] = {} for invoke in psyir.walk(AlgorithmInvokeCall): kernels[id(invoke)] = {} for kern in invoke.walk(KernelFunctor): @@ -386,6 +406,7 @@ def generate(filename: str, sys.exit(1) # Raise to Kernel PSyIR + kern_trans: Transformation if api in GOCEAN_API_NAMES: kern_trans = RaisePSyIR2GOceanKernTrans(kern.symbol.name) kern_trans.apply(kernel_psyir) @@ -398,6 +419,7 @@ def generate(filename: str, kernels[id(invoke)][id(kern)] = kernel_psyir # Transform 'invoke' calls into calls to PSy-layer subroutines + invoke_trans: Transformation if api in GOCEAN_API_NAMES: invoke_trans = GOceanAlgInvoke2PSyCallTrans() else: # api in LFRIC_API_NAMES @@ -426,8 +448,10 @@ def generate(filename: str, if script_name is not None: # Call the optimisation script for psy-layer optimisations - recipe, _, _ = load_script(script_name) - recipe(psy.container.root) + trans_func, _, _, kwargs = load_script(script_name, kwargs_str) + # recipe is always defined, otherwise an exception is raised + assert trans_func + trans_func(psy.container.root, **kwargs) # TODO issue #1618 remove Alg class and tests from PSyclone if api in LFRIC_API_NAMES and not LFRIC_TESTING: @@ -441,7 +465,7 @@ def generate(filename: str, return alg_gen, psy.gen -def main(arguments): +def main(arguments: list[str]) -> None: ''' Parses and checks the command line arguments, calls the generate function if all is well, catches any errors and outputs the @@ -449,7 +473,6 @@ def main(arguments): :param arguments: the list of command-line arguments that PSyclone has been invoked with. - :type arguments: List[str] ''' # pylint: disable=too-many-statements,too-many-branches @@ -470,8 +493,11 @@ def main(arguments): help='display version information') parser.add_argument('-c', '--config', help='config file with ' 'PSyclone specific options') - parser.add_argument('-s', '--script', help='filename of a PSyclone' - ' optimisation recipe') + scripts = parser.add_argument_group("Transformation scripts") + scripts.add_argument('-s', '--script', + help='filename of a PSyclone optimisation recipe') + scripts.add_argument('--script-kwargs', help='Keyword arguments for the ' + 'transformation script.') parser.add_argument( '--enable-cache', action="store_true", default=False, help='whether to enable caching of imported module dependencies (if ' @@ -609,6 +635,7 @@ def main(arguments): args = parser.parse_args(arguments) # Set the logging system up. + handler: logging.Handler if args.log_file: handler = logging.FileHandler(args.log_file, mode="a", encoding="utf-8") @@ -732,7 +759,8 @@ def main(arguments): if not args.psykal_dsl: code_transformation_mode( input_file=args.filename, - recipe_file=args.script, + script_name=args.script, + kwargs_str=args.script_kwargs, output_file=args.o, keep_comments=args.keep_comments, keep_directives=args.keep_directives, @@ -764,6 +792,7 @@ def main(arguments): args.filename, api=api, kernel_paths=args.directory, script_name=args.script, + kwargs_str=args.script_kwargs, line_length=(args.limit == 'all'), distributed_memory=args.dist_mem, kern_out_path=kern_out_path, @@ -817,12 +846,13 @@ def main(arguments): print(f"Generated psy layer code:\n{psy_str}") -def check_psyir(psyir, filename): +def check_psyir(psyir: FileContainer, + filename: str) -> None: '''Check the supplied psyir to make sure that it contains a single program or module. :param psyir: the psyir to check. - :type psyir: py:class:`psyclone.psyir.nodes.FileContainer` + :param filename: filename to use in error messages. :raises GenerationError: if the algorithm file contains \ multiple modules or programs. @@ -845,7 +875,8 @@ def check_psyir(psyir, filename): f"found '{type(psyir.children[0]).__name__}'.") -def add_builtins_use(fp2_tree, name): +def add_builtins_use(fp2_tree: Fortran2003.Program, + name: str) -> None: '''Modify the fparser2 tree adding a 'use ' so that builtin kernel functors do not appear to be undeclared. @@ -880,23 +911,26 @@ def add_builtins_use(fp2_tree, name): spec_part.children.insert(0, use_stmt) -def code_transformation_mode(input_file, recipe_file, output_file, - keep_comments: bool, keep_directives: bool, +def code_transformation_mode(input_file: str, + script_name: str, + output_file: str, + keep_comments: bool, + keep_directives: bool, keep_conditional_openmp_statements: bool, - free_form: bool = True, line_length="off"): - ''' Process the input_file with the recipe_file instructions and - store it in the output_file. + kwargs_str: Optional[str] = None, + free_form: bool = True, + line_length="off"): + ''' + Process the input_file with the transformations script specified in + `script_name` and store it in the output_file. Note: there is some duplicated logic in the PSyKAl path, we could attempt to merge them when adopting the LFRIC_TESTING PATH and removing the previous way. :param input_file: the given input file. - :type input_file: str | os.PathLike - :param recipe_file: the given transformation recipe file. - :type input_file: Optional[str | os.PathLike] + :param script_name: the given transformation recipe file. :param output_file: the output file where to store the resulting code. - :type output_file: Optional[str | os.PathLike] :param keep_comments: whether to keep comments from the original source. :param keep_directives: whether to keep directives from the original source. @@ -910,11 +944,13 @@ def code_transformation_mode(input_file, recipe_file, output_file, ''' logger = logging.getLogger(__name__) - # Load recipe file - if recipe_file: - trans_recipe, files_to_skip, resolve_mods = load_script(recipe_file) + # Load script file + if script_name: + (trans_recipe, files_to_skip, + resolve_mods, kwargs) = load_script(script_name, kwargs_str) else: - trans_recipe, files_to_skip, resolve_mods = (None, [], False) + trans_recipe, files_to_skip, resolve_mods, kwargs = (None, [], False, + {}) _, filename = os.path.split(input_file) if filename not in files_to_skip: @@ -950,7 +986,7 @@ def code_transformation_mode(input_file, recipe_file, output_file, # Modify file if trans_recipe: - trans_recipe(psyir) + trans_recipe(psyir, **kwargs) # Add profiling if automatic profiling has been requested for routine in psyir.walk(Routine): diff --git a/src/psyclone/tests/generator_test.py b/src/psyclone/tests/generator_test.py index 4e993b4620..cf44e712ad 100644 --- a/src/psyclone/tests/generator_test.py +++ b/src/psyclone/tests/generator_test.py @@ -2110,3 +2110,147 @@ def test_config_overwrite() -> None: main([filename, "--config-opts", "DOES_NOT_EXIST=27"]) assert ("Attempt to overwrite unknown configuration option: " "'DOES_NOT_EXIST=27'" in str(err.value)) + + +@pytest.mark.parametrize("kwargs, out", [("", "{}"), + ("'a':1", "{'a': 1}"), + ("'b': {1: 2}", "{'b': {1: 2}}"), + ("'l': [1,2]", "{'l': [1, 2]}"), + ("a:1", "{'a': 1}"), + ("b: {1: 2}", "{'b': {1: 2}}"), + ("l: [1,2]", "{'l': [1, 2]}") + + ]) +def test_script_arguments_transmute(kwargs, out, tmp_path, capsys): + """Tests that script arguments are received as expected with the + transmute approach. This test creates a dummy script that prints the + arguments, which we check for using capsys + + :param kwargs: the input string for the command line + :param out: the expected output of the print statement (which can be + different from the input if the keyword argument is not a string). + """ + recipe = ''' +def trans(psyir, **kwargs): + print("XARGS:", kwargs) + ''' + script_path = tmp_path / "print_args_transmute.py" + script_path.write_text(recipe) + + inputfile = Path(get_base_path("nemo")) / "afunction.f90" + outputfile = tmp_path / "output.f90" + main([str(inputfile), "-s", str(script_path), + "--script-kwargs", kwargs, + "-o", str(outputfile)]) + stdout, _ = capsys.readouterr() + assert f"ARGS: {out}" in stdout + + +@pytest.mark.parametrize("kwargs, out", [("", "{}"), + ("'a':1", "{'a': 1}"), + ("'b': {1: 2}", "{'b': {1: 2}}"), + ("'l': [1,2]", "{'l': [1, 2]}"), + ("a:1", "{'a': 1}"), + ("b: {1: 2}", "{'b': {1: 2}}"), + ("l: [1,2]", "{'l': [1, 2]}") + ]) +def test_script_arguments_lfric_testing(kwargs, out, tmp_path, capsys, + monkeypatch): + """Tests that script arguments are received as expected using the + LFRic API. This test creates a dummy script that prints the arguments + for trans and trans_alg, which we check for. This uses LFRIC_TESTING, + which will also call trans_alg (which by LFRic otherwise would not do). + + :param kwargs: the input string for the command line + :param out: the expected output of the print statement (which can be + different from the input if the keyword argument is not a string). + """ + monkeypatch.setattr(generator, "LFRIC_TESTING", True) + + recipe = ''' +def trans(psyir, **kwargs): + print("trans args:", kwargs) + +def trans_alg(psyir, **kwargs): + print("trans_alg args:", kwargs) + ''' + script_path = tmp_path / "print_args_lfric_testing.py" + script_path.write_text(recipe) + + inputfile = Path(get_base_path("lfric")) / "1_single_invoke.f90" + psy_file = tmp_path / "psy.f90" + alg_file = tmp_path / "alg.f90" + main([str(inputfile), "-s", str(script_path), "--psykal-dsl", "lfric", + "--script-kwargs", kwargs, + "-opsy", str(psy_file), + "-oalg", str(alg_file)]) + stdout, _ = capsys.readouterr() + assert f"trans args: {out}" in stdout + assert f"trans_alg args: {out}" in stdout + + +@pytest.mark.parametrize("kwargs, out", [("", "{}"), + ("'a':1", "{'a': 1}"), + ("'b': {1: 2}", "{'b': {1: 2}}"), + ("'l': [1,2]", "{'l': [1, 2]}"), + ("a:1", "{'a': 1}"), + ("b: {1: 2}", "{'b': {1: 2}}"), + ("l: [1,2]", "{'l': [1, 2]}") + ]) +def test_script_arguments_lfric_default(kwargs, out, tmp_path, capsys): + """Tests that script arguments are received as expected using the + LFRic API. This test creates a dummy script that prints the arguments + for trans and trans_alg, which we check for. This uses default + LFRic handling, which does not call trans_alg. + + :param kwargs: the input string for the command line + :param out: the expected output of the print statement (which can be + different from the input if the keyword argument is not a string). + """ + + recipe = ''' +def trans(psyir, **kwargs): + print("trans args:", kwargs) + +def trans_alg(psyir, **kwargs): + print("trans_alg args:", kwargs) + ''' + script_path = tmp_path / "print_args_lfric_default.py" + script_path.write_text(recipe) + + inputfile = Path(get_base_path("lfric")) / "1_single_invoke.f90" + psy_file = tmp_path / "psy.f90" + alg_file = tmp_path / "alg.f90" + main([str(inputfile), "-s", str(script_path), "--psykal-dsl", "lfric", + "--script-kwargs", kwargs, + "-opsy", str(psy_file), + "-oalg", str(alg_file)]) + stdout, _ = capsys.readouterr() + assert f"trans args: {out}" in stdout + # Default LFRic API does not call trans_alg!!! This line is here to + # fail once we switch LFRic over. + assert f"trans_alg args: {out}" not in stdout + + +@pytest.mark.parametrize("kwargs", ["1", "'a'", "[1,2]", "{1:2}", + "a=1"]) +def test_script_arguments_errors(kwargs, tmp_path): + """Tests that script arguments are received as expected. This + test creates a dummy script that prints the arguments, which + we check for + :param kwargs: the input string for the command line + """ + recipe = ''' +def trans(psyir, **kwargs): + print("YARGS:", kwargs) + ''' + script_path = tmp_path / "print_args_errors.py" + script_path.write_text(recipe) + + inputfile = Path(get_base_path("lfric")) / "1_single_invoke.f90" + outputfile = tmp_path / "output.f90" + with pytest.raises(ValueError) as err: + main([str(inputfile), "-s", str(script_path), + "--script-kwargs", kwargs, + "-o", str(outputfile)]) + assert "Invalid syntax for keyword arguments" in str(err.value) diff --git a/src/psyclone/tests/utils_test.py b/src/psyclone/tests/utils_test.py index 06d2dbb122..73251fedf7 100644 --- a/src/psyclone/tests/utils_test.py +++ b/src/psyclone/tests/utils_test.py @@ -36,17 +36,17 @@ '''This module implements tests for the generic utility functions.''' import inspect -import pytest import sys - from typing import Union +import pytest + from psyclone.errors import InternalError from psyclone.transformations import Transformation from psyclone.utils import ( - within_virtual_env, a_or_an, + a_or_an, parse_kwargs, transformation_documentation_wrapper, - stringify_annotation, + stringify_annotation, within_virtual_env, ) @@ -472,3 +472,38 @@ def apply(self, node, opt1: bool = False, opt2=None, :type opt3: int :param int opt3: (Option provided for SubTrans2) opt3 docstring.""" assert correct in BaseTrans.apply.__doc__ + + +@pytest.mark.parametrize("kwargs, expected", + [("", {}), + ("'a':1", {'a': 1}), + ("'b': {1: 2}", {'b': {1: 2}}), + ("'l': [1,2]", {'l': [1, 2]}), + ("a:1", {'a': 1}), + ("b: {1: 2}", {'b': {1: 2}}), + ("l: [1,2]", {'l': [1, 2]}), + ]) +def test_parse_kwargs(kwargs, expected): + """ + Test that the parsing function for user-specific script options + work as expected. + + :param kwargs: the input string for the command line + """ + result = parse_kwargs(kwargs) + assert result == expected + + +@pytest.mark.parametrize("kwargs", ["1", "'a'", "[1,2]", "{1:2}", + "a=1", 123]) +def test_parse_kwargs_errors(kwargs): + """ + Test that the parsing function for user-specific script options + work as expected. + + :param kwargs: the input string for the command line + """ + with pytest.raises(ValueError) as err: + parse_kwargs(kwargs) + + assert "Invalid syntax for keyword arguments" in str(err.value) diff --git a/src/psyclone/utils.py b/src/psyclone/utils.py index a822d56fa2..ffcfc917f6 100644 --- a/src/psyclone/utils.py +++ b/src/psyclone/utils.py @@ -37,9 +37,11 @@ '''This module provides generic utility functions.''' -from typing import Type, TYPE_CHECKING, Union +import ast from collections import OrderedDict import sys +from typing import Any, Type, TYPE_CHECKING, Union + from psyclone.errors import InternalError from psyclone.docstring_parser import ( DocstringData, ReturnsData @@ -234,3 +236,82 @@ def wrapper(cls): return wrapper(*args) else: return wrapper + + +# ---------------------------------------------------------------------------- +def parse_kwargs(kwargs: str) -> dict[str, Any]: + """ + This function safely parses a user string provided on the command line + using '--kwargs ...` into a python dictionary. It especially simplifies + the syntax for the user by not requiring the keys to be escaped, e.g. + --kwargs "'a':1,'b':2" and --kwargs "a:1,b:2" will both work as expected. + + This is done by using Python'kwargs ast parser, then adding a separate + transformation step that replaces keys that are an ast.Name + with an ast.Constant, then finally calling literal_eval to + create the dictionary. + + :param kwargs: the string to parse. + + :raises ValueError: if the string cannot be converted to a kwargs-style + dictionary. + """ + + # Parse as an expression. Note that various ast functions can + # raise different exceptions, so we catch all exceptions and + # re-raise them as a ValueError + try: + # Make it look like a dict literal + wrapped = "{" + kwargs.strip().rstrip(",") + "}" + expr = ast.parse(wrapped, mode="eval") + + # Convert bare-name keys to string keys + transformer = NameKeysToStr() + expr = transformer.visit(expr) + # This call will update line-number, column, ... information in + # the modified tree, since the newly created nodes won't have + # this information + ast.fix_missing_locations(expr) + + # Safely evaluate literals/containers + result = ast.literal_eval(expr) + # pylint: disable=broad-exception-caught + except Exception: + # This will trigger an exception in the next statement + result = None + + if not isinstance(result, dict): + raise ValueError(f"Invalid syntax for keyword arguments '{kwargs}' ") + + return result + + +class NameKeysToStr(ast.NodeTransformer): + """ + This is a helper class to convert dictionary keys that are + an ast.Name (i.e. not a string) into an ast.Constant (a string). + + It will effectively change `{a:1}` to `{'a':1}` + + :param node: the dictionary node. + """ + + # pylint: disable=invalid-name + def visit_Dict(self, node: ast.Dict): + """Function to replace non-string keys in a dictionary + with strings representing the same name. + + :param node: the dictionary node in Python AST. + """ + + # Transform keys: Name(...) -> Constant("name") + new_keys = [] + for k in node.keys: + if isinstance(k, ast.Name): + new_keys.append(ast.Constant(k.id)) + else: + new_keys.append(self.visit(k)) + node.keys = new_keys + # Still visit values normally + node.values = [self.visit(v) for v in node.values] + return node diff --git a/tutorial/training/transformation/3.8-code-creation/Makefile b/tutorial/training/transformation/3.8-code-creation/Makefile index 0e435b7bbc..526abcd28b 100644 --- a/tutorial/training/transformation/3.8-code-creation/Makefile +++ b/tutorial/training/transformation/3.8-code-creation/Makefile @@ -11,19 +11,20 @@ SRC_F90 = gol.f90 compute_die_mod.f90 compute_born_mod.f90 \ output_field_mod.f90 time_step_mod.f90 OBJ = $(SRC_F90:%.f90=%.o) $(PROCESSED_SRC:%.f90=%.o) +PARSE_STRING?=False transform: $(PROCESSED_SRC) test: @# First test the code creation using the PSyIR API $(MAKE) clean - PARSE_STRING="" $(MAKE) combine_mod.processed.f90 | grep "Creating tree" + $(MAKE) PARSE_STRING=False combine_mod.processed.f90 | grep "Creating tree" @# Count the number of j-loops, which should be 2 cat combine_mod.processed.f90 | grep -c 'j = ystart, ystop' | grep 2 @# Then test the code creation by parsing a string $(MAKE) clean - PARSE_STRING="yes" $(MAKE) combine_mod.processed.f90 | grep "Parsing string" + $(MAKE) PARSE_STRING=True combine_mod.processed.f90 | grep "Parsing string" @# Count the number of j-loops, which should be 2 cat combine_mod.processed.f90 | grep -c 'j = ystart, ystop' | grep 2 @@ -37,7 +38,8 @@ $(EXE): $(OBJ) .precious: $(PROCESSED_SRC) combine_mod.processed.f90: combine_mod.f90 $(SCRIPT) - $(PSYCLONE) -s ./$(SCRIPT) -o combine_mod.processed.f90 combine_mod.f90 + $(PSYCLONE) -s ./$(SCRIPT) --script-kwargs "parse_string:$(PARSE_STRING)" \ + -o combine_mod.processed.f90 combine_mod.f90 %.o: %.f90 $(F90) -c $(F90FLAGS) $< diff --git a/tutorial/training/transformation/3.8-code-creation/README.md b/tutorial/training/transformation/3.8-code-creation/README.md index e2f1ef3b63..383f1fee27 100644 --- a/tutorial/training/transformation/3.8-code-creation/README.md +++ b/tutorial/training/transformation/3.8-code-creation/README.md @@ -69,10 +69,10 @@ There are two ways of creating the if condition: approach, you don't need to worry about converting the start- and stop-expressions, you just use a copy of the tree. -The ``add_if.py`` template supports both ways, and queries the -environment variable ``PARSE_STRING`` to decide which one to use. Just -set ``PARSE_STRING`` to a non-empty string to use string-parsing, -otherwise (also as default) it will use the PSyIR subtree creation method. +The ``add_if.py`` template supports both ways, and it takes a keyword +argument to select which one to use. Use the PSyclone command line +option ``--scripts-kwargs "parse_string: True`` (or ``False``, which +is also the default) to select which version to use. Next, create the ``if``-statement (this time only using the PSyIR tree assembly, which feels more natural and faster than creating the full diff --git a/tutorial/training/transformation/3.8-code-creation/add_if.py b/tutorial/training/transformation/3.8-code-creation/add_if.py index 915851e2ba..4e913ca21d 100755 --- a/tutorial/training/transformation/3.8-code-creation/add_if.py +++ b/tutorial/training/transformation/3.8-code-creation/add_if.py @@ -39,8 +39,6 @@ of a loop, depending on iteration count. ''' -import os - from psyclone.psyir.frontend.fortran import FortranReader from psyclone.psyir.backend.fortran import FortranWriter @@ -50,7 +48,7 @@ Literal, Loop, Routine, Schedule) -def trans(psyir: FileContainer) -> None: +def trans(psyir: FileContainer, parse_string: False) -> None: ''' Create two versions of a loop, depending in iteration count. @@ -74,7 +72,7 @@ def trans(psyir: FileContainer) -> None: # Use an environment variable to select which one you # want to use - if os.environ.get("PARSE_STRING", False): + if parse_string: # Option 1: Create expression by parsing a Fortran string: writer = FortranWriter() # TODO: Create a string with the Fortran condition diff --git a/tutorial/training/transformation/3.8-code-creation/solution/Makefile b/tutorial/training/transformation/3.8-code-creation/solution/Makefile index 60884a72fb..a9597bb0ea 100644 --- a/tutorial/training/transformation/3.8-code-creation/solution/Makefile +++ b/tutorial/training/transformation/3.8-code-creation/solution/Makefile @@ -11,19 +11,20 @@ SRC_F90 = gol.f90 compute_die_mod.f90 compute_born_mod.f90 \ output_field_mod.f90 time_step_mod.f90 OBJ = $(SRC_F90:%.f90=%.o) $(PROCESSED_SRC:%.f90=%.o) +PARSE_STRING?=False transform: $(PROCESSED_SRC) test: @# First test the code creation using the PSyIR API $(MAKE) clean - PARSE_STRING="" $(MAKE) combine_mod.processed.f90 | grep "Creating tree" + $(MAKE) PARSE_STRING=False combine_mod.processed.f90 | grep "Creating tree" @# Count the number of j-loops, which should be 2 cat combine_mod.processed.f90 | grep -c 'j = ystart, ystop' | grep 2 @# Then test the code creation by parsing a string $(MAKE) clean - PARSE_STRING="yes" $(MAKE) combine_mod.processed.f90 | grep "Parsing string" + $(MAKE) PARSE_STRING=True combine_mod.processed.f90 | grep "Parsing string" @# Count the number of j-loops, which should be 2 cat combine_mod.processed.f90 | grep -c 'j = ystart, ystop' | grep 2 @@ -37,7 +38,8 @@ $(EXE): $(OBJ) .precious: $(PROCESSED_SRC) combine_mod.processed.f90: combine_mod.f90 $(SCRIPT) - $(PSYCLONE) -s ./$(SCRIPT) -o combine_mod.processed.f90 combine_mod.f90 + $(PSYCLONE) -s ./$(SCRIPT) --script-kwargs "parse_string:$(PARSE_STRING)" \ + -o combine_mod.processed.f90 combine_mod.f90 %.o: %.f90 $(F90) -c $(F90FLAGS) $< diff --git a/tutorial/training/transformation/3.8-code-creation/solution/add_if.py b/tutorial/training/transformation/3.8-code-creation/solution/add_if.py index 7c1e0c01e3..58da45f6b4 100755 --- a/tutorial/training/transformation/3.8-code-creation/solution/add_if.py +++ b/tutorial/training/transformation/3.8-code-creation/solution/add_if.py @@ -38,8 +38,6 @@ of a loop, depending on iteration count. ''' -import os - from psyclone.psyir.frontend.fortran import FortranReader from psyclone.psyir.backend.fortran import FortranWriter @@ -49,7 +47,7 @@ Literal, Loop, Routine, Schedule) -def trans(psyir: FileContainer) -> None: +def trans(psyir: FileContainer, parse_string: False) -> None: ''' Create two versions of a loop, depending in iteration count. @@ -71,7 +69,7 @@ def trans(psyir: FileContainer) -> None: # 1. Parsing a Fortran expression given as string # 2. Creating the tree representation using the PSyIR API - if os.environ.get("PARSE_STRING", False): + if parse_string: # Option 1: Create expression by parsing a Fortran string: writer = FortranWriter() expr_str = (f"{writer(outer_loop.stop_expr)} - "