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

[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);
----------------
jeanPerier wrote:

Why not placing the whole repacking under `if (isPresent)` and just have the fir.if return the fir.box?

This would produce clearer/simpler IR and code in my opinion (the whole non optional repacking code generation can be in its own member function).

Note that it is OK to do an alloca inside a fir.if an use its address outside of it because fir.if is is not an alloca "owner" (it does not have the `AutomaticAllocationScope` trait). We are already relying on this.

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