[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/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")
+static void genACC(Fortran::lower::AbstractConverter &absConv,
+                   Fortran::lower::pft::Evaluation &eval,
+                   const Fortran::parser::OpenACCLoopConstruct &loopConstruct) {
+
+  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();
+    llvm::ArrayRef<mlir::Type> argTy;
+    mlir::ValueRange range;
+    // Temporarly set to default 0 as operands are not generated yet.
+    llvm::SmallVector<int32_t, 2> operandSegmentSizes(/*Size=*/2,
+                                                      /*Value=*/0);
+    auto loopOp =
+        firOpBuilder.create<mlir::acc::LoopOp>(currentLocation, argTy, range);
+    loopOp.setAttr(mlir::acc::LoopOp::getOperandSegmentSizeAttr(),
+                   firOpBuilder.getI32VectorAttr(operandSegmentSizes));
+
+    firOpBuilder.createBlock(&loopOp.getRegion());
+    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);
+  }
+}
 
 void Fortran::lower::genOpenACCConstruct(
-    Fortran::lower::AbstractConverter &absConv,
+    Fortran::lower::AbstractConverter &converter,
     Fortran::lower::pft::Evaluation &eval,
-    const Fortran::parser::OpenACCConstruct &accConstruct) {
-
+    const Fortran::parser::OpenACCConstruct &acc) {
   std::visit(
-      common::visitors{
+      Fortran::common::visitors{
           [&](const Fortran::parser::OpenACCBlockConstruct &blockConstruct) {
             TODO();
           },
           [&](const Fortran::parser::OpenACCCombinedConstruct
                   &combinedConstruct) { TODO(); },
           [&](const Fortran::parser::OpenACCLoopConstruct &loopConstruct) {
-            TODO();
+            genACC(converter, eval, loopConstruct);
           },
           [&](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);
+      acc.u);
 }

flang/unittests/Lower/CMakeLists.txt

 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 set(LIBS
   MLIRLLVMIR
   ${dialect_libs}
 )
 
-add_flang_unittest(FlangLoweringOpenMPTests
+add_flang_unittest(FlangLoweringDirectivesTests
   OpenMPLoweringTest.cpp
 )
-target_link_libraries(FlangLoweringOpenMPTests
+target_link_libraries(FlangLoweringDirectivesTests
   PRIVATE
   ${LIBS})