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

I think the best would be for fir.pack_array to propagate the bounds from its input fir.box.

I guess the issue is that there is no easy tool/helper like `readExtents` to extract that for you in FIR, maybe you could move and expose `genLboundsAndExtentsFromBox` to FIRBuilder.cpp.

The issue I see with the fir.shift is that it is extra complexity to the operation a bit outside of the operation scope. Plus, if there is no repacking, then returning the original fir.box would be treacherous because its bounds would not be guaranteed to be the one of the fir.shift.

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