[flang] Code generation for fir.pack/unpack_array. (#132080)

The code generation relies on `ShallowCopyDirect` runtime
to copy data between the original and the temporary arrays
(both directions). The allocations are done by the compiler
generated code. The heap allocations could have been passed
to `ShallowCopy` runtime, but I decided to expose the allocations
so that the temporary descriptor passed to `ShallowCopyDirect`
has `nocapture` - maybe this will be better for LLVM optimizations.

Author: vzakhari

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -268,6 +268,40 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
                       mlir::ValueRange lenParams = {},
                       llvm::ArrayRef<mlir::NamedAttribute> attrs = {});
 
+  /// Sample genDeclare callback for createArrayTemp() below.
+  /// It creates fir.declare operation using the given operands.
+  /// \p memref is the base of the allocated temporary,
+  /// which may be !fir.ref<!fir.array<>> or !fir.ref<!fir.box/class<>>.
+  static mlir::Value genTempDeclareOp(fir::FirOpBuilder &builder,
+                                      mlir::Location loc, mlir::Value memref,
+                                      llvm::StringRef name, mlir::Value shape,
+                                      llvm::ArrayRef<mlir::Value> typeParams,
+                                      fir::FortranVariableFlagsAttr attrs);
+
+  /// Create a temporary array with the given \p arrayType,
+  /// \p shape, \p extents and \p typeParams. An optional
+  /// \p polymorphicMold specifies the entity which dynamic type
+  /// has to be used for the allocation.
+  /// \p genDeclare callback generates a declare operation
+  /// for the created temporary. FIR passes may use genTempDeclareOp()
+  /// function above that creates fir.declare.
+  /// HLFIR passes may provide their own callback that generates
+  /// hlfir.declare. Some passes may provide a callback that
+  /// just passes through the base of the temporary.
+  /// If \p useStack is true, the function will try to do the allocation
+  /// in stack memory (which is not always possible currently).
+  /// The first return value is the base of the temporary object,
+  /// which may be !fir.ref<!fir.array<>> or !fir.ref<!fir.box/class<>>.
+  /// The second return value is true, if the actual allocation
+  /// was done in heap memory.
+  std::pair<mlir::Value, bool>
+  createArrayTemp(mlir::Location loc, fir::SequenceType arrayType,
+                  mlir::Value shape, llvm::ArrayRef<mlir::Value> extents,
+                  llvm::ArrayRef<mlir::Value> typeParams,
+                  const std::function<decltype(genTempDeclareOp)> &genDeclare,
+                  mlir::Value polymorphicMold, bool useStack = false,
+                  llvm::StringRef tmpName = ".tmp.array");
+
   /// Create an LLVM stack save intrinsic op. Returns the saved stack pointer.
   /// The stack address space is fetched from the data layout of the current
   /// module.
@@ -596,6 +630,15 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
     return result;
   }
 
+  /// Compare two pointer-like values using the given predicate.
+  mlir::Value genPtrCompare(mlir::Location loc,
+                            mlir::arith::CmpIPredicate predicate,
+                            mlir::Value ptr1, mlir::Value ptr2) {
+    ptr1 = createConvert(loc, getIndexType(), ptr1);
+    ptr2 = createConvert(loc, getIndexType(), ptr2);
+    return create<mlir::arith::CmpIOp>(loc, predicate, ptr1, ptr2);
+  }
+
 private:
   /// Set attributes (e.g. FastMathAttr) to \p op operation
   /// based on the current attributes setting.
@@ -850,6 +893,17 @@ llvm::SmallVector<mlir::Value> deduceOptimalExtents(mlir::ValueRange extents1,
 ///   %result1 = arith.select %p4, %c0, %e1 : index
 llvm::SmallVector<mlir::Value> updateRuntimeExtentsForEmptyArrays(
     fir::FirOpBuilder &builder, mlir::Location loc, mlir::ValueRange extents);
+
+/// Given \p box of type fir::BaseBoxType representing an array,
+/// the function generates code to fetch the lower bounds,
+/// the extents and the strides from the box. The values are returned via
+/// \p lbounds, \p extents and \p strides.
+void genDimInfoFromBox(fir::FirOpBuilder &builder, mlir::Location loc,
+                       mlir::Value box,
+                       llvm::SmallVectorImpl<mlir::Value> *lbounds,
+                       llvm::SmallVectorImpl<mlir::Value> *extents,
+                       llvm::SmallVectorImpl<mlir::Value> *strides);
+
 } // namespace fir::factory
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

flang/include/flang/Optimizer/CodeGen/CGPasses.td

@@ -99,4 +99,15 @@ def BoxedProcedurePass : Pass<"boxed-procedure", "mlir::ModuleOp"> {
   ];
 }
 
