[Flang][OpenMP] Add lowering from PFT to new MapEntry and Bounds operations and tie them to relevant Target operations

This patch builds on top of a prior patch in review which adds a new map
and bounds operation by modifying the OpenMP PFT lowering to support
these operations and generate them from the PFT.

A significant amount of the support for the Bounds operation is borrowed
from OpenACC's own current implementation and lowering, just ported
over to OpenMP.

The patch also adds very preliminary/initial support for lowering to
a new Capture attribute, which is stored on the new Map Operation,
which helps the later lowering from OpenMP -> LLVM IR by indicating
how a map argument should be handled. This capture type will
influence how a map argument is accessed on device and passed by
the host (different load/store handling etc.). It is reflective of a
similar piece of information stored in the Clang AST which performs a
similar role.

As well as some minor adjustments to how the map type (map bitshift
which dictates to the runtime how it should handle an argument) is
generated to further support more use-cases for future patches that
build on this work.

Finally it adds the map entry operation creation and tying it to the relevant
target operations as well as the addition of some new tests and alteration
of previous tests to support the new changes.

Depends on D158732

reviewers: kiranchandramohan, TIFitis, clementval, razvanlupusoru

Differential Revision: https://reviews.llvm.org/D158734

Author: agozillon

flang/include/flang/Lower/OpenMP.h

