[flang-commits] [flang] [flang] add AllocationPolicy attribute to module and use it in InlineHLFIRCopy (PR #222013)

[via flang-commits]

================
@@ -38,23 +45,81 @@ static llvm::cl::opt<bool> noInlineHLFIRCopy(
     llvm::cl::init(false));
 
 namespace {
+/// Everything needed to compute the constant byte size of a buffer, gathered
+/// once by the pass since it is module-level information.
+struct SizeContext {
+  std::optional<mlir::DataLayout> dataLayout;
+  std::optional<fir::KindMapping> kindMap;
+};
+
+/// Gather the module level information needed to compute buffer sizes. Without
+/// a data layout no size can be computed, and all the buffers are then left on
+/// the heap.
+static SizeContext getSizeContext(mlir::Operation *op) {
+  auto module = mlir::dyn_cast<mlir::ModuleOp>(op);
+  if (!module)
+    module = op->getParentOfType<mlir::ModuleOp>();
+  if (!module)
+    return SizeContext{std::nullopt, std::nullopt};
+  return SizeContext{fir::support::getOrSetMLIRDataLayout(
+                         module, /*allowDefaultLayout=*/false),
+                     fir::getKindMapping(module)};
+}
+
 class InlineCopyInConversion : public mlir::OpRewritePattern<hlfir::CopyInOp> {
 public:
-  using mlir::OpRewritePattern<hlfir::CopyInOp>::OpRewritePattern;
+  InlineCopyInConversion(mlir::MLIRContext *context,
+                         const fir::AllocationPolicy &policy,
+                         const SizeContext &sizeContext)
+      : mlir::OpRewritePattern<hlfir::CopyInOp>(context), policy(policy),
+        sizeContext(sizeContext) {}
 
   llvm::LogicalResult
   matchAndRewrite(hlfir::CopyInOp copyIn,
                   mlir::PatternRewriter &rewriter) const override;
+
+private:
+  /// Return true if the copy-in buffer of type \p sequenceType should be
+  /// allocated on the stack rather than on the heap.
+  bool shouldUseStack(mlir::Location loc, mlir::Type sequenceType) const;
+
+  fir::AllocationPolicy policy;
+  const SizeContext &sizeContext;
 };
 
+bool InlineCopyInConversion::shouldUseStack(mlir::Location loc,
+                                            mlir::Type sequenceType) const {
+  // Only buffers with a compile-time constant size are considered. A buffer
+  // with a runtime size would need stack save/restore to avoid growing the
+  // stack when the copy-in is inside a loop. There is also little to gain: for
+  // a big buffer the element-per-element copy costs much more than the
+  // allocation itself.
+  if (fir::hasDynamicSize(sequenceType))
+    return false;
+  if (!sizeContext.dataLayout || !sizeContext.kindMap)
+    return false;
+  auto sizeAndAlignment = fir::getTypeSizeAndAlignment(
+      loc, sequenceType, *sizeContext.dataLayout, *sizeContext.kindMap);
+  if (!sizeAndAlignment)
+    return false;
+
+  fir::PendingAllocationInfo info;
+  info.isTemporary = true;
+  info.isDynamic = false;
+  info.byteSize = static_cast<std::int64_t>(sizeAndAlignment->first);
+  // The per-function stack budget is not tracked here: the
+  // allocation-placement pass sees the fir.alloca generated below and can
+  // still move it back to the heap if the budget turns out to be exceeded.
+  return fir::shouldAllocateOnStack(info, policy, /*stackBytesUsed=*/0);
----------------
jeanPerier wrote:

It is also used in   `bool wantStack = fir::shouldAllocateOnStack(info, policy, stackBytesUsed);`  [here in fir::decideAllocationPlacement](https://github.com/jeanPerier/llvm-project/blob/780461d8f609834cd5cf6e4c02622c07c0b47f96/flang/lib/Optimizer/Support/AllocationPolicy.cpp#L80) in AllocationPolicy.cpp.

fir::decideAllocationPlacement is used in the AllocationPlacement pass with an actual threshold (for now quite high at 4Gb).

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