[flang][OpenMP] Support for Collapse

[Leporacanthicus]

Summary

Convert Fortran parse-tree into MLIR for collapse-clause.

Includes simple Fortran to LLVM-IR test, with auto-generated
check-lines (some of which have been edited by hand).

[kiranchandramohan]

Looks OK. A few Nit comments.

flang/lib/Lower/OpenMP.cpp

140

Nit: I think you are introducing an auto here. The previous state (without the auto) is better.

408

Did you accidentally remove this assert. I think bounds are required for a worksharing loop. This might be a change from fir-dev.

453

Nit: I prefer the previous formatting.

466

Nit: 2->3

[peixin]

LGTM except for @kiranchandramohan 's comments.

flang/lib/Lower/Bridge.cpp

1423

nit

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

5

Nit: Can you write these in one line and remove the check-prefix?

! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s

[shraiysh]

A few nits. Thanks for sending this.

flang/lib/Lower/OpenMP.cpp

138

It would be nice to know what the argument "args" refers to, it is very generic. Can we maybe please add a description of arguments over this function as a comment?

409

Can we please move this up? Right now there is (almost) a very nice separation of clauses in the upper half of the file and constructs in the lower half.

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

53

nit: end do?

[shraiysh]

If D125282 gets merged before this, please add tests for simd too.

[peixin]

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

53

end do is totally the same as enddo for Fortran code. Both of them are commonly used in HPC workloads.

[Leporacanthicus]

Update for review comments

[Leporacanthicus]

flang/lib/Lower/OpenMP.cpp

138

I have added a comment. I'm not sure the comment is THAT helpful, because I only know roughly the overall purpose of this function in the first place - I just changed arg to args and loop over the vector instead of

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

53

I agree with both comments: end do and enddo1 both works. I prefer end do - not sure why it wasn't in the first place - I have changed it - all other tests in Lower/OpenMP use that form.

There are enddo in Semantics tests.

