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

[Slava Zakharin via flang-commits]

================
@@ -0,0 +1,423 @@
+//===-- LowerRepackArrays.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
+//
+//===----------------------------------------------------------------------===//
+/// \file
+/// This pass expands fir.pack_array and fir.unpack_array operations
+/// into sequences of other FIR operations and Fortran runtime calls.
+/// This pass is using structured control flow FIR operations such
+/// as fir.if, so its placement in the pipeline should guarantee
+/// further lowering of these operations.
+///
+/// A fir.pack_array operation is converted into a sequence of checks
+/// identifying whether an array needs to be copied into a contiguous
+/// temporary. When the checks pass, a new memory allocation is done
+/// for the temporary array (in either stack or heap memory).
+/// If `fir.pack_array` does not have no_copy attribute, then
+/// the original array is shallow-copied into the temporary.
+///
+/// A fir.unpack_array operations is converted into a check
+/// of whether the original and the temporary arrays are different
+/// memory. When the check passes, the temporary array might be
+/// shallow-copied into the original array, and then the temporary
+/// array is deallocated (if it was allocated in stack memory,
+/// then there is no explicit deallocation).
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/CodeGen/CodeGen.h"
+
+#include "flang/Optimizer/Builder/BoxValue.h"
+#include "flang/Optimizer/Builder/Character.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/MutableBox.h"
+#include "flang/Optimizer/Builder/Runtime/Allocatable.h"
+#include "flang/Optimizer/Builder/Runtime/Transformational.h"
+#include "flang/Optimizer/Builder/Todo.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Support/DataLayout.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace fir {
+#define GEN_PASS_DEF_LOWERREPACKARRAYSPASS
+#include "flang/Optimizer/CodeGen/CGPasses.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "lower-repack-arrays"
+
+namespace {
+class RepackArrayConversion {
+public:
+  RepackArrayConversion(std::optional<mlir::DataLayout> dataLayout)
+      : dataLayout(dataLayout) {}
+
+protected:
+  std::optional<mlir::DataLayout> dataLayout;
+
+  static bool canAllocateTempOnStack(mlir::Value box);
+};
+
+class PackArrayConversion : public mlir::OpRewritePattern<fir::PackArrayOp>,
+                            RepackArrayConversion {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  PackArrayConversion(mlir::MLIRContext *context,
+                      std::optional<mlir::DataLayout> dataLayout)
+      : OpRewritePattern(context), RepackArrayConversion(dataLayout) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::PackArrayOp op,
+                  mlir::PatternRewriter &rewriter) const override;
+
+private:
+  static constexpr llvm::StringRef bufferName = ".repacked";
+
+  static mlir::Value allocateTempBuffer(fir::FirOpBuilder &builder,
+                                        mlir::Location loc, bool useStack,
+                                        mlir::Value origBox,
+                                        llvm::ArrayRef<mlir::Value> extents,
+                                        llvm::ArrayRef<mlir::Value> typeParams);
+};
+
+class UnpackArrayConversion : public mlir::OpRewritePattern<fir::UnpackArrayOp>,
+                              RepackArrayConversion {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  UnpackArrayConversion(mlir::MLIRContext *context,
+                        std::optional<mlir::DataLayout> dataLayout)
+      : OpRewritePattern(context), RepackArrayConversion(dataLayout) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::UnpackArrayOp op,
+                  mlir::PatternRewriter &rewriter) const override;
+};
+} // anonymous namespace
+
+bool RepackArrayConversion::canAllocateTempOnStack(mlir::Value box) {
+  return !fir::isPolymorphicType(box.getType());
+}
+
+mlir::LogicalResult
+PackArrayConversion::matchAndRewrite(fir::PackArrayOp op,
+                                     mlir::PatternRewriter &rewriter) const {
+  mlir::Location loc = op.getLoc();
+  fir::FirOpBuilder builder(rewriter, op.getOperation());
+  if (op.getMaxSize() || op.getMaxElementSize() || op.getMinStride())
+    TODO(loc, "fir.pack_array with constraints");
+  if (op.getHeuristics() != fir::PackArrayHeuristics::None)
+    TODO(loc, "fir.pack_array with heuristics");
+
+  mlir::Value box = op.getArray();
+  llvm::SmallVector<mlir::Value> typeParams(op.getTypeparams().begin(),
+                                            op.getTypeparams().end());
+  // TODO: set non-default lower bounds on fir.pack_array,
+  // so that we can preserve lower bounds in the temporary box.
+  fir::BoxValue boxValue(box, /*lbounds=*/{}, typeParams);
+  mlir::Type boxType = boxValue.getBoxTy();
+  unsigned rank = boxValue.rank();
+  mlir::Type indexType = builder.getIndexType();
+  mlir::Value zero = fir::factory::createZeroValue(builder, loc, indexType);
+
+  // Fetch the extents from the box, and see if the array
+  // is not empty.
+  // If the type params are not explicitly provided, then we must also
+  // fetch the type parameters from the box.
+  //
+  // bool isNotEmpty;
+  // vector<int64_t> extents;
+  // if (IsPresent(box) && !IsContiguous[UpTo](box[, 1])) {
+  //   isNotEmpty = box->base_addr != null;
+  //   extents = SHAPE(box);
----------------
vzakhari wrote:

This is a good point, Jean! Though I am a little bit nervous about placing the alloca under an IF :)
I've been running my testing with the heap allocations so far, and it worked pretty well.  I am going to do what you suggest and test it with using the stack allocations. We'll see what pops up.

https://github.com/llvm/llvm-project/pull/132080