[flang][OpenMP] Common lowering flow for atomic update #69866
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
NimishMishra
requested review from
clementval,
kiranchandramohan,
raghavendhra,
agozillon,
DominikAdamski,
shraiysh and
kiranktp
llvmbot
added
flang
|
@llvm/pr-subscribers-flang-fir-hlfir @llvm/pr-subscribers-flang-openmp Author: None (NimishMishra) ChangesOffers a common lowering flow for scalar/non-scalar atomic variables. Fixes #68384 TODOs:
Full diff: https://github.com/llvm/llvm-project/pull/69866.diff 2 Files Affected:
diff --git a/flang/lib/Lower/DirectivesCommon.h b/flang/lib/Lower/DirectivesCommon.h
index ed44598bc925212..56d79de2d31995d 100644
--- a/flang/lib/Lower/DirectivesCommon.h
+++ b/flang/lib/Lower/DirectivesCommon.h
@@ -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) {
+ // 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,
@@ -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};
@@ -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,
+ convertRhs);
+ } 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>(
+ currentLocation, mlir::arith::CmpIPredicate::eq, val, convertRhs);
+ } else if (std::get_if<Fortran::parser::Expr::NEQV>(&assignmentStmtExpr.u)) {
+ op = firOpBuilder.create<mlir::arith::CmpIOp>(
+ currentLocation, mlir::arith::CmpIPredicate::ne, val, convertRhs);
}
-
- tempOp.erase();
+ firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, op);
}
/// Processes an atomic construct with write clause.
diff --git a/flang/test/Lower/OpenMP/common-atomic-lowering.f90 b/flang/test/Lower/OpenMP/common-atomic-lowering.f90
new file mode 100644
index 000000000000000..7da30243e676c00
--- /dev/null
+++ b/flang/test/Lower/OpenMP/common-atomic-lowering.f90
@@ -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
|
| convertRhs); | ||
| } else if (std::get_if<Fortran::parser::Expr::Divide>( | ||
| &assignmentStmtExpr.u)) { | ||
| op = firOpBuilder.create<mlir::arith::DivUIOp>(currentLocation, val, |
There was a problem hiding this comment.
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.
|
Does your code also handle arrays with descriptor? e.g. |
|
Thanks for creating this PR. Having this here (vs a branch) helps everyone to provide suggestions. I think you should expand the summary a bit to describe the context and how this solution was reached. I will try that briefly below. The The initial lowering reused the existing lowering code. This was achieved by binding the block argument of the atomic update region as the value corresponding to the symbol in the lowering for the body of the atomic region. But this flow ran into a few difficulties with:
For handling issue (1) above, we could think of two approaches:
Note: We also discussed this a bit in the following thread: https://flang-compiler.slack.com/archives/C01PY03PP9P/p1694494832259789 |
| 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) { |
There was a problem hiding this comment.
Would the check mis-handle cases like the following?
arr(1) = arr(2) + arr(1)
| } | ||
|
|
||
| tempOp.erase(); | ||
| firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, op); |
There was a problem hiding this comment.
Should there be an if to decide between OpenMP and OpenACC Yields?
|
The five tests below are broken with this patch. Can you look at them? |
clementval
reviewed
Oct 23, 2023
| 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>( |
There was a problem hiding this comment.
Don't you need a SelectOp after the comparison?
clementval
reviewed
Oct 23, 2023
| op = firOpBuilder.create<mlir::arith::CmpIOp>( | ||
| currentLocation, mlir::arith::CmpIPredicate::eq, val, convertRhs); | ||
| } else if (std::get_if<Fortran::parser::Expr::NEQV>(&assignmentStmtExpr.u)) { | ||
| op = firOpBuilder.create<mlir::arith::CmpIOp>( |
jeanPerier
reviewed
Oct 23, 2023
There was a problem hiding this comment.
Thanks for the patch, and thanks for the summary @kiranchandramohan.
I think it would still be better to add the lowering infrastructure to have a finer control over expression lowering and allow mapping pre-lowered mlir::Value sub-expression to their evaluate::Expr representation so that later lowering can use the value without lowering the sub-expression again.
I think this may have more use cases than this and will avoid duplicating expression lowering.
I made a draft proposal in #69944
| } else if (std::get_if<Fortran::parser::Expr::Divide>( | ||
| &assignmentStmtExpr.u)) { | ||
| op = firOpBuilder.create<mlir::arith::DivUIOp>(currentLocation, val, | ||
| convertRhs); |
There was a problem hiding this comment.
Don't you need to make a disinction between x / expr and expr / x here ?
jeanPerier
added a commit
to jeanPerier/llvm-project
that referenced
this pull request
Oct 24, 2023
OpenACC/OpenMP atomic lowering needs a finer control over expression lowering. This patch allows mapping evaluate::Expr<T> to mlir::Value so that any subsequent expression lowering will use these values when an operand is a mapped Expr<T>. This is an alternative to llvm#69866 From which I took the test. The same test as in llvm#69866 are failing because the "non atomic part" is now out of the atomic.update op, which in some case causing verification failure because this is generated in the middle of an omp.atomic.capture. I did not try fixing these failures. My patch is about the lowering infrastructure and how to use it rather than the OpenMP semantics. Co-authored-by: Nimish Mishra <[email protected]>
jeanPerier
added a commit
that referenced
this pull request
Oct 25, 2023
OpenACC/OpenMP atomic lowering needs a finer control over expression lowering. This patch allows mapping evaluate::Expr<T> to mlir::Value so that any subsequent expression lowering will use these values when an operand is a mapped Expr<T>. This is an alternative to #69866 From which I took the test and some of the logic to extract the non-atomic sub-expression. --------- Co-authored-by: Nimish Mishra <[email protected]>
|
Abandoning as an alternative lowering infrastructure is in place. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Offers a common lowering flow for scalar/non-scalar atomic variables. Fixes #68384
Atomic updates in OpenMP work on a well-defined restricted subset of expressions as given below. This patch proposes to do a custom lowering that:
-> Lowers
expr-> Generates the address of
x-> And custom-constructs the expression
TODOs:
Note: This modifies the common code shared with OpenACC.