[flang][OpenMP] Lowering support for lastprivate on unstructured sections construct

[NimishMishra]

Summary

This patch adds lowering support for lastprivate privatization on unstructured construct. Generic lastprivate lowering on sections involves creating a scf.if operation on the lexically last section block and updating the lastprivate variable inside it. However, this control flow is not needed if the section construct is lowered as an unstructured construct. Hence, this patch modifies the lowering control flow to not emit an unnecessary scf.if if the section block is lowered as an unstructured construct.

Depends on D133686

[NimishMishra]

Rebase with main

[NimishMishra]

Fix CI fails

[kiranchandramohan]

Could you add a description in the summary or add a comment in the changed code on what this change does and how it fixes the issue?

[kiranchandramohan]

LG.

While looking at this patch I saw that there is some firstprivatisation that is happening with the lastprivate clause in those section constructs that are not the last section. I am not sure why this happens. This might not be a bug since the private values can have any value, but it is definitely unnecessary instructions.

flang/lib/Lower/OpenMP.cpp

63–64

Nit: Insert braces in the else part.

190–191

Nit: Add a comment on what is the location of lastPrivOp. Is it the end of the last block in the lexically last section construct?

Do you know why the same code does not work for the structured case? I believe we inserted the scf.if operation to avoid some issue, but it is not actually required.

[NimishMishra]

In D143077#4281689, @kiranchandramohan wrote:

LG.

While looking at this patch I saw that there is some firstprivatisation that is happening with the lastprivate clause in those section constructs that are not the last section. I am not sure why this happens. This might not be a bug since the private values can have any value, but it is definitely unnecessary instructions.

Thanks.

I was just wondering, are you referring to the following IR generated by the following test case. In this case, shouldn't the "lexically last section block" apply on the lastprivate case? Like isn't firstprivate applicable on all section blocks? If that is so, that is why firstprivate code is being produced in each section block, while lastprivate code is just in the final block.

If this is wrong though, let me know and I will fix it on priority.

!$omp sections firstprivate(x) lastprivate(x)
      !$omp section
            x = x * 10
      !$omp section
            x = x + 1
!$omp end sections

and the generated IR

omp.sections   {
      omp.section {
        %1 = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
        %2 = fir.load %0 : !fir.ref<i32>
        fir.store %2 to %1 : !fir.ref<i32>
        omp.barrier
        %3 = fir.load %0 : !fir.ref<i32>
        fir.store %3 to %1 : !fir.ref<i32>
        %c10_i32 = arith.constant 10 : i32
        %4 = fir.load %1 : !fir.ref<i32>
        %5 = arith.muli %c10_i32, %4 : i32
        fir.store %5 to %1 : !fir.ref<i32>
        omp.terminator
      }

omp.section {
        %1 = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
        %2 = fir.load %0 : !fir.ref<i32>
        fir.store %2 to %1 : !fir.ref<i32>
        omp.barrier
        %true = arith.constant true
        %3 = fir.load %1 : !fir.ref<i32>
        %c1_i32 = arith.constant 1 : i32
        %4 = arith.addi %3, %c1_i32 : i32
        fir.store %4 to %1 : !fir.ref<i32>
        scf.if %true {
          %5 = fir.load %1 : !fir.ref<i32>
          fir.store %5 to %0 : !fir.ref<i32>
        }
        omp.terminator
      }
omp.terminator
    }

[kiranchandramohan]

The issue need not block this patch.

I was referring to tests like the following:

subroutine lastprivate()
  integer :: x
  !$omp sections lastprivate(x)
  !$omp section
      x = x * 10
  !$omp section
      x = x + 10
  !$omp end sections
end subroutine

that generates IR like the following:

func.func @_QPlastprivate() {
  %0 = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFlastprivateEx"}
  omp.sections   {
    omp.section {
      %1 = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
      %2 = fir.load %0 : !fir.ref<i32>
      fir.store %2 to %1 : !fir.ref<i32>
      %c10_i32 = arith.constant 10 : i32
      %3 = fir.load %1 : !fir.ref<i32>
      %4 = arith.muli %c10_i32, %3 : i32
      fir.store %4 to %1 : !fir.ref<i32>
      omp.terminator
    }
    omp.section {
      %1 = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
      %true = arith.constant true
      %2 = fir.load %1 : !fir.ref<i32>
      %c10_i32 = arith.constant 10 : i32
      %3 = arith.addi %2, %c10_i32 : i32
      fir.store %3 to %1 : !fir.ref<i32>
      fir.if %true {
        %4 = fir.load %1 : !fir.ref<i32>
        fir.store %4 to %0 : !fir.ref<i32>
      }
      omp.terminator
    }
    omp.terminator
  }
  return
}

Note the following IR for the first section that seems to do a firstprivatisation although firstprivate is not specified.

%1 = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
%2 = fir.load %0 : !fir.ref<i32>
fir.store %2 to %1 : !fir.ref<i32>

flang/lib/Lower/OpenMP.cpp

