[flang][OpenMP] Privatize locally destroyed values in do concurent (#112)

Collects values that are local to a loop: "loop-local values". A loop-local
value is one that is used exclusively inside the loop but allocated outside
of it. This usually corresponds to temporary values that are used inside the
loop body for initialzing other variables for example. For a "loop-local" 
value within a loop's scope, localizes that value within the scope of the 
parallel region the loop maps to.

Author: ergawy

flang/docs/DoConcurrentConversionToOpenMP.md

@@ -234,9 +234,16 @@ see the "Data environment" section below.
 By default, variables that are used inside a `do concurernt` loop nest are
 either treated as `shared` in case of mapping to `host`, or mapped into the
 `target` region using a `map` clause in case of mapping to `device`. The only
-exception to this is the loop's iteration variable(s) (IV) of **perfect** loop
-nest. In that case, for each IV, we allocate a local copy as shown the by the
-mapping examples above.
+exceptions to this are:
+  1. the loop's iteration variable(s) (IV) of **perfect** loop nests. In that
+     case, for each IV, we allocate a local copy as shown the by the mapping
+     examples above.
+  1. any values that are from allocations outside the loop nest and used
+     exclusively inside of it. In such cases, a local privatized
+     value is created in the OpenMP region to prevent multiple teams of threads
+     from accessing and destroying the same memory block which causes runtime
+     issues. For an example of such cases, see
+     `flang/test/Transforms/DoConcurrent/locally_destroyed_temp.f90`.
 
 #### Non-perfectly-nested loops' IVs
 

flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp

@@ -15,7 +15,10 @@
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/IRMapping.h"
@@ -468,6 +471,61 @@ void sinkLoopIVArgs(mlir::ConversionPatternRewriter &rewriter,
     ++idx;
   }
 }
+
+/// Collects values that are local to a loop: "loop-local values". A loop-local
+/// value is one that is used exclusively inside the loop but allocated outside
+/// of it. This usually corresponds to temporary values that are used inside the
+/// loop body for initialzing other variables for example.
+///
+/// \param [in] doLoop - the loop within which the function searches for values
+/// used exclusively inside.
+///
+/// \param [out] locals - the list of loop-local values detected for \p doLoop.
+static void collectLoopLocalValues(fir::DoLoopOp doLoop,
+                                   llvm::SetVector<mlir::Value> &locals) {
+  doLoop.walk([&](mlir::Operation *op) {
+    for (mlir::Value operand : op->getOperands()) {
+      if (locals.contains(operand))
+        continue;
+
+      bool isLocal = true;
+
+      if (!mlir::isa_and_present<fir::AllocaOp>(operand.getDefiningOp()))
+        continue;
+
+      // Values defined inside the loop are not interesting since they do not
+      // need to be localized.
+      if (doLoop->isAncestor(operand.getDefiningOp()))
+        continue;
+
+      for (auto *user : operand.getUsers()) {
+        if (!doLoop->isAncestor(user)) {
+          isLocal = false;
+          break;
+        }
+      }
+
+      if (isLocal)
+        locals.insert(operand);
+    }
+  });
+}
+
+/// For a "loop-local" value \p local within a loop's scope, localizes that
+/// value within the scope of the parallel region the loop maps to. Towards that
+/// end, this function moves the allocation of \p local within \p allocRegion.
+///
+/// \param local - the value used exclusively within a loop's scope (see
+/// collectLoopLocalValues).
+///
+/// \param allocRegion - the parallel region where \p local's allocation will be
+/// privatized.
+///
+/// \param rewriter - builder used for updating \p allocRegion.
+static void localizeLoopLocalValue(mlir::Value local, mlir::Region &allocRegion,
+                            mlir::ConversionPatternRewriter &rewriter) {
+  rewriter.moveOpBefore(local.getDefiningOp(), &allocRegion.front().front());
+}
 } // namespace looputils
 
 class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
@@ -519,9 +577,13 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
     bool hasRemainingNestedLoops =
         failed(looputils::collectLoopNest(doLoop, loopNest));
 
+    mlir::IRMapping mapper;
+
+    llvm::SetVector<mlir::Value> locals;
+    looputils::collectLoopLocalValues(loopNest.back().first, locals);
+
     looputils::sinkLoopIVArgs(rewriter, loopNest);
 
-    mlir::IRMapping mapper;
     mlir::omp::TargetOp targetOp;
     mlir::omp::LoopNestClauseOps loopNestClauseOps;
 
