[flang][openacc] Lower loop with collapse clause

[clementval]

Summary

Lower OpenACCLoopConstruct and the collapse clause to the OpenACC acc.loop operation in MLIR.
This patch refelcts what can be upstream from PR flang-compiler/f18-llvm-project#410

[clementval]

Update to reflect change in PR

[DavidTruby]

Is it possible to add lit tests based on the generated mlir?

[clementval]

In D86423#2249060, @DavidTruby wrote:

Is it possible to add lit tests based on the generated mlir?

Not without the Bridge. Tests are present in fir-dev and will be upstream asap.

[clementval]

This patch is now part of D87389

flang/include/flang/Lower/OpenACC.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_LOWER_OPENACC_H
 #define FORTRAN_LOWER_OPENACC_H
 
 namespace Fortran {
 namespace parser {
-struct OpenACCConstruct;
+struct OpenACCLoopConstruct;
 } // namespace parser
 
 namespace lower {
 
 class AbstractConverter;
 
 namespace pft {
 struct Evaluation;
 } // namespace pft
 
-void genOpenACCConstruct(AbstractConverter &, pft::Evaluation &,
-                         const parser::OpenACCConstruct &);
+void genACC(AbstractConverter &, pft::Evaluation &,
+            const parser::OpenACCLoopConstruct &);
 
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENACC_H

flang/include/flang/Optimizer/Dialect/FIRDialect.h

 };
 
 /// Register the dialect with MLIR
 inline void registerFIR() {
   // we want to register exactly once
   [[maybe_unused]] static bool init_once = [] {
     mlir::registerDialect<mlir::AffineDialect>();
     mlir::registerDialect<mlir::LLVM::LLVMDialect>();
+    mlir::registerDialect<mlir::acc::OpenACCDialect>();
     mlir::registerDialect<mlir::omp::OpenMPDialect>();
     mlir::registerDialect<mlir::scf::SCFDialect>();
     mlir::registerDialect<mlir::StandardOpsDialect>();
     mlir::registerDialect<mlir::vector::VectorDialect>();
     mlir::registerDialect<FIROpsDialect>();
     return true;
   }();
 }

flang/lib/Lower/OpenACC.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Lower/OpenACC.h"
 #include "flang/Lower/Bridge.h"
 #include "flang/Lower/FIRBuilder.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Parser/parse-tree.h"
+#include "flang/Semantics/tools.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "llvm/Frontend/OpenACC/ACC.h.inc"
 
 #define TODO() llvm_unreachable("not yet implemented")
 
-void Fortran::lower::genOpenACCConstruct(
+void Fortran::lower::genACC(
     Fortran::lower::AbstractConverter &absConv,
     Fortran::lower::pft::Evaluation &eval,
-    const Fortran::parser::OpenACCConstruct &accConstruct) {
+    const Fortran::parser::OpenACCLoopConstruct &loopConstruct) {
 
-  std::visit(
-      common::visitors{
-          [&](const Fortran::parser::OpenACCBlockConstruct &blockConstruct) {
-            TODO();
-          },
-          [&](const Fortran::parser::OpenACCCombinedConstruct
-                  &combinedConstruct) { TODO(); },
-          [&](const Fortran::parser::OpenACCLoopConstruct &loopConstruct) {
-            TODO();
-          },
-          [&](const Fortran::parser::OpenACCStandaloneConstruct
-                  &standaloneConstruct) { TODO(); },
-          [&](const Fortran::parser::OpenACCRoutineConstruct
-                  &routineConstruct) { TODO(); },
-          [&](const Fortran::parser::OpenACCCacheConstruct &cacheConstruct) {
-            TODO();
-          },
-          [&](const Fortran::parser::OpenACCWaitConstruct &waitConstruct) {
-            TODO();
-          },
-          [&](const Fortran::parser::OpenACCAtomicConstruct &atomicConstruct) {
-            TODO();
-          },
-      },
-      accConstruct.u);
+  const auto &beginLoopDirective =
+      std::get<Fortran::parser::AccBeginLoopDirective>(loopConstruct.t);
+  const auto &loopDirective =
+      std::get<Fortran::parser::AccLoopDirective>(beginLoopDirective.t);
+
+  if (loopDirective.v == llvm::acc::ACCD_loop) {
+    auto &firOpBuilder = absConv.getFirOpBuilder();
+    auto currentLocation = absConv.getCurrentLocation();
+    auto insertPt = firOpBuilder.saveInsertionPoint();
+    llvm::ArrayRef<mlir::Type> argTy;
+    mlir::ValueRange range;
+    llvm::SmallVector<int32_t, 2> operandSegmentSizes(2 /*Size=*/,
+                                                      0 /*Value=*/);
+    auto loopOp =
+        firOpBuilder.create<mlir::acc::LoopOp>(currentLocation, argTy, range);
+    loopOp.setAttr(mlir::acc::LoopOp::getOperandSegmentSizeAttr(),
+                   firOpBuilder.getI32VectorAttr(operandSegmentSizes));
+    loopOp.getRegion().push_back(new Block{});
+    auto &block = loopOp.getRegion().back();
+    firOpBuilder.setInsertionPointToStart(&block);
+    // ensure the block is well-formed.
+    firOpBuilder.create<mlir::acc::YieldOp>(currentLocation);
+
+    // Add attribute extracted from clauses.
+    const auto &accClauseList =
+        std::get<Fortran::parser::AccClauseList>(beginLoopDirective.t);
+    for (const auto &clause : accClauseList.v) {
+      if (const auto *collapseClause{
+              std::get_if<Fortran::parser::AccClause::Collapse>(&clause.u)}) {
+
+        const auto *expr = Fortran::semantics::GetExpr(collapseClause->v);
+        const auto collapseValue = Fortran::evaluate::ToInt64(*expr);
+        if (collapseValue.has_value()) {
+          loopOp.setAttr(mlir::acc::LoopOp::getCollapseAttrName(),
+                         firOpBuilder.getI64IntegerAttr(collapseValue.value()));
+        }
+      }
+    }
+
+    // Place the insertion point to the start of the first block.
+    firOpBuilder.setInsertionPointToStart(&block);
+  }
 }