def trans(psy):
'''PSyclone transformation script for the Dynamo0.3 API to optimise
the matvec kernel for many-core CPUs. For the moment simply find
the first matvec kernel in the example, transform the matmul
intrinsic to equivalant inline code and then print out its PSyIR
representation and output it as Fortran using the PSyIR Fortran
back-end.
'''
matmul2code_trans = Matmul2CodeTrans()
fortran_writer = FortranWriter()
for invoke in psy.invokes.invoke_list:
schedule = invoke.schedule
for kernel in schedule.coded_kernels():
if kernel.name.lower() == "matrix_vector_kernel_code":
kernel_schedule = kernel.get_kernel_schedule()
# Replace matmul with inline code
for bin_op in kernel_schedule.walk(BinaryOperation):
if bin_op.operator is BinaryOperation.Operator.MATMUL:
matmul2code_trans.apply(bin_op)
# Future optimisations will go here.
kernel_schedule.view()
result = fortran_writer(kernel_schedule)
print(result)
# Abort after the first matrix vector kernel for the
# time being.
print("Aborting to view the modifications to the matrix "
"vector kernel")
sys.exit()
return psy
def test_apply2(tmpdir):
'''Test that the matmul2code apply method produces the expected
PSyIR. We use the Fortran backend to help provide the test for
correctness. This example includes extra indices for the vector
and matrix arrays with additional indices being literals.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
matmul.children[0].children[2] = Literal("1", INTEGER_TYPE)
matmul.children[1].children[1] = Literal("2", INTEGER_TYPE)
trans.apply(matmul)
writer = FortranWriter()
result = writer(matmul.root)
assert ("subroutine my_kern()\n"
" integer, parameter :: idx = 3\n"
" real, dimension(5,10,15) :: x\n"
" real, dimension(10,20) :: y\n"
" real, dimension(10) :: result\n"
" integer :: i\n"
" integer :: j\n"
"\n"
" do i = 1, 5, 1\n"
" result(i)=0.0\n"
" do j = 1, 10, 1\n"
" result(i)=result(i) + x(i,j,1) * y(j,2)\n"
" enddo\n"
" enddo\n"
"\n"
"end subroutine my_kern" in result)
assert Compile(tmpdir).string_compiles(result)
def trans(psy):
'''PSyclone transformation script for the Dynamo0.3 API to optimise
the matvec kernel for many-core CPUs. This is currently limited to
running on the scaled_matrix_vector_code kernel but should work
more generally. Any matmul calls are replaced with inline matric
vector code.
:param psy: a PSyclone PSy object which captures the algorithm and \
kernel information required by PSyclone.
:type psy: subclass of :py:class:`psyclone.psyGen.PSy`
'''
matmul2code_trans = Matmul2CodeTrans()
for invoke in psy.invokes.invoke_list:
schedule = invoke.schedule
for kernel in schedule.coded_kernels():
if kernel.name.lower() == "scaled_matrix_vector_code":
kernel.modified = True
kernel_schedule = kernel.get_kernel_schedule()
# Replace matmul with inline code
for bin_op in kernel_schedule.walk(BinaryOperation):
if bin_op.operator is BinaryOperation.Operator.MATMUL:
matmul2code_trans.apply(bin_op)
kernel_schedule.view()
return psy
def test_initialise():
'''Check that the str and name methods work as expected.
'''
trans = Matmul2CodeTrans()
assert (str(trans) == "Convert the PSyIR MATMUL intrinsic to equivalent "
"PSyIR code.")
assert trans.name == "Matmul2CodeTrans"
def test_validate14():
'''Check that the Matmul2Code validate method returns without any
exceptions when the supplied node is a MATMUL binary operation
that obeys the required rules and constraints.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
trans.validate(matmul)
def test_validate1():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is the wrong type.
'''
trans = Matmul2CodeTrans()
with pytest.raises(TransformationError) as excinfo:
trans.validate(None)
assert ("Error in Matmul2CodeTrans transformation. The supplied node "
"argument is not a MATMUL operator, found 'NoneType'."
in str(excinfo.value))
def test_validate2():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a binary operation but not a
MATMUL.
'''
trans = Matmul2CodeTrans()
with pytest.raises(TransformationError) as excinfo:
trans.validate(BinaryOperation.create(
BinaryOperation.Operator.ADD, Literal("1.0", REAL_TYPE),
Literal("1.0", REAL_TYPE)))
assert ("Transformation Error: Error in Matmul2CodeTrans transformation. "
"The supplied node operator is invalid, found 'Operator.ADD'."
in str(excinfo.value))
def test_validate12():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
the first dimension of its second (vector) argument is not a full
range.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
vector = matmul.children[1]
vector.children[0] = Literal("1", INTEGER_TYPE)
with pytest.raises(NotImplementedError) as excinfo:
trans.validate(matmul)
assert ("To use matmul2code_trans on matmul, index 0 of the 2nd (vector) "
"argument 'x' must be a full range." in str(excinfo.value))
def test_validate10():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
the first two dimensions of its first (matrix) argument are not
full ranges.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
matrix = matmul.children[0]
matrix.children[0] = Literal("1", INTEGER_TYPE)
with pytest.raises(NotImplementedError) as excinfo:
trans.validate(matmul)
assert ("To use matmul2code_trans on matmul, indices 0 and 1 of the "
"1st (matrix) argument 'x' must be full ranges."
in str(excinfo.value))
def test_validate7():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
its first (matrix) argument has fewer than 2 dimensions.
'''
trans = Matmul2CodeTrans()
array_type = ArrayType(REAL_TYPE, [10])
array = Reference(DataSymbol("x", array_type))
matmul = BinaryOperation.create(BinaryOperation.Operator.MATMUL, array,
array)
_ = Assignment.create(array, matmul)
with pytest.raises(TransformationError) as excinfo:
trans.validate(matmul)
assert ("Transformation Error: Expected 1st child of a MATMUL "
"BinaryOperation to be a matrix with at least 2 dimensions, "
"but found '1'." in str(excinfo.value))
def test_validate6():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
either or both of its arguments are references to datasymbols that
are not arrays.
'''
trans = Matmul2CodeTrans()
scalar = Reference(DataSymbol("x", REAL_TYPE))
matmul = BinaryOperation.create(BinaryOperation.Operator.MATMUL, scalar,
scalar)
_ = Assignment.create(scalar, matmul)
with pytest.raises(TransformationError) as excinfo:
trans.validate(matmul)
assert ("Transformation Error: Expected children of a MATMUL "
"BinaryOperation to be references to arrays, but found "
"'DataSymbol', 'DataSymbol'." in str(excinfo.value))
def test_validate13():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
the second (or higher) dimension of the second (vector) argument is
indexed via a range.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
vector = matmul.children[1]
my_range = vector.children[0]
vector.children[1] = my_range
with pytest.raises(NotImplementedError) as excinfo:
trans.validate(matmul)
assert ("To use matmul2code_trans on matmul, only the first index of the "
"2nd (vector) argument is permitted to be a Range but found "
"Range at index 1." in str(excinfo.value))
def test_validate11():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
the third (or higher) dimension of the first (matrix) argument is
indexed via a range.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
matrix = matmul.children[0]
my_range = matrix.children[0]
matrix.children[2] = my_range
with pytest.raises(NotImplementedError) as excinfo:
trans.validate(matmul)
assert ("To use matmul2code_trans on matmul, only the first two indices "
"of the 1st (matrix) argument are permitted to be Ranges but "
"found Range at index 2." in str(excinfo.value))
def test_apply3(tmpdir):
'''Test that the matmul2code apply method produces the expected
PSyIR. We use the Fortran backend to help provide the test for
correctness. This example includes the array and vector being
passed with no index information.
'''
trans = Matmul2CodeTrans()
matmul = create_matmul()
root = matmul.root
matrix = matmul.children[0]
lhs_vector = matrix.parent.parent.lhs
matrix_symbol = matrix.symbol
matmul.children[0] = Reference(matrix_symbol)
matrix_symbol.datatype._shape = [
Literal("10", INTEGER_TYPE),
Literal("20", INTEGER_TYPE)
]
rhs_vector = matmul.children[1]
rhs_vector_symbol = rhs_vector.symbol
rhs_vector_symbol.datatype._shape = [Literal("20", INTEGER_TYPE)]
matmul.children[1] = Reference(rhs_vector_symbol)
lhs_vector_symbol = lhs_vector.symbol
lhs_vector_symbol._shape = [Literal("10", INTEGER_TYPE)]
lhs_vector.replace_with(Reference(lhs_vector_symbol))
trans.apply(matmul)
writer = FortranWriter()
result = writer(root)
assert ("subroutine my_kern()\n"
" integer, parameter :: idx = 3\n"
" real, dimension(10,20) :: x\n"
" real, dimension(20) :: y\n"
" real, dimension(10) :: result\n"
" integer :: i\n"
" integer :: j\n"
"\n"
" do i = 1, 10, 1\n"
" result(i) = 0.0\n"
" do j = 1, 20, 1\n"
" result(i) = result(i) + x(i,j) * y(j)\n"
" enddo\n"
" enddo\n"
"\n"
"end subroutine my_kern" in result)
assert Compile(tmpdir).string_compiles(result)
def test_validate3():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
doesn't have an assignment as an ancestor.
'''
trans = Matmul2CodeTrans()
vector_type = ArrayType(REAL_TYPE, [10])
array_type = ArrayType(REAL_TYPE, [10, 10])
vector = Reference(DataSymbol("x", vector_type))
array = Reference(DataSymbol("y", array_type))
matmul = BinaryOperation.create(BinaryOperation.Operator.MATMUL, array,
vector)
with pytest.raises(TransformationError) as excinfo:
trans.validate(matmul)
assert ("Transformation Error: Error in Matmul2CodeTrans transformation. "
"This transformation requires the operator to be part of an "
"assignment statement, but no such assignment was found."
in str(excinfo.value))
def test_validate5():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
either or both arguments are not references.
'''
trans = Matmul2CodeTrans()
array_type = ArrayType(REAL_TYPE, [10])
array = Array.create(DataSymbol("x", array_type),
[Literal("10", INTEGER_TYPE)])
mult = BinaryOperation.create(BinaryOperation.Operator.MUL, array, array)
matmul = BinaryOperation.create(BinaryOperation.Operator.MATMUL, mult,
mult)
_ = Assignment.create(array, matmul)
with pytest.raises(TransformationError) as excinfo:
trans.validate(matmul)
assert ("Expected children of a MATMUL BinaryOperation to be references, "
"but found 'BinaryOperation', 'BinaryOperation'."
in str(excinfo.value))
def test_validate4():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
it is not the only operation on the RHS of an assignment.
'''
trans = Matmul2CodeTrans()
vector_type = ArrayType(REAL_TYPE, [10])
array_type = ArrayType(REAL_TYPE, [10, 10])
vector = Reference(DataSymbol("x", vector_type))
array = Reference(DataSymbol("y", array_type))
matmul = BinaryOperation.create(BinaryOperation.Operator.MATMUL, array,
vector)
rhs = BinaryOperation.create(BinaryOperation.Operator.MUL, matmul, vector)
_ = Assignment.create(array, rhs)
with pytest.raises(TransformationError) as excinfo:
trans.validate(matmul)
assert ("Transformation Error: Matmul2CodeTrans only supports the "
"transformation of a MATMUL operation when it is the sole "
"operation on the rhs of an assignment." in str(excinfo.value))
def test_validate9():
'''Check that the Matmul2Code validate method raises the expected
exception when the supplied node is a MATMUL binary operation but
its second (vector) argument is a reference to a vector with
greater than 1 dimension.
'''
trans = Matmul2CodeTrans()
array_type = ArrayType(REAL_TYPE, [10, 10])
array = Reference(DataSymbol("x", array_type))
vector_type = ArrayType(REAL_TYPE, [10, 10, 10])
vector = Reference(DataSymbol("y", vector_type))
matmul = BinaryOperation.create(BinaryOperation.Operator.MATMUL, array,
vector)
_ = Assignment.create(array, matmul)
with pytest.raises(TransformationError) as excinfo:
trans.validate(matmul)
assert ("Transformation Error: Expected 2nd child of a MATMUL "
"BinaryOperation to have 1 dimension, but found '3'."
in str(excinfo.value))
def test_apply4(tmpdir):
'''Test that the matmul2code apply method produces the expected
PSyIR. We use the Fortran backend to help provide the test for
correctness. This example make the lhs be the same array as the
second operand of the matmul (the vector in this case).
'''
trans = Matmul2CodeTrans()
one = Literal("1", INTEGER_TYPE)
five = Literal("5", INTEGER_TYPE)
matmul = create_matmul()
root = matmul.root
assignment = matmul.parent
vector = assignment.scope.symbol_table.lookup("y")
assignment.children[0] = ArrayReference.create(
vector, [Range.create(one, five, one.copy()),
one.copy()])
trans.apply(matmul)
writer = FortranWriter()
result = writer(root)
assert ("subroutine my_kern()\n"
" integer, parameter :: idx = 3\n"
" real, dimension(5,10,15) :: x\n"
" real, dimension(10,20) :: y\n"
" real, dimension(10) :: result\n"
" integer :: i\n"
" integer :: j\n"
"\n"
" do i = 1, 5, 1\n"
" y(i,1) = 0.0\n"
" do j = 1, 10, 1\n"
" y(i,1) = y(i,1) + x(i,j,idx) * y(j,idx)\n"
" enddo\n"
" enddo\n"
"\n"
"end subroutine my_kern" in result)
assert Compile(tmpdir).string_compiles(result)
