Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Lowering for the Matmul Intrinsic #931

Merged
merged 2 commits into from
Jul 29, 2021
Merged

Lowering for the Matmul Intrinsic #931

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
019b009
Lowering for the Matmul Intrinsic
kiranchandramohan Jul 20, 2021
dbe480d
Fixes to address TODOs and review comments
kiranchandramohan Jul 22, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions flang/include/flang/Lower/TransformationalRuntime.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@ void genCshiftVector(FirOpBuilder &builder, mlir::Location loc,
mlir::Value resultBox, mlir::Value arrayBox,
mlir::Value shiftBox);

void genMatmul(Fortran::lower::FirOpBuilder &builder, mlir::Location loc,
mlir::Value matrixABox, mlir::Value matrixBBox,
mlir::Value resultBox);

void genReshape(FirOpBuilder &builder, mlir::Location loc,
mlir::Value resultBox, mlir::Value sourceBox,
mlir::Value shapeBox, mlir::Value padBox, mlir::Value orderBox);
Expand Down
35 changes: 34 additions & 1 deletion flang/lib/Lower/IntrinsicCall.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,7 @@ struct IntrinsicLibrary {
mlir::Value genIshftc(mlir::Type, llvm::ArrayRef<mlir::Value>);
fir::ExtendedValue genLen(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
fir::ExtendedValue genLenTrim(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
fir::ExtendedValue genMatmul(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
fir::ExtendedValue genMaxloc(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
fir::ExtendedValue genMaxval(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
mlir::Value genMerge(mlir::Type, llvm::ArrayRef<mlir::Value>);
Expand Down Expand Up @@ -703,6 +704,10 @@ static constexpr IntrinsicHandler handlers[]{
{"lgt", &I::genCharacterCompare<mlir::CmpIPredicate::sgt>},
{"lle", &I::genCharacterCompare<mlir::CmpIPredicate::sle>},
{"llt", &I::genCharacterCompare<mlir::CmpIPredicate::slt>},
{"matmul",
&I::genMatmul,
{{{"matrix_a", asAddr}, {"matrix_b", asAddr}}},
/*isElemental=*/false},
{"max", &I::genExtremum<Extremum::Max, ExtremumBehavior::MinMaxss>},
{"maxloc",
&I::genMaxloc,
Expand Down Expand Up @@ -2427,6 +2432,34 @@ IntrinsicLibrary::genCharacterCompare(mlir::Type type,
fir::getBase(args[1]), fir::getLen(args[1]));
}

// MATMUL
fir::ExtendedValue
IntrinsicLibrary::genMatmul(mlir::Type resultType,
llvm::ArrayRef<fir::ExtendedValue> args) {
assert(args.size() == 2);

// Handle required matmul arguments
fir::BoxValue matrixTmpA = builder.createBox(loc, args[0]);
mlir::Value matrixA = fir::getBase(matrixTmpA);
fir::BoxValue matrixTmpB = builder.createBox(loc, args[1]);
mlir::Value matrixB = fir::getBase(matrixTmpB);
unsigned resultRank =
(matrixTmpA.rank() == 1 || matrixTmpB.rank() == 1) ? 1 : 2;

// Create mutable fir.box to be passed to the runtime for the result.
auto resultArrayType = builder.getVarLenSeqTy(resultType, resultRank);
auto resultMutableBox =
Fortran::lower::createTempMutableBox(builder, loc, resultArrayType);
auto resultIrBox =
Fortran::lower::getMutableIRBox(builder, loc, resultMutableBox);
// Call runtime. The runtime is allocating the result.
Fortran::lower::genMatmul(builder, loc, resultIrBox, matrixA, matrixB);
// Read result from mutable fir.box and add it to the list of temps to be
// finalized by the StatementContext.
return readAndAddCleanUp(resultMutableBox, resultType,
"unexpected result for MATMUL");
}

// MERGE
mlir::Value IntrinsicLibrary::genMerge(mlir::Type,
llvm::ArrayRef<mlir::Value> args) {
Expand Down Expand Up @@ -2653,7 +2686,7 @@ IntrinsicLibrary::genReshape(mlir::Type resultType,
// Handle shape argument
auto shape = builder.createBox(loc, args[1]);
fir::BoxValue shapeTmp = shape;
auto shapeRank = shapeTmp.rank();
[[maybe_unused]] auto shapeRank = shapeTmp.rank();
assert(shapeRank == 1);
auto shapeTy = shape.getType();
auto shapeArrTy = fir::dyn_cast_ptrOrBoxEleTy(shapeTy);
Expand Down
16 changes: 16 additions & 0 deletions flang/lib/Lower/TransformationalRuntime.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//

#include "flang/Lower/TransformationalRuntime.h"
#include "../../runtime/matmul.h"
#include "../../runtime/transformational.h"
#include "RTBuilder.h"
#include "flang/Lower/Bridge.h"
Expand Down Expand Up @@ -52,6 +53,21 @@ void Fortran::lower::genCshiftVector(Fortran::lower::FirOpBuilder &builder,
builder.create<fir::CallOp>(loc, cshiftFunc, args);
}

/// Generate call to Matmul intrinsic runtime routine.
void Fortran::lower::genMatmul(Fortran::lower::FirOpBuilder &builder,
mlir::Location loc, mlir::Value resultBox,
mlir::Value matrixABox, mlir::Value matrixBBox) {
auto func = Fortran::lower::getRuntimeFunc<mkRTKey(Matmul)>(loc, builder);
auto fTy = func.getType();
auto sourceFile = Fortran::lower::locationToFilename(builder, loc);
auto sourceLine =
Fortran::lower::locationToLineNo(builder, loc, fTy.getInput(4));
auto args =
Fortran::lower::createArguments(builder, loc, fTy, resultBox, matrixABox,
matrixBBox, sourceFile, sourceLine);
builder.create<fir::CallOp>(loc, func, args);
}

/// Generate call to Reshape intrinsic runtime routine.
void Fortran::lower::genReshape(Fortran::lower::FirOpBuilder &builder,
mlir::Location loc, mlir::Value resultBox,
Expand Down
65 changes: 65 additions & 0 deletions flang/test/Lower/intrinsic-procedures.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -906,6 +906,71 @@ subroutine lge_test
print*, llt(c1, c2)
end

! MATMUL
! CHECK-LABEL: matmul_test
! CHECK-SAME: (%[[X:.*]]: !fir.ref<!fir.array<3x1xf32>>, %[[Y:.*]]: !fir.ref<!fir.array<1x3xf32>>, %[[Z:.*]]: !fir.ref<!fir.array<2x2xf32>>)
! CHECK: %[[RESULT_BOX_ADDR:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?x?xf32>>> {uniq_name = ""}
! CHECK: %[[C3:.*]] = constant 3 : index
! CHECK: %[[C1:.*]] = constant 1 : index
! CHECK: %[[C1_0:.*]] = constant 1 : index
! CHECK: %[[C3_1:.*]] = constant 3 : index
! CHECK: %[[Z_BOX:.*]] = fir.array_load %[[Z]]({{.*}}) : (!fir.ref<!fir.array<2x2xf32>>, !fir.shape<2>) -> !fir.array<2x2xf32>
! CHECK: %[[X_SHAPE:.*]] = fir.shape %[[C3]], %[[C1]] : (index, index) -> !fir.shape<2>
! CHECK: %[[X_BOX:.*]] = fir.embox %[[X]](%[[X_SHAPE]]) : (!fir.ref<!fir.array<3x1xf32>>, !fir.shape<2>) -> !fir.box<!fir.array<3x1xf32>>
! CHECK: %[[Y_SHAPE:.*]] = fir.shape %[[C1_0]], %[[C3_1]] : (index, index) -> !fir.shape<2>
! CHECK: %[[Y_BOX:.*]] = fir.embox %[[Y]](%[[Y_SHAPE]]) : (!fir.ref<!fir.array<1x3xf32>>, !fir.shape<2>) -> !fir.box<!fir.array<1x3xf32>>
! CHECK: %[[ZERO_INIT:.*]] = fir.zero_bits !fir.heap<!fir.array<?x?xf32>>
! CHECK: %[[C0:.*]] = constant 0 : index
! CHECK: %[[RESULT_SHAPE:.*]] = fir.shape %[[C0]], %[[C0]] : (index, index) -> !fir.shape<2>
! CHECK: %[[RESULT_BOX_VAL:.*]] = fir.embox %[[ZERO_INIT]](%[[RESULT_SHAPE]]) : (!fir.heap<!fir.array<?x?xf32>>, !fir.shape<2>) -> !fir.box<!fir.heap<!fir.array<?x?xf32>>>
! CHECK: fir.store %[[RESULT_BOX_VAL]] to %[[RESULT_BOX_ADDR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>
! CHECK: %[[RESULT_BOX_ADDR_RUNTIME:.*]] = fir.convert %[[RESULT_BOX_ADDR]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>) -> !fir.ref<!fir.box<none>>
! CHECK: %[[X_BOX_RUNTIME:.*]] = fir.convert %[[X_BOX]] : (!fir.box<!fir.array<3x1xf32>>) -> !fir.box<none>
! CHECK: %[[Y_BOX_RUNTIME:.*]] = fir.convert %[[Y_BOX]] : (!fir.box<!fir.array<1x3xf32>>) -> !fir.box<none>
! CHECK: {{.*}}fir.call @_FortranAMatmul(%[[RESULT_BOX_ADDR_RUNTIME]], %[[X_BOX_RUNTIME]], %[[Y_BOX_RUNTIME]], {{.*}}, {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> none
! CHECK: %[[RESULT_BOX:.*]] = fir.load %[[RESULT_BOX_ADDR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>
! CHECK: %[[RESULT_TMP:.*]] = fir.box_addr %[[RESULT_BOX]] : (!fir.box<!fir.heap<!fir.array<?x?xf32>>>) -> !fir.heap<!fir.array<?x?xf32>>
! CHECK: %[[Z_COPY_FROM_RESULT:.*]] = fir.do_loop
! CHECK: {{.*}}fir.array_fetch
! CHECK: {{.*}}fir.array_update
! CHECK: fir.result
! CHECK: }
! CHECK: fir.array_merge_store %[[Z_BOX]], %[[Z_COPY_FROM_RESULT]] to %[[Z]] : !fir.array<2x2xf32>, !fir.array<2x2xf32>, !fir.ref<!fir.array<2x2xf32>>
! CHECK: fir.freemem %[[RESULT_TMP]] : !fir.heap<!fir.array<?x?xf32>>
subroutine matmul_test(x,y,z)
real :: x(3,1), y(1,3), z(2,2)
z = matmul(x,y)
end subroutine

! CHECK-LABEL: matmul_test2
! CHECK-SAME: (%[[X_BOX:.*]]: !fir.box<!fir.array<?x?x!fir.logical<4>>>, %[[Y_BOX:.*]]: !fir.box<!fir.array<?x!fir.logical<4>>>, %[[Z_BOX:.*]]: !fir.box<!fir.array<?x!fir.logical<4>>>)
!CHECK: %[[RESULT_BOX_ADDR:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?x!fir.logical<4>>>> {uniq_name = ""}
!CHECK: %[[Z:.*]] = fir.array_load %[[Z_BOX]] : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> !fir.array<?x!fir.logical<4>>
!CHECK: %[[ZERO_INIT:.*]] = fir.zero_bits !fir.heap<!fir.array<?x!fir.logical<4>>>
!CHECK: %[[C0:.*]] = constant 0 : index
!CHECK: %[[RESULT_SHAPE:.*]] = fir.shape %[[C0]] : (index) -> !fir.shape<1>
!CHECK: %[[RESULT_BOX:.*]] = fir.embox %[[ZERO_INIT]](%[[RESULT_SHAPE]]) : (!fir.heap<!fir.array<?x!fir.logical<4>>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<?x!fir.logical<4>>>>
!CHECK: fir.store %[[RESULT_BOX]] to %[[RESULT_BOX_ADDR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.logical<4>>>>>
!CHECK: %[[RESULT_BOX_RUNTIME:.*]] = fir.convert %[[RESULT_BOX_ADDR]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.logical<4>>>>>) -> !fir.ref<!fir.box<none>>
!CHECK: %[[X_BOX_RUNTIME:.*]] = fir.convert %[[X_BOX]] : (!fir.box<!fir.array<?x?x!fir.logical<4>>>) -> !fir.box<none>
!CHECK: %[[Y_BOX_RUNTIME:.*]] = fir.convert %[[Y_BOX]] : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> !fir.box<none>
!CHECK: {{.*}}fir.call @_FortranAMatmul(%[[RESULT_BOX_RUNTIME]], %[[X_BOX_RUNTIME]], %[[Y_BOX_RUNTIME]], {{.*}}, {{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.ref<i8>, i32) -> none
!CHECK: %[[RESULT_BOX:.*]] = fir.load %[[RESULT_BOX_ADDR]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?x!fir.logical<4>>>>>
!CHECK: %[[RESULT_TMP:.*]] = fir.box_addr %[[RESULT_BOX]] : (!fir.box<!fir.heap<!fir.array<?x!fir.logical<4>>>>) -> !fir.heap<!fir.array<?x!fir.logical<4>>>
!CHECK: %[[Z_COPY_FROM_RESULT:.*]] = fir.do_loop
!CHECK: {{.*}}fir.array_fetch
!CHECK: {{.*}}fir.array_update
!CHECK: fir.result
!CHECK: }
!CHECK: fir.array_merge_store %[[Z]], %[[Z_COPY_FROM_RESULT]] to %[[Z_BOX]] : !fir.array<?x!fir.logical<4>>, !fir.array<?x!fir.logical<4>>, !fir.box<!fir.array<?x!fir.logical<4>>>
!CHECK: fir.freemem %[[RESULT_TMP]] : !fir.heap<!fir.array<?x!fir.logical<4>>>
subroutine matmul_test2(X, Y, Z)
logical :: X(:,:)
logical :: Y(:)
logical :: Z(:)
Z = matmul(X, Y)
end subroutine

! MAXLOC
! CHECK-LABEL: maxloc_test
! CHECK-SAME: %[[arg0:.*]]: !fir.box<!fir.array<?xi32>>,
Expand Down