[flang] loop interchange from kruse

Author: ronlieb

flang/include/flang/Lower/AbstractConverter.h

@@ -401,6 +401,10 @@ class AbstractConverter {
 
   virtual mlir::StateStack &getStateStack() = 0;
 
+  virtual void
+  genPermutatedLoops(llvm::ArrayRef<Fortran::lower::pft::Evaluation *> doStmts,
+                     Fortran::lower::pft::Evaluation *innermostDo) = 0;
+
 private:
   /// Options controlling lowering behavior.
   const Fortran::lower::LoweringOptions &loweringOptions;

flang/include/flang/Lower/OpenMP/Utils.h

@@ -187,6 +187,11 @@ void collectTileSizesFromOpenMPConstruct(
     llvm::SmallVectorImpl<int64_t> &tileSizes,
     Fortran::semantics::SemanticsContext &semaCtx);
 
+void collectPermutationFromOpenMPConstruct(
+    const parser::OpenMPConstruct *ompCons,
+    llvm::SmallVectorImpl<int64_t> &permuation,
+    Fortran::semantics::SemanticsContext &semaCtx);
+
 } // namespace omp
 } // namespace lower
 } // namespace Fortran

flang/lib/Lower/Allocatable.cpp

@@ -37,6 +37,7 @@
 #include "flang/Runtime/pointer.h"
 #include "flang/Semantics/tools.h"
 #include "flang/Semantics/type.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
 #include "llvm/Support/CommandLine.h"
 

flang/lib/Lower/Bridge.cpp

@@ -2268,6 +2268,124 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     // so no clean-up needs to be generated for these entities.
   }
 