+def LowerRepackArraysPass : Pass<"lower-repack-arrays", "mlir::ModuleOp"> {
+  let summary = "Convert fir.pack/unpack_array to other FIR operations";
+  let description = [{
+    Convert fir.pack/unpack_array operations to other FIR operations
+    and Fortran runtime calls that implement the semantics
+    of packing/unpacking.
+  }];
+  let dependentDialects = ["fir::FIROpsDialect", "mlir::arith::ArithDialect",
+                           "mlir::func::FuncDialect"];
+}
+
 #endif // FORTRAN_OPTIMIZER_CODEGEN_FIR_PASSES

flang/include/flang/Optimizer/CodeGen/CodeGen.h

@@ -26,6 +26,7 @@ struct NameUniquer;
 #define GEN_PASS_DECL_CODEGENREWRITE
 #define GEN_PASS_DECL_TARGETREWRITEPASS
 #define GEN_PASS_DECL_BOXEDPROCEDUREPASS
+#define GEN_PASS_DECL_LOWERREPACKARRAYSPASS
 #include "flang/Optimizer/CodeGen/CGPasses.h.inc"
 
 /// FIR to LLVM translation pass options.

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -11,6 +11,7 @@
 #include "flang/Optimizer/Builder/Character.h"
 #include "flang/Optimizer/Builder/Complex.h"
 #include "flang/Optimizer/Builder/MutableBox.h"
+#include "flang/Optimizer/Builder/Runtime/Allocatable.h"
 #include "flang/Optimizer/Builder/Runtime/Assign.h"
 #include "flang/Optimizer/Builder/Runtime/Derived.h"
 #include "flang/Optimizer/Builder/Todo.h"
@@ -362,6 +363,72 @@ mlir::Value fir::FirOpBuilder::createHeapTemporary(
                                  name, dynamicLength, dynamicShape, attrs);
 }
 
+std::pair<mlir::Value, bool> fir::FirOpBuilder::createArrayTemp(
+    mlir::Location loc, fir::SequenceType arrayType, mlir::Value shape,
+    llvm::ArrayRef<mlir::Value> extents, llvm::ArrayRef<mlir::Value> typeParams,
+    const std::function<decltype(FirOpBuilder::genTempDeclareOp)> &genDeclare,
+    mlir::Value polymorphicMold, bool useStack, llvm::StringRef tmpName) {
+  if (polymorphicMold) {
+    // Create *allocated* polymorphic temporary using the dynamic type
+    // of the mold and the provided shape/extents. The created temporary
+    // array will be written element per element, that is why it has to be
+    // allocated.
+    mlir::Type boxHeapType = fir::HeapType::get(arrayType);
+    mlir::Value alloc = fir::factory::genNullBoxStorage(
+        *this, loc, fir::ClassType::get(boxHeapType));
+    fir::FortranVariableFlagsAttr declAttrs =
+        fir::FortranVariableFlagsAttr::get(
+            getContext(), fir::FortranVariableFlagsEnum::allocatable);
+
+    mlir::Value base = genDeclare(*this, loc, alloc, tmpName,
+                                  /*shape=*/nullptr, typeParams, declAttrs);
+
+    int rank = extents.size();
+    fir::runtime::genAllocatableApplyMold(*this, loc, alloc, polymorphicMold,
+                                          rank);
+    if (!extents.empty()) {
+      mlir::Type idxTy = getIndexType();
+      mlir::Value one = createIntegerConstant(loc, idxTy, 1);
+      unsigned dim = 0;
+      for (mlir::Value extent : extents) {
+        mlir::Value dimIndex = createIntegerConstant(loc, idxTy, dim++);
+        fir::runtime::genAllocatableSetBounds(*this, loc, alloc, dimIndex, one,
+                                              extent);
+      }
+    }
+    if (!typeParams.empty()) {
+      // We should call AllocatableSetDerivedLength() here.
+      // TODO: does the mold provide the length parameters or
+      // the operation itself or should they be in sync?
+      TODO(loc, "polymorphic type with length parameters");
+    }
+    fir::runtime::genAllocatableAllocate(*this, loc, alloc);
+
+    return {base, /*isHeapAllocation=*/true};
+  }
+  mlir::Value allocmem;
+  if (useStack)
+    allocmem = createTemporary(loc, arrayType, tmpName, extents, typeParams);
+  else
+    allocmem =
+        createHeapTemporary(loc, arrayType, tmpName, extents, typeParams);
+  mlir::Value base = genDeclare(*this, loc, allocmem, tmpName, shape,
+                                typeParams, fir::FortranVariableFlagsAttr{});
+  return {base, !useStack};
+}
+
+mlir::Value fir::FirOpBuilder::genTempDeclareOp(
+    fir::FirOpBuilder &builder, mlir::Location loc, mlir::Value memref,
+    llvm::StringRef name, mlir::Value shape,
+    llvm::ArrayRef<mlir::Value> typeParams,
+    fir::FortranVariableFlagsAttr fortranAttrs) {
+  auto nameAttr = mlir::StringAttr::get(builder.getContext(), name);
+  return builder.create<fir::DeclareOp>(loc, memref.getType(), memref, shape,
+                                        typeParams,
+                                        /*dummy_scope=*/nullptr, nameAttr,
+                                        fortranAttrs, cuf::DataAttributeAttr{});
+}
+
 mlir::Value fir::FirOpBuilder::genStackSave(mlir::Location loc) {
   mlir::Type voidPtr = mlir::LLVM::LLVMPointerType::get(
       getContext(), fir::factory::getAllocaAddressSpace(&getDataLayout()));
@@ -1825,3 +1892,29 @@ llvm::SmallVector<mlir::Value> fir::factory::updateRuntimeExtentsForEmptyArrays(
   }
   return newExtents;
 }
