[flang][OpenMP] Fix firstprivate bug

[arnamoy10]

Summary

In case where the bound(s) of a workshare loop use(s) firstprivate var(s), currently, that use is not updated with the created clone. It still uses the shared variable. This patch fixes that.

[kiranchandramohan]

Not requesting a change: But have you considered performing privatisation before creating the loop bounds for the worksharing loop?

flang/lib/Lower/Bridge.cpp

515–521

There are functions called replaceAllUsesWith. Are they not sufficient for the purpose here?

[peixin]

In D127137#3565831, @kiranchandramohan wrote:

Not requesting a change: But have you considered performing privatisation before creating the loop bounds for the worksharing loop?

+1 It's better to use copyHostAssociateVar unconditionally for all the cases and handle special usage elsewhere.

[shraiysh]

But have you considered performing privatisation before creating the loop bounds for the worksharing loop?

Earlier, I thought about doing it this way but realized that privatization is tied to body creation and both run in a single call to createBodyOfOp. Separating them would require a little more change to the IR but I think that would be a step in the right direction. If we carefully separate privatization from body creation, then we will automatically handle this bug.

[peixin]

In D127137#3566056, @shraiysh wrote:

But have you considered performing privatisation before creating the loop bounds for the worksharing loop?

Earlier, I thought about doing it this way but realized that privatization is tied to body creation and both run in a single call to createBodyOfOp. Separating them would require a little more change to the IR but I think that would be a step in the right direction. If we carefully separate privatization from body creation, then we will automatically handle this bug.

I think the problem is not if the privatization is tied to the body creation or not. The problem is getting the loop bounds values in worksharing-loop use the value of the host symbol instead of the associated symbol. There should be no problem with privatisation. For example, the lowering of copyin clause (D127468) does not have this problem since I explicitly use the host and associated symbols for both threadprivate and copyin variables.

flang/lib/Lower/Bridge.cpp

515–521

If you want to do this here. Please make sure it won't crash the usage of this function elsewhere such as lowering for copyin clause D127468.

[arnamoy10]

Rebasing. Also did not attempt to move this before loop-bounds creation (will not be a simple fix IMHO), adding a TODO

[peixin]

flang/lib/Lower/Bridge.cpp

475

Can you fix the comment to describe that this is esepcially for loop bounds since loop bounds used old values which need to be fixed by using the new copied value?

478

Remove the TODO. The code cannot be moved there.

483

[arnamoy10]

Rebase, update comment and code to reflect reviewer's comments.

[peixin]

LGTM

[kiranchandramohan]

flang/lib/Lower/Bridge.cpp

478

Can we add a TODO to call privatization before the loop operation is created as a directly submitted NFC patch?

flang/lib/Lower/Bridge.cpp

fir::ExtendedValue exv = hexv.match(
         },
         [&](const auto &) -> fir::ExtendedValue {
           mlir::Value temp =
               allocate(fir::factory::getExtents(loc, *builder, hexv),
                        fir::factory::getTypeParams(loc, *builder, hexv));
           return fir::substBase(hexv, temp);
         });
 
+    // Replace all uses of the original with the clone/copy,
+    // esepcially for loop bounds (that uses the variable being privatised)
+    // since loop bounds use old values that need to be fixed by using the
+    // new copied value.
+    // Not able to use replaceAllUsesWith() because uses outside
+    // the loop body should not use the clone.
+    mlir::Region &curRegion = getFirOpBuilder().getRegion();
+    mlir::Value oldVal = fir::getBase(hexv);
+    mlir::Value cloneVal = fir::getBase(exv);
+    for (auto &oper : curRegion.getOps()) {
+      for (unsigned int ii = 0; ii < oper.getNumOperands(); ++ii) {
+        if (oper.getOperand(ii) == oldVal) {
+          oper.setOperand(ii, cloneVal);
+        }
+      }
+    }
     return bindIfNewSymbol(sym, exv);
   }
 
   void
   copyHostAssociateVar(const Fortran::semantics::Symbol &sym) override final {
     // 1) Fetch the original copy of the variable.
     assert(sym.has<Fortran::semantics::HostAssocDetails>() &&
            "No host-association found");
     const Fortran::semantics::Symbol &hsym = sym.GetUltimate();
     Fortran::lower::SymbolBox hsb = lookupOneLevelUpSymbol(hsym);
     assert(hsb && "Host symbol box not found");
     fir::ExtendedValue hexv = getExtendedValue(hsb);
 
     // 2) Fetch the copied one that will mask the original.
     Fortran::lower::SymbolBox sb = shallowLookupSymbol(sym);
     assert(sb && "Host-associated symbol box not found");
     assert(hsb.getAddr() != sb.getAddr() &&
            "Host and associated symbol boxes are the same");
     fir::ExtendedValue exv = getExtendedValue(sb);
 
     // 3) Perform the assignment.
     builder->setInsertionPointAfter(fir::getBase(exv).getDefiningOp());
     mlir::Location loc = genLocation(sym.name());
     mlir::Type symType = genType(sym);
     if (auto seqTy = symType.dyn_cast<fir::SequenceType>()) {
       Fortran::lower::StatementContext stmtCtx;
       Fortran::lower::createSomeArrayAssignment(*this, exv, hexv, localSymbols,
                                                 stmtCtx);
       stmtCtx.finalize();
     } else if (hexv.getBoxOf<fir::CharBoxValue>()) {
       fir::factory::CharacterExprHelper{*builder, loc}.createAssign(exv, hexv);
     } else if (hexv.getBoxOf<fir::MutableBoxValue>()) {
       TODO(loc, "firstprivatisation of allocatable variables");
     } else {
       auto loadVal = builder->create<fir::LoadOp>(loc, fir::getBase(hexv));
       builder->create<fir::StoreOp>(loc, loadVal, fir::getBase(exv));
     }
   }
 

