[flang-commits] [flang] [flang] Wire allocation-placement into the optimizer pipeline (default off) - memory passes unification [3/5] (PR #210745)

[via flang-commits]

https://github.com/jeanPerier updated https://github.com/llvm/llvm-project/pull/210745

>From 922d6396f261735e92f3ab6168de616280f38278 Mon Sep 17 00:00:00 2001
From: root <jperier at nvidia.com>
Date: Mon, 20 Jul 2026 08:40:07 -0700
Subject: [PATCH] [flang] Wire allocation-placement into the optimizer pipeline
 (experimental)

Add a hidden -enable-allocation-placement flag that, when set, replaces the
stack-arrays and memory-allocation-opt passes in the default FIR optimizer
pipeline with the unified allocation-placement pass. The flag is off by
default, so the legacy passes remain the default path and behavior is
unchanged.

When enabled, the pass runs with its default byte-size thresholds; the
-fstack-arrays strategy is forwarded through the new stackArrays option.
---
 .../flang/Optimizer/Passes/CommandLineOpts.h  | 12 +++++++++++
 .../flang/Optimizer/Passes/Pipelines.h        |  2 ++
 .../lib/Optimizer/Passes/CommandLineOpts.cpp  | 16 ++++++++++++++
 flang/lib/Optimizer/Passes/Pipelines.cpp      | 12 ++++++++++-
 .../Transforms/AllocationPlacement.cpp        |  4 ++--
 .../Fir/allocation-placement-pipeline.fir     | 21 +++++++++++++++++++
 6 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 flang/test/Fir/allocation-placement-pipeline.fir

diff --git a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
index 074f9c0a98e51..d42ce8b762e63 100644
--- a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
+++ b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
@@ -57,6 +57,18 @@ extern llvm::cl::opt<bool> disableFirMao;
 extern llvm::cl::opt<bool> enableFirLICM;
 extern llvm::cl::opt<bool> useOldAliasTags;
 
+/// Experimental option to replace the stack-arrays and memory-allocation-opt
+/// passes with the unified allocation-placement pass.
+extern llvm::cl::opt<bool> enableAllocationPlacement;
+
+/// Constant-size arrays up to this many bytes are considered "small" and placed
+/// on the stack by the allocation-placement pass.
+extern llvm::cl::opt<std::size_t> allocationPlacementSmallArraySize;
+
+/// Per-function budget (in bytes) for small arrays placed on the stack by the
+/// allocation-placement pass.
+extern llvm::cl::opt<std::size_t> allocationPlacementStackLimit;
+
 /// CodeGen Passes
 extern llvm::cl::opt<bool> disableCodeGenRewrite;
 extern llvm::cl::opt<bool> disableTargetRewrite;
diff --git a/flang/include/flang/Optimizer/Passes/Pipelines.h b/flang/include/flang/Optimizer/Passes/Pipelines.h
index a824f4cd2050d..e8ead893dd089 100644
--- a/flang/include/flang/Optimizer/Passes/Pipelines.h
+++ b/flang/include/flang/Optimizer/Passes/Pipelines.h
@@ -83,6 +83,8 @@ void addCfgConversionPass(mlir::PassManager &pm,
 
 void addMemoryAllocationOpt(mlir::PassManager &pm);
 
+void addAllocationPlacement(mlir::PassManager &pm, bool stackArrays);
+
 void addCodeGenRewritePass(mlir::PassManager &pm, bool preserveDeclare);
 
 void addTargetRewritePass(mlir::PassManager &pm);
diff --git a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
index bb61916fcfcbb..d2af7dd5ecd64 100644
--- a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
+++ b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
@@ -61,6 +61,22 @@ cl::opt<bool> useOldAliasTags(
              "the FIR alias tags pass"),
     cl::init(false), cl::Hidden);
 EnableOption(FirLICM, "fir-licm", "FIR loop invariant code motion");
+EnableOption(AllocationPlacement, "allocation-placement",
+             "unified array allocation placement (experimental; replaces "
+             "stack-arrays and memory-allocation-opt)");
+
+cl::opt<std::size_t> allocationPlacementSmallArraySize(
+    "allocation-placement-small-array-size",
+    cl::desc("constant-size arrays up to <size> bytes are placed on the stack "
+             "by the allocation-placement pass"),
+    cl::init(64), cl::Hidden);
+
+cl::opt<std::size_t> allocationPlacementStackLimit(
+    "allocation-placement-stack-limit",
+    cl::desc(
+        "per-function budget in bytes for small arrays placed on the stack "
+        "by the allocation-placement pass"),
+    cl::init(4ull * 1024 * 1024), cl::Hidden);
 
 /// CodeGen Passes
 DisableOption(CodeGenRewrite, "codegen-rewrite", "rewrite FIR for codegen");
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index 1ab59b0a53a6b..f40d99aa5a66c 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -73,6 +73,14 @@ void addMemoryAllocationOpt(mlir::PassManager &pm) {
   });
 }
 
