[flang] Add stack reclaim pass to reclaim allocas in loop (llvm#95309)

Some passes in the flang pipeline are creating `fir.alloca` operation
like `hlfir.concat`. When these allocas are located in a loop, the stack
can quickly be used too much leading to segfaults.

This behavior can be seen in
https://github.com/jacobwilliams/json-fortran/blob/master/src/tests/jf_test_36.F90

This patch insert a call to LLVM stacksave/stackrestore in the body of
the loop to reclaim the alloca in its scope.

This PR is an alternative implementation to llvm#95173

Author: clementval

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

@@ -49,6 +49,7 @@ namespace fir {
 #define GEN_PASS_DECL_OPENACCDATAOPERANDCONVERSION
 #define GEN_PASS_DECL_ADDDEBUGINFO
 #define GEN_PASS_DECL_STACKARRAYS
+#define GEN_PASS_DECL_STACKRECLAIM
 #define GEN_PASS_DECL_LOOPVERSIONING
 #define GEN_PASS_DECL_ADDALIASTAGS
 #define GEN_PASS_DECL_OMPMAPINFOFINALIZATIONPASS

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -260,6 +260,15 @@ def StackArrays : Pass<"stack-arrays", "mlir::ModuleOp"> {
   let dependentDialects = [ "fir::FIROpsDialect" ];
 }
 
+def StackReclaim : Pass<"stack-reclaim"> {
+  let summary = "Insert stacksave/stackrestore in region with allocas";
+  let description = [{
+    Insert stacksave/stackrestore in loop region to reclaim alloca done in its
+    scope.
+  }];
+  let dependentDialects = [ "mlir::LLVM::LLVMDialect" ];
+}
+
 def AddAliasTags : Pass<"fir-add-alias-tags", "mlir::ModuleOp"> {
   let summary = "Add tbaa tags to operations that implement FirAliasAnalysisOpInterface";
   let description = [{

flang/include/flang/Tools/CLOptions.inc

@@ -295,6 +295,7 @@ inline void createDefaultFIROptimizerPassPipeline(
   if (pc.AliasAnalysis && !disableFirAliasTags && !useOldAliasTags)
     pm.addPass(fir::createAddAliasTags());
 
+  addNestedPassToAllTopLevelOperations(pm, fir::createStackReclaim);
   // convert control flow to CFG form
   fir::addCfgConversionPass(pm, pc);
   pm.addPass(mlir::createConvertSCFToCFPass());

flang/lib/Optimizer/Transforms/CMakeLists.txt

@@ -21,6 +21,7 @@ add_flang_library(FIRTransforms
   OMPFunctionFiltering.cpp
   OMPMapInfoFinalization.cpp
   OMPMarkDeclareTarget.cpp
+  StackReclaim.cpp
   VScaleAttr.cpp
   FunctionAttr.cpp
   DebugTypeGenerator.cpp

flang/lib/Optimizer/Transforms/StackReclaim.cpp

@@ -0,0 +1,52 @@
+//===- StackReclaim.cpp -- Insert stacksave/stackrestore in region --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Common/Fortran.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Pass/Pass.h"
+
+namespace fir {
+#define GEN_PASS_DEF_STACKRECLAIM
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+using namespace mlir;
+
+namespace {
+
+class StackReclaimPass : public fir::impl::StackReclaimBase<StackReclaimPass> {
+public:
+  using StackReclaimBase<StackReclaimPass>::StackReclaimBase;
+
+  void runOnOperation() override;
+};
+} // namespace
+
+void StackReclaimPass::runOnOperation() {
+  auto *op = getOperation();
+  auto *context = &getContext();
+  mlir::OpBuilder builder(context);
+  mlir::Type voidPtr = mlir::LLVM::LLVMPointerType::get(context);
+
+  op->walk([&](fir::DoLoopOp loopOp) {
+    mlir::Location loc = loopOp.getLoc();
+
+    if (!loopOp.getRegion().getOps<fir::AllocaOp>().empty()) {
+      builder.setInsertionPointToStart(&loopOp.getRegion().front());
+      auto stackSaveOp = builder.create<LLVM::StackSaveOp>(loc, voidPtr);
+
+      auto *terminator = loopOp.getRegion().back().getTerminator();
+      builder.setInsertionPoint(terminator);
+      builder.create<LLVM::StackRestoreOp>(loc, stackSaveOp);
+    }
+  });
+}

flang/test/Driver/bbc-mlir-pass-pipeline.f90

@@ -50,12 +50,16 @@
 
 ! CHECK-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! CHECK-NEXT: 'fir.global' Pipeline
+! CHECK-NEXT:   StackReclaim
 ! CHECK-NEXT:   CFGConversion
 ! CHECK-NEXT: 'func.func' Pipeline
+! CHECK-NEXT:   StackReclaim
 ! CHECK-NEXT:   CFGConversion
 ! CHECK-NEXT: 'omp.declare_reduction' Pipeline
+! CHECK-NEXT:   StackReclaim
 ! CHECK-NEXT:   CFGConversion
 ! CHECK-NEXT: 'omp.private' Pipeline
+! CHECK-NEXT:   StackReclaim
 ! CHECK-NEXT:   CFGConversion
 
 ! CHECK-NEXT: SCFToControlFlow

flang/test/Driver/mlir-debug-pass-pipeline.f90

@@ -77,12 +77,16 @@
 
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! ALL-NEXT:   'fir.global' Pipeline
+! ALL-NEXT:     StackReclaim
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT:   'func.func' Pipeline
+! ALL-NEXT:     StackReclaim
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT:   'omp.declare_reduction' Pipeline
+! ALL-NEXT:     StackReclaim
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT:   'omp.private' Pipeline
+! ALL-NEXT:     StackReclaim
 ! ALL-NEXT:     CFGConversion
 ! ALL-NEXT: SCFToControlFlow
 ! ALL-NEXT: Canonicalizer

flang/test/Driver/mlir-pass-pipeline.f90

@@ -85,12 +85,16 @@
 
 ! ALL-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 ! ALL-NEXT:    'fir.global' Pipeline
+! ALL-NEXT:      StackReclaim
 ! ALL-NEXT:      CFGConversion
 ! ALL-NEXT:    'func.func' Pipeline
+! ALL-NEXT:      StackReclaim
 ! ALL-NEXT:      CFGConversion
 ! ALL-NEXT:   'omp.declare_reduction' Pipeline
+! ALL-NEXT:      StackReclaim
 ! ALL-NEXT:      CFGConversion
 ! ALL-NEXT:   'omp.private' Pipeline
+! ALL-NEXT:      StackReclaim
 ! ALL-NEXT:      CFGConversion
 
 ! ALL-NEXT: SCFToControlFlow

flang/test/Fir/basic-program.fir

@@ -85,12 +85,16 @@ func.func @_QQmain() {
 
 // PASSES-NEXT: Pipeline Collection : ['fir.global', 'func.func', 'omp.declare_reduction', 'omp.private']
 // PASSES-NEXT: 'fir.global' Pipeline
+// PASSES-NEXT:   StackReclaim
 // PASSES-NEXT:   CFGConversion
 // PASSES-NEXT: 'func.func' Pipeline
+// PASSES-NEXT:   StackReclaim
 // PASSES-NEXT:   CFGConversion
 // PASSES-NEXT: 'omp.declare_reduction' Pipeline
+// PASSES-NEXT:   StackReclaim
 // PASSES-NEXT:   CFGConversion
 // PASSES-NEXT: 'omp.private' Pipeline
+// PASSES-NEXT:   StackReclaim
 // PASSES-NEXT:   CFGConversion
 
 // PASSES-NEXT: SCFToControlFlow

flang/test/Transforms/stack-reclaime.fir

@@ -0,0 +1,14 @@
+// RUN: fir-opt --split-input-file --stack-reclaim %s | FileCheck %s
+
+func.func @alloca_in_loop(%lb : index, %ub : index, %step : index, %b : i1, %addr : !fir.ref<index>) {
+  fir.do_loop %iv = %lb to %ub step %step unordered {
+    %0 = fir.alloca !fir.box<!fir.heap<!fir.char<1,?>>>
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @alloca_in_loop
+// CHECK: fir.do_loop
+// CHECK: %[[STACKPTR:.*]] = llvm.intr.stacksave : !llvm.ptr
+// CHECK: %{{.*}} = fir.alloca !fir.box<!fir.heap<!fir.char<1,?>>>
+// CHECK: llvm.intr.stackrestore %0 : !llvm.ptr