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)