+
+void fir::factory::genDimInfoFromBox(
+    fir::FirOpBuilder &builder, mlir::Location loc, mlir::Value box,
+    llvm::SmallVectorImpl<mlir::Value> *lbounds,
+    llvm::SmallVectorImpl<mlir::Value> *extents,
+    llvm::SmallVectorImpl<mlir::Value> *strides) {
+  auto boxType = mlir::dyn_cast<fir::BaseBoxType>(box.getType());
+  assert(boxType && "must be a box");
+  if (!lbounds && !extents && !strides)
+    return;
+
+  unsigned rank = fir::getBoxRank(boxType);
+  assert(rank != 0 && "must be an array of known rank");
+  mlir::Type idxTy = builder.getIndexType();
+  for (unsigned i = 0; i < rank; ++i) {
+    mlir::Value dim = builder.createIntegerConstant(loc, idxTy, i);
+    auto dimInfo =
+        builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy, box, dim);
+    if (lbounds)
+      lbounds->push_back(dimInfo.getLowerBound());
+    if (extents)
+      extents->push_back(dimInfo.getExtent());
+    if (strides)
+      strides->push_back(dimInfo.getByteStride());
+  }
+}

flang/lib/Optimizer/Builder/HLFIRTools.cpp

@@ -95,24 +95,6 @@ getExplicitLbounds(fir::FortranVariableOpInterface var) {
   return {};
 }
 
-static void
-genLboundsAndExtentsFromBox(mlir::Location loc, fir::FirOpBuilder &builder,
-                            hlfir::Entity boxEntity,
-                            llvm::SmallVectorImpl<mlir::Value> &lbounds,
-                            llvm::SmallVectorImpl<mlir::Value> *extents) {
-  assert(mlir::isa<fir::BaseBoxType>(boxEntity.getType()) && "must be a box");
-  mlir::Type idxTy = builder.getIndexType();
-  const int rank = boxEntity.getRank();
-  for (int i = 0; i < rank; ++i) {
-    mlir::Value dim = builder.createIntegerConstant(loc, idxTy, i);
-    auto dimInfo = builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy,
-                                                  boxEntity, dim);
-    lbounds.push_back(dimInfo.getLowerBound());
-    if (extents)
-      extents->push_back(dimInfo.getExtent());
-  }
-}
-
 static llvm::SmallVector<mlir::Value>
 getNonDefaultLowerBounds(mlir::Location loc, fir::FirOpBuilder &builder,
                          hlfir::Entity entity) {
@@ -128,8 +110,8 @@ getNonDefaultLowerBounds(mlir::Location loc, fir::FirOpBuilder &builder,
   if (entity.isMutableBox())
     entity = hlfir::derefPointersAndAllocatables(loc, builder, entity);
   llvm::SmallVector<mlir::Value> lowerBounds;
-  genLboundsAndExtentsFromBox(loc, builder, entity, lowerBounds,
-                              /*extents=*/nullptr);
+  fir::factory::genDimInfoFromBox(builder, loc, entity, &lowerBounds,
+                                  /*extents=*/nullptr, /*strides=*/nullptr);
   return lowerBounds;
 }
 
@@ -1149,8 +1131,8 @@ static fir::ExtendedValue translateVariableToExtendedValue(
       variable.mayHaveNonDefaultLowerBounds()) {
     // This special case avoids generating two sets of identical
     // fir.box_dim to get both the lower bounds and extents.
-    genLboundsAndExtentsFromBox(loc, builder, variable, nonDefaultLbounds,
-                                &extents);
+    fir::factory::genDimInfoFromBox(builder, loc, variable, &nonDefaultLbounds,
+                                    &extents, /*strides=*/nullptr);
   } else {
     extents = getVariableExtents(loc, builder, variable);
     nonDefaultLbounds = getNonDefaultLowerBounds(loc, builder, variable);

flang/lib/Optimizer/CodeGen/CMakeLists.txt

@@ -4,6 +4,7 @@ add_flang_library(FIRCodeGen
   CodeGen.cpp
   CodeGenOpenMP.cpp
   FIROpPatterns.cpp
+  LowerRepackArrays.cpp
   PreCGRewrite.cpp
   TBAABuilder.cpp
   Target.cpp