[Flang][OpenMP] Use fir.if instead of scf.if in lastprivate lowering

[kiranchandramohan]

Summary

For finding the last iteration of a loop, or the last section an
if condition is generated. Using scf::if can cause some lowering
issues since the code contained inside it can have branches. Using
fir::if instead ensures that the fir::if is lowered into branches
along with any code contained inside that can generate branches.

Fixes #62458.

[NimishMishra]

LGTM. Thank you for the patch. I was not aware of this problem with scf.if.

flang/lib/Lower/OpenMP.cpp

 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/StatementContext.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Semantics/tools.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
 using namespace mlir;
 
 int64_t Fortran::lower::getCollapseValue(
     const Fortran::parser::OmpClauseList &clauseList) {
   for (const auto &clause : clauseList.v) {
     if (const auto &collapseClause =
if (std::get_if<Fortran::parser::OmpClause::Lastprivate>(&clause.u)) {
           // shall be generated for the same:
           //
           // omp.sections lastprivate(...) {
           //  omp.section {...}
           //  omp.section {...}
           //  omp.section {
           //      fir.allocate for `private`/`firstprivate`
           //      <More operations here>
-          //      scf.if %true {
+          //      fir.if %true {
           //          ^%lpv_update_blk
           //      }
           //  }
           // }
           //
           // To keep code consistency while handling privatization
-          // through this control flow, add a `scf.if` operation
+          // through this control flow, add a `fir.if` operation
           // that always evaluates to true, in order to create
           // a dedicated sub-region in `omp.section` where
           // lastprivate FIR can reside. Later canonicalizations
           // will optimize away this operation.
 
-          mlir::scf::IfOp ifOp = firOpBuilder.create<mlir::scf::IfOp>(
+          auto ifOp = firOpBuilder.create<fir::IfOp>(
               op->getLoc(),
               firOpBuilder.createIntegerConstant(
                   op->getLoc(), firOpBuilder.getIntegerType(1), 0x1),
               /*else*/ false);
           firOpBuilder.setInsertionPointToStart(&ifOp.getThenRegion().front());
 
           const Fortran::parser::OpenMPConstruct *parentOmpConstruct =
               eval.parentConstruct->getIf<Fortran::parser::OpenMPConstruct>();
if (std::get_if<Fortran::parser::OmpClause::Lastprivate>(&clause.u)) {
 
         // Update the original variable just before exiting the worksharing
         // loop. Conversion as follows:
         //
         //                       omp.wsloop {
         // omp.wsloop {            ...
         //    ...                  store
         //    store       ===>     %cmp = llvm.icmp "eq" %iv %ub
-        //    omp.yield            scf.if %cmp {
+        //    omp.yield            fir.if %cmp {
         // }                         ^%lpv_update_blk:
         //                         }
         //                         omp.yield
         //                       }
         //
 
         // Only generate the compare once in presence of multiple LastPrivate
         // clauses.
         if (!cmpCreated) {
           cmpOp = firOpBuilder.create<mlir::arith::CmpIOp>(
               op->getLoc(), mlir::arith::CmpIPredicate::eq,
               op->getRegion(0).front().getArguments()[0],
               mlir::dyn_cast<mlir::omp::WsLoopOp>(op).getUpperBound()[0]);
         }
-        mlir::scf::IfOp ifOp = firOpBuilder.create<mlir::scf::IfOp>(
-            op->getLoc(), cmpOp, /*else*/ false);
+        auto ifOp =
+            firOpBuilder.create<fir::IfOp>(op->getLoc(), cmpOp, /*else*/ false);
         firOpBuilder.setInsertionPointToStart(&ifOp.getThenRegion().front());
         lastPrivIP = firOpBuilder.saveInsertionPoint();
       } else {
         TODO(converter.getCurrentLocation(),
              "lastprivate clause in constructs other than "
              "simd/worksharing-loop");
       }
     }

flang/test/Lower/OpenMP/parallel-lastprivate-clause-scalar.f90

 !CHECK-NEXT: %[[CALL_BEGIN_IO:.*]] = fir.call @_FortranAioBeginExternalListOutput(%[[NEG_ONE]], %[[CVT0]], %[[CNST]]) {{.*}}: (i32, !fir.ref<i8>, i32) -> !fir.ref<i8>
 !CHECK-NEXT: %[[CVT_0_1:.*]] = fir.convert %[[ARG1_PVT]] 
 !CHECK-NEXT: %[[CVT_0_2:.*]] = fir.convert %[[FIVE]]
 !CHECK-NEXT: %[[CALL_OP_ASCII:.*]] = fir.call @_FortranAioOutputAscii(%[[CALL_BEGIN_IO]], %[[CVT_0_1]], %[[CVT_0_2]])
 !CHECK-NEXT: %[[CALL_END_IO:.*]] = fir.call @_FortranAioEndIoStatement(%[[CALL_BEGIN_IO]])
 
 ! Testing last iteration check
 !CHECK-NEXT: %[[IV_CMP:.*]] = arith.cmpi eq, %[[INDX_WS]]
-!CHECK: scf.if %[[IV_CMP]] {
+!CHECK: fir.if %[[IV_CMP]] {
 
 ! Testing lastprivate val update
 !CHECK-DAG: %[[CVT:.*]] = fir.convert %[[ARG1_REF]] : (!fir.ref<!fir.char<1,5>>) -> !fir.ref<i8>
 !CHECK-DAG: %[[CVT1:.*]] = fir.convert %[[ARG1_PVT]] : (!fir.ref<!fir.char<1,5>>) -> !fir.ref<i8>
 !CHECK-DAG: fir.call @llvm.memmove.p0.p0.i64(%[[CVT]], %[[CVT1]]{{.*}})
 !CHECK-DAG: } 
 !CHECK-DAG: omp.yield
 

 
 !CHECK: func @_QPlastprivate_int(%[[ARG1:.*]]: !fir.ref<i32> {fir.bindc_name = "arg1"}) {
 !CHECK-DAG: omp.parallel  {
 !CHECK-DAG: %[[CLONE:.*]] = fir.alloca i32 {bindc_name = "arg1"
 !CHECK: omp.wsloop for (%[[INDX_WS:.*]]) : {{.*}} {
 
 ! Testing last iteration check
 !CHECK-DAG: %[[IV_CMP:.*]] = arith.cmpi eq, %[[INDX_WS]]
-!CHECK-DAG: scf.if %[[IV_CMP]] {
+!CHECK-DAG: fir.if %[[IV_CMP]] {
 
 ! Testing lastprivate val update
 !CHECK-NEXT: %[[CLONE_LD:.*]] = fir.load %[[CLONE]] : !fir.ref<i32>
 !CHECK-NEXT: fir.store %[[CLONE_LD]] to %[[ARG1]] : !fir.ref<i32>
 !CHECK-DAG: }
 !CHECK-DAG: omp.yield
 
 subroutine lastprivate_int(arg1)

 !CHECK: func.func @_QPmult_lastprivate_int(%[[ARG1:.*]]: !fir.ref<i32> {fir.bindc_name = "arg1"}, %[[ARG2:.*]]: !fir.ref<i32> {fir.bindc_name = "arg2"}) {
 !CHECK: omp.parallel  {
 !CHECK-DAG: %[[CLONE1:.*]] = fir.alloca i32 {bindc_name = "arg1"
 !CHECK-DAG: %[[CLONE2:.*]] = fir.alloca i32 {bindc_name = "arg2"
 !CHECK: omp.wsloop for (%[[INDX_WS:.*]]) : {{.*}} {
 
 ! Testing last iteration check
 !CHECK: %[[IV_CMP1:.*]] = arith.cmpi eq, %[[INDX_WS]]
-!CHECK-NEXT: scf.if %[[IV_CMP1]] {
+!CHECK-NEXT: fir.if %[[IV_CMP1]] {
 ! Testing lastprivate val update
 !CHECK-DAG: %[[CLONE_LD1:.*]] = fir.load %[[CLONE1]] : !fir.ref<i32>
 !CHECK-DAG: fir.store %[[CLONE_LD1]] to %[[ARG1]] : !fir.ref<i32>
 !CHECK-DAG: %[[CLONE_LD2:.*]] = fir.load %[[CLONE2]] : !fir.ref<i32>
 !CHECK-DAG: fir.store %[[CLONE_LD2]] to %[[ARG2]] : !fir.ref<i32>
 !CHECK: }
 !CHECK: omp.yield
 

 !CHECK: func.func @_QPmult_lastprivate_int2(%[[ARG1:.*]]: !fir.ref<i32> {fir.bindc_name = "arg1"}, %[[ARG2:.*]]: !fir.ref<i32> {fir.bindc_name = "arg2"}) {
 !CHECK: omp.parallel  {
 !CHECK-DAG: %[[CLONE1:.*]] = fir.alloca i32 {bindc_name = "arg1"
 !CHECK-DAG: %[[CLONE2:.*]] = fir.alloca i32 {bindc_name = "arg2"
 !CHECK: omp.wsloop for (%[[INDX_WS:.*]]) : {{.*}} {
 
 !Testing last iteration check
 !CHECK: %[[IV_CMP1:.*]] = arith.cmpi eq, %[[INDX_WS]]
-!CHECK-NEXT: scf.if %[[IV_CMP1]] {
+!CHECK-NEXT: fir.if %[[IV_CMP1]] {
 !Testing lastprivate val update
 !CHECK-DAG: %[[CLONE_LD2:.*]] = fir.load %[[CLONE2]] : !fir.ref<i32>
 !CHECK-DAG: fir.store %[[CLONE_LD2]] to %[[ARG2]] : !fir.ref<i32>
 !CHECK-DAG: %[[CLONE_LD1:.*]] = fir.load %[[CLONE1]] : !fir.ref<i32>
 !CHECK-DAG: fir.store %[[CLONE_LD1]] to %[[ARG1]] : !fir.ref<i32>
 !CHECK: }
 !CHECK: omp.yield
 

 !CHECK-DAG: fir.store %[[FPV_LD]] to %[[CLONE1]] : !fir.ref<i32>
 ! Lastprivate Allocation
 !CHECK-DAG: %[[CLONE2:.*]] = fir.alloca i32 {bindc_name = "arg2"
 !CHECK-NOT: omp.barrier
 !CHECK: omp.wsloop for (%[[INDX_WS:.*]]) : {{.*}} {
 
 ! Testing last iteration check
 !CHECK: %[[IV_CMP1:.*]] = arith.cmpi eq, %[[INDX_WS]]
-!CHECK-NEXT: scf.if %[[IV_CMP1]] {
+!CHECK-NEXT: fir.if %[[IV_CMP1]] {
 ! Testing lastprivate val update
 !CHECK-NEXT: %[[CLONE_LD:.*]] = fir.load %[[CLONE2]] : !fir.ref<i32>
 !CHECK-NEXT: fir.store %[[CLONE_LD]] to %[[ARG2]] : !fir.ref<i32>
 !CHECK-NEXT: }
 !CHECK-NEXT: omp.yield
 
 subroutine firstpriv_lastpriv_int(arg1, arg2)
         integer :: arg1, arg2

 ! Firstprivate update
 !CHECK-NEXT: %[[CLONE1:.*]] = fir.alloca i32 {bindc_name = "arg1"
 !CHECK-NEXT: %[[FPV_LD:.*]] = fir.load %[[ARG1]] : !fir.ref<i32>
 !CHECK-NEXT: fir.store %[[FPV_LD]] to %[[CLONE1]] : !fir.ref<i32>
 !CHECK-NEXT: omp.barrier
 !CHECK: omp.wsloop for (%[[INDX_WS:.*]]) : {{.*}} {
 ! Testing last iteration check
 !CHECK: %[[IV_CMP1:.*]] = arith.cmpi eq, %[[INDX_WS]]
-!CHECK-NEXT: scf.if %[[IV_CMP1]] {
+!CHECK-NEXT: fir.if %[[IV_CMP1]] {
 ! Testing lastprivate val update
 !CHECK-NEXT: %[[CLONE_LD:.*]] = fir.load %[[CLONE1]] : !fir.ref<i32>
 !CHECK-NEXT: fir.store %[[CLONE_LD]] to %[[ARG1]] : !fir.ref<i32>
 !CHECK-NEXT: }
 !CHECK-NEXT: omp.yield
 
 subroutine firstpriv_lastpriv_int2(arg1)
         integer :: arg1

flang/test/Lower/OpenMP/sections.f90

!CHECK: }
             x = x * 10
 !CHECK: omp.section {
 !CHECK: %[[PRIVATE_X:.*]] = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
 !CHECK: %[[true:.*]] = arith.constant true
 !CHECK: %[[temp:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: %[[const:.*]] = arith.constant 1 : i32
 !CHECK: %[[result:.*]] = arith.addi %[[temp]], %[[const]] : i32
 !CHECK: fir.store %[[result]] to %[[PRIVATE_X]] : !fir.ref<i32>
-!CHECK: scf.if %[[true]] {
+!CHECK: fir.if %[[true]] {
 !CHECK: %[[temp:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: fir.store %[[temp]] to %[[X]] : !fir.ref<i32>
 !CHECK: }
 !CHECK: omp.terminator
 !CHECK: }
         !$omp section
             x = x + 1
 !CHECK: omp.terminator

 !CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK: fir.store %[[temp]] to %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: omp.barrier
 !CHECK: %[[true:.*]] = arith.constant true
 !CHECK: %[[temp:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: %[[const:.*]] = arith.constant 1 : i32
 !CHECK: %[[result:.*]] = arith.addi %[[temp]], %[[const]] : i32
 !CHECK: fir.store %[[result]] to %[[PRIVATE_X]] : !fir.ref<i32>
-!CHECK: scf.if %true {
+!CHECK: fir.if %true {
 !CHECK: %[[temp:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: fir.store %[[temp]] to %[[X]] : !fir.ref<i32>
 !CHECK: }
 !CHECK: omp.terminator
 !CHECK: }
         !$omp section
             x = x + 1
 !CHECK: omp.terminator

 !CHECK: %[[temp:.*]] = fir.load %[[X]] : !fir.ref<i32>
 !CHECK: fir.store %[[temp]] to %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: omp.barrier
 !CHECK: %[[true:.*]] = arith.constant true
 !CHECK: %[[temp:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: %[[const:.*]] = arith.constant 1 : i32
 !CHECK: %[[result:.*]] = arith.addi %[[temp]], %[[const]] : i32
 !CHECK: fir.store %[[result]] to %[[PRIVATE_X]] : !fir.ref<i32>
-!CHECK: scf.if %true {
+!CHECK: fir.if %true {
 !CHECK: %[[temp:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
 !CHECK: fir.store %[[temp]] to %[[X]] : !fir.ref<i32>
 !CHECK: omp.barrier
 !CHECK: }
 !CHECK: omp.terminator
 !CHECK: }
         !$omp section
             x = x + 1
 !CHECK: omp.terminator
 !CHECK: }
     !$omp end sections nowait
 end subroutine