Initial port of functionality from #971 adapted for PSyTran

psytran/clauses.py

@@ -4,22 +4,28 @@
 # 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,
 )
 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

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))