+  void
+  genPermutatedLoops(llvm::ArrayRef<Fortran::lower::pft::Evaluation *> doStmts,
+                     Fortran::lower::pft::Evaluation *innermostDo) override {
+    // Fortran::lower::pft::Evaluation &eval = getEval();
+    // bool unstructuredContext = eval.lowerAsUnstructured();
+
+    llvm::SmallVector<mlir::Block *> headerBlocks;
+    llvm::SmallVector<IncrementLoopNestInfo, 1> loopInfos;
+
+    auto enterLoop = [&](Fortran::lower::pft::Evaluation &eval) {
+      bool unstructuredContext = eval.lowerAsUnstructured();
+
+      // Collect loop nest information.
+      // Generate begin loop code directly for infinite and while loops.
+      Fortran::lower::pft::Evaluation &doStmtEval =
+          eval.getFirstNestedEvaluation();
+      auto *doStmt = doStmtEval.getIf<Fortran::parser::NonLabelDoStmt>();
+      const auto &loopControl =
+          std::get<std::optional<Fortran::parser::LoopControl>>(doStmt->t);
+      mlir::Block *preheaderBlock = doStmtEval.block;
+      mlir::Block *beginBlock =
+          preheaderBlock ? preheaderBlock : builder->getBlock();
+      auto createNextBeginBlock = [&]() {
+        // Step beginBlock through unstructured preheader, header, and mask
+        // blocks, created in outermost to innermost order.
+        return beginBlock = beginBlock->splitBlock(beginBlock->end());
+      };
+      mlir::Block *headerBlock =
+          unstructuredContext ? createNextBeginBlock() : nullptr;
+      headerBlocks.push_back(headerBlock);
+      mlir::Block *bodyBlock = doStmtEval.lexicalSuccessor->block;
+      mlir::Block *exitBlock = doStmtEval.parentConstruct->constructExit->block;
+      IncrementLoopNestInfo &incrementLoopNestInfo = loopInfos.emplace_back();
+      const Fortran::parser::ScalarLogicalExpr *whileCondition = nullptr;
+      bool infiniteLoop = !loopControl.has_value();
+      if (infiniteLoop) {
+        assert(unstructuredContext && "infinite loop must be unstructured");
+        startBlock(headerBlock);
+      } else if ((whileCondition =
+                      std::get_if<Fortran::parser::ScalarLogicalExpr>(
+                          &loopControl->u))) {
+        assert(unstructuredContext && "while loop must be unstructured");
+        maybeStartBlock(preheaderBlock); // no block or empty block
+        startBlock(headerBlock);
+        genConditionalBranch(*whileCondition, bodyBlock, exitBlock);
+      } else if (const auto *bounds =
+                     std::get_if<Fortran::parser::LoopControl::Bounds>(
+                         &loopControl->u)) {
+        // Non-concurrent increment loop.
+        IncrementLoopInfo &info = incrementLoopNestInfo.emplace_back(
+            *bounds->name.thing.symbol, bounds->lower, bounds->upper,
+            bounds->step);
+        if (unstructuredContext) {
+          maybeStartBlock(preheaderBlock);
+          info.hasRealControl = info.loopVariableSym->GetType()->IsNumeric(
+              Fortran::common::TypeCategory::Real);
+          info.headerBlock = headerBlock;
+          info.bodyBlock = bodyBlock;
+          info.exitBlock = exitBlock;
+        }
+      } else {
+        llvm_unreachable("Cannot permute DO CONCURRENT");
+      }
+
+      // Increment loop begin code. (Infinite/while code was already generated.)
+      if (!infiniteLoop && !whileCondition)
+        genFIRIncrementLoopBegin(incrementLoopNestInfo, doStmtEval.dirs);
+    };
+
+    auto leaveLoop = [&](Fortran::lower::pft::Evaluation &eval,
+                         mlir::Block *headerBlock,
+                         IncrementLoopNestInfo &incrementLoopNestInfo) {
+      bool unstructuredContext = eval.lowerAsUnstructured();
+
+      Fortran::lower::pft::Evaluation &doStmtEval =
+          eval.getFirstNestedEvaluation();
+      auto *doStmt = doStmtEval.getIf<Fortran::parser::NonLabelDoStmt>();
+
+      const auto &loopControl =
+          std::get<std::optional<Fortran::parser::LoopControl>>(doStmt->t);
+      bool infiniteLoop = !loopControl.has_value();
+      const Fortran::parser::ScalarLogicalExpr *whileCondition =
+          std::get_if<Fortran::parser::ScalarLogicalExpr>(&loopControl->u);
+
+      auto iter = std::prev(eval.getNestedEvaluations().end());
+
+      // An EndDoStmt in unstructured code may start a new block.
+      Fortran::lower::pft::Evaluation &endDoEval = *iter;
+      assert(endDoEval.getIf<Fortran::parser::EndDoStmt>() && "no enddo stmt");
+      if (unstructuredContext)
+        maybeStartBlock(endDoEval.block);
+
+      // Loop end code.
+      if (infiniteLoop || whileCondition)
+        genBranch(headerBlock);
+      else
+        genFIRIncrementLoopEnd(incrementLoopNestInfo);
+
+      // This call may generate a branch in some contexts.
+      genFIR(endDoEval, unstructuredContext);
+    };
+
+    for (auto l : doStmts)
+      enterLoop(*l);
+
+    // Loop body code.
+    bool innermostUnstructuredContext = innermostDo->lowerAsUnstructured();
+
+    auto iter = innermostDo->getNestedEvaluations().begin();
+    for (auto end = --innermostDo->getNestedEvaluations().end(); iter != end;
+         ++iter)
+      genFIR(*iter, innermostUnstructuredContext);
+
+    for (auto &&[l, headerBlock, li] :
+         llvm::zip_equal(doStmts, headerBlocks, loopInfos))
+      leaveLoop(*l, headerBlock, li);
+  }
+
   void attachInlineAttributes(
       mlir::Operation &op,
       const llvm::ArrayRef<const Fortran::parser::CompilerDirective *> &dirs) {

flang/lib/Lower/OpenMP/Decomposer.cpp

@@ -98,14 +98,49 @@ ConstructQueue buildConstructQueue(
   return decompose.output;
 }
 
+// from clang
+// There is a copy in check-omp-loops.cpp
+static bool isOpenMPLoopTransformationDirective(llvm::omp::Directive dir) {
+  switch (dir) {
+  // TODO case llvm::omp::Directive::OMPD_flatten:
+  case llvm::omp::Directive::OMPD_fuse:
+  case llvm::omp::Directive::OMPD_interchange:
+    // case llvm::omp::Directive::OMPD_nothing:
+  case llvm::omp::Directive::OMPD_reverse:
+  // TODO case llvm::omp::Directive::OMPD_split:
+  case llvm::omp::Directive::OMPD_stripe:
+  case llvm::omp::Directive::OMPD_tile:
+  case llvm::omp::Directive::OMPD_unroll:
+    return true;
+  default:
+    return false;
+  }
+}
+
+llvm::iterator_range<ConstructQueue::const_iterator> getNonTransformQueue(
+    llvm::iterator_range<ConstructQueue::const_iterator> range) {
+  // remove trailing loop transformations
+  auto b = range.begin();
+  auto e = range.end();
+  while (e != b) {
+    auto e2 = e - 1;
+    if (!isOpenMPLoopTransformationDirective(e2->id))
+      break;
+    e = e2;
+  }
+
+  return llvm::make_range(b, e);
+}
+
 bool matchLeafSequence(ConstructQueue::const_iterator item,
                        const ConstructQueue &queue,
                        llvm::omp::Directive directive) {
   llvm::ArrayRef<llvm::omp::Directive> leafDirs =
       llvm::omp::getLeafConstructsOrSelf(directive);
 
-  for (auto [dir, leaf] :
-       llvm::zip_longest(leafDirs, llvm::make_range(item, queue.end()))) {
+  for (auto [dir, leaf] : llvm::zip_longest(
+           leafDirs,
+           getNonTransformQueue(llvm::make_range(item, queue.end())))) {
     if (!dir.has_value() || !leaf.has_value())
       return false;
 

flang/lib/Lower/OpenMP/Decomposer.h

@@ -57,6 +57,10 @@ bool isLastItemInQueue(ConstructQueue::const_iterator item,
 bool matchLeafSequence(ConstructQueue::const_iterator item,
                        const ConstructQueue &queue,
                        llvm::omp::Directive directive);
+
+llvm::iterator_range<ConstructQueue::const_iterator> getNonTransformQueue(
+    llvm::iterator_range<ConstructQueue::const_iterator> range);
+
 } // namespace Fortran::lower::omp
 
 #endif // FORTRAN_LOWER_OPENMP_DECOMPOSER_H