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

[via llvm-branch-commits]

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

>From 11640c8923e02979ef511199c620e4a2f748898e 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 882f02032a3b8..dc874937fa45b 100644
--- a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
+++ b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
@@ -58,6 +58,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 8d867612d405c..007cfe306e547 100644
--- a/flang/include/flang/Optimizer/Passes/Pipelines.h
+++ b/flang/include/flang/Optimizer/Passes/Pipelines.h
@@ -85,6 +85,8 @@ void addAVC(mlir::PassManager &pm, const llvm::OptimizationLevel &optLevel);
 
 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 d461c1b9757b5..b9b34c9bb3e87 100644
--- a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
+++ b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
@@ -62,6 +62,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 879a224824119..f51e55aeb92c5 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -80,6 +80,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;
@@ -223,7 +231,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
+}