@@ -416,7 +416,8 @@ static inline void genAtomicUpdateStatement(
416
416
Fortran::lower::AbstractConverter &converter, mlir::Value lhsAddr,
417
417
mlir::Type varType, const Fortran::parser::Variable &assignmentStmtVariable,
418
418
const Fortran::parser::Expr &assignmentStmtExpr, mlir::Location loc,
419
- mlir::Operation *atomicCaptureOp = nullptr ) {
419
+ mlir::Operation *atomicCaptureOp = nullptr ,
420
+ Fortran::lower::StatementContext *atomicCaptureStmtCtx = nullptr ) {
420
421
// Generate `atomic.update` operation for atomic assignment statements
421
422
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder ();
422
423
mlir::Location currentLocation = converter.getCurrentLocation ();
@@ -496,15 +497,24 @@ static inline void genAtomicUpdateStatement(
496
497
},
497
498
assignmentStmtExpr.u );
498
499
Fortran::lower::StatementContext nonAtomicStmtCtx;
500
+ Fortran::lower::StatementContext *stmtCtxPtr = &nonAtomicStmtCtx;
499
501
if (!nonAtomicSubExprs.empty ()) {
500
502
// Generate non atomic part before all the atomic operations.
501
503
auto insertionPoint = firOpBuilder.saveInsertionPoint ();
502
- if (atomicCaptureOp)
504
+ if (atomicCaptureOp) {
505
+ assert (atomicCaptureStmtCtx && " must specify statement context" );
503
506
firOpBuilder.setInsertionPoint (atomicCaptureOp);
507
+ // Any clean-ups associated with the expression lowering
508
+ // must also be generated outside of the atomic update operation
509
+ // and after the atomic capture operation.
510
+ // The atomicCaptureStmtCtx will be finalized at the end
511
+ // of the atomic capture operation generation.
512
+ stmtCtxPtr = atomicCaptureStmtCtx;
513
+ }
504
514
mlir::Value nonAtomicVal;
505
515
for (auto *nonAtomicSubExpr : nonAtomicSubExprs) {
506
516
nonAtomicVal = fir::getBase (converter.genExprValue (
507
- currentLocation, *nonAtomicSubExpr, nonAtomicStmtCtx ));
517
+ currentLocation, *nonAtomicSubExpr, *stmtCtxPtr ));
508
518
exprValueOverrides.try_emplace (nonAtomicSubExpr, nonAtomicVal);
509
519
}
510
520
if (atomicCaptureOp)
@@ -652,7 +662,7 @@ void genAtomicCapture(Fortran::lower::AbstractConverter &converter,
652
662
genAtomicCaptureStatement (converter, stmt2LHSArg, stmt1LHSArg,
653
663
elementType, loc);
654
664
genAtomicUpdateStatement (converter, stmt2LHSArg, stmt2VarType, stmt2Var,
655
- stmt2Expr, loc, atomicCaptureOp);
665
+ stmt2Expr, loc, atomicCaptureOp, &stmtCtx );
656
666
} else {
657
667
// Atomic capture construct is of the form [capture-stmt, write-stmt]
658
668
firOpBuilder.setInsertionPoint (atomicCaptureOp);
@@ -672,13 +682,15 @@ void genAtomicCapture(Fortran::lower::AbstractConverter &converter,
672
682
*Fortran::semantics::GetExpr (stmt2Expr);
673
683
mlir::Type elementType = converter.genType (fromExpr);
674
684
genAtomicUpdateStatement (converter, stmt1LHSArg, stmt1VarType, stmt1Var,
675
- stmt1Expr, loc, atomicCaptureOp);
685
+ stmt1Expr, loc, atomicCaptureOp, &stmtCtx );
676
686
genAtomicCaptureStatement (converter, stmt1LHSArg, stmt2LHSArg, elementType,
677
687
loc);
678
688
}
679
689
firOpBuilder.setInsertionPointToEnd (&block);
680
690
firOpBuilder.create <mlir::acc::TerminatorOp>(loc);
681
- firOpBuilder.setInsertionPointToStart (&block);
691
+ // The clean-ups associated with the statements inside the capture
692
+ // construct must be generated after the AtomicCaptureOp.
693
+ firOpBuilder.setInsertionPointAfter (atomicCaptureOp);
682
694
}
683
695
684
696
template <typename Op>
0 commit comments