def gen_code(self, parent):
''' Generate the Fortran source for this loop '''
# Our schedule holds the names to use for the loop bounds.
# Climb up the tree looking for our enclosing Schedule
schedule = self.ancestor(GOSchedule)
if schedule is None or not isinstance(schedule, GOSchedule):
raise GenerationError("Internal error: cannot find parent"
" GOSchedule for this Do loop")
# Walk down the tree looking for a kernel so that we can
# look-up what index-offset convention we are to use
go_kernels = self.walk(self.children, GOKern)
if len(go_kernels) == 0:
raise GenerationError("Internal error: cannot find the "
"GOcean Kernel enclosed by this loop")
index_offset = go_kernels[0].index_offset
if schedule.const_loop_bounds and \
index_offset not in SUPPORTED_OFFSETS:
raise GenerationError("Constant bounds generation"
" not implemented for a grid offset "
"of {0}. Supported offsets are {1}".format(
index_offset, SUPPORTED_OFFSETS))
# Check that all kernels enclosed by this loop expect the same
# grid offset
for kernel in go_kernels:
if kernel.index_offset != index_offset:
raise GenerationError("All Kernels must expect the same "
"grid offset but kernel {0} has offset "
"{1} which does not match {2}".format(
kernel.name, kernel.index_offset,
index_offset))
# Generate the upper and lower loop bounds
self._start = self._lower_bound()
self._stop = self._upper_bound()
Loop.gen_code(self, parent)
def gen_code(self, parent):
if self.field_space == "every":
from psyclone.f2pygen import DeclGen
dim_var = DeclGen(parent, datatype="INTEGER",
entity_decls=[self._variable_name])
parent.add(dim_var)
# loop bounds
self._start = "1"
if self._loop_type == "inner":
index = "1"
elif self._loop_type == "outer":
index = "2"
self._stop = ("SIZE(" + self.field_name + ", " +
index + ")")
else: # one of our spaces so use values provided by the infrastructure
# loop bounds
if self._loop_type == "inner":
self._start = self.field_space + "%istart"
self._stop = self.field_space + "%istop"
elif self._loop_type == "outer":
self._start = self.field_space + "%jstart"
self._stop = self.field_space + "%jstop"
Loop.gen_code(self, parent)
def __init__(self, parent=None,
topology_name="", loop_type=""):
Loop.__init__(self, parent=parent,
valid_loop_types=["inner", "outer"])
self.loop_type = loop_type
if self._loop_type == "inner":
self._variable_name = "i"
elif self._loop_type == "outer":
self._variable_name = "j"
def gen_code(self, parent):
''' Work out the appropriate loop bounds and then call the base
class to generate the code '''
self._start = "1"
if self._loop_type == "colours":
self._stop = "ncolour"
elif self._loop_type == "colour":
self._stop = "ncp_ncolour(colour)"
else:
self._stop = self.field_name + "%get_ncell()"
Loop.gen_code(self, parent)
def gen_code(self, parent):
''' Work out the appropriate loop bounds and then call the base
class to generate the code '''
self.start_expr = Literal("1", parent=self)
if self._loop_type == "colours":
self.stop_expr = Reference("ncolour", parent=self)
elif self._loop_type == "colour":
self.stop_expr = ArrayReference("ncp_ncolour", parent=self)
self.stop_expr.addchild(Reference("colour"), parent=self.stop_expr)
else:
self.stop_expr = Reference(self.field_name + "%get_ncell()",
parent=self)
Loop.gen_code(self, parent)
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":
self._variable_name = "colour"
elif self._loop_type == "colour":
self._variable_name = "cell"
else:
self._variable_name = "cell"
def __init__(self, parent=None,
topology_name="", loop_type=""):
Loop.__init__(self, parent=parent,
valid_loop_types=["inner", "outer"])
self.loop_type = loop_type
if self._loop_type == "inner":
self._variable_name = "i"
elif self._loop_type == "outer":
self._variable_name = "j"
# Pre-initialise the Loop children # TODO: See issue #440
self.addchild(Literal("NOT_INITIALISED", parent=self)) # start
self.addchild(Literal("NOT_INITIALISED", parent=self)) # stop
self.addchild(Literal("1", parent=self)) # step
self.addchild(Schedule(parent=self)) # loop body
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":
self._variable_name = "colour"
elif self._loop_type == "colour":
self._variable_name = "cell"
else:
self._variable_name = "cell"
# Pre-initialise the Loop children # TODO: See issue #440
self.addchild(Literal("NOT_INITIALISED", parent=self)) # start
self.addchild(Literal("NOT_INITIALISED", parent=self)) # stop
self.addchild(Literal("1", parent=self)) # step
self.addchild(Schedule(parent=self)) # loop body
def __init__(self, ast, parent=None):
from fparser.two.Fortran2003 import Loop_Control
Loop.__init__(self, parent=parent, valid_loop_types=VALID_LOOP_TYPES)
NemoFparser2ASTProcessor.__init__(self)
# Keep a ptr to the corresponding node in the AST
self._ast = ast
# Get the loop variable
ctrl = walk_ast(ast.content, [Loop_Control])
# Second element of items member of Loop Control is itself a tuple
# containing:
# Loop variable, [start value expression, end value expression, step
# expression]
# Loop variable will be an instance of Fortran2003.Name
loop_var = str(ctrl[0].items[1][0])
self._variable_name = str(loop_var)
# Identify the type of loop
if self._variable_name in NEMO_LOOP_TYPE_MAPPING:
self.loop_type = NEMO_LOOP_TYPE_MAPPING[self._variable_name]
else:
self.loop_type = "unknown"
# Get the loop limits. These are given in a list which is the second
# element of a tuple which is itself the second element of the items
# tuple:
# (None, (Name('jk'), [Int_Literal_Constant('1', None), Name('jpk'),
# Int_Literal_Constant('1', None)]), None)
limits_list = ctrl[0].items[1][1]
self._start = str(limits_list[0])
self._stop = str(limits_list[1])
if len(limits_list) == 3:
self._step = str(limits_list[2])
else:
# Default loop increment is 1
self._step = "1"
# Is this loop body a kernel?
if NemoKern.match(self._ast):
self.addchild(NemoKern(self._ast, parent=self))
return
# It's not - walk on down the AST...
self.process_nodes(self, self._ast.content, self._ast)
def gen_code(self, parent):
if self.field_space == "every":
from psyclone.f2pygen import DeclGen
from psyclone.psyGen import BinaryOperation, Reference
dim_var = DeclGen(parent, datatype="INTEGER",
entity_decls=[self._variable_name])
parent.add(dim_var)
# Update start loop bound
self.start_expr = Literal("1", parent=self)
# Update stop loop bound
if self._loop_type == "inner":
index = "1"
elif self._loop_type == "outer":
index = "2"
self.stop_expr = BinaryOperation(BinaryOperation.Operator.SIZE,
parent=self)
self.stop_expr.addchild(Reference(self.field_name,
parent=self.stop_expr))
self.stop_expr.addchild(Literal(index, parent=self.stop_expr))
else: # one of our spaces so use values provided by the infrastructure
# loop bounds
# TODO: Issue 440. Implement derive types in PSyIR
if self._loop_type == "inner":
self.start_expr = Reference(
self.field_space + "%istart", parent=self)
self.stop_expr = Reference(
self.field_space + "%istop", parent=self)
elif self._loop_type == "outer":
self.start_expr = Reference(
self.field_space + "%jstart", parent=self)
self.stop_expr = Reference(
self.field_space + "%jstop", parent=self)
Loop.gen_code(self, parent)
def gen_code(self, parent):
if self.field_space == "every":
from psyclone.f2pygen import DeclGen
dim_var = DeclGen(parent, datatype="INTEGER",
entity_decls=[self._variable_name])
parent.add(dim_var)
# Update start loop bound
self.start_expr = Literal("1", parent=self)
# Update stop loop bound
if self._loop_type == "inner":
index = "1"
elif self._loop_type == "outer":
index = "2"
# TODO: Issue 440. Implement SIZE intrinsic in PSyIR
self.stop_expr = Literal("SIZE(" + self.field_name + "," +
index + ")", parent=self)
else: # one of our spaces so use values provided by the infrastructure
# loop bounds
# TODO: Issue 440. Implement derive types in PSyIR
if self._loop_type == "inner":
self.start_expr = Reference(
self.field_space + "%istart", parent=self)
self.stop_expr = Reference(
self.field_space + "%istop", parent=self)
elif self._loop_type == "outer":
self.start_expr = Reference(
self.field_space + "%jstart", parent=self)
self.stop_expr = Reference(
self.field_space + "%jstop", parent=self)
Loop.gen_code(self, parent)
def test_profile_basic(capsys):
'''Check basic functionality: node names, schedule view.
'''
from psyclone.psyGen import colored, SCHEDULE_COLOUR_MAP
Profiler.set_options([Profiler.INVOKES])
_, invoke = get_invoke("test11_different_iterates_over_one_invoke.f90",
"gocean1.0",
idx=0)
assert isinstance(invoke.schedule.children[0], ProfileNode)
invoke.schedule.view()
out, _ = capsys.readouterr()
gsched = colored("GOInvokeSchedule", SCHEDULE_COLOUR_MAP["Schedule"])
loop = Loop().coloured_text
profile = invoke.schedule.children[0].coloured_text
# Do one test based on schedule view, to make sure colouring
# and indentation is correct
expected = (gsched + "[invoke='invoke_0', Constant loop bounds=True]\n"
" " + profile + "\n"
" " + loop + "[type='outer', field_space='go_cv', "
"it_space='go_internal_pts']\n")
assert expected in out
prt = ProfileRegionTrans()
# Insert a profile call between outer and inner loop.
# This tests that we find the subroutine node even
# if it is not the immediate parent.
new_sched, _ = prt.apply(
invoke.schedule.children[0].children[0].children[0])
new_sched_str = str(new_sched)
correct = ("""GOInvokeSchedule(Constant loop bounds=True):
ProfileStart[var=profile]
GOLoop[id:'', variable:'j', loop_type:'outer']
Literal[value:'2']
Literal[value:'jstop-1']
Literal[value:'1']
Schedule:
GOLoop[id:'', variable:'i', loop_type:'inner']
Literal[value:'2']
Literal[value:'istop']
Literal[value:'1']
Schedule:
kern call: compute_cv_code
End Schedule
End GOLoop
End Schedule
End GOLoop
GOLoop[id:'', variable:'j', loop_type:'outer']
Literal[value:'1']
Literal[value:'jstop+1']
Literal[value:'1']
Schedule:
GOLoop[id:'', variable:'i', loop_type:'inner']
Literal[value:'1']
Literal[value:'istop+1']
Literal[value:'1']
Schedule:
kern call: bc_ssh_code
End Schedule
End GOLoop
End Schedule
End GOLoop
ProfileEnd
End Schedule""")
assert correct in new_sched_str
Profiler.set_options(None)
def test_psy_data_trans_basic(capsys):
'''Check basic functionality: node names, schedule view.
'''
_, invoke = get_invoke("test11_different_iterates_over_one_invoke.f90",
"gocean1.0",
idx=0,
dist_mem=False)
schedule = invoke.schedule
# This test expects constant loop bounds
schedule._const_loop_bounds = True
data_trans = PSyDataTrans()
assert "Create a sub-tree of the PSyIR that has a node of type " \
"PSyDataNode at its root" in str(data_trans)
assert data_trans.name == "PSyDataTrans"
data_trans.apply(schedule)
assert isinstance(invoke.schedule[0], PSyDataNode)
schedule.view()
out, _ = capsys.readouterr()
gsched = colored("GOInvokeSchedule", GOInvokeSchedule._colour)
sched = colored("Schedule", Schedule._colour)
loop = Loop().coloured_name(True)
data = invoke.schedule[0].coloured_name(True)
# Do one test based on schedule view, to make sure colouring
# and indentation is correct
expected = (gsched + "[invoke='invoke_0', Constant loop bounds=True]\n"
" 0: " + data + "[]\n"
" " + sched + "[]\n"
" 0: " + loop +
"[type='outer', field_space='go_cv', "
"it_space='go_internal_pts']\n")
assert expected in out
# Insert a DataTrans call between outer and inner loop.
# This tests that we find the subroutine node even
# if it is not the immediate parent.
data_trans.apply(invoke.schedule[0].psy_data_body[0].loop_body[0])
new_sched_str = str(invoke.schedule)
correct = ("""GOInvokeSchedule[invoke='invoke_0', \
Constant loop bounds=True]:
PSyDataStart[var=psy_data]
GOLoop[id:'', variable:'j', loop_type:'outer']
Literal[value:'2', Scalar<INTEGER, UNDEFINED>]
Literal[value:'jstop-1', Scalar<INTEGER, UNDEFINED>]
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Schedule:
PSyDataStart[var=psy_data_1]
GOLoop[id:'', variable:'i', loop_type:'inner']
Literal[value:'2', Scalar<INTEGER, UNDEFINED>]
Literal[value:'istop', Scalar<INTEGER, UNDEFINED>]
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Schedule:
kern call: compute_cv_code
End Schedule
End GOLoop
PSyDataEnd[var=psy_data_1]
End Schedule
End GOLoop
GOLoop[id:'', variable:'j', loop_type:'outer']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Literal[value:'jstop+1', Scalar<INTEGER, UNDEFINED>]
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Schedule:
GOLoop[id:'', variable:'i', loop_type:'inner']
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Literal[value:'istop+1', Scalar<INTEGER, UNDEFINED>]
Literal[value:'1', Scalar<INTEGER, UNDEFINED>]
Schedule:
kern call: bc_ssh_code
End Schedule
End GOLoop
End Schedule
End GOLoop
PSyDataEnd[var=psy_data]
End Schedule""")
assert correct in new_sched_str
def test_profile_basic(capsys):
'''Check basic functionality: node names, schedule view.
'''
Profiler.set_options([Profiler.INVOKES])
_, invoke = get_invoke("test11_different_iterates_over_one_invoke.f90",
"gocean1.0", idx=0)
assert isinstance(invoke.schedule.children[0], ProfileNode)
invoke.schedule.view()
out, _ = capsys.readouterr()
coloured_schedule = GOSchedule([]).coloured_text
coloured_loop = Loop().coloured_text
coloured_kern = GOKern().coloured_text
coloured_profile = invoke.schedule.children[0].coloured_text
# Do one test based on schedule view, to make sure colouring
# and indentation is correct
correct = (
'''{0}[invoke='invoke_0',Constant loop bounds=True]
{3}
{1}[type='outer',field_space='go_cv',it_space='go_internal_pts']
{1}[type='inner',field_space='go_cv',it_space='go_internal_pts']
{2} compute_cv_code(cv_fld,p_fld,v_fld) '''
'''[module_inline=False]
{1}[type='outer',field_space='go_ct',it_space='go_all_pts']
{1}[type='inner',field_space='go_ct',it_space='go_all_pts']
{2} bc_ssh_code(ncycle,p_fld,tmask) '''
'''[module_inline=False]'''.format(coloured_schedule, coloured_loop,
coloured_kern, coloured_profile)
)
assert correct in out
prt = ProfileRegionTrans()
# Insert a profile call between outer and inner loop.
# This tests that we find the subroutine node even
# if it is not the immediate parent.
new_sched, _ = prt.apply(invoke.schedule.children[0]
.children[0].children[0])
new_sched_str = str(new_sched)
correct = ("""GOSchedule(Constant loop bounds=True):
ProfileStart[var=profile]
Loop[]: j= lower=2,jstop-1,1
ProfileStart[var=profile_1]
Loop[]: i= lower=2,istop,1
kern call: compute_cv_code
EndLoop
ProfileEnd
EndLoop
Loop[]: j= lower=1,jstop+1,1
Loop[]: i= lower=1,istop+1,1
kern call: bc_ssh_code
EndLoop
EndLoop
ProfileEnd
End Schedule""")
assert correct in new_sched_str
Profiler.set_options(None)
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 __init__(self, parent=None, topology_name="", loop_type=""):
Loop.__init__(self, parent=parent, valid_loop_types=VALID_LOOP_TYPES)
self.loop_type = loop_type
# We set the loop variable name in the constructor so that it is
# available when we're determining which vars should be OpenMP
# PRIVATE (which is done *before* code generation is performed)
if self.loop_type == "inner":
self._variable_name = "i"
elif self.loop_type == "outer":
self._variable_name = "j"
else:
raise GenerationError(
"Invalid loop type of '{0}'. Expected one of {1}".format(
self._loop_type, VALID_LOOP_TYPES))
# Create a dictionary to simplify the business of looking-up
# loop bounds
self._bounds_lookup = {}
for grid_offset in SUPPORTED_OFFSETS:
self._bounds_lookup[grid_offset] = {}
for gridpt_type in VALID_FIELD_GRID_TYPES:
self._bounds_lookup[grid_offset][gridpt_type] = {}
for itspace in VALID_ITERATES_OVER:
self._bounds_lookup[grid_offset][gridpt_type][itspace] = {}
# Loop bounds for a mesh with NE offset
self._bounds_lookup['offset_ne']['ct']['all_pts'] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': "+1"}}
self._bounds_lookup['offset_ne']['ct']['internal_pts'] = \
{'inner': {'start': "2", 'stop': ""},
'outer': {'start': "2", 'stop': ""}}
self._bounds_lookup['offset_ne']['cu']['all_pts'] = \
{'inner': {'start': "1", 'stop': ""},
'outer': {'start': "1", 'stop': "+1"}}
self._bounds_lookup['offset_ne']['cu']['internal_pts'] = \
{'inner': {'start': "2", 'stop': "-1"},
'outer': {'start': "2", 'stop': ""}}
self._bounds_lookup['offset_ne']['cv']['all_pts'] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': ""}}
self._bounds_lookup['offset_ne']['cv']['internal_pts'] = \
{'inner': {'start': "2", 'stop': ""},
'outer': {'start': "2", 'stop': "-1"}}
self._bounds_lookup['offset_ne']['cf']['all_pts'] = \
{'inner': {'start': "1", 'stop': ""},
'outer': {'start': "1", 'stop': ""}}
self._bounds_lookup['offset_ne']['cf']['internal_pts'] = \
{'inner': {'start': "1", 'stop': "-1"},
'outer': {'start': "1", 'stop': "-1"}}
# Loop bounds for a mesh with SE offset
self._bounds_lookup['offset_sw']['ct']['all_pts'] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': "+1"}}
self._bounds_lookup['offset_sw']['ct']['internal_pts'] = \
{'inner': {'start': "2", 'stop': ""},
'outer': {'start': "2", 'stop': ""}}
self._bounds_lookup['offset_sw']['cu']['all_pts'] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': "+1"}}
self._bounds_lookup['offset_sw']['cu']['internal_pts'] = \
{'inner': {'start': "2", 'stop': "+1"},
'outer': {'start': "2", 'stop': ""}}
self._bounds_lookup['offset_sw']['cv']['all_pts'] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': "+1"}}
self._bounds_lookup['offset_sw']['cv']['internal_pts'] = \
{'inner': {'start': "2", 'stop': ""},
'outer': {'start': "2", 'stop': "+1"}}
self._bounds_lookup['offset_sw']['cf']['all_pts'] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': "+1"}}
self._bounds_lookup['offset_sw']['cf']['internal_pts'] = \
{'inner': {'start': "2", 'stop': "+1"},
'outer': {'start': "2", 'stop': "+1"}}
# For offset 'any'
for gridpt_type in VALID_FIELD_GRID_TYPES:
for itspace in VALID_ITERATES_OVER:
self._bounds_lookup['offset_any'][gridpt_type][itspace] = \
{'inner': {'start': "1", 'stop': ""},
'outer': {'start': "1", 'stop': ""}}
# For 'every' grid-point type
for offset in SUPPORTED_OFFSETS:
for itspace in VALID_ITERATES_OVER:
self._bounds_lookup[offset]['every'][itspace] = \
{'inner': {'start': "1", 'stop': "+1"},
'outer': {'start': "1", 'stop': "+1"}}
def test_fusetrans_error_not_same_parent():
''' Check that we reject attempts to fuse loops which don't share the
same parent '''
from psyclone.psyGen import Loop, Schedule, Literal
from psyclone.transformations import LoopFuseTrans, TransformationError
sch1 = Schedule()
sch2 = Schedule()
loop1 = Loop(variable_name="i", parent=sch1)
loop2 = Loop(variable_name="j", parent=sch2)
sch1.addchild(loop1)
sch2.addchild(loop2)
loop1.addchild(Literal("1", parent=loop1)) # start
loop1.addchild(Literal("10", parent=loop1)) # stop
loop1.addchild(Literal("1", parent=loop1)) # step
loop1.addchild(Schedule(parent=loop1)) # loop body
loop2.addchild(Literal("1", parent=loop2)) # start
loop2.addchild(Literal("10", parent=loop2)) # stop
loop2.addchild(Literal("1", parent=loop2)) # step
loop2.addchild(Schedule(parent=loop2)) # loop body
fuse = LoopFuseTrans()
# Try to fuse loops with different parents
with pytest.raises(TransformationError) as err:
fuse.validate(loop1, loop2)
assert "Error in LoopFuse transformation. Loops do not have the " \
"same parent" in str(err.value)
def test_fusetrans_error_incomplete():
''' Check that we reject attempts to fuse loops which are incomplete. '''
from psyclone.psyGen import Loop, Schedule, Literal, Return
from psyclone.transformations import LoopFuseTrans, TransformationError
sch = Schedule()
loop1 = Loop(variable_name="i", parent=sch)
loop2 = Loop(variable_name="j", parent=sch)
sch.addchild(loop1)
sch.addchild(loop2)
fuse = LoopFuseTrans()
# Check first loop
with pytest.raises(TransformationError) as err:
fuse.validate(loop1, loop2)
assert "Error in LoopFuse transformation. The first loop does not have " \
"4 children." in str(err.value)
loop1.addchild(Literal("start", parent=loop1))
loop1.addchild(Literal("stop", parent=loop1))
loop1.addchild(Literal("step", parent=loop1))
loop1.addchild(Schedule(parent=loop1))
loop1.loop_body.addchild(Return(parent=loop1.loop_body))
# Check second loop
with pytest.raises(TransformationError) as err:
fuse.validate(loop1, loop2)
assert "Error in LoopFuse transformation. The second loop does not have " \
"4 children." in str(err.value)
loop2.addchild(Literal("start", parent=loop2))
loop2.addchild(Literal("stop", parent=loop2))
loop2.addchild(Literal("step", parent=loop2))
loop2.addchild(Schedule(parent=loop2))
loop2.loop_body.addchild(Return(parent=loop2.loop_body))
# Validation should now pass
fuse.validate(loop1, loop2)
def __init__(self, ast, parent=None):
Loop.__init__(self, parent=parent, valid_loop_types=VALID_LOOP_TYPES)
# Keep a ptr to the corresponding node in the AST
self._ast = ast
def __init__(self, ast, parent=None):
# pylint: disable=super-init-not-called, non-parent-init-called
valid_loop_types = Config.get().api_conf("nemo").get_valid_loop_types()
Loop.__init__(self, parent=parent, valid_loop_types=valid_loop_types)
# Keep a ptr to the corresponding node in the AST
self._ast = ast