[flang][acc] Implement MappableType interfaces for fir.box and fir.ar…

…ray (llvm#122495)

The newly introduced MappableType interface in `acc` dialect was
primarily intended to allow variables with non-materialized storage to
be used in acc data clauses (previously everything was required to be
`pointer-like`). One motivator for this was `fir.box` since it is
possible to be passed to functions without a wrapping `fir.ref` and also
it can be generated directly via operations like `fir.embox` - and
unlike other variable representations in FIR, the underlying storage for
it does not get materialized until LLVM codegen.

The new interface is being attached to both `fir.box` and `fir.array`.
Strictly speaking, attaching to the latter is primarily for consistency
since the MappableType interface requires implementation of utilities to
compute byte size - and it made sense that a
`fir.box<fir.array<10xi32>>` and `fir.array<10xi32>` would have a
consistently computable size. This decision may be revisited as
MappableType interface evolves.

The new interface attachments are made in a new library named
`FIROpenACCSupport`. The reason for this is to avoid circular
dependencies since the implementation of this library is reusing code
from lowering of OpenACC. More specifically, the types are defined in
`FIRDialect` and `FortranLower` depends on it. Thus we cannot attach
these interfaces in `FIRDialect`.

flang/include/flang/Optimizer/OpenACC/FIROpenACCTypeInterfaces.h

@@ -0,0 +1,43 @@
+//===- FIROpenACCTypeInterfaces.h -------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains external dialect interfaces for FIR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FLANG_OPTIMIZER_OPENACC_FIROPENACCTYPEINTERFACES_H_
+#define FLANG_OPTIMIZER_OPENACC_FIROPENACCTYPEINTERFACES_H_
+
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+
+namespace fir::acc {
+
+template <typename T>
+struct OpenACCMappableModel
+    : public mlir::acc::MappableType::ExternalModel<OpenACCMappableModel<T>,
+                                                    T> {
+  mlir::TypedValue<mlir::acc::PointerLikeType> getVarPtr(::mlir::Type type,
+                                                         mlir::Value var) const;
+
+  std::optional<llvm::TypeSize>
+  getSizeInBytes(mlir::Type type, mlir::Value var, mlir::ValueRange accBounds,
+                 const mlir::DataLayout &dataLayout) const;
+
+  std::optional<int64_t>
+  getOffsetInBytes(mlir::Type type, mlir::Value var, mlir::ValueRange accBounds,
+                   const mlir::DataLayout &dataLayout) const;
+
+  llvm::SmallVector<mlir::Value>
+  generateAccBounds(mlir::Type type, mlir::Value var,
+                    mlir::OpBuilder &builder) const;
+};
+
+} // namespace fir::acc
+
+#endif // FLANG_OPTIMIZER_OPENACC_FIROPENACCTYPEINTERFACES_H_

flang/include/flang/Optimizer/OpenACC/RegisterOpenACCExtensions.h

@@ -0,0 +1,22 @@
+//===- RegisterOpenACCExtensions.h - OpenACC Extension Registration --===--===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FLANG_OPTIMIZER_OPENACC_REGISTEROPENACCEXTENSIONS_H_
+#define FLANG_OPTIMIZER_OPENACC_REGISTEROPENACCEXTENSIONS_H_
+
+namespace mlir {
+class DialectRegistry;
+} // namespace mlir
+
+namespace fir::acc {
+
+void registerOpenACCExtensions(mlir::DialectRegistry &registry);
+
+} // namespace fir::acc
+
+#endif // FLANG_OPTIMIZER_OPENACC_REGISTEROPENACCEXTENSIONS_H_

flang/include/flang/Optimizer/Support/InitFIR.h

@@ -17,6 +17,7 @@
 #include "flang/Optimizer/Dialect/CUF/CUFToLLVMIRTranslation.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
+#include "flang/Optimizer/OpenACC/RegisterOpenACCExtensions.h"
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
@@ -63,6 +64,7 @@ inline void addFIRExtensions(mlir::DialectRegistry &registry,
     addFIRInlinerExtension(registry);
   addFIRToLLVMIRExtension(registry);
   cuf::registerCUFDialectTranslation(registry);
+  fir::acc::registerOpenACCExtensions(registry);
 }
 
 inline void loadNonCodegenDialects(mlir::MLIRContext &context) {

flang/lib/Frontend/CMakeLists.txt

@@ -39,6 +39,7 @@ add_flang_library(flangFrontend
   HLFIRDialect
   HLFIRTransforms
   flangPasses
+  FIROpenACCSupport
   FlangOpenMPTransforms
   MLIRTransforms
   MLIRBuiltinToLLVMIRTranslation

flang/lib/Optimizer/CMakeLists.txt

@@ -3,6 +3,7 @@ add_subdirectory(Builder)
 add_subdirectory(CodeGen)
 add_subdirectory(Dialect)
 add_subdirectory(HLFIR)
+add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(Passes)
 add_subdirectory(Support)

flang/lib/Optimizer/Dialect/CMakeLists.txt

@@ -6,8 +6,8 @@ add_flang_library(FIRDialect
   FIRDialect.cpp
   FIROps.cpp
   FIRType.cpp
-  FortranVariableInterface.cpp
   FirAliasTagOpInterface.cpp
+  FortranVariableInterface.cpp
   Inliner.cpp
 
   DEPENDS

flang/lib/Optimizer/OpenACC/CMakeLists.txt

@@ -0,0 +1,22 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+
+add_flang_library(FIROpenACCSupport
+  FIROpenACCTypeInterfaces.cpp
+  RegisterOpenACCExtensions.cpp
+
+  DEPENDS
+  FIRBuilder
+  FIRDialect
+  FIRDialectSupport
+  FIRSupport
+  HLFIRDialect
+  MLIROpenACCDialect
+
+  LINK_LIBS
+  FIRBuilder
+  FIRDialect
+  FIRDialectSupport
+  FIRSupport
+  HLFIRDialect
+  MLIROpenACCDialect
+)

flang/lib/Optimizer/OpenACC/FIROpenACCTypeInterfaces.cpp

@@ -0,0 +1,227 @@
+//===-- FIROpenACCTypeInterfaces.cpp --------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implementation of external dialect interfaces for FIR.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/OpenACC/FIROpenACCTypeInterfaces.h"
+#include "flang/Lower/DirectivesCommon.h"
+#include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/HLFIRTools.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
+#include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Support/LLVM.h"
+
+namespace fir::acc {
+
+static mlir::TypedValue<mlir::acc::PointerLikeType>
+getPtrFromVar(mlir::Value var) {
+  if (auto ptr =
+          mlir::dyn_cast<mlir::TypedValue<mlir::acc::PointerLikeType>>(var))
+    return ptr;
+
+  if (auto load = mlir::dyn_cast_if_present<fir::LoadOp>(var.getDefiningOp())) {
+    // All FIR reference types implement the PointerLikeType interface.
+    return mlir::cast<mlir::TypedValue<mlir::acc::PointerLikeType>>(
+        load.getMemref());
+  }
+
+  return {};
+}
+
+template <>
+mlir::TypedValue<mlir::acc::PointerLikeType>
+OpenACCMappableModel<fir::SequenceType>::getVarPtr(mlir::Type type,
+                                                   mlir::Value var) const {
+  return getPtrFromVar(var);
+}
+
+template <>
+mlir::TypedValue<mlir::acc::PointerLikeType>
+OpenACCMappableModel<fir::BaseBoxType>::getVarPtr(mlir::Type type,
+                                                  mlir::Value var) const {
+  return getPtrFromVar(var);
+}
+
+template <>
+std::optional<llvm::TypeSize>
+OpenACCMappableModel<fir::SequenceType>::getSizeInBytes(
+    mlir::Type type, mlir::Value var, mlir::ValueRange accBounds,
+    const mlir::DataLayout &dataLayout) const {
+  // TODO: Bounds operation affect the total size - add support to take them
+  // into account.
+  if (!accBounds.empty())
+    return {};
+
+  // Dynamic extents or unknown ranks generally do not have compile-time
+  // computable dimensions.
+  auto seqType = mlir::cast<fir::SequenceType>(type);
+  if (seqType.hasDynamicExtents() || seqType.hasUnknownShape())
+    return {};
+
+  // Attempt to find an operation that a lookup for KindMapping can be done
+  // from.
+  mlir::Operation *kindMapSrcOp = var.getDefiningOp();
+  if (!kindMapSrcOp) {
+    kindMapSrcOp = var.getParentRegion()->getParentOp();
+    if (!kindMapSrcOp)
+      return {};
+  }
+  auto kindMap = fir::getKindMapping(kindMapSrcOp);
+
+  auto sizeAndAlignment =
+      fir::getTypeSizeAndAlignment(var.getLoc(), type, dataLayout, kindMap);
+  if (!sizeAndAlignment.has_value())
+    return {};
+
+  return {llvm::TypeSize::getFixed(sizeAndAlignment->first)};
+}
+
+template <>
+std::optional<llvm::TypeSize>
+OpenACCMappableModel<fir::BaseBoxType>::getSizeInBytes(
+    mlir::Type type, mlir::Value var, mlir::ValueRange accBounds,
+    const mlir::DataLayout &dataLayout) const {
+  // If we have a box value instead of box reference, the intent is to
+  // get the size of the data not the box itself.
+  if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(var.getType())) {
+    if (auto mappableTy = mlir::dyn_cast<mlir::acc::MappableType>(
+            fir::unwrapRefType(boxTy.getEleTy()))) {
+      return mappableTy.getSizeInBytes(var, accBounds, dataLayout);
+    }
+  }
+  // Size for boxes is not computable until it gets materialized.
+  return {};
+}
+
+template <>
+std::optional<int64_t>
+OpenACCMappableModel<fir::SequenceType>::getOffsetInBytes(
+    mlir::Type type, mlir::Value var, mlir::ValueRange accBounds,
+    const mlir::DataLayout &dataLayout) const {
+  // TODO: Bounds operation affect the offset- add support to take them
+  // into account.
+  if (!accBounds.empty())
+    return {};
+
+  // Dynamic extents (aka descriptor-based arrays) - may have a offset.
+  // For example, a negative stride may mean a negative offset to compute the
+  // start of array.
+  auto seqType = mlir::cast<fir::SequenceType>(type);
+  if (seqType.hasDynamicExtents() || seqType.hasUnknownShape())
+    return {};
+
+  // We have non-dynamic extents - but if for some reason the size is not
+  // computable - assume offset is not either. Otherwise, it is an offset of
+  // zero.
+  if (getSizeInBytes(type, var, accBounds, dataLayout).has_value()) {
+    return {0};
+  }
+  return {};
+}
+
+template <>
+std::optional<int64_t> OpenACCMappableModel<fir::BaseBoxType>::getOffsetInBytes(
+    mlir::Type type, mlir::Value var, mlir::ValueRange accBounds,
+    const mlir::DataLayout &dataLayout) const {
+  // If we have a box value instead of box reference, the intent is to
+  // get the offset of the data not the offset of the box itself.
+  if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(var.getType())) {
+    if (auto mappableTy = mlir::dyn_cast<mlir::acc::MappableType>(
+            fir::unwrapRefType(boxTy.getEleTy()))) {
+      return mappableTy.getOffsetInBytes(var, accBounds, dataLayout);
+    }
+  }
+  // Until boxes get materialized, the offset is not evident because it is
+  // relative to the pointer being held.
+  return {};
+}
+
+template <>
+llvm::SmallVector<mlir::Value>
+OpenACCMappableModel<fir::SequenceType>::generateAccBounds(
+    mlir::Type type, mlir::Value var, mlir::OpBuilder &builder) const {
+  assert((mlir::isa<mlir::acc::PointerLikeType>(var.getType()) ||
+          mlir::isa<mlir::acc::MappableType>(var.getType())) &&
+         "must be pointer-like or mappable");
+
+  fir::FirOpBuilder firBuilder(builder, var.getDefiningOp());
+  auto seqType = mlir::cast<fir::SequenceType>(type);
+  mlir::Location loc = var.getLoc();
+
+  mlir::Value varPtr =
+      mlir::isa<mlir::acc::PointerLikeType>(var.getType())
+          ? var
+          : mlir::cast<mlir::acc::MappableType>(var.getType()).getVarPtr(var);
+
+  if (seqType.hasDynamicExtents() || seqType.hasUnknownShape()) {
+    if (auto boxAddr =
+            mlir::dyn_cast_if_present<fir::BoxAddrOp>(varPtr.getDefiningOp())) {
+      mlir::Value box = boxAddr.getVal();
+      auto res =
+          hlfir::translateToExtendedValue(loc, firBuilder, hlfir::Entity(box));
+      fir::ExtendedValue exv = res.first;
+      mlir::Value boxRef = box;
+      if (auto boxPtr = getPtrFromVar(box)) {
+        boxRef = boxPtr;
+      }
+      // TODO: Handle Fortran optional.
+      const mlir::Value isPresent;
+      Fortran::lower::AddrAndBoundsInfo info(box, boxRef, isPresent,
+                                             box.getType());
+      return Fortran::lower::genBoundsOpsFromBox<mlir::acc::DataBoundsOp,
+                                                 mlir::acc::DataBoundsType>(
+          firBuilder, loc, exv, info);
+    }
+    assert(false && "array with unknown dimension expected to have descriptor");
+    return {};
+  }
+
+  // TODO: Detect assumed-size case.
+  const bool isAssumedSize = false;
+  auto valToCheck = varPtr;
+  if (auto boxAddr =
+          mlir::dyn_cast_if_present<fir::BoxAddrOp>(varPtr.getDefiningOp())) {
+    valToCheck = boxAddr.getVal();
+  }
+  auto res = hlfir::translateToExtendedValue(loc, firBuilder,
+                                             hlfir::Entity(valToCheck));
+  fir::ExtendedValue exv = res.first;
+  return Fortran::lower::genBaseBoundsOps<mlir::acc::DataBoundsOp,
+                                          mlir::acc::DataBoundsType>(
+      firBuilder, loc, exv,
+      /*isAssumedSize=*/isAssumedSize);
+}
+
+template <>
+llvm::SmallVector<mlir::Value>
+OpenACCMappableModel<fir::BaseBoxType>::generateAccBounds(
+    mlir::Type type, mlir::Value var, mlir::OpBuilder &builder) const {
+  // If we have a box value instead of box reference, the intent is to
+  // get the bounds of the data not the bounds of the box itself.
+  if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(var.getType())) {
+    if (auto mappableTy = mlir::dyn_cast<mlir::acc::MappableType>(
+            fir::unwrapRefType(boxTy.getEleTy()))) {
+      mlir::Value data = builder.create<fir::BoxAddrOp>(var.getLoc(), var);
+      return mappableTy.generateAccBounds(data, builder);
+    }
+  }
+  // Box references are not arrays - thus generating acc.bounds does not make
+  // sense.
+  return {};
+}
+
+} // namespace fir::acc

flang/lib/Optimizer/OpenACC/RegisterOpenACCExtensions.cpp

@@ -0,0 +1,28 @@
+//===-- RegisterOpenACCExtensions.cpp -------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Registration for OpenACC extensions as applied to FIR dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/OpenACC/RegisterOpenACCExtensions.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/OpenACC/FIROpenACCTypeInterfaces.h"
+
+namespace fir::acc {
+void registerOpenACCExtensions(mlir::DialectRegistry &registry) {
+  registry.addExtension(+[](mlir::MLIRContext *ctx,
+                            fir::FIROpsDialect *dialect) {
+    fir::SequenceType::attachInterface<OpenACCMappableModel<fir::SequenceType>>(
+        *ctx);
+    fir::BoxType::attachInterface<OpenACCMappableModel<fir::BaseBoxType>>(*ctx);
+  });
+}
+
+} // namespace fir::acc