@@ -541,8 +603,13 @@ class DoConcurrentConversion : public mlir::OpConversionPattern<fir::DoLoopOp> {
       genDistributeOp(doLoop.getLoc(), rewriter);
     }
 
-    genParallelOp(doLoop.getLoc(), rewriter, loopNest, mapper,
-                  loopNestClauseOps);
+    mlir::omp::ParallelOp parallelOp = genParallelOp(
+        doLoop.getLoc(), rewriter, loopNest, mapper, loopNestClauseOps);
+
+    for (mlir::Value local : locals)
+      looputils::localizeLoopLocalValue(local, parallelOp.getRegion(),
+                                        rewriter);
+
     mlir::omp::LoopNestOp ompLoopNest =
         genWsLoopOp(rewriter, loopNest.back().first, mapper, loopNestClauseOps);
 
@@ -919,9 +986,10 @@ class DoConcurrentConversionPass
         context, mapTo == fir::omp::DoConcurrentMappingKind::DCMK_Device,
         concurrentLoopsToSkip);
     mlir::ConversionTarget target(*context);
-    target.addLegalDialect<fir::FIROpsDialect, hlfir::hlfirDialect,
-                           mlir::arith::ArithDialect, mlir::func::FuncDialect,
-                           mlir::omp::OpenMPDialect>();
+    target.addLegalDialect<
+        fir::FIROpsDialect, hlfir::hlfirDialect, mlir::arith::ArithDialect,
+        mlir::func::FuncDialect, mlir::omp::OpenMPDialect,
+        mlir::cf::ControlFlowDialect, mlir::math::MathDialect>();
 
     target.addDynamicallyLegalOp<fir::DoLoopOp>([&](fir::DoLoopOp op) {
       return !op.getUnordered() || concurrentLoopsToSkip.contains(op);

flang/test/Transforms/DoConcurrent/locally_destroyed_temp.f90

@@ -0,0 +1,66 @@
+! Tests that locally destroyed values in a `do concurrent` loop are properly
+! handled. Locally destroyed values are those values for which the Fortran runtime
+! calls `@_FortranADestroy` inside the loops body. If these values are allocated
+! outside the loop, and the loop is mapped to OpenMP, then a runtime error would
+! occur due to multiple teams trying to access the same allocation.
+
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fdo-concurrent-parallel=host %s -o - \
+! RUN:   | FileCheck %s
+
+module struct_mod
+    type test_struct
+        integer, allocatable :: x_
+    end type
+
+    interface test_struct
+        pure module function construct_from_components(x) result(struct)
+            implicit none
+            integer, intent(in) :: x
+            type(test_struct) struct
+        end function
+    end interface
+end module
+
+submodule(struct_mod) struct_sub
+    implicit none
+
+contains
+    module procedure construct_from_components
+        struct%x_ = x
+    end procedure
+end submodule struct_sub
+
+program main
+    use struct_mod, only : test_struct
+
+    implicit none
+    type(test_struct), dimension(10) :: a
+    integer :: i
+    integer :: total
+
+    do concurrent (i=1:10)
+        a(i) = test_struct(i)
+    end do
+
+    do i=1,10
+        total = total + a(i)%x_
+    end do
+
+    print *, "total =", total
+end program main
+
+! CHECK: omp.parallel {
+! CHECK:   %[[LOCAL_TEMP:.*]] = fir.alloca !fir.type<_QMstruct_modTtest_struct{x_:!fir.box<!fir.heap<i32>>}> {bindc_name = ".result"}
+! CHECK:   omp.wsloop {
+! CHECK:     omp.loop_nest {{.*}} {
+! CHECK:       %[[TEMP_VAL:.*]] = fir.call @_QMstruct_modPconstruct_from_components
+! CHECK:       fir.save_result %[[TEMP_VAL]] to %[[LOCAL_TEMP]]
+! CHECK:       %[[EMBOXED_LOCAL:.*]] = fir.embox %[[LOCAL_TEMP]]
+! CHECK:       %[[CONVERTED_LOCAL:.*]] = fir.convert %[[EMBOXED_LOCAL]]
+! CHECK:       fir.call @_FortranADestroy(%[[CONVERTED_LOCAL]])
+! CHECK:       omp.yield
+! CHECK:     }
+! CHECK:     omp.terminator
+! CHECK:   }
+! CHECK:   omp.terminator
+! CHECK: }