Merge pull request #1098 from schweitzpgi/ch-where

Implement the correct splitting of FORALL and array expression loops

Author: schweitzpgi

flang/include/flang/Lower/ConvertExpr.h

@@ -17,7 +17,7 @@
 #ifndef FORTRAN_LOWER_CONVERTEXPR_H
 #define FORTRAN_LOWER_CONVERTEXPR_H
 
-#include "flang/Evaluate/shape.h"
+#include "flang/Evaluate/expression.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 
@@ -169,10 +169,24 @@ createSomeArrayTempValue(AbstractConverter &converter,
                          const evaluate::Expr<evaluate::SomeType> &expr,
                          SymMap &symMap, StatementContext &stmtCtx);
 
+// Lambda to reload the dynamically allocated pointers to a lazy buffer and its
+// extents. This is used to introduce these ssa-values in a place that will
+// dominate any/all subsequent uses after the loop that created the lazy buffer.
+using LoadLazyBufferLambda =
+    std::function<std::pair<fir::ExtendedValue, mlir::Value>(
+        fir::FirOpBuilder &)>;
+
+// Creating a lazy array temporary returns a pair of values. The first is an
+// extended value which is a pointer to the buffer, of array type, with the
+// appropriate dynamic extents. The second argument is a continuation to reload
+// the buffer at some future point in the code gen.
+using CreateLazyArrayResult =
+    std::pair<fir::ExtendedValue, LoadLazyBufferLambda>;
+
 /// Like createSomeArrayTempValue, but the temporary buffer is allocated lazily
 /// (inside the loops instead of before the loops). This can be useful if a
 /// loop's bounds are functions of other loop indices, for example.
-fir::ExtendedValue
+CreateLazyArrayResult
 createLazyArrayTempValue(AbstractConverter &converter,
                          const evaluate::Expr<evaluate::SomeType> &expr,
                          mlir::Value var, mlir::Value shapeBuffer,

flang/lib/Lower/Bridge.cpp

@@ -1235,6 +1235,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         mlir::Value by;
         if (outermost) {
           assert(headerIndex < lows.size());
+          if (headerIndex == 0)
+            explicitIterSpace.resetInnerArgs();
           lb = lows[headerIndex];
           ub = highs[headerIndex];
           by = steps[headerIndex++];
@@ -1261,7 +1263,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         forceControlVariableBinding(ctrlVar, lp.getInductionVar());
         loops.push_back(lp);
       }
-      explicitIterSpace.setOuterLoop(loops[0]);
+      if (outermost)
+         explicitIterSpace.setOuterLoop(loops[0]);
       if (const auto &mask =
               std::get<std::optional<Fortran::parser::ScalarLogicalExpr>>(
                   header.t);
@@ -2730,13 +2733,14 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     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);
-      builder->genIfThen(loc, cmp)
-          .genThen([&]() { builder->create<fir::FreeMemOp>(loc, load); })
-          .end();
-    });
+    explicitIterSpace.outermostContext().attachCleanup(
+        [builder = this->builder, loc, var]() {
+          auto load = builder->create<fir::LoadOp>(loc, var);
+          auto cmp = builder->genIsNotNull(loc, load);
+          builder->genIfThen(loc, cmp)
+              .genThen([&]() { builder->create<fir::FreeMemOp>(loc, load); })
+              .end();
+        });
   }
 
   //===--------------------------------------------------------------------===//

flang/lib/Lower/ConvertExpr.cpp

@@ -255,11 +255,13 @@ class ExplicitIterSpace {
       innerArgs.push_back(arg);
   }
 
-  void setOuterLoop(fir::DoLoopOp loop) {
-    if (!outerLoop.hasValue())
-      outerLoop = loop;
-  }
+  /// Reset the outermost `array_load` arguments to the loop nest.
+  void resetInnerArgs() { innerArgs = initialArgs; }
+
+  /// Capture the current outermost loop.
+  void setOuterLoop(fir::DoLoopOp loop) { outerLoop = loop; }
 