flang/test/Lower/OpenMP/omp-parallel-wsloop-firstpriv.f90

+! This test checks lowering of OpenMP parallel DO, with the loop bound being
+! a firstprivate variable
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+! CHECK: func @_QPomp_do_firstprivate(%[[ARG0:.*]]: !fir.ref<i32> {fir.bindc_name = "a"}) 
+subroutine omp_do_firstprivate(a)
+  integer::a
+  integer::n
+  n = a+1
+  !$omp parallel do firstprivate(a)
+  ! CHECK:  omp.parallel {
+  ! CHECK-NEXT: %[[CLONE:.*]] = fir.alloca i32 {bindc_name = "a", pinned
+  ! CHECK-NEXT: %[[LD:.*]] = fir.load %[[ARG0]] : !fir.ref<i32>
+  ! CHECK-NEXT: fir.store %[[LD]] to %[[CLONE]] : !fir.ref<i32>
+  ! CHECK-NEXT: omp.barrier
+  ! CHECK-NEXT: %[[REF:.*]] = fir.alloca i32 {adapt.valuebyref, pinned}
+  ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
+  ! CHECK-NEXT: %[[UB:.*]] = fir.load %[[CLONE]] : !fir.ref<i32>
+  ! CHECK-NEXT: %[[STEP:.*]] = arith.constant 1 : i32
+  ! CHECK-NEXT: omp.wsloop   for  (%[[ARG1:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]])
+  ! CHECK-NEXT: fir.store %[[ARG1]] to %[[REF]] : !fir.ref<i32>
+  ! CHECK-NEXT: fir.call @_QPfoo(%[[REF]], %[[CLONE]]) : (!fir.ref<i32>, !fir.ref<i32>) -> ()
+  ! CHECK-NEXT: omp.yield
+    do i=1, a
+      call foo(i, a)
+    end do
+  !$omp end parallel do
+  !CHECK: fir.call @_QPbar(%[[ARG0]]) : (!fir.ref<i32>) -> ()
+  call bar(a)
+end subroutine omp_do_firstprivate
+
+! CHECK: func @_QPomp_do_firstprivate2(%[[ARG0:.*]]: !fir.ref<i32> {fir.bindc_name = "a"}, %[[ARG1:.*]]: !fir.ref<i32> {fir.bindc_name = "n"}) 
+subroutine omp_do_firstprivate2(a, n)
+  integer::a
+  integer::n
+  n = a+1
+  !$omp parallel do firstprivate(a, n)
+  ! CHECK:  omp.parallel {
+  ! CHECK-NEXT: %[[CLONE:.*]] = fir.alloca i32 {bindc_name = "a", pinned
+  ! CHECK-NEXT: %[[LD:.*]] = fir.load %[[ARG0]] : !fir.ref<i32>
+  ! CHECK-NEXT: fir.store %[[LD]] to %[[CLONE]] : !fir.ref<i32>
+  ! CHECK-NEXT: %[[CLONE1:.*]] = fir.alloca i32 {bindc_name = "n", pinned
+  ! CHECK-NEXT: %[[LD1:.*]] = fir.load %[[ARG1]] : !fir.ref<i32>
+  ! CHECK-NEXT: fir.store %[[LD1]] to %[[CLONE1]] : !fir.ref<i32>
+  ! CHECK-NEXT: omp.barrier
+  ! CHECK-NEXT: %[[REF:.*]] = fir.alloca i32 {adapt.valuebyref, pinned}
+
+
+  ! CHECK: %[[LB:.*]] = fir.load %[[CLONE]] : !fir.ref<i32>
+  ! CHECK-NEXT: %[[UB:.*]] = fir.load %[[CLONE1]] : !fir.ref<i32>
+  ! CHECK-NEXT: %[[STEP:.*]] = arith.constant 1 : i32
+  ! CHECK-NEXT: omp.wsloop   for  (%[[ARG2:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) inclusive step (%[[STEP]])
+  ! CHECK-NEXT: fir.store %[[ARG2]] to %[[REF]] : !fir.ref<i32>
+  ! CHECK-NEXT: fir.call @_QPfoo(%[[REF]], %[[CLONE]]) : (!fir.ref<i32>, !fir.ref<i32>) -> ()
+  ! CHECK-NEXT: omp.yield
+    do i= a, n
+      call foo(i, a)
+    end do
+  !$omp end parallel do
+  !CHECK: fir.call @_QPbar(%[[ARG1]]) : (!fir.ref<i32>) -> ()
+  call bar(n)
+end subroutine omp_do_firstprivate2