Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
''' 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
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
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"
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)
def loops(self):
'''Return all loops currently in this schedule.'''
from psyclone.psyir.nodes import Loop
return self.walk(Loop)
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)
@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
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
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
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)