Fix evaluation semantics of FORALL constructs per 10.2.4.2.4.

When lowering of FORALL was refactored to split the work between the
bridge and array expression lowering, the individual statements in a
forall construct body were no longer being lowered one at a time. These
changes add back the proper lowering.

Changes include:

  - Factoring the rhs array base analysis to be by statement.
  - Fix bug with finalization of context stack.
  - Fix bug with rerunning the analysis.
  - Merge aspects of cleanup code gen.
  - Restructure loop nest lowering such that each assignment resides in
    its own copy of a loop nest.
  - Add a lazy shape buffer. Thread lazy mask buffers through the loop
    nest so that cached results are available.
  - Regenerate checks for tests.
  - Fixes for p9.

Make forall-2.f90 source file name independent.

Fix for k6 test, among others. Using incorrect statement context.

Fix bug exposed by m7.

Fix m7 bug.

review comments

Author: schweitzpgi

flang/include/flang/Lower/ConvertExpr.h

@@ -176,8 +176,8 @@ createSomeArrayTempValue(AbstractConverter &converter,
 fir::ExtendedValue
 createLazyArrayTempValue(AbstractConverter &converter,
                          const evaluate::Expr<evaluate::SomeType> &expr,
-                         mlir::Value var, SymMap &symMap,
-                         StatementContext &stmtCtx);
+                         mlir::Value var, mlir::Value shapeBuffer,
+                         SymMap &symMap, StatementContext &stmtCtx);
 
 /// Lower an array expression to a value of type box. The expression must be a
 /// variable.

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -248,10 +248,10 @@ class FirOpBuilder : public mlir::OpBuilder {
   }
 
   /// Construct one of the two forms of shape op from an array box.