@@ -36,6 +36,7 @@ struct OmpClauseList;
 
 namespace semantics {
 class Symbol;
+class SemanticsContext;
 } // namespace semantics
 
 namespace lower {
@@ -51,8 +52,8 @@ struct Variable;
 void genOpenMPTerminator(fir::FirOpBuilder &, mlir::Operation *,
                          mlir::Location);
 
-void genOpenMPConstruct(AbstractConverter &, pft::Evaluation &,
-                        const parser::OpenMPConstruct &);
+void genOpenMPConstruct(AbstractConverter &, semantics::SemanticsContext &,
+                        pft::Evaluation &, const parser::OpenMPConstruct &);
 void genOpenMPDeclarativeConstruct(AbstractConverter &, pft::Evaluation &,
                                    const parser::OpenMPDeclarativeConstruct &);
 int64_t getCollapseValue(const Fortran::parser::OmpClauseList &clauseList);

flang/lib/Lower/Bridge.cpp

@@ -2320,7 +2320,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   void genFIR(const Fortran::parser::OpenMPConstruct &omp) {
     mlir::OpBuilder::InsertPoint insertPt = builder->saveInsertionPoint();
     localSymbols.pushScope();
-    genOpenMPConstruct(*this, getEval(), omp);
+    genOpenMPConstruct(*this, bridge.getSemanticsContext(), getEval(), omp);
 
     const Fortran::parser::OpenMPLoopConstruct *ompLoop =
         std::get_if<Fortran::parser::OpenMPLoopConstruct>(&omp.u);

flang/lib/Lower/DirectivesCommon.h

@@ -12,20 +12,23 @@
 ///
 /// A location to place directive utilities shared across multiple lowering
 /// files, e.g. utilities shared in OpenMP and OpenACC. The header file can
-/// be used for both declarations and templated/inline implementations.
+/// be used for both declarations and templated/inline implementations
 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_LOWER_DIRECTIVES_COMMON_H
 #define FORTRAN_LOWER_DIRECTIVES_COMMON_H
 
 #include "flang/Common/idioms.h"
+#include "flang/Evaluate/tools.h"
+#include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/Bridge.h"
 #include "flang/Lower/ConvertExpr.h"
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/OpenACC.h"
 #include "flang/Lower/OpenMP.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/StatementContext.h"
+#include "flang/Lower/Support/Utils.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
@@ -36,7 +39,9 @@
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Value.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
+#include <list>
 #include <type_traits>
 
 namespace Fortran {
@@ -611,6 +616,309 @@ void createEmptyRegionBlocks(
   }
 }
 
+inline mlir::Value
+getDataOperandBaseAddr(Fortran::lower::AbstractConverter &converter,
+                       fir::FirOpBuilder &builder,
+                       Fortran::lower::SymbolRef sym, mlir::Location loc) {
+  mlir::Value symAddr = converter.getSymbolAddress(sym);
+  // TODO: Might need revisiting to handle for non-shared clauses
+  if (!symAddr) {
+    if (const auto *details =
+            sym->detailsIf<Fortran::semantics::HostAssocDetails>())
+      symAddr = converter.getSymbolAddress(details->symbol());
+  }
+
+  if (!symAddr)
+    llvm::report_fatal_error("could not retrieve symbol address");
+
+  if (auto boxTy =
+          fir::unwrapRefType(symAddr.getType()).dyn_cast<fir::BaseBoxType>()) {
+    if (boxTy.getEleTy().isa<fir::RecordType>())
+      TODO(loc, "derived type");
+
+    // Load the box when baseAddr is a `fir.ref<fir.box<T>>` or a
+    // `fir.ref<fir.class<T>>` type.
+    if (symAddr.getType().isa<fir::ReferenceType>())
+      return builder.create<fir::LoadOp>(loc, symAddr);
+  }
+  return symAddr;
+}
+
+/// Generate the bounds operation from the descriptor information.
+template <typename BoundsOp, typename BoundsType>
+llvm::SmallVector<mlir::Value>
+genBoundsOpsFromBox(fir::FirOpBuilder &builder, mlir::Location loc,
+                    Fortran::lower::AbstractConverter &converter,
+                    fir::ExtendedValue dataExv, mlir::Value box) {
+  llvm::SmallVector<mlir::Value> bounds;
+  mlir::Type idxTy = builder.getIndexType();
+  mlir::Type boundTy = builder.getType<BoundsType>();
+  mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+  assert(box.getType().isa<fir::BaseBoxType>() &&
+         "expect fir.box or fir.class");
+  for (unsigned dim = 0; dim < dataExv.rank(); ++dim) {
+    mlir::Value d = builder.createIntegerConstant(loc, idxTy, dim);
+    mlir::Value baseLb =
+        fir::factory::readLowerBound(builder, loc, dataExv, dim, one);
+    auto dimInfo =
+        builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy, box, d);
+    mlir::Value lb = builder.createIntegerConstant(loc, idxTy, 0);
+    mlir::Value ub =
+        builder.create<mlir::arith::SubIOp>(loc, dimInfo.getExtent(), one);
+    mlir::Value bound =
+        builder.create<BoundsOp>(loc, boundTy, lb, ub, mlir::Value(),
+                                 dimInfo.getByteStride(), true, baseLb);
+    bounds.push_back(bound);
+  }
+  return bounds;
+}
+
+/// Generate bounds operation for base array without any subscripts
+/// provided.
+template <typename BoundsOp, typename BoundsType>
+llvm::SmallVector<mlir::Value>
+genBaseBoundsOps(fir::FirOpBuilder &builder, mlir::Location loc,
+                 Fortran::lower::AbstractConverter &converter,
+                 fir::ExtendedValue dataExv, mlir::Value baseAddr) {
+  mlir::Type idxTy = builder.getIndexType();
+  mlir::Type boundTy = builder.getType<BoundsType>();
+  llvm::SmallVector<mlir::Value> bounds;
+
+  if (dataExv.rank() == 0)
+    return bounds;
+
+  mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+  for (std::size_t dim = 0; dim < dataExv.rank(); ++dim) {
+    mlir::Value baseLb =
+        fir::factory::readLowerBound(builder, loc, dataExv, dim, one);
+    mlir::Value ext = fir::factory::readExtent(builder, loc, dataExv, dim);
+    mlir::Value lb = builder.createIntegerConstant(loc, idxTy, 0);
+
+    // ub = extent - 1
+    mlir::Value ub = builder.create<mlir::arith::SubIOp>(loc, ext, one);
+    mlir::Value bound =
+        builder.create<BoundsOp>(loc, boundTy, lb, ub, ext, one, false, baseLb);
+    bounds.push_back(bound);
+  }
+  return bounds;
+}
+
+/// Generate bounds operations for an array section when subscripts are
+/// provided.
+template <typename BoundsOp, typename BoundsType>
+llvm::SmallVector<mlir::Value>
+genBoundsOps(fir::FirOpBuilder &builder, mlir::Location loc,
+             Fortran::lower::AbstractConverter &converter,
+             Fortran::lower::StatementContext &stmtCtx,
+             const std::list<Fortran::parser::SectionSubscript> &subscripts,
+             std::stringstream &asFortran, fir::ExtendedValue &dataExv,
+             mlir::Value baseAddr) {
+  int dimension = 0;
+  mlir::Type idxTy = builder.getIndexType();
+  mlir::Type boundTy = builder.getType<BoundsType>();
+  llvm::SmallVector<mlir::Value> bounds;
+
+  mlir::Value zero = builder.createIntegerConstant(loc, idxTy, 0);
+  mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+  for (const auto &subscript : subscripts) {
+    if (const auto *triplet{
+            std::get_if<Fortran::parser::SubscriptTriplet>(&subscript.u)}) {
+      if (dimension != 0)
+        asFortran << ',';
+      mlir::Value lbound, ubound, extent;
+      std::optional<std::int64_t> lval, uval;
+      mlir::Value baseLb =
+          fir::factory::readLowerBound(builder, loc, dataExv, dimension, one);
+      bool defaultLb = baseLb == one;
+      mlir::Value stride = one;
+      bool strideInBytes = false;
+
+      if (fir::unwrapRefType(baseAddr.getType()).isa<fir::BaseBoxType>()) {
+        mlir::Value d = builder.createIntegerConstant(loc, idxTy, dimension);
+        auto dimInfo = builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy,
+                                                      baseAddr, d);
+        stride = dimInfo.getByteStride();
+        strideInBytes = true;
+      }
+
+      const auto &lower{std::get<0>(triplet->t)};
+      if (lower) {
+        lval = Fortran::semantics::GetIntValue(lower);
+        if (lval) {
+          if (defaultLb) {
+            lbound = builder.createIntegerConstant(loc, idxTy, *lval - 1);
+          } else {
+            mlir::Value lb = builder.createIntegerConstant(loc, idxTy, *lval);
+            lbound = builder.create<mlir::arith::SubIOp>(loc, lb, baseLb);
+          }
+          asFortran << *lval;
+        } else {
+          const Fortran::lower::SomeExpr *lexpr =
+              Fortran::semantics::GetExpr(*lower);
+          mlir::Value lb =
+              fir::getBase(converter.genExprValue(loc, *lexpr, stmtCtx));
+          lb = builder.createConvert(loc, baseLb.getType(), lb);
+          lbound = builder.create<mlir::arith::SubIOp>(loc, lb, baseLb);
+          asFortran << lexpr->AsFortran();
+        }
+      } else {
+        lbound = defaultLb ? zero : baseLb;
+      }
+      asFortran << ':';
+      const auto &upper{std::get<1>(triplet->t)};
+      if (upper) {
+        uval = Fortran::semantics::GetIntValue(upper);
+        if (uval) {
+          if (defaultLb) {
+            ubound = builder.createIntegerConstant(loc, idxTy, *uval - 1);
+          } else {
+            mlir::Value ub = builder.createIntegerConstant(loc, idxTy, *uval);
+            ubound = builder.create<mlir::arith::SubIOp>(loc, ub, baseLb);
+          }
+          asFortran << *uval;
+        } else {
+          const Fortran::lower::SomeExpr *uexpr =
+              Fortran::semantics::GetExpr(*upper);
+          mlir::Value ub =
+              fir::getBase(converter.genExprValue(loc, *uexpr, stmtCtx));
+          ub = builder.createConvert(loc, baseLb.getType(), ub);
+          ubound = builder.create<mlir::arith::SubIOp>(loc, ub, baseLb);
+          asFortran << uexpr->AsFortran();
+        }
+      }
+      if (lower && upper) {
+        if (lval && uval && *uval < *lval) {
+          mlir::emitError(loc, "zero sized array section");
+          break;
+        } else if (std::get<2>(triplet->t)) {
+          const auto &strideExpr{std::get<2>(triplet->t)};
+          if (strideExpr) {
+            mlir::emitError(loc, "stride cannot be specified on "
+                                 "an OpenMP array section");
+            break;
+          }
+        }
+      }
+      // ub = baseLb + extent - 1
+      if (!ubound) {
+        mlir::Value ext =
+            fir::factory::readExtent(builder, loc, dataExv, dimension);
+        mlir::Value lbExt =
+            builder.create<mlir::arith::AddIOp>(loc, ext, baseLb);
+        ubound = builder.create<mlir::arith::SubIOp>(loc, lbExt, one);
+      }
+      mlir::Value bound = builder.create<BoundsOp>(
+          loc, boundTy, lbound, ubound, extent, stride, strideInBytes, baseLb);
+      bounds.push_back(bound);
+      ++dimension;
+    }
+  }
+  return bounds;
+}
+
+template <typename ObjectType, typename BoundsOp, typename BoundsType>
+mlir::Value gatherDataOperandAddrAndBounds(
+    Fortran::lower::AbstractConverter &converter, fir::FirOpBuilder &builder,
+    Fortran::semantics::SemanticsContext &semanticsContext,
+    Fortran::lower::StatementContext &stmtCtx, const ObjectType &object,
+    mlir::Location operandLocation, std::stringstream &asFortran,
+    llvm::SmallVector<mlir::Value> &bounds) {
+  mlir::Value baseAddr;
+
+  std::visit(
+      Fortran::common::visitors{
+          [&](const Fortran::parser::Designator &designator) {
+            if (auto expr{Fortran::semantics::AnalyzeExpr(semanticsContext,
+                                                          designator)}) {
+              if ((*expr).Rank() > 0 &&
+                  Fortran::parser::Unwrap<Fortran::parser::ArrayElement>(
+                      designator)) {
+                const auto *arrayElement =
+                    Fortran::parser::Unwrap<Fortran::parser::ArrayElement>(
+                        designator);
+                const auto *dataRef =
+                    std::get_if<Fortran::parser::DataRef>(&designator.u);
+                fir::ExtendedValue dataExv;
+                if (Fortran::parser::Unwrap<
+                        Fortran::parser::StructureComponent>(
+                        arrayElement->base)) {
+                  auto exprBase = Fortran::semantics::AnalyzeExpr(
+                      semanticsContext, arrayElement->base);
+                  dataExv = converter.genExprAddr(operandLocation, *exprBase,
+                                                  stmtCtx);
+                  baseAddr = fir::getBase(dataExv);
+                  asFortran << (*exprBase).AsFortran();
+                } else {
+                  const Fortran::parser::Name &name =
+                      Fortran::parser::GetLastName(*dataRef);
+                  baseAddr = getDataOperandBaseAddr(
+                      converter, builder, *name.symbol, operandLocation);
+                  dataExv = converter.getSymbolExtendedValue(*name.symbol);
+                  asFortran << name.ToString();
+                }
+
+                if (!arrayElement->subscripts.empty()) {
+                  asFortran << '(';
+                  bounds = genBoundsOps<BoundsType, BoundsOp>(
+                      builder, operandLocation, converter, stmtCtx,
+                      arrayElement->subscripts, asFortran, dataExv, baseAddr);
+                }
+                asFortran << ')';
+              } else if (Fortran::parser::Unwrap<
+                             Fortran::parser::StructureComponent>(designator)) {
+                fir::ExtendedValue compExv =
+                    converter.genExprAddr(operandLocation, *expr, stmtCtx);
+                baseAddr = fir::getBase(compExv);
+                if (fir::unwrapRefType(baseAddr.getType())
+                        .isa<fir::SequenceType>())
+                  bounds = genBaseBoundsOps<BoundsType, BoundsOp>(
+                      builder, operandLocation, converter, compExv, baseAddr);
+                asFortran << (*expr).AsFortran();
+
+                // If the component is an allocatable or pointer the result of
+                // genExprAddr will be the result of a fir.box_addr operation.
+                // Retrieve the box so we handle it like other descriptor.
+                if (auto boxAddrOp = mlir::dyn_cast_or_null<fir::BoxAddrOp>(
+                        baseAddr.getDefiningOp())) {
+                  baseAddr = boxAddrOp.getVal();
+                  bounds = genBoundsOpsFromBox<BoundsType, BoundsOp>(
+                      builder, operandLocation, converter, compExv, baseAddr);
+                }
+              } else {
+                // Scalar or full array.
+                if (const auto *dataRef{
+                        std::get_if<Fortran::parser::DataRef>(&designator.u)}) {
+                  const Fortran::parser::Name &name =
+                      Fortran::parser::GetLastName(*dataRef);
+                  fir::ExtendedValue dataExv =
+                      converter.getSymbolExtendedValue(*name.symbol);
+                  baseAddr = getDataOperandBaseAddr(
+                      converter, builder, *name.symbol, operandLocation);
+                  if (fir::unwrapRefType(baseAddr.getType())
+                          .isa<fir::BaseBoxType>())
+                    bounds = genBoundsOpsFromBox<BoundsType, BoundsOp>(
+                        builder, operandLocation, converter, dataExv, baseAddr);
+                  if (fir::unwrapRefType(baseAddr.getType())
+                          .isa<fir::SequenceType>())
+                    bounds = genBaseBoundsOps<BoundsType, BoundsOp>(
+                        builder, operandLocation, converter, dataExv, baseAddr);
+                  asFortran << name.ToString();
+                } else { // Unsupported
+                  llvm::report_fatal_error(
+                      "Unsupported type of OpenACC operand");
+                }
+              }
+            }
+          },
+          [&](const Fortran::parser::Name &name) {
+            baseAddr = getDataOperandBaseAddr(converter, builder, *name.symbol,
+                                              operandLocation);
+            asFortran << name.ToString();
+          }},
+      object.u);
+  return baseAddr;
+}
+
 } // namespace lower
 } // namespace Fortran