+void addAllocationPlacement(mlir::PassManager &pm, bool stackArrays) {
+  fir::AllocationPlacementOptions options;
+  options.stackArrays = stackArrays;
+  options.smallArrayThresholdBytes = allocationPlacementSmallArraySize;
+  options.totalStackLimitBytes = allocationPlacementStackLimit;
+  pm.addPass(fir::createAllocationPlacement(options));
+}
+
 void addCodeGenRewritePass(mlir::PassManager &pm, bool preserveDeclare) {
   fir::CodeGenRewriteOptions options;
   options.preserveDeclare = preserveDeclare;
@@ -215,7 +223,9 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
 
   pm.addPass(mlir::createCSEPass());
 
-  if (pc.StackArrays)
+  if (enableAllocationPlacement)
+    fir::addAllocationPlacement(pm, pc.StackArrays);
+  else if (pc.StackArrays)
     pm.addPass(fir::createStackArrays());
   else
     fir::addMemoryAllocationOpt(pm);
diff --git a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
index 162cb5fbb21fe..d070376dd5637 100644
--- a/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
+++ b/flang/lib/Optimizer/Transforms/AllocationPlacement.cpp
@@ -288,10 +288,10 @@ void AllocationPlacementPass::runOnOperation() {
     mlir::MLIRContext &context = getContext();
     mlir::RewritePatternSet patterns(&context);
     mlir::GreedyRewriteConfig config;
+    // Prevent the pattern driver from merging blocks; otherwise use the default
+    // configuration (folding on) as the legacy stack-arrays pass did.
     config.setRegionSimplificationLevel(
         mlir::GreedySimplifyRegionLevel::Disabled);
-    config.setStrictness(mlir::GreedyRewriteStrictness::ExistingAndNewOps);
-    config.enableFolding(false);
     patterns.insert<fir::AllocMemConversion>(&context, *candidateOps, dl,
                                              kindMap);
     if (mlir::failed(mlir::applyOpPatternsGreedily(
diff --git a/flang/test/Fir/allocation-placement-pipeline.fir b/flang/test/Fir/allocation-placement-pipeline.fir
new file mode 100644
index 0000000000000..e8aa904b2d40d
--- /dev/null
+++ b/flang/test/Fir/allocation-placement-pipeline.fir
@@ -0,0 +1,21 @@
+// Test that the experimental -enable-allocation-placement flag replaces the
+// legacy stack-arrays / memory-allocation-opt passes with the unified
+// allocation-placement pass in the default optimizer pipeline. By default the
+// legacy passes are still used.
+
+// RUN: tco -enable-allocation-placement %s --mlir-pass-statistics --mlir-pass-statistics-display=pipeline 2>&1 | FileCheck %s --check-prefix=NEW
+// RUN: tco %s --mlir-pass-statistics --mlir-pass-statistics-display=pipeline 2>&1 | FileCheck %s --check-prefix=OLD
+
+// REQUIRES: asserts
+
+// NEW-NOT: MemoryAllocationOpt
+// NEW: AllocationPlacement
+// NEW-NOT: MemoryAllocationOpt
+
+// OLD-NOT: AllocationPlacement
+// OLD: MemoryAllocationOpt
+// OLD-NOT: AllocationPlacement
+
+func.func @_QPfoo() {
+  return
+}