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

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.

Assisted-by: AI

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;

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);

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");

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);

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(

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
+}