[Total in flang/test enddo: 205, end do: 1132 - no, I'm not volunteering to unify this, I just wanted to see if one was more common than the other]

[peixin]

LGTM

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

10

BTW, which script file did you use to generate these FIR checks?

[kiranchandramohan]

It seems you missed a comment. Otherwise LGTM.

flang/lib/Lower/OpenMP.cpp

408

Don't know whether you missed this comment or prefer not to change.

The OpenMP standard says the following, so i think the assert is better.
"The do-loop cannot be a DO WHILE or a DO loop without loop control."

[shraiysh]

LGTM with comments addressed. Thanks!

flang/lib/Lower/OpenMP.cpp

408

+1. Maybe graceful error out during semantic checks for this (not necessarily in this patch) plus an assert here should be added.

[Leporacanthicus]

In D125302#3512417, @shraiysh wrote:

LGTM with comments addressed. Thanks!

Yes, I will fix this. I'll try to get done on Monday - been working on some other things today.

[peixin]

flang/lib/Lower/OpenMP.cpp

408

Can you add this back, rebase and land this patch?

@kiranchandramohan @Leporacanthicus Can we promote the priority of the upstream work? I compared fir-dev and LLVM main. We have 5 patches not upstreamed for OpenMP: this one, D125465, D125740, https://github.com/flang-compiler/f18-llvm-project/pull/1559 (@kiranchandramohan Can you also upstream this?), schedule chunk/modifier (@Leporacanthicus Can you upstream it?) For this one (https://github.com/flang-compiler/f18-llvm-project/pull/1178), it is not in a hurry to upstream it since it also has some works in FIR and it won't affect the following refactor work.

The current framework in OpenMP.cpp does not organize the clauses very well. Continuing adding the support for other constructs can easily bring in some bugs. OpenMP Loop constructs (https://github.com/llvm/llvm-project/blob/5646d82885f64f1f886d84aff5422697f355f6c4/flang/lib/Parser/openmp-parsers.cpp#L261-L294) and Block constructs do need to be refactored. After the above upstream work finishes, @shraiysh proposes one good patch (D121583) to reorganize the clauses.

All in all, the more other additional code merged, the harder the refactor work will be.

[kiranchandramohan]

flang/lib/Lower/OpenMP.cpp

408

We were originally planning to push forward with upstreaming this week. Unfortunately, something else came up. @Leporacanthicus is on leave today.
I will get to the remaining upstreaming and patch reviews from Thursday. Apologies for the delay.

[peixin]

flang/lib/Lower/OpenMP.cpp

408

OK. Thanks.

[Leporacanthicus]

Replace if with assert, as it was in earlier code (before the loop was added)

[Leporacanthicus]

Add assert for bounds back in, remove if-statment. This puts the code in
the new loop back to how the code was before the loop.

[Leporacanthicus]

flang/lib/Lower/OpenMP.cpp

408

Maybe graceful error out during semantic checks for this (not necessarily in this patch) plus an assert here should be added.

There is some checking in the Sema that produces an error when you try to do this:

i = 0
!$OMP DO
do
   if (i .eq. 100) exit
   i = i + 1
end do
!$OMP END DO

With other errors if you use do while (...) [Although when I did both, it only showed errors for the latter - I'm guessing because it does some random order evaluation and checking and quits on first error... Either way, we should not hit this assert. I have added one - and removed the if-statement, since it's redundant [and if we crash on bounds being NULL in non-debug, it's probably better than some kind weirder code-generation failure or bug somewhere else]

flang/include/flang/Lower/OpenMP.h

 //===-- Lower/OpenMP.h -- lower Open MP directives --------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_LOWER_OPENMP_H
 #define FORTRAN_LOWER_OPENMP_H
 
+#include <cinttypes>
+
 namespace Fortran {
 namespace parser {
 struct OpenMPConstruct;
 struct OpenMPDeclarativeConstruct;
+struct OmpEndLoopDirective;
+struct OmpClauseList;
 } // namespace parser
 
 namespace lower {
 
 class AbstractConverter;
 
 namespace pft {
 struct Evaluation;
 } // namespace pft
 
 void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
                         const parser::OpenMPConstruct &);
 void genOpenMPDeclarativeConstruct(AbstractConverter &, pft::Evaluation &,
                                    const parser::OpenMPDeclarativeConstruct &);
+int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);
 
 } // namespace lower
 } // namespace Fortran
 
 #endif // FORTRAN_LOWER_OPENMP_H

flang/lib/Lower/Bridge.cpp

void genFIR(const Fortran::parser::OpenACCDeclarativeConstruct &accDecl) {
     for (Fortran::lower::pft::Evaluation &e : getEval().getNestedEvaluations())
       genFIR(e);
     builder->restoreInsertionPoint(insertPt);
   }
 
   void genFIR(const Fortran::parser::OpenMPConstruct &omp) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     localSymbols.pushScope();
-    Fortran::lower::genOpenMPConstruct(*this, getEval(), omp);
+    genOpenMPConstruct(*this, getEval(), omp);
+
+    const Fortran::parser::OpenMPLoopConstruct *ompLoop =
+        std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u);
 
     // If loop is part of an OpenMP Construct then the OpenMP dialect
     // workshare loop operation has already been created. Only the
     // body needs to be created here and the do_loop can be skipped.
-    Fortran::lower::pft::Evaluation *curEval =
-        std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u)
-            ? &getEval().getFirstNestedEvaluation()
-            : &getEval();
+    // Skip the number of collapsed loops, which is 1 when there is a
+    // no collapse requested.
+
+    Fortran::lower::pft::Evaluation *curEval = &getEval();
+    if (ompLoop) {
+      const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
+          std::get<Fortran::parser::OmpBeginLoopDirective>(ompLoop->t).t);
+      int64_t collapseValue =
+          Fortran::lower::getCollapseValue(wsLoopOpClauseList);
+
+      curEval = &curEval->getFirstNestedEvaluation();
+      for (int64_t i = 1; i < collapseValue; i++) {
+        curEval = &*std::next(curEval->getNestedEvaluations().begin());
+      }
+    }
 
     for (Fortran::lower::pft::Evaluation &e : curEval->getNestedEvaluations())
       genFIR(e);
     localSymbols.popScope();
     builder->restoreInsertionPoint(insertPt);
   }
 
   void genFIR(const Fortran::parser::OpenMPDeclarativeConstruct &ompDecl) {

flang/lib/Lower/OpenMP.cpp

 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Semantics/tools.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
 using namespace mlir;
 
+int64_t Fortran::lower::getCollapseValue(
+    const Fortran::parser::OmpClauseList &clauseList) {
+  for (const auto &clause : clauseList.v) {
+    if (const auto &collapseClause =
+            std::get_if<Fortran::parser::OmpClause::Collapse>(&clause.u)) {
+      const auto *expr = Fortran::semantics::GetExpr(collapseClause->v);
+      return Fortran::evaluate::ToInt64(*expr).value();
+    }
+  }
+  return 1;
+}
+
 static const Fortran::parser::Name *
 getDesignatorNameIfDataRef(const Fortran::parser::Designator &designator) {
   const auto *dataRef = std::get_if<Fortran::parser::DataRef>(&designator.u);
   return dataRef ? std::get_if<Fortran::parser::Name>(&dataRef->u) : nullptr;
 }
 
 template <typename T>
 static void createPrivateVarSyms(Fortran::lower::AbstractConverter &converter,
std::visit(Fortran::common::visitors{
                    },
                    [&](const Fortran::parser::Name &name) {
                      addOperands(*name.symbol);
                    }},
                ompObject.u);
   }
 }
 
