[flang] Add MemoryAllocation pass to the pipeline

[clementval]

Summary

Add the MemoryAllocation pass into the pipeline. Add
the possibilty to pass the options directly within the tool (tco).

This patch is part of the upstreaming effort from fir-dev branch.

[jeanPerier]

LGTM

flang/include/flang/Optimizer/Transforms/Passes.h

 std::unique_ptr<mlir::Pass> createAffineDemotionPass();
 std::unique_ptr<mlir::Pass> createArrayValueCopyPass();
 std::unique_ptr<mlir::Pass> createFirToCfgPass();
 std::unique_ptr<mlir::Pass> createCharacterConversionPass();
 std::unique_ptr<mlir::Pass> createExternalNameConversionPass();
 std::unique_ptr<mlir::Pass> createMemDataFlowOptPass();
 std::unique_ptr<mlir::Pass> createPromoteToAffinePass();
 std::unique_ptr<mlir::Pass> createMemoryAllocationPass();
+std::unique_ptr<mlir::Pass>
+createMemoryAllocationPass(bool dynOnHeap, std::size_t maxStackSize);
 
 // declarative passes
 #define GEN_PASS_REGISTRATION
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_TRANSFORMS_PASSES_H

flang/include/flang/Tools/CLOptions.inc

addNestedPassConditionally<mlir::FuncOp>(
       pm, disableCfgConversion, fir::createFirToCfgPass);
 }
 
 inline void addAVC(mlir::PassManager &pm) {
   addNestedPassConditionally<mlir::FuncOp>(
       pm, disableFirAvc, fir::createArrayValueCopyPass);
 }
 
+inline void addMemoryAllocationOpt(mlir::PassManager &pm) {
+  addNestedPassConditionally<mlir::FuncOp>(pm, disableFirMao, [&]() {
+    return fir::createMemoryAllocationPass(
+        dynamicArrayStackToHeapAllocation, arrayStackAllocationThreshold);
+  });
+}
+
 #if !defined(FLANG_EXCLUDE_CODEGEN)
 inline void addCodeGenRewritePass(mlir::PassManager &pm) {
   addPassConditionally(
       pm, disableCodeGenRewrite, fir::createFirCodeGenRewritePass);
 }
 
 inline void addTargetRewritePass(mlir::PassManager &pm) {
   addPassConditionally(pm, disableTargetRewrite, []() {

 /// \param pm - MLIR pass manager that will hold the pipeline definition
 inline void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm) {
   // simplify the IR
   mlir::GreedyRewriteConfig config;
   config.enableRegionSimplification = false;
   fir::addAVC(pm);
   pm.addNestedPass<mlir::FuncOp>(fir::createCharacterConversionPass());
   pm.addPass(mlir::createCanonicalizerPass(config));
+  fir::addMemoryAllocationOpt(pm);
 
   // The default inliner pass adds the canonicalizer pass with the default
   // configuration. Create the inliner pass with tco config.
   llvm::StringMap<mlir::OpPassManager> pipelines;
   pm.addPass(
       mlir::createInlinerPass(pipelines, defaultFlangInlinerOptPipeline));
   pm.addPass(mlir::createCSEPass());
 

flang/lib/Optimizer/Transforms/MemoryAllocation.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #define DEBUG_TYPE "flang-memory-allocation-opt"
 
 // Number of elements in an array does not determine where it is allocated.
-static constexpr std::size_t UnlimitedArraySize = ~static_cast<std::size_t>(0);
+static constexpr std::size_t unlimitedArraySize = ~static_cast<std::size_t>(0);
 
 namespace {
 struct MemoryAllocationOptions {
   // Always move dynamic array allocations to the heap. This may result in more
   // heap fragmentation, so may impact performance negatively.
   bool dynamicArrayOnHeap = false;
 
   // Number of elements in array threshold for moving to heap. In environments
   // with limited stack size, moving large arrays to the heap can avoid running
   // out of stack space.
-  std::size_t maxStackArraySize = UnlimitedArraySize;
+  std::size_t maxStackArraySize = unlimitedArraySize;
 };
 
 class ReturnAnalysis {
 public:
   ReturnAnalysis(mlir::Operation *op) {
     if (auto func = mlir::dyn_cast<mlir::FuncOp>(op))
       for (mlir::Block &block : func)
         for (mlir::Operation &i : block)

 /// Currently, implements two conversions from stack to heap allocation.
 ///   1. If a stack allocation is an array larger than some threshold value
 ///      make it a heap allocation.
 ///   2. If a stack allocation is an array with a runtime evaluated size make
 ///      it a heap allocation.
 class MemoryAllocationOpt
     : public fir::MemoryAllocationOptBase<MemoryAllocationOpt> {
 public:
+  MemoryAllocationOpt() {
+    // Set options with default values. (See Passes.td.) Note that the
+    // command-line options, e.g. dynamicArrayOnHeap,  are not set yet.
+    options = {dynamicArrayOnHeap, maxStackArraySize};
+  }
+
+  MemoryAllocationOpt(bool dynOnHeap, std::size_t maxStackSize) {
+    // Set options with default values. (See Passes.td.)
+    options = {dynOnHeap, maxStackSize};
+  }
+
+  /// Override `options` if command-line options have been set.
+  inline void useCommandLineOptions() {
+    if (dynamicArrayOnHeap)
+      options.dynamicArrayOnHeap = dynamicArrayOnHeap;
+    if (maxStackArraySize != unlimitedArraySize)
+      options.maxStackArraySize = maxStackArraySize;
+  }
+
   void runOnOperation() override {
     auto *context = &getContext();
     auto func = getOperation();
     mlir::OwningRewritePatternList patterns(context);
     mlir::ConversionTarget target(*context);
-    MemoryAllocationOptions options = {dynamicArrayOnHeap.getValue(),
-                                       maxStackArraySize.getValue()};
+
+    useCommandLineOptions();
+    LLVM_DEBUG(llvm::dbgs()
+               << "dynamic arrays on heap: " << options.dynamicArrayOnHeap
+               << "\nmaximum number of elements of array on stack: "
+               << options.maxStackArraySize << '\n');
 
     // If func is a declaration, skip it.
     if (func.empty())
       return;
 
     const auto &analysis = getAnalysis<ReturnAnalysis>();
 
     target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithmeticDialect,
                            mlir::StandardOpsDialect>();
     target.addDynamicallyLegalOp<fir::AllocaOp>([&](fir::AllocaOp alloca) {
       return keepStackAllocation(alloca, &func.front(), options);
     });
 
     patterns.insert<AllocaOpConversion>(context, analysis.getReturns(func));
     if (mlir::failed(
             mlir::applyPartialConversion(func, target, std::move(patterns)))) {
       mlir::emitError(func.getLoc(),
                       "error in memory allocation optimization\n");
       signalPassFailure();
     }
   }
+
+private:
+  MemoryAllocationOptions options;
 };
 } // namespace
 
 std::unique_ptr<mlir::Pass> fir::createMemoryAllocationPass() {
   return std::make_unique<MemoryAllocationOpt>();
 }
+
+std::unique_ptr<mlir::Pass>
+fir::createMemoryAllocationPass(bool dynOnHeap, std::size_t maxStackSize) {
+  return std::make_unique<MemoryAllocationOpt>(dynOnHeap, maxStackSize);
+}