Skip to content

[flang][OpenMP] Common lowering flow for atomic update #69866

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

Closed
wants to merge 1 commit into from
Closed

[flang][OpenMP] Common lowering flow for atomic update #69866

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
188 changes: 86 additions & 102 deletions flang/lib/Lower/DirectivesCommon.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,76 +204,62 @@ static inline void genOmpAccAtomicUpdateStatement(
// Generate `omp.atomic.update` operation for atomic assignment statements
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
mlir::Location currentLocation = converter.getCurrentLocation();

const auto *varDesignator =
std::get_if<Fortran::common::Indirection<Fortran::parser::Designator>>(
&assignmentStmtVariable.u);
assert(varDesignator && "Variable designator for atomic update assignment "
"statement does not exist");
const Fortran::parser::Name *name =
Fortran::semantics::getDesignatorNameIfDataRef(varDesignator->value());
if (!name)
TODO(converter.getCurrentLocation(),
"Array references as atomic update variable");
assert(name && name->symbol &&
"No symbol attached to atomic update variable");
if (Fortran::semantics::IsAllocatableOrPointer(name->symbol->GetUltimate()))
converter.bindSymbol(*name->symbol, lhsAddr);

// Lowering is in two steps :
// subroutine sb
// integer :: a, b
// !$omp atomic update
// a = a + b
// end subroutine
//
// 1. Lower to scf.execute_region_op
//
// func.func @_QPsb() {
// %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFsbEa"}
// %1 = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFsbEb"}
// %2 = scf.execute_region -> i32 {
// %3 = fir.load %0 : !fir.ref<i32>
// %4 = fir.load %1 : !fir.ref<i32>
// %5 = arith.addi %3, %4 : i32
// scf.yield %5 : i32
// }
// return
// }
auto tempOp =
firOpBuilder.create<mlir::scf::ExecuteRegionOp>(currentLocation, varType);
firOpBuilder.createBlock(&tempOp.getRegion());
mlir::Block &block = tempOp.getRegion().back();
firOpBuilder.setInsertionPointToEnd(&block);
Fortran::lower::StatementContext stmtCtx;
mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
*Fortran::semantics::GetExpr(assignmentStmtExpr), stmtCtx));
mlir::Value convertResult =
firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
// Insert the terminator: YieldOp.
firOpBuilder.create<mlir::scf::YieldOp>(currentLocation, convertResult);
firOpBuilder.setInsertionPointToStart(&block);

// 2. Create the omp.atomic.update Operation using the Operations in the
// temporary scf.execute_region Operation.
//
// func.func @_QPsb() {
// %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFsbEa"}
// %1 = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFsbEb"}
// %2 = fir.load %1 : !fir.ref<i32>
// omp.atomic.update %0 : !fir.ref<i32> {
// ^bb0(%arg0: i32):
// %3 = fir.load %1 : !fir.ref<i32>
// %4 = arith.addi %arg0, %3 : i32
// omp.yield(%3 : i32)
// }
// return
// }
mlir::Value updateVar = converter.getSymbolAddress(*name->symbol);
if (auto decl = updateVar.getDefiningOp<hlfir::DeclareOp>())
updateVar = decl.getBase();

firOpBuilder.setInsertionPointAfter(tempOp);
mlir::Value convertRhs = nullptr;

auto lowerExpression = [&](const auto &intrinsicBinaryExpr) {
const auto &variableName{assignmentStmtVariable.GetSource().ToString()};
const auto &exprLeft{std::get<0>(intrinsicBinaryExpr.t)};
if (exprLeft.value().source.ToString() == variableName) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would the check mis-handle cases like the following?

arr(1) = arr(2) + arr(1)

// Update statement is of form `x = x op expr`
const auto &exprToLower{std::get<1>(intrinsicBinaryExpr.t)};
mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
*Fortran::semantics::GetExpr(exprToLower), stmtCtx));
convertRhs =
firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
} else {
// Update statement is of form `x = expr op x`
const auto &exprToLower{std::get<0>(intrinsicBinaryExpr.t)};
mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
*Fortran::semantics::GetExpr(exprToLower), stmtCtx));
convertRhs =
firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
}
};
Fortran::common::visit(
Fortran::common::visitors{
[&](const common::Indirection<parser::FunctionReference> &x) {
TODO(converter.getCurrentLocation(),
"Not yet implemented: intrinsic procedure in atomic update "
"expressions");
},
[&](const Fortran::parser::Expr::Add &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::Subtract &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::Multiply &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::Divide &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::AND &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::OR &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::EQV &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const Fortran::parser::Expr::NEQV &intrinsicBinaryExpr) {
lowerExpression(intrinsicBinaryExpr);
},
[&](const auto &) {},
},
assignmentStmtExpr.u);

