How to use the psyclone.psyir.nodes.Loop function in PSyclone

To help you get started, we’ve selected a few PSyclone examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github stfc / PSyclone / examples / lfric / eg3 / colouring_and_omp.py View on Github external
''' PSyclone transformation script for the dynamo0p3 api to apply
    colouring and OpenMP generically.'''
    ctrans = Dynamo0p3ColourTrans()
    otrans = DynamoOMPParallelLoopTrans()

    # Loop over all of the Invokes in the PSy object
    for invoke in psy.invokes.invoke_list:

        print("Transforming invoke '"+invoke.name+"'...")
        schedule = invoke.schedule

        # Colour all of the loops over cells unless they are on
        # discontinuous spaces
        cschedule = schedule
        for child in schedule.children:
            if isinstance(child, Loop) \
               and child.field_space.orig_name \
               not in FunctionSpace.VALID_DISCONTINUOUS_NAMES \
               and child.iteration_space == "cells":
                cschedule, _ = ctrans.apply(child)
        # Then apply OpenMP to each of the colour loops
        schedule = cschedule
        for child in schedule.children:
            if isinstance(child, Loop):
                if child.loop_type == "colours":
                    schedule, _ = otrans.apply(child.loop_body[0])
                else:
                    schedule, _ = otrans.apply(child)

        schedule.view()
        invoke.schedule = schedule