-  mlir::Value consShape(mlir::Location loc, const fir::AbstractArrayBox &arr);
-  mlir::Value consShape(mlir::Location loc, llvm::ArrayRef<mlir::Value> shift,
+  mlir::Value genShape(mlir::Location loc, const fir::AbstractArrayBox &arr);
+  mlir::Value genShape(mlir::Location loc, llvm::ArrayRef<mlir::Value> shift,
                         llvm::ArrayRef<mlir::Value> exts);
-  mlir::Value consShape(mlir::Location loc, llvm::ArrayRef<mlir::Value> exts);
+  mlir::Value genShape(mlir::Location loc, llvm::ArrayRef<mlir::Value> exts);
 
   /// Create one of the shape ops given an extended value. For a boxed value,
   /// this may create a `fir.shift` op.

flang/lib/Lower/Bridge.cpp

@@ -1211,80 +1211,123 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   /// Process a concurrent header for a FORALL. (Concurrent headers for DO
   /// CONCURRENT loops are lowered elsewhere.)
   void genFIR(const Fortran::parser::ConcurrentHeader &header) {
-    // Create our iteration space from the header spec.
-    localSymbols.pushScope();
-    auto idxTy = builder->getIndexType();
-    auto loc = toLocation();
-    llvm::SmallVector<fir::DoLoopOp> loops;
-    for (auto &ctrl :
-         std::get<std::list<Fortran::parser::ConcurrentControl>>(header.t)) {
-      const auto *ctrlVar = std::get<Fortran::parser::Name>(ctrl.t).symbol;
-      const auto *lo = Fortran::semantics::GetExpr(std::get<1>(ctrl.t));
-      const auto *hi = Fortran::semantics::GetExpr(std::get<2>(ctrl.t));
-      auto &optStep =
-          std::get<std::optional<Fortran::parser::ScalarIntExpr>>(ctrl.t);
-      auto lb = builder->createConvert(
-          loc, idxTy,
-          fir::getBase(genExprValue(*lo, explicitIterSpace.stmtContext())));
-      auto ub = builder->createConvert(
-          loc, idxTy,
-          fir::getBase(genExprValue(*hi, explicitIterSpace.stmtContext())));
-      auto by = optStep.has_value()
-                    ? builder->createConvert(
-                          loc, idxTy,
-                          fir::getBase(genExprValue(
-                              *Fortran::semantics::GetExpr(*optStep),
-                              explicitIterSpace.stmtContext())))
-                    : builder->createIntegerConstant(loc, idxTy, 1);
-      auto lp = builder->create<fir::DoLoopOp>(
-          loc, lb, ub, by, /*unordered=*/true,
-          /*finalCount=*/false, explicitIterSpace.getInnerArgs());
-      if (!loops.empty())
-        builder->create<fir::ResultOp>(loc, lp.getResults());
-      explicitIterSpace.setInnerArgs(lp.getRegionIterArgs());
-      builder->setInsertionPointToStart(lp.getBody());
-      forceControlVariableBinding(ctrlVar, lp.getInductionVar());
-      loops.push_back(lp);
-    }
-    explicitIterSpace.setOuterLoop(loops[0]);
-    if (const auto &mask =
-            std::get<std::optional<Fortran::parser::ScalarLogicalExpr>>(
-                header.t);
-        mask.has_value()) {
-      auto i1Ty = builder->getI1Type();
-      auto maskExv = genExprValue(*Fortran::semantics::GetExpr(mask.value()),
-                                  explicitIterSpace.stmtContext());
-      auto cond = builder->createConvert(loc, i1Ty, fir::getBase(maskExv));
-      auto ifOp = builder->create<fir::IfOp>(
-          loc, explicitIterSpace.innerArgTypes(), cond,
-          /*withElseRegion=*/true);
-      builder->create<fir::ResultOp>(loc, ifOp.getResults());
-      builder->setInsertionPointToStart(&ifOp.elseRegion().front());
-      builder->create<fir::ResultOp>(loc, explicitIterSpace.getInnerArgs());
-      builder->setInsertionPointToStart(&ifOp.thenRegion().front());
+    llvm::SmallVector<mlir::Value> lows;
+    llvm::SmallVector<mlir::Value> highs;
+    llvm::SmallVector<mlir::Value> steps;
+    if (explicitIterSpace.isOutermostForall()) {
+      // For the outermost forall, we evaluate the bounds expressions once.
+      // Contrastingly, if this forall is nested, the bounds expressions are
+      // assumed to be pure, possibly dependent on outer concurrent control
+      // variables, possibly variant with respect to arguments, and will be
+      // re-evaluated.
+      auto loc = toLocation();
+      auto idxTy = builder->getIndexType();
+      auto &stmtCtx = explicitIterSpace.stmtContext();
+      auto lowerExpr = [&](auto &e) {
+        return fir::getBase(genExprValue(e, stmtCtx));
+      };
+      for (auto &ctrl :
+           std::get<std::list<Fortran::parser::ConcurrentControl>>(header.t)) {
+        const auto *lo = Fortran::semantics::GetExpr(std::get<1>(ctrl.t));
+        const auto *hi = Fortran::semantics::GetExpr(std::get<2>(ctrl.t));
+        auto &optStep =
+            std::get<std::optional<Fortran::parser::ScalarIntExpr>>(ctrl.t);
+        lows.push_back(builder->createConvert(loc, idxTy, lowerExpr(*lo)));
+        highs.push_back(builder->createConvert(loc, idxTy, lowerExpr(*hi)));
+        steps.push_back(
+            optStep.has_value()
+                ? builder->createConvert(
+                      loc, idxTy,
+                      lowerExpr(*Fortran::semantics::GetExpr(*optStep)))
+                : builder->createIntegerConstant(loc, idxTy, 1));
+      }
     }
+    auto lambda = [&, lows, highs, steps]() {
+      // Create our iteration space from the header spec.
+      auto loc = toLocation();
+      auto idxTy = builder->getIndexType();
+      llvm::SmallVector<fir::DoLoopOp> loops;
+      auto &stmtCtx = explicitIterSpace.stmtContext();
+      auto lowerExpr = [&](auto &e) {
+        return fir::getBase(genExprValue(e, stmtCtx));
+      };
+      const auto outermost = !lows.empty();
+      std::size_t headerIndex = 0;
+      for (auto &ctrl :
+           std::get<std::list<Fortran::parser::ConcurrentControl>>(header.t)) {
+        const auto *ctrlVar = std::get<Fortran::parser::Name>(ctrl.t).symbol;
+        mlir::Value lb;
+        mlir::Value ub;
+        mlir::Value by;
+        if (outermost) {
+          assert(headerIndex < lows.size());
+          lb = lows[headerIndex];
+          ub = highs[headerIndex];
+          by = steps[headerIndex++];
+        } else {
+          const auto *lo = Fortran::semantics::GetExpr(std::get<1>(ctrl.t));
+          const auto *hi = Fortran::semantics::GetExpr(std::get<2>(ctrl.t));
+          auto &optStep =
+              std::get<std::optional<Fortran::parser::ScalarIntExpr>>(ctrl.t);
+          lb = builder->createConvert(loc, idxTy, lowerExpr(*lo));
+          ub = builder->createConvert(loc, idxTy, lowerExpr(*hi));
+          by = optStep.has_value()
+                   ? builder->createConvert(
+                         loc, idxTy,
+                         lowerExpr(*Fortran::semantics::GetExpr(*optStep)))
+                   : builder->createIntegerConstant(loc, idxTy, 1);
+        }
+        auto lp = builder->create<fir::DoLoopOp>(
+            loc, lb, ub, by, /*unordered=*/true,
+            /*finalCount=*/false, explicitIterSpace.getInnerArgs());
+        if (!loops.empty() || !outermost)
+          builder->create<fir::ResultOp>(loc, lp.getResults());
+        explicitIterSpace.setInnerArgs(lp.getRegionIterArgs());
+        builder->setInsertionPointToStart(lp.getBody());
+        forceControlVariableBinding(ctrlVar, lp.getInductionVar());
+        loops.push_back(lp);
+      }
+      explicitIterSpace.setOuterLoop(loops[0]);
+      if (const auto &mask =
+              std::get<std::optional<Fortran::parser::ScalarLogicalExpr>>(
+                  header.t);
+          mask.has_value()) {
+        auto i1Ty = builder->getI1Type();
+        auto maskExv =
+            genExprValue(*Fortran::semantics::GetExpr(mask.value()), stmtCtx);
+        auto cond = builder->createConvert(loc, i1Ty, fir::getBase(maskExv));
+        auto ifOp = builder->create<fir::IfOp>(
+            loc, explicitIterSpace.innerArgTypes(), cond,
+            /*withElseRegion=*/true);
+        builder->create<fir::ResultOp>(loc, ifOp.getResults());
+        builder->setInsertionPointToStart(&ifOp.elseRegion().front());
+        builder->create<fir::ResultOp>(loc, explicitIterSpace.getInnerArgs());
+        builder->setInsertionPointToStart(&ifOp.thenRegion().front());
+      }
+    };
+    // Push the lambda to gen the loop nest context.
+    explicitIterSpace.pushLoopNest(lambda);
   }
 
   void genFIR(const Fortran::parser::ForallAssignmentStmt &stmt) {
     std::visit([&](const auto &x) { genFIR(x); }, stmt.u);
   }
 
   void genFIR(const Fortran::parser::EndForallStmt &) {
-    explicitIterSpace.finalize();
     cleanupExplicitSpace();
   }
 
   template <typename A>
   void prepareExplicitSpace(const A &forall) {
-    analyzeExplicitSpace(forall);
+    if (!explicitIterSpace.isActive())
+      analyzeExplicitSpace(forall);
+    localSymbols.pushScope();
     explicitIterSpace.enter();
-    Fortran::lower::createArrayLoads(*this, explicitIterSpace, localSymbols);
   }
 
   /// Cleanup all the FORALL context information when we exit.
   void cleanupExplicitSpace() {
-    Fortran::lower::createArrayMergeStores(*this, explicitIterSpace);
-    explicitIterSpace.conditionalCleanup();
+    explicitIterSpace.leave();
     localSymbols.popScope();
   }
 
@@ -1824,6 +1867,10 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   void genAssignment(const Fortran::evaluate::Assignment &assign) {
     Fortran::lower::StatementContext stmtCtx;
     auto loc = toLocation();
+    if (explicitIterationSpace()) {
+      Fortran::lower::createArrayLoads(*this, explicitIterSpace, localSymbols);
+      explicitIterSpace.genLoopNest();
+    }
     std::visit(
         Fortran::common::visitors{
             // [1] Plain old assignment.
@@ -1920,7 +1967,10 @@ class FirConverter : public Fortran::lower::AbstractConverter {
               if (implicitIterationSpace())
                 TODO(loc, "user defined assignment within WHERE");
               Fortran::semantics::SomeExpr expr{procRef};
-              createFIRExpr(toLocation(), &expr, stmtCtx);
+              createFIRExpr(toLocation(), &expr,
+                            explicitIterationSpace()
+                                ? explicitIterSpace.stmtContext()
+                                : stmtCtx);
             },
 
             // [3] Pointer assignment with possibly empty bounds-spec. R1035: a
@@ -1981,6 +2031,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
             },
         },
         assign.u);
+    if (explicitIterationSpace())
+      Fortran::lower::createArrayMergeStores(*this, explicitIterSpace);
   }
 
   void genFIR(const Fortran::parser::WhereConstruct &c) {
@@ -2563,6 +2615,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   void analyzeExplicitSpace(const Fortran::evaluate::Assignment *assign) {
     analyzeExplicitSpace</*LHS=*/true>(assign->lhs);
     analyzeExplicitSpace(assign->rhs);
+    explicitIterSpace.endAssign();
   }
   void analyzeExplicitSpace(const Fortran::parser::ForallAssignmentStmt &stmt) {
     std::visit([&](const auto &s) { analyzeExplicitSpace(s); }, stmt.u);
@@ -2693,7 +2746,11 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     auto var = builder->createTemporary(loc, ty);
     auto nil = builder->createNullConstant(loc, ty);
     builder->create<fir::StoreOp>(loc, nil, var);
-    implicitIterSpace.addMaskVariable(exp, var);
+    auto shTy = fir::HeapType::get(builder->getIndexType());
+    auto shape = builder->createTemporary(loc, shTy);
+    auto nilSh = builder->createNullConstant(loc, shTy);
+    builder->create<fir::StoreOp>(loc, nilSh, shape);
+    implicitIterSpace.addMaskVariable(exp, var, shape);
     explicitIterSpace.outermostContext().attachCleanup([=]() {
       auto load = builder->create<fir::LoadOp>(loc, var);
       auto cmp = builder->genIsNotNull(loc, load);