[flang-commits] [flang] [flang] Insert stacksave/stackrestore when alloca are present in loops (PR #95173)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

<details>
<summary>Changes</summary>

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. 

note: this would require another solution for unstructured loop that are not lowered to fir.do_loop. 

---
Full diff: https://github.com/llvm/llvm-project/pull/95173.diff


2 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp (+19) 
- (modified) flang/test/Fir/loop01.fir (+19) 


``````````diff
diff --git a/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp b/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp
index a233e7fbdcd1e..e75803d9571cb 100644
--- a/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp
+++ b/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp
@@ -6,6 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Builder/LowLevelIntrinsics.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
@@ -51,6 +53,7 @@ class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> {
   matchAndRewrite(DoLoopOp loop,
                   mlir::PatternRewriter &rewriter) const override {
     auto loc = loop.getLoc();
+    bool hasAllocas = !loop.getBody()->getOps<fir::AllocaOp>().empty();
     mlir::arith::IntegerOverflowFlags flags{};
     if (setNSW)
       flags = bitEnumSet(flags, mlir::arith::IntegerOverflowFlags::nsw);
@@ -72,6 +75,22 @@ class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> {
         rewriter.splitBlock(conditionalBlock, conditionalBlock->begin());
     auto *lastBlock = &loop.getRegion().back();
 
+    // Insert stacksave/stackrestore if there is fir.alloca operation in the
+    // loop.
+    if (hasAllocas) {
+      auto mod = loop.getOperation()->getParentOfType<mlir::ModuleOp>();
+      fir::FirOpBuilder builder(rewriter, mod);
+      builder.setInsertionPointToStart(firstBlock);
+      mlir::func::FuncOp stackSave = fir::factory::getLlvmStackSave(builder);
+      mlir::func::FuncOp stackRestore =
+          fir::factory::getLlvmStackRestore(builder);
+      mlir::Value stackPtr =
+          builder.create<fir::CallOp>(loc, stackSave).getResult(0);
+      auto *terminator = lastBlock->getTerminator();
+      builder.setInsertionPoint(terminator);
+      builder.create<fir::CallOp>(loc, stackRestore, stackPtr);
+    }
+
     // Move the blocks from the DoLoopOp between initBlock and endBlock
     rewriter.inlineRegionBefore(loop.getRegion(), endBlock);
 
diff --git a/flang/test/Fir/loop01.fir b/flang/test/Fir/loop01.fir
index c1cbb522c378c..55de3bd67b32b 100644
--- a/flang/test/Fir/loop01.fir
+++ b/flang/test/Fir/loop01.fir
@@ -542,3 +542,22 @@ func.func @y5(%lo : index, %up : index) -> index {
 // NSW:           fir.call @f3(%[[VAL_7]]) : (i16) -> ()
 // NSW:           return %[[VAL_5]] : index
 // NSW:         }
+
+// -----
+
+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: ^bb1
+// CHECK: ^bb2
+// CHECK:   %[[STACKPTR:.*]] = fir.call @llvm.stacksave.p0() : () -> !fir.ref<i8>
+// CHECK:   fir.alloca !fir.box<!fir.heap<!fir.char<1,?>>>
+// CHECK:   fir.call @llvm.stackrestore.p0(%[[STACKPTR]]) : (!fir.ref<i8>) -> ()
+// CHECK:   cf.br ^bb1
+// CHECK: ^bb3:
+// CHECK:   return

``````````

</details>


https://github.com/llvm/llvm-project/pull/95173