std::visit(Fortran::common::visitors{
                            getDesignatorNameIfDataRef(designator)) {
                      sym = name->symbol;
                    }
                  },
                  [&](const Fortran::parser::Name &name) { sym = name.symbol; }},
              ompObject.u);
   return sym;
 }
 
 template <typename Op>
 static void privatizeSymbol(
     Op &op, Fortran::lower::AbstractConverter &converter,
     const Fortran::semantics::Symbol *sym,
     [[maybe_unused]] mlir::OpBuilder::InsertPoint *lastPrivIP = nullptr) {
   // Privatization for symbols which are pre-determined (like loop index
   // variables) happen separately, for everything else privatize here.
   if (sym->test(Fortran::semantics::Symbol::Flag::OmpPreDetermined))
     return;
if (const auto &privateClause =
           // To keep code consistency while handling privatization
           // through this control flow, add a `scf.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.
 
           omp::SectionOp *sectionOp = dyn_cast<omp::SectionOp>(&op);
+          if (!eval.lowerAsUnstructured()) {
             mlir::scf::IfOp ifOp = firOpBuilder.create<mlir::scf::IfOp>(
                 sectionOp->getLoc(),
                 firOpBuilder.createIntegerConstant(
                     sectionOp->getLoc(), firOpBuilder.getIntegerType(1), 0x1),
                 /*else*/ false);
-          firOpBuilder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+            firOpBuilder.setInsertionPointToStart(
+                &ifOp.getThenRegion().front());
 
             const Fortran::parser::OpenMPConstruct *parentOmpConstruct =
                 eval.parentConstruct->getIf<Fortran::parser::OpenMPConstruct>();
             assert(parentOmpConstruct &&
                    "Expected a valid enclosing OpenMP construct");
             const Fortran::parser::OpenMPSectionsConstruct *sectionsConstruct =
                 std::get_if<Fortran::parser::OpenMPSectionsConstruct>(
                     &parentOmpConstruct->u);
             assert(sectionsConstruct &&
                    "Expected an enclosing omp.sections construct");
             const Fortran::parser::OmpClauseList &sectionsEndClauseList =
                 std::get<Fortran::parser::OmpClauseList>(
                     std::get<Fortran::parser::OmpEndSectionsDirective>(
                         sectionsConstruct->t)
                         .t);
             for (const Fortran::parser::OmpClause &otherClause :
                  sectionsEndClauseList.v)
-            if (std::get_if<Fortran::parser::OmpClause::Nowait>(&otherClause.u))
+              if (std::get_if<Fortran::parser::OmpClause::Nowait>(
+                      &otherClause.u))
                 // Emit implicit barrier to synchronize threads and avoid data
                 // races on post-update of lastprivate variables when `nowait`
                 // clause is present.
                 firOpBuilder.create<mlir::omp::BarrierOp>(
                     converter.getCurrentLocation());
-          firOpBuilder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+            firOpBuilder.setInsertionPointToStart(
+                &ifOp.getThenRegion().front());
             lastPrivIP = firOpBuilder.saveInsertionPoint();
             firOpBuilder.setInsertionPoint(ifOp);
             insPt = firOpBuilder.saveInsertionPoint();
+          } else {
+            mlir::OpBuilder::InsertPoint unstructuredSectionsIP =
+                firOpBuilder.saveInsertionPoint();
+            firOpBuilder.setInsertionPointToStart(
+                &sectionOp->getRegion().back());
+            lastPrivIP = firOpBuilder.saveInsertionPoint();
+            firOpBuilder.restoreInsertionPoint(unstructuredSectionsIP);
+          }
         }
       } else if (std::is_same_v<Op, omp::WsLoopOp>) {
         omp::WsLoopOp *wsLoopOp = dyn_cast<omp::WsLoopOp>(&op);
         mlir::Operation *lastOper =
             wsLoopOp->getRegion().back().getTerminator();
         firOpBuilder.setInsertionPoint(lastOper);
 
         // Our goal here is to introduce the following control flow
if (const auto &defaultClause =
         collectSymbols(Fortran::semantics::Symbol::Flag::OmpPrivate);
       else if (defaultClause->v.v ==
                Fortran::parser::OmpDefaultClause::Type::Firstprivate)
         collectSymbols(Fortran::semantics::Symbol::Flag::OmpFirstPrivate);
     }
   }
 
   bool needBarrier = false;
-  if (mlir::isa<mlir::omp::SectionOp>(op))
+  if (mlir::isa<mlir::omp::SectionOp>(op)) {
+    if (!eval.lowerAsUnstructured())
       firOpBuilder.setInsertionPointToStart(&op.getRegion().back());
-  else
+  } else
     firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
   for (auto sym : privatizedSymbols) {
     privatizeSymbol(op, converter, sym, &lastPrivIP);
     if (sym->test(Fortran::semantics::Symbol::Flag::OmpFirstPrivate) &&
         sym->test(Fortran::semantics::Symbol::Flag::OmpLastPrivate))
       needBarrier = true;
   }
 

flang/test/Lower/OpenMP/sections.f90

 !CHECK: }
 !CHECK: omp.terminator
 !CHECK: }
         !$omp section
             x = x + 1
 !CHECK: omp.terminator
 !CHECK: }
     !$omp end sections nowait
+
+!CHECK: omp.sections {
+!CHECK: omp.section {
+!CHECK: %[[PRIVATE_X:.*]] = fir.alloca i32 {bindc_name = "x", pinned, uniq_name = "_QFlastprivateEx"}
+!CHECK: cf.br ^bb1
+!CHECK: ^bb1:  // pred: ^bb0
+!CHECK: %[[INNER_PRIVATE_X:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
+!CHECK: %[[const:.*]] = arith.constant 1 : i32
+!CHECK: %[[result:.*]] = arith.addi %[[INNER_PRIVATE_X]], %[[const]] : i32
+!CHECK: fir.store %[[result]] to %[[PRIVATE_X]] : !fir.ref<i32>
+!CHECK: %[[loaded_value:.*]] = fir.load %[[PRIVATE_X]] : !fir.ref<i32>
+!CHECK: fir.store %[[loaded_value]] to %[[X]] : !fir.ref<i32>
+!CHECK: omp.terminator
+!CHECK: }
+!CHECK: }
+    !$omp sections lastprivate(x)
+        !$omp section
+            goto 30
+        30  x = x + 1
+    !$omp end sections
 end subroutine