diff --git a/psytran/clauses.py b/psytran/clauses.py index b74513c..83462ac 100644 --- a/psytran/clauses.py +++ b/psytran/clauses.py @@ -4,10 +4,14 @@ # See LICENSE in the root of the repository for full licensing details. r""" -This module implements functions for querying whether :py:class:`Node`\s have -OpenACC clauses associated with them, as well as for applying such clauses. +This module implements functions for querying clauses. +This includes whether :py:class:`Node`\s have OpenACC clauses +associated with them, as well as for applying such clauses, or a workaround +for the firstprivate issues on 3.1 version of PSyclone. """ +from psyclone.psyir import nodes +from psyclone.psyir import symbols from psytran.directives import ( has_loop_directive, _check_directive, @@ -15,11 +19,13 @@ from psytran.family import get_ancestors from psytran.loop import _check_loop + __all__ = [ "has_seq_clause", "has_gang_clause", "has_vector_clause", "has_collapse_clause", + "first_priv_red_init", ] @@ -96,3 +102,43 @@ def has_collapse_clause(loop): continue return collapse > i return False + + +def first_priv_red_init(node_target, init_scalars): + ''' + Add redundant initialisation before a Node, generally a Loop, where + a OMP clause has firstprivate added by PSyclone. + Software stack version of psyclone is adding firstprivates which fail + with CCE. + This is mostly fixed with PSyclone release 3.2, however that fix + may still have unforeseen edge-cases. + + Parameters + ---------- + Node : Target Node to reference from, adds redundant initialisation before. + str list: List of str variable indexes to reference against. + + Returns + ---------- + None : Note the tree has been modified + ''' + # Ensure scalars that may be emitted as FIRSTPRIVATE have a value + parent = node_target.parent + insert_at = parent.children.index(node_target) + for nm in init_scalars: # e.g., ("jdir", "k") + try: + sym = node_target.scope.symbol_table.lookup(nm) + # ensure character variables are initialised with CHARACTER_TYPE + # rather than UnsupportedFortranType + if isinstance(sym.datatype, symbols.UnsupportedFortranType): + init = nodes.Assignment.create( + nodes.Reference(sym), + nodes.Literal("", symbols.CHARACTER_TYPE)) + else: + init = nodes.Assignment.create( + nodes.Reference(sym), + nodes.Literal("0", sym.datatype)) + parent.children.insert(insert_at, init) + insert_at += 1 + except KeyError: + continue diff --git a/psytran/convert.py b/psytran/convert.py index 8aa2228..fff1a27 100644 --- a/psytran/convert.py +++ b/psytran/convert.py @@ -4,8 +4,10 @@ # See LICENSE in the root of the repository for full licensing details. """ -This module provides functions for converting the array notation used in a -:py:class:`Schedule`. +This module provides functions for converting aspects of the tree including +symbols. This includes the array notation used in a +:py:class:`Schedule`, and a method for converting an n_threads reference +into a library call. """ from psyclone.psyir import nodes @@ -36,3 +38,55 @@ def convert_array_notation(schedule): trans.Reference2ArrayRangeTrans().apply(reference) except TransformationError: # pragma: no cover pass + + +def replace_n_threads(psyir, n_threads_var_name): + ''' + If a scheme would use omp_get_max_threads() to determine + how work is divided, it will often be done so through + an omp clause. PSyclone will remove this. + We will therefore need to be able to add it to the source. + With a given variable name for n_threads know by the developer + in the source, replace its initialisation (often to 1) with + omp_get_max_threads(). + + Parameters + ---------- + psyir object : Uses whole psyir representation + n_threads_var_name str : The name of the variable in the + Scheme. Set relative to the scheme. + + Returns + ---------- + None : Note the tree has been modified + ''' + + imported_lib = False + # Walk the schedules + for shed in psyir.walk(nodes.Schedule): + # Walk the Assignments + for assign in shed.walk(nodes.Assignment): + # If the assignment has a lhs and is a reference... + if isinstance(assign.lhs, nodes.Reference): + # and that lhs name is n_threads_var_name + if assign.lhs.name == n_threads_var_name: + # Do this once, but only if needed + if imported_lib is False: + # Get the symbol table of the current schedule + symtab = shed.symbol_table + # Set up the omp_library symbol + omp_lib = symtab.find_or_create( + "omp_lib", + symbol_type=symbols.ContainerSymbol) + # Set up the omp_get_max_threads symbol + omp_get_max_threads = symtab.find_or_create( + "omp_get_max_threads", + symbol_type=symbols.RoutineSymbol, + # Import the reference + interface=symbols.ImportInterface(omp_lib)) + imported_lib = True + # Replace the rhs of the reference with the + # omp_get_max_threads symbol + assign.rhs.replace_with( + # pylint: disable=possibly-used-before-assignment + nodes.Call.create(omp_get_max_threads))