+  /// Sets the inner loop argument at position \p offset to \p val.
   void setInnerArg(size_t offset, mlir::Value val) {
     assert(offset < innerArgs.size());
     innerArgs[offset] = val;
@@ -385,6 +387,7 @@ class ExplicitIterSpace {
   // Assignment statement context (inside the loop nest).
   StatementContext stmtCtx;
   llvm::SmallVector<mlir::Value> innerArgs;
+  llvm::SmallVector<mlir::Value> initialArgs;
   llvm::Optional<fir::DoLoopOp> outerLoop;
   llvm::Optional<std::function<void(fir::FirOpBuilder &)>> loopCleanup;
   std::size_t forallContextOpen = 0;

flang/lib/Lower/IterationSpace.h

@@ -71,7 +71,8 @@ Fortran::lower::operator<<(llvm::raw_ostream &os,
   for (auto i : llvm::enumerate(symMap.symbolMapStack)) {
     os << " level " << i.index() << "<{\n";
     for (auto iter : i.value())
-      os << "  symbol [" << *iter.first << "] ->\n    " << iter.second;
+      os << "  symbol @" << (void *)iter.first << " [" << *iter.first
+         << "] ->\n    " << iter.second;
     os << " }>\n";
   }
   return os;

flang/lib/Lower/SymbolMap.cpp

@@ -96,23 +96,23 @@ subroutine slice_with_explicit_iters
   ! CHECK:           %[[VAL_25:.*]] = divi_signed %[[VAL_24]], %[[VAL_19]] : index
   ! CHECK:           %[[VAL_26:.*]] = cmpi sgt, %[[VAL_25]], %[[VAL_20]] : index
   ! CHECK:           %[[VAL_27:.*]] = select %[[VAL_26]], %[[VAL_25]], %[[VAL_20]] : index
-  ! CHECK:           %[[VAL_28:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
-  ! CHECK:           %[[VAL_29:.*]] = constant 1 : index
-  ! CHECK:           %[[VAL_30:.*]] = constant 0 : index
-  ! CHECK:           %[[VAL_31:.*]] = subi %[[VAL_27]], %[[VAL_29]] : index
-  ! CHECK:           %[[VAL_32:.*]] = fir.do_loop %[[VAL_33:.*]] = %[[VAL_30]] to %[[VAL_31]] step %[[VAL_29]] unordered iter_args(%[[VAL_34:.*]] = %[[VAL_10]]) -> (!fir.array<10x10xi32>) {
+  ! CHECK:           %[[VAL_28:.*]] = constant 1 : index
+  ! CHECK:           %[[VAL_29:.*]] = constant 0 : index
+  ! CHECK:           %[[VAL_30:.*]] = subi %[[VAL_27]], %[[VAL_28]] : index
+  ! CHECK:           %[[VAL_31:.*]] = fir.do_loop %[[VAL_32:.*]] = %[[VAL_29]] to %[[VAL_30]] step %[[VAL_28]] unordered iter_args(%[[VAL_33:.*]] = %[[VAL_13]]) -> (!fir.array<10x10xi32>) {
+  ! CHECK:             %[[VAL_34:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
   ! CHECK:             %[[VAL_35:.*]] = constant 0 : i32
-  ! CHECK:             %[[VAL_36:.*]] = subi %[[VAL_35]], %[[VAL_28]] : i32
+  ! CHECK:             %[[VAL_36:.*]] = subi %[[VAL_35]], %[[VAL_34]] : i32
   ! CHECK:             %[[VAL_37:.*]] = constant 1 : i64
   ! CHECK:             %[[VAL_38:.*]] = fir.convert %[[VAL_37]] : (i64) -> index
   ! CHECK:             %[[VAL_39:.*]] = constant 1 : i64
   ! CHECK:             %[[VAL_40:.*]] = fir.convert %[[VAL_39]] : (i64) -> index
-  ! CHECK:             %[[VAL_41:.*]] = muli %[[VAL_33]], %[[VAL_40]] : index
+  ! CHECK:             %[[VAL_41:.*]] = muli %[[VAL_32]], %[[VAL_40]] : index
   ! CHECK:             %[[VAL_42:.*]] = addi %[[VAL_38]], %[[VAL_41]] : index
   ! CHECK:             %[[VAL_43:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
   ! CHECK:             %[[VAL_44:.*]] = fir.convert %[[VAL_43]] : (i32) -> i64
   ! CHECK:             %[[VAL_45:.*]] = fir.convert %[[VAL_44]] : (i64) -> index
-  ! CHECK:             %[[VAL_46:.*]] = fir.array_update %[[VAL_13]], %[[VAL_36]], %[[VAL_42]], %[[VAL_45]] {Fortran.offsets} : (!fir.array<10x10xi32>, i32, index, index) -> !fir.array<10x10xi32>
+  ! CHECK:             %[[VAL_46:.*]] = fir.array_update %[[VAL_33]], %[[VAL_36]], %[[VAL_42]], %[[VAL_45]] {Fortran.offsets} : (!fir.array<10x10xi32>, i32, index, index) -> !fir.array<10x10xi32>
   ! CHECK:             fir.result %[[VAL_46]] : !fir.array<10x10xi32>
   ! CHECK:           }
   ! CHECK:           fir.result %[[VAL_47:.*]] : !fir.array<10x10xi32>