+/// Create the body (block) for an OpenMP Operation.
+///
+/// \param [in]    op - the operation the body belongs to.
+/// \param [inout] converter - converter to use for the clauses.
+/// \param [in]    loc - location in source code.
+/// \oaran [in]    clauses - list of clauses to process.
+/// \param [in]    args - block arguments (induction variable[s]) for the
+////                      region.
+/// \param [in]    outerCombined - is this an outer operation - prevents
+///                                privatization.
 template <typename Op>
 static void
 createBodyOfOp(Op &op, Fortran::lower::AbstractConverter &converter,
                mlir::Location &loc,
                const Fortran::parser::OmpClauseList *clauses = nullptr,
-               const Fortran::semantics::Symbol *arg = nullptr,
+               const SmallVector<const Fortran::semantics::Symbol *> &args = {},
                bool outerCombined = false) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-  // If an argument for the region is provided then create the block with that
-  // argument. Also update the symbol's address with the mlir argument value.
-  // e.g. For loops the argument is the induction variable. And all further
+  // If arguments for the region are provided then create the block with those
+  // arguments. Also update the symbol's address with the mlir argument values.
+  // e.g. For loops the arguments are the induction variable. And all further
   // uses of the induction variable should use this mlir value.
-  if (arg) {
-    firOpBuilder.createBlock(&op.getRegion(), {}, {converter.genType(*arg)},
-                             {loc});
-    converter.bindSymbol(*arg, op.getRegion().front().getArgument(0));
+  if (args.size()) {
+    SmallVector<Type> tiv;
+    SmallVector<Location> locs;
+    int argIndex = 0;
+    for (auto &arg : args) {
+      tiv.push_back(converter.genType(*arg));
+      locs.push_back(loc);
+    }
+    firOpBuilder.createBlock(&op.getRegion(), {}, tiv, locs);
+    for (auto &arg : args) {
+      fir::ExtendedValue exval = op.getRegion().front().getArgument(argIndex);
+      converter.bindSymbol(*arg, exval);
+      argIndex++;
+    }
   } else {
     firOpBuilder.createBlock(&op.getRegion());
   }
   auto &block = op.getRegion().back();
   firOpBuilder.setInsertionPointToStart(&block);
 
   // Insert the terminator.
   if constexpr (std::is_same_v<Op, omp::WsLoopOp>) {
auto taskOp = firOpBuilder.create<mlir::omp::TaskOp>(
         /*in_reductions=*/nullptr, priorityClauseOperand, allocateOperands,
         allocatorOperands);
     createBodyOfOp(taskOp, converter, currentLocation, &opClauseList);
   } else {
     TODO(converter.getCurrentLocation(), "Unhandled block directive");
   }
 }
 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
       linearStepVars, reductionVars;
   mlir::Value scheduleChunkClauseOperand;
   mlir::Attribute scheduleClauseOperand, collapseClauseOperand,
       noWaitClauseOperand, orderedClauseOperand, orderClauseOperand;
   const auto &wsLoopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
       std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
   if (llvm::omp::OMPD_do !=
       std::get<Fortran::parser::OmpLoopDirective>(
           std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
           .v) {
     TODO(converter.getCurrentLocation(), "Combined worksharing loop construct");
   }
 
-  Fortran::lower::pft::Evaluation *doConstructEval =
-      &eval.getFirstNestedEvaluation();
+  int64_t collapseValue = Fortran::lower::getCollapseValue(wsLoopOpClauseList);
+
+  // Collect the loops to collapse.
+  auto *doConstructEval = &eval.getFirstNestedEvaluation();
 
-  Fortran::lower::pft::Evaluation *doLoop =
-      &doConstructEval->getFirstNestedEvaluation();
+  SmallVector<const Fortran::semantics::Symbol *> iv;
+  do {
+    auto *doLoop = &doConstructEval->getFirstNestedEvaluation();
     auto *doStmt = doLoop->getIf<Fortran::parser::NonLabelDoStmt>();
     assert(doStmt && "Expected do loop to be in the nested evaluation");
     const auto &loopControl =
         std::get<std::optional<Fortran::parser::LoopControl>>(doStmt->t);
     const Fortran::parser::LoopControl::Bounds *bounds =
         std::get_if<Fortran::parser::LoopControl::Bounds>(&loopControl->u);
     assert(bounds && "Expected bounds for worksharing do loop");
-  Fortran::semantics::Symbol *iv = nullptr;
     Fortran::lower::StatementContext stmtCtx;
     lowerBound.push_back(fir::getBase(converter.genExprValue(
         *Fortran::semantics::GetExpr(bounds->lower), stmtCtx)));
     upperBound.push_back(fir::getBase(converter.genExprValue(
         *Fortran::semantics::GetExpr(bounds->upper), stmtCtx)));
     if (bounds->step) {
       step.push_back(fir::getBase(converter.genExprValue(
           *Fortran::semantics::GetExpr(bounds->step), stmtCtx)));
     } else { // If `step` is not present, assume it as `1`.
       step.push_back(firOpBuilder.createIntegerConstant(
           currentLocation, firOpBuilder.getIntegerType(32), 1));
     }
-  iv = bounds->name.thing.symbol;
+    iv.push_back(bounds->name.thing.symbol);
+
+    collapseValue--;
+    doConstructEval =
+        &*std::next(doConstructEval->getNestedEvaluations().begin());
+  } while (collapseValue > 0);
 
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order
-  // 3. collapse
-  // 4. schedule (with chunk)
+  // 3. schedule (with chunk)
   auto wsLoopOp = firOpBuilder.create<mlir::omp::WsLoopOp>(
       currentLocation, lowerBound, upperBound, step, linearVars, linearStepVars,
       reductionVars, /*reductions=*/nullptr,
       scheduleClauseOperand.dyn_cast_or_null<omp::ClauseScheduleKindAttr>(),
       scheduleChunkClauseOperand, /*schedule_modifiers=*/nullptr,
       /*simd_modifier=*/nullptr,
       collapseClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       noWaitClauseOperand.dyn_cast_or_null<UnitAttr>(),
if (const auto &orderedClause =
         const auto *expr = Fortran::semantics::GetExpr(orderedClause->v);
         const std::optional<std::int64_t> orderedClauseValue =
             Fortran::evaluate::ToInt64(*expr);
         wsLoopOp.ordered_valAttr(
             firOpBuilder.getI64IntegerAttr(*orderedClauseValue));
       } else {
         wsLoopOp.ordered_valAttr(firOpBuilder.getI64IntegerAttr(0));
       }
+    } else if (const auto &collapseClause =
+                   std::get_if<Fortran::parser::OmpClause::Collapse>(
+                       &clause.u)) {
+      const auto *expr = Fortran::semantics::GetExpr(collapseClause->v);
+      const std::optional<std::int64_t> collapseValue =
+          Fortran::evaluate::ToInt64(*expr);
+      wsLoopOp.collapse_valAttr(firOpBuilder.getI64IntegerAttr(*collapseValue));
     } else if (const auto &scheduleClause =
                    std::get_if<Fortran::parser::OmpClause::Schedule>(
                        &clause.u)) {
       mlir::MLIRContext *context = firOpBuilder.getContext();
       const auto &scheduleType = scheduleClause->v;
       const auto &scheduleKind =
           std::get<Fortran::parser::OmpScheduleClause::ScheduleType>(
               scheduleType.t);

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

+! This test checks lowering of OpenMP DO Directive(Worksharing) with collapse.
+
+! RUN: bbc -fopenmp -emit-fir %s -o - | FileCheck %s
+
+program wsloop_collapse
+  integer :: i, j, k
+  integer :: a, b, c
+  integer :: x
+! CHECK:         %[[VAL_0:.*]] = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFEa"}
+! CHECK:         %[[VAL_1:.*]] = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFEb"}
+! CHECK:         %[[VAL_2:.*]] = fir.alloca i32 {bindc_name = "c", uniq_name = "_QFEc"}
+! CHECK:         %[[VAL_3:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFEi"}
+! CHECK:         %[[VAL_4:.*]] = fir.alloca i32 {bindc_name = "j", uniq_name = "_QFEj"}
+! CHECK:         %[[VAL_5:.*]] = fir.alloca i32 {bindc_name = "k", uniq_name = "_QFEk"}
+! CHECK:         %[[VAL_6:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFEx"}
+  a=3
+! CHECK:         %[[VAL_7:.*]] = arith.constant 3 : i32
+! CHECK:         fir.store %[[VAL_7]] to %[[VAL_0]] : !fir.ref<i32>
+  b=2
+! CHECK:         %[[VAL_8:.*]] = arith.constant 2 : i32
+! CHECK:         fir.store %[[VAL_8]] to %[[VAL_1]] : !fir.ref<i32>
+  c=5
+! CHECK:         %[[VAL_9:.*]] = arith.constant 5 : i32
+! CHECK:         fir.store %[[VAL_9]] to %[[VAL_2]] : !fir.ref<i32>
+  x=0
+! CHECK:         %[[VAL_10:.*]] = arith.constant 0 : i32
+! CHECK:         fir.store %[[VAL_10]] to %[[VAL_6]] : !fir.ref<i32>
+
+  !$omp do collapse(3)
+! CHECK:           %[[VAL_20:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_21:.*]] = fir.load %[[VAL_0]] : !fir.ref<i32>
+! CHECK:           %[[VAL_22:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_23:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_24:.*]] = fir.load %[[VAL_1]] : !fir.ref<i32>
+! CHECK:           %[[VAL_25:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_26:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_27:.*]] = fir.load %[[VAL_2]] : !fir.ref<i32>
+! CHECK:           %[[VAL_28:.*]] = arith.constant 1 : i32
+  do i = 1, a
+     do j= 1, b
+        do k = 1, c
+! CHECK:           omp.wsloop collapse(3) for (%[[VAL_9:.*]], %[[VAL_10:.*]], %[[VAL_11:.*]]) : i32 = (%[[VAL_20]], %[[VAL_23]], %[[VAL_26]]) to (%[[VAL_21]], %[[VAL_24]], %[[VAL_27]]) inclusive step (%[[VAL_22]], %[[VAL_25]], %[[VAL_28]]) {
+! CHECK:             %[[VAL_12:.*]] = fir.load %[[VAL_6]] : !fir.ref<i32>
+! CHECK:             %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_9]] : i32
+! CHECK:             %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_10]] : i32
+! CHECK:             %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_11]] : i32
+! CHECK:             fir.store %[[VAL_15]] to %[[VAL_6]] : !fir.ref<i32>
+! CHECK:             omp.yield
+! CHECK:           }
+           x = x + i + j + k
+        end do
+     end do
+  end do
+  !$omp end do
+! CHECK:         return
+! CHECK:       }
+end program wsloop_collapse