mlir::Operation *atomicUpdateOp = nullptr;
if constexpr (std::is_same<AtomicListT,
Expand All @@ -289,10 +275,10 @@ static inline void genOmpAccAtomicUpdateStatement(
genOmpAtomicHintAndMemoryOrderClauses(converter, *rightHandClauseList,
hint, memoryOrder);
atomicUpdateOp = firOpBuilder.create<mlir::omp::AtomicUpdateOp>(
currentLocation, updateVar, hint, memoryOrder);
currentLocation, lhsAddr, hint, memoryOrder);
} else {
atomicUpdateOp = firOpBuilder.create<mlir::acc::AtomicUpdateOp>(
currentLocation, updateVar);
currentLocation, lhsAddr);
}

llvm::SmallVector<mlir::Type> varTys = {varType};
Expand All @@ -301,38 +287,36 @@ static inline void genOmpAccAtomicUpdateStatement(
mlir::Value val =
fir::getBase(atomicUpdateOp->getRegion(0).front().getArgument(0));

llvm::SmallVector<mlir::Operation *> ops;
for (mlir::Operation &op : tempOp.getRegion().getOps())
ops.push_back(&op);

// SCF Yield is converted to OMP Yield. All other operations are copied
for (mlir::Operation *op : ops) {
if (auto y = mlir::dyn_cast<mlir::scf::YieldOp>(op)) {
firOpBuilder.setInsertionPointToEnd(
&atomicUpdateOp->getRegion(0).front());
if constexpr (std::is_same<AtomicListT,
Fortran::parser::OmpAtomicClauseList>()) {
firOpBuilder.create<mlir::omp::YieldOp>(currentLocation,
y.getResults());
} else {
firOpBuilder.create<mlir::acc::YieldOp>(currentLocation,
y.getResults());
}
op->erase();
} else {
op->remove();
atomicUpdateOp->getRegion(0).front().push_back(op);
}
}

// Remove the load and replace all uses of load with the block argument
for (mlir::Operation &op : atomicUpdateOp->getRegion(0).getOps()) {
fir::LoadOp y = mlir::dyn_cast<fir::LoadOp>(&op);
if (y && y.getMemref() == updateVar)
y.getRes().replaceAllUsesWith(val);
mlir::Value op = nullptr;
if (std::get_if<Fortran::parser::Expr::Add>(&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::AddIOp>(currentLocation, val,
convertRhs);
} else if (std::get_if<Fortran::parser::Expr::Subtract>(
&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::SubIOp>(currentLocation, val,
convertRhs);
} else if (std::get_if<Fortran::parser::Expr::Multiply>(
&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::MulIOp>(currentLocation, val,
convertRhs);
} else if (std::get_if<Fortran::parser::Expr::Divide>(
&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::DivUIOp>(currentLocation, val,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can see from your commit message that deciding type of divide is a TODO. And probably deciding between int and float ops is also a TODO. That said, this whole section which selects the operation seems a bit brittle to me - it feels like selection of appropriate operation should be delegated (and be consistent) with the rest of FIR lowering which handles these expressions.

convertRhs);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't you need to make a disinction between x / expr and expr / x here ?

} else if (std::get_if<Fortran::parser::Expr::AND>(&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::AndIOp>(currentLocation, val,
convertRhs);
} else if (std::get_if<Fortran::parser::Expr::OR>(&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::OrIOp>(currentLocation, val,
convertRhs);
} else if (std::get_if<Fortran::parser::Expr::EQV>(&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::CmpIOp>(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't you need a SelectOp after the comparison?

currentLocation, mlir::arith::CmpIPredicate::eq, val, convertRhs);
} else if (std::get_if<Fortran::parser::Expr::NEQV>(&assignmentStmtExpr.u)) {
op = firOpBuilder.create<mlir::arith::CmpIOp>(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Select?

currentLocation, mlir::arith::CmpIPredicate::ne, val, convertRhs);
}

tempOp.erase();
firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, op);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should there be an if to decide between OpenMP and OpenACC Yields?

}

/// Processes an atomic construct with write clause.
Expand Down
74 changes: 74 additions & 0 deletions flang/test/Lower/OpenMP/common-atomic-lowering.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s

!CHECK: func.func @_QQmain() attributes {fir.bindc_name = "sample"} {
!CHECK: %[[val_0:.*]] = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFEa"}
!CHECK: %[[val_1:.*]]:2 = hlfir.declare %[[val_0]] {uniq_name = "_QFEa"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
!CHECK: %[[val_2:.*]] = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFEb"}
!CHECK: %[[val_3:.*]]:2 = hlfir.declare %[[val_2]] {uniq_name = "_QFEb"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
!CHECK: %[[val_4:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
!CHECK: %[[val_5:.*]]:2 = hlfir.declare %[[val_4]] {uniq_name = "_QFEx"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
!CHECK: %[[val_c5:.*]] = arith.constant 5 : index
!CHECK: %[[val_6:.*]] = fir.alloca !fir.array<5xi32> {bindc_name = "y", uniq_name = "_QFEy"}
!CHECK: %[[val_7:.*]] = fir.shape %[[val_c5]] : (index) -> !fir.shape<1>
!CHECK: %[[val_8:.*]]:2 = hlfir.declare %[[val_6]](%[[val_7]]) {uniq_name = "_QFEy"} : (!fir.ref<!fir.array<5xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<5xi32>>, !fir.ref<!fir.array<5xi32>>)
!CHECK: %[[val_c2:.*]] = arith.constant 2 : index
!CHECK: %[[val_9:.*]] = hlfir.designate %[[val_8]]#0 (%[[val_c2]]) : (!fir.ref<!fir.array<5xi32>>, index) -> !fir.ref<i32>
!CHECK: %[[val_c8:.*]] = arith.constant 8 : i32
!CHECK: %[[val_10:.*]] = fir.load %[[val_5]]#0 : !fir.ref<i32>
!CHECK: %[[val_11:.*]] = arith.addi %[[val_c8]], %[[val_10]] : i32
!CHECK: %[[val_12:.*]] = hlfir.no_reassoc %[[val_11]] : i32
!CHECK: omp.atomic.update %[[val_9]] : !fir.ref<i32> {
!CHECK: ^bb0(%[[ARG:.*]]: i32):
!CHECK: %[[val_18:.*]] = arith.muli %[[ARG]], %[[val_12]] : i32
!CHECK: omp.yield(%[[val_18]] : i32)
!CHECK: }
!CHECK: %[[val_c2_0:.*]] = arith.constant 2 : index
!CHECK: %[[val_13:.*]] = hlfir.designate %[[val_8]]#0 (%[[val_c2_0]]) : (!fir.ref<!fir.array<5xi32>>, index) -> !fir.ref<i32>
!CHECK: %[[val_c8_1:.*]] = arith.constant 8 : i32
!CHECK: omp.atomic.update %[[val_13:.*]] : !fir.ref<i32> {
!CHECK: ^bb0(%[[ARG:.*]]: i32):
!CHECK: %[[val_18:.*]] = arith.divui %[[ARG]], %[[val_c8_1]] : i32
!CHECK: omp.yield(%[[val_18]] : i32)
!CHECK: }
!CHECK: %[[val_c8_2:.*]] = arith.constant 8 : i32
!CHECK: %[[val_c4:.*]] = arith.constant 4 : index
!CHECK: %[[val_14:.*]] = hlfir.designate %[[val_8]]#0 (%[[val_c4]]) : (!fir.ref<!fir.array<5xi32>>, index) -> !fir.ref<i32>
!CHECK: %[[val_15:.*]] = fir.load %[[val_14]] : !fir.ref<i32>
!CHECK: %[[val_16:.*]] = arith.addi %[[val_c8_2]], %[[val_15]] : i32
!CHECK: %[[val_17:.*]] = hlfir.no_reassoc %[[val_16]] : i32
!CHECK: omp.atomic.update %[[val_5]]#1 : !fir.ref<i32> {
!CHECK: ^bb0(%[[ARG:.*]]: i32):
!CHECK: %[[val_18:.*]] = arith.addi %[[ARG]], %[[val_17]] : i32
!CHECK: omp.yield(%[[val_18]] : i32)
!CHECK: }
!CHECK: %[[val_c8_3:.*]] = arith.constant 8 : i32
!CHECK: omp.atomic.update %[[val_5]]#1 : !fir.ref<i32> {
!CHECK: ^bb0(%[[ARG]]: i32):
!CHECK: %[[val_18:.*]] = arith.subi %[[ARG]], %[[val_c8_3]] : i32
!CHECK: omp.yield(%[[val_18]] : i32)
!CHECK: }
!CHECK: return
!CHECK: }
program sample

integer :: x
integer, dimension(5) :: y
integer :: a, b

!$omp atomic update
y(2) = (8 + x) * y(2)
!$omp end atomic

!$omp atomic update
y(2) = y(2) / 8
!$omp end atomic

!$omp atomic update
x = (8 + y(4)) + x
!$omp end atomic

!$omp atomic update
x = 8 - x
!$omp end atomic

end program sample