def test_name():
'''Check that the name property of the NemoAllArrayRange2LoopTrans
class returns the expected value.
'''
assert (
NemoAllArrayRange2LoopTrans().name == "NemoAllArrayRange2LoopTrans")
def test_apply_multi_assignments():
'''Check that the apply() method can be called for multiple assigments
with and without ranges.
'''
_, invoke_info = get_invoke("array_syntax.f90", api=API, idx=0)
schedule = invoke_info.schedule
trans = NemoAllArrayRange2LoopTrans()
for assignment in schedule.walk(Assignment):
trans.apply(assignment)
writer = FortranWriter()
result = writer(schedule)
expected = ("do jk = 1, jpk, 1\n"
" do jj = 1, jpj, 1\n"
" do ji = 1, jpi, 1\n"
" zftv(ji,jj,jk) = 0.0e0\n"
" enddo\n"
" enddo\n"
"enddo\n"
"if (l_ptr) then\n"
" call dia_ptr_hst(jn, ''ldf'', -zftv(:,:,:))\n"
"end if\n"
"call dia_ptr_hst(jn, ''ldf'', -zftv(:,:,:))\n"
"do jj = 1, jpj, 1\n"
" do ji = 1, jpi, 1\n"
" zftu(ji,jj,1) = 1.0e0\n"
" enddo\n"
"enddo\n"
"do jj = 1, jpj, 1\n"
" do ji = 1, jpi, 1\n"
" tmask(ji,jj) = jpi\n"
" enddo\n"
"enddo\n")
assert expected in result
def test_str():
'''Test that the str of an instance of the NemoAllArrayRange2LoopTrans
class returns the expected value.
'''
assert (str(NemoAllArrayRange2LoopTrans()) == "Convert all array ranges "
"in a PSyIR assignment into PSyIR NemoLoops.")
def test_apply_calls_validate():
'''Check that the apply() method calls the validate method.'''
trans = NemoAllArrayRange2LoopTrans()
with pytest.raises(TransformationError) as info:
trans.apply(None)
assert ("Error in NemoAllArrayRange2LoopTrans transformation. The "
"supplied node argument should be a PSyIR Assignment, but "
"found 'NoneType'." in str(info.value))
def trans(psy):
'''Transformation routine for use with PSyclone. Applies the PSyIR2SIR
transform to the supplied invokes after replacing any ABS, SIGN or
MIN intrinsics with equivalent code. This is done because the SIR
does not support intrinsics.
:param psy: the PSy object which this script will transform.
:type psy: :py:class:`psyclone.psyGen.PSy`
:returns: the transformed PSy object.
:rtype: :py:class:`psyclone.psyGen.PSy`
'''
abs_trans = Abs2CodeTrans()
sign_trans = Sign2CodeTrans()
min_trans = Min2CodeTrans()
nemo_loop_trans = NemoAllArrayRange2LoopTrans()
sir_writer = SIRWriter()
fortran_writer = FortranWriter()
# For each Invoke write out the SIR representation of the
# schedule. Note, there is no algorithm layer in the NEMO API so
# the invokes represent all of the original code.
for invoke in psy.invokes.invoke_list:
schedule = invoke.schedule
for assignment in schedule.walk(Assignment):
nemo_loop_trans.apply(assignment)
for kernel in schedule.walk(NemoKern):
# The NEMO api currently has no symbol table so create one
# to allow the generation of new variables. Note, this
# does not guarantee unique names as we don't know any of
# the existing names (so generated names could clash).
symbol_table = SymbolTable()
kernel_schedule = kernel.get_kernel_schedule()
for oper in kernel_schedule.walk(Operation):
if oper.operator == UnaryOperation.Operator.ABS:
# Apply ABS transformation
abs_trans.apply(oper, symbol_table)
elif oper.operator == BinaryOperation.Operator.SIGN:
# Apply SIGN transformation
sign_trans.apply(oper, symbol_table)
elif oper.operator in [
BinaryOperation.Operator.MIN,
NaryOperation.Operator.MIN
]:
# Apply (2-n arg) MIN transformation
min_trans.apply(oper, symbol_table)
kern = fortran_writer(schedule)
print(kern)
kern = sir_writer(schedule)
print(kern)
return psy
def test_validate_assignment():
'''Check that the validate method raises an exception if the supplied
argument is not an Assignment node.
'''
trans = NemoAllArrayRange2LoopTrans()
with pytest.raises(TransformationError) as info:
trans.validate(None)
assert ("Error in NemoAllArrayRange2LoopTrans transformation. The "
"supplied node argument should be a PSyIR Assignment, but "
"found 'NoneType'." in str(info.value))
def test_transform_apply_mixed_implicit_do():
'''Check that the PSyIR is transformed as expected for a lat,lon,levs
loop with some of its indices accessed using array notation and
some using explicit loops. The resultant Fortran code is used to
confirm the transformation has worked correctly.
'''
_, invoke_info = get_invoke("explicit_over_implicit.f90", api=API, idx=0)
schedule = invoke_info.schedule
assignment = schedule[0].loop_body[0]
trans = NemoAllArrayRange2LoopTrans()
trans.apply(assignment)
writer = FortranWriter()
result = writer(schedule)
expected = ("do jk = 1, jpk, 1\n"
" do jj = 1, jpj, 1\n"
" do ji = 1, jpi, 1\n"
" umask(ji,jj,jk) = vmask(ji,jj,jk) + 1.0\n"
" enddo\n"
" enddo\n"
"enddo")
assert expected in result
'''
from __future__ import print_function
from psyclone.domain.nemo.transformations import NemoAllArrayRange2LoopTrans
from psyclone.errors import InternalError
from psyclone.psyir.nodes import Assignment, CodeBlock, Call, Literal, Loop, \
ACCLoopDirective
from psyclone.transformations import ACCLoopTrans, TransformationError, \
ACCKernelsTrans
COLLAPSE_LOOPS = False
EXPAND_IMPLICIT_LOOPS = False
# Get the PSyclone transformations we will use
ARRAY_RANGE_TRANS = NemoAllArrayRange2LoopTrans()
def trans(psy):
'''A PSyclone-script compliant transformation function. Applies
OpenACC 'kernels' and 'data' directives to NEMO code.
:param psy: The PSy layer object to apply transformations to.
:type psy: :py:class:`psyclone.psyGen.PSy`
'''
print("Invokes found:\n{0}\n".format("\n".join(
[str(name) for name in psy.invokes.names])))
for invoke in psy.invokes.invoke_list:
def test_transform():
'''Check that it is possible to create an instance of
NemoArrayRange2LoopTrans and that it is a Transformation.
'''
assert isinstance(NemoAllArrayRange2LoopTrans(), Transformation)
def make_sir_compliant(schedule):
'''
Applies various transformations to the supplied schedule to replace any
features that cannot be represented in SIR with alternative forms:
1. Converts any accesses of individual array elements into 1-trip loops.
2. Transforms array assignments into loops.
3. Replaces any ABS, SIGN, MIN or MAX intrinsics with equivalent PSyIR.
4. Hoists any loop-invariant assignments out of loops over levels.
:param schedule: the schedule to transform.
:type schedule: :py:class:`psyclone.psyir.nodes.Schedule`
'''
abs_trans = Abs2CodeTrans()
sign_trans = Sign2CodeTrans()
min_trans = Min2CodeTrans()
max_trans = Max2CodeTrans()
array_range_trans = NemoAllArrayRange2LoopTrans()
array_access_trans = NemoAllArrayAccess2LoopTrans()
hoist_trans = HoistTrans()
# Transform any single index accesses in array assignments
# (e.g. a(1)) into 1-trip loops.
for assignment in schedule.walk(Assignment):
array_access_trans.apply(assignment)
# Transform any array assignments (Fortran ':' notation) into loops.
for assignment in schedule.walk(Assignment):
array_range_trans.apply(assignment)
for kernel in schedule.walk(NemoKern):
kernel_schedule = kernel.get_kernel_schedule()
for oper in kernel_schedule.walk(Operation):
if oper.operator == UnaryOperation.Operator.ABS:
# Apply ABS transformation
abs_trans.apply(oper)
elif oper.operator == BinaryOperation.Operator.SIGN:
# Apply SIGN transformation
sign_trans.apply(oper)
elif oper.operator in [BinaryOperation.Operator.MIN,
NaryOperation.Operator.MIN]:
# Apply (2-n arg) MIN transformation
min_trans.apply(oper)
elif oper.operator in [BinaryOperation.Operator.MAX,
NaryOperation.Operator.MAX]:
# Apply (2-n arg) MAX transformation
max_trans.apply(oper)
# Remove any loop invariant assignments inside k-loops to make
# them perfectly nested. At the moment this transformation
# does not perform any dependence analysis validation so could
# move code that should not be moved, see issue
# #1387. However, it is known that it is safe do apply this
# transformation to this particular code
# (tra_adv_compute.F90).
for loop in schedule.loops():
# outermost only
if loop.loop_type == "levels":
for child in loop.loop_body[:]:
if isinstance(child, Assignment):
hoist_trans.apply(child)