github stfc / PSyclone / src / psyclone / psyir / tools / dependency_tools.py View on Github external
has a certain type (e.g. 'latitude'), and optional tests for
        nested loops (to avoid parallelising single loops which will likely
        result in a slowdown due to thread synchronisation costs). This
        function is used by can_loop_be_parallelised() and can of course be
        overwritten by the user to implement different suitability criteria.

        :param loop: the loop to test.
        :type loop: :py:class:`psyclone.psyir.nodes.Loop`
        :param bool only_nested_loops: true (default) if only nested loops\
                                        should be considered.

        :returns: true if the loop fulfills the requirements.
        :rtype: bool
        '''
        if only_nested_loops:
            all_loops = loop.walk(Loop)
            if len(all_loops) == 1:
                self._add_info("Not a nested loop.")
                return False

        if loop.loop_type not in self._loop_types_to_parallelise:
            self._add_info("Loop has wrong loop type '{0}'.".
                           format(loop.loop_type))
            return False
        return True
github stfc / PSyclone / src / psyclone / dynamo0p1.py View on Github external
my_typedecl = TypeDeclGen(invoke_sub, datatype="field_type",
                                  entity_decls=self.psy_unique_var_names,
                                  intent="inout")
        invoke_sub.add(my_typedecl)


class DynInvokeSchedule(InvokeSchedule):
    ''' The Dynamo specific InvokeSchedule sub-class. This passes the Dynamo
        specific loop and infrastructure classes to the base class so it
        creates the ones we require. '''
    def __init__(self, arg, reserved_names=None):
        InvokeSchedule.__init__(self, DynKernCallFactory,
                                DynBuiltInCallFactory, arg, reserved_names)


class DynLoop(Loop):
    ''' The Dynamo specific Loop class. This passes the Dynamo specific
        loop information to the base class so it creates the one we require.
        Creates Dynamo specific loop bounds when the code is being generated.
    '''
    def __init__(self, parent=None, loop_type=""):
        Loop.__init__(self, parent=parent,
                      valid_loop_types=["", "colours", "colour"])
        self.loop_type = loop_type

        # Work out the variable name from  the loop type
        if self._loop_type == "colours":
            tag = "colours_loop_idx"
            suggested_name = "colour"
        elif self._loop_type == "colour":
            tag = "colour_loop_idx"
            suggested_name = "cell"
github stfc / PSyclone / src / psyclone / nemo.py View on Github external
def __init__(self, parent=None, variable_name=''):
        valid_loop_types = Config.get().api_conf("nemo").get_valid_loop_types()
        Loop.__init__(self, parent=parent,
                      variable_name=variable_name,
                      valid_loop_types=valid_loop_types)
github stfc / PSyclone / src / psyclone / psyir / nodes / node.py View on Github external
def loops(self):
        '''Return all loops currently in this schedule.'''
        from psyclone.psyir.nodes import Loop
        return self.walk(Loop)
github stfc / PSyclone / src / psyclone / psyir / frontend / fparser2.py View on Github external
def _create_loop(self, parent, variable_name):
        '''
        Create a Loop instance. This is done outside _do_construct_handler
        because some APIs may want to instantiate a specialised Loop.

        :param parent: the parent of the node.
        :type parent: :py:class:`psyclone.psyir.nodes.Node`
        :param str variable_name: name of the iteration variable.

        :return: a new Loop instance.
        :rtype: :py:class:`psyclone.psyir.nodes.Loop`

        '''
        return Loop(parent=parent, variable_name=variable_name)
github stfc / PSyclone / src / psyclone / gocean1p0.py View on Github external
    @property
    def const_loop_bounds(self):
        ''' Returns True if constant loop bounds are enabled for this
        schedule. Returns False otherwise. '''
        return self._const_loop_bounds

    @const_loop_bounds.setter
    def const_loop_bounds(self, obj):
        ''' Set whether the InvokeSchedule will use constant loop bounds or
        will look them up from the field object for every loop '''
        self._const_loop_bounds = obj


# pylint: disable=too-many-instance-attributes
class GOLoop(Loop):
    ''' The GOcean specific Loop class. This passes the GOcean specific
        single loop information to the base class so it creates the one we
        require. Adds a GOcean specific setBounds method which tells the loop
        what to iterate over. Need to harmonise with the topology_name method
        in the Dynamo api. '''

    _bounds_lookup = {}

    def __init__(self, parent=None,
                 topology_name="", loop_type=""):
        '''Constructs a GOLoop instance.
        :param parent: Optional parent node (default None).
        :type parent: :py:class:`psyclone.psyGen.node`
        :param str topology_name: Optional opology of the loop (unused atm).
        :param str loop_type: Loop type - must be 'inner' or 'outer'.'''
        # pylint: disable=unused-argument
github stfc / PSyclone / src / psyclone / psyir / transformations / extract_trans.py View on Github external
After applying the transformation the Nodes marked for extraction are
    children of the ExtractNode.
    Nodes to extract can be individual constructs within an Invoke (e.g.
    Loops containing a Kernel or BuiltIn call) or entire Invokes. This
    functionality does not support distributed memory.

    :param node_class: The Node class of which an instance will be inserted \
        into the tree (defaults to ExtractNode), but can be any derived class.
    :type node_class: :py:class:`psyclone.psyir.nodes.ExtractNode` or \
        derived class

    '''
    from psyclone.psyir import nodes
    from psyclone import psyGen
    # The types of node that this transformation can enclose
    valid_node_types = (nodes.Loop, psyGen.Kern, psyGen.BuiltIn,
                        psyGen.Directive, nodes.Literal, nodes.Reference)

    def __init__(self, node_class=ExtractNode):
        super(ExtractTrans, self).__init__(node_class=node_class)

    def __str__(self):
        return ("Create a sub-tree of the PSyIR that has ExtractNode "
                "at its root.")

    @property
    def name(self):
        ''' Returns the name of this transformation as a string.'''
        return "ExtractTrans"

    def validate(self, node_list, options=None):
        '''Performs validation checks specific to extract-based
github stfc / PSyclone / src / psyclone / psyir / frontend / fparser2.py View on Github external
integer_type = default_integer_type()
        # Now create a loop nest of depth `rank`
        new_parent = parent
        for idx in range(rank, 0, -1):

            # TODO #630 this creation of a new symbol really needs to account
            # for all of the symbols in the hierarchy of symbol tables. Since
            # we don't yet have that functionality we just add everything to
            # the top-level symbol table.
            loop_vars[idx-1] = symbol_table.new_symbol_name(
                "widx{0}".format(idx))

            symbol_table.add(DataSymbol(loop_vars[idx-1], integer_type))

            loop = Loop(parent=new_parent, variable_name=loop_vars[idx-1],
                        annotations=annotations)
            # Point to the original WHERE statement in the parse tree.
            loop.ast = node
            # Add loop lower bound
            loop.addchild(Literal("1", integer_type, parent=loop))
            # Add loop upper bound - we use the SIZE operator to query the
            # extent of the current array dimension
            size_node = BinaryOperation(BinaryOperation.Operator.SIZE,
                                        parent=loop)
            loop.addchild(size_node)
            symbol = size_node.find_or_create_symbol(arrays[0].name)

            size_node.addchild(Reference(symbol, parent=size_node))
            size_node.addchild(Literal(str(idx), integer_type,
                                       parent=size_node))
            # Add loop increment
github stfc / PSyclone / examples / gocean / eg4 / acc_transform.py View on Github external
from psyclone.transformations import ACCParallelTrans, \
        ACCEnterDataTrans, ACCLoopTrans, ACCRoutineTrans
    ptrans = ACCParallelTrans()
    ltrans = ACCLoopTrans()
    dtrans = ACCEnterDataTrans()
    ktrans = ACCRoutineTrans()

    invoke = psy.invokes.get('invoke_0')
    schedule = invoke.schedule
    schedule.view()

    # Apply the OpenACC Loop transformation to *every* loop
    # nest in the schedule
    from psyclone.psyir.nodes import Loop
    for child in schedule.children:
        if isinstance(child, Loop):
            newschedule, _ = ltrans.apply(child, {"collapse": 2})
            schedule = newschedule

    # Put all of the loops in a single parallel region
    newschedule, _ = ptrans.apply(schedule.children)

    # Add an enter-data directive
    newschedule, _ = dtrans.apply(schedule)

    # Put an 'acc routine' directive inside each kernel
    for kern in schedule.coded_kernels():
        _, _ = ktrans.apply(kern)
        # Ideally we would module-inline the kernel here (to save having to
        # rely on the compiler to do it) but this does not currently work
        # for the fparser2 AST (issue #229).
        # _, _ = itrans.apply(kern)