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

clementval · 2024-06-12T20:48:51Z

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 #95173

llvmbot · 2024-06-12T20:49:24Z

@llvm/pr-subscribers-flang-driver

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

Changes

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 #95173

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

10 Files Affected:

(modified) flang/include/flang/Optimizer/Transforms/Passes.h (+1)
(modified) flang/include/flang/Optimizer/Transforms/Passes.td (+9)
(modified) flang/include/flang/Tools/CLOptions.inc (+1)
(modified) flang/lib/Optimizer/Transforms/CMakeLists.txt (+1)
(added) flang/lib/Optimizer/Transforms/StackReclaim.cpp (+56)
(modified) flang/test/Driver/bbc-mlir-pass-pipeline.f90 (+4)
(modified) flang/test/Driver/mlir-debug-pass-pipeline.f90 (+4)
(modified) flang/test/Driver/mlir-pass-pipeline.f90 (+4)
(modified) flang/test/Fir/basic-program.fir (+4)
(added) flang/test/Transforms/stack-reclaime.fir (+14)

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h
index 9fa819e2bf502..1ca1539e76fc6 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.h
+++ b/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
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 7a3baca4c19da..27aee5650e75d 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/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 = [{
diff --git a/flang/include/flang/Tools/CLOptions.inc b/flang/include/flang/Tools/CLOptions.inc
index 2a0cfc04aa350..df396e04b2a76 100644
--- a/flang/include/flang/Tools/CLOptions.inc
+++ b/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());
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 5ef930fdb2c2f..149afdf601c93 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/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
diff --git a/flang/lib/Optimizer/Transforms/StackReclaim.cpp b/flang/lib/Optimizer/Transforms/StackReclaim.cpp
new file mode 100644
index 0000000000000..541311aaa183e
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/StackReclaim.cpp
@@ -0,0 +1,56 @@
+//===- 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([&](mlir::Operation *op) {
+    if (!mlir::isa<fir::DoLoopOp>(op))
+      return;
+
+    auto loopOp = mlir::dyn_cast<fir::DoLoopOp>(op);
+    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);
+    }
+  });
+}
diff --git a/flang/test/Driver/bbc-mlir-pass-pipeline.f90 b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
index c94b98c7c5805..5520d750e2ce1 100644
--- a/flang/test/Driver/bbc-mlir-pass-pipeline.f90
+++ b/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
diff --git a/flang/test/Driver/mlir-debug-pass-pipeline.f90 b/flang/test/Driver/mlir-debug-pass-pipeline.f90
index 49b1f8c5c3134..6e9846fa422e5 100644
--- a/flang/test/Driver/mlir-debug-pass-pipeline.f90
+++ b/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
diff --git a/flang/test/Driver/mlir-pass-pipeline.f90 b/flang/test/Driver/mlir-pass-pipeline.f90
index 8e1a3d43edd1c..db4551e93fe64 100644
--- a/flang/test/Driver/mlir-pass-pipeline.f90
+++ b/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
diff --git a/flang/test/Fir/basic-program.fir b/flang/test/Fir/basic-program.fir
index dd184d99cb809..7bbfd709b0aaf 100644
--- a/flang/test/Fir/basic-program.fir
+++ b/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
diff --git a/flang/test/Transforms/stack-reclaime.fir b/flang/test/Transforms/stack-reclaime.fir
new file mode 100644
index 0000000000000..b53cc96035751
--- /dev/null
+++ b/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

llvmbot · 2024-06-12T20:49:24Z

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

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

Changes

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 #95173

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

10 Files Affected:

(modified) flang/include/flang/Optimizer/Transforms/Passes.h (+1)
(modified) flang/include/flang/Optimizer/Transforms/Passes.td (+9)
(modified) flang/include/flang/Tools/CLOptions.inc (+1)
(modified) flang/lib/Optimizer/Transforms/CMakeLists.txt (+1)
(added) flang/lib/Optimizer/Transforms/StackReclaim.cpp (+56)
(modified) flang/test/Driver/bbc-mlir-pass-pipeline.f90 (+4)
(modified) flang/test/Driver/mlir-debug-pass-pipeline.f90 (+4)
(modified) flang/test/Driver/mlir-pass-pipeline.f90 (+4)
(modified) flang/test/Fir/basic-program.fir (+4)
(added) flang/test/Transforms/stack-reclaime.fir (+14)

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h
index 9fa819e2bf502..1ca1539e76fc6 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.h
+++ b/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
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 7a3baca4c19da..27aee5650e75d 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/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 = [{
diff --git a/flang/include/flang/Tools/CLOptions.inc b/flang/include/flang/Tools/CLOptions.inc
index 2a0cfc04aa350..df396e04b2a76 100644
--- a/flang/include/flang/Tools/CLOptions.inc
+++ b/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());
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 5ef930fdb2c2f..149afdf601c93 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/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
diff --git a/flang/lib/Optimizer/Transforms/StackReclaim.cpp b/flang/lib/Optimizer/Transforms/StackReclaim.cpp
new file mode 100644
index 0000000000000..541311aaa183e
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/StackReclaim.cpp
@@ -0,0 +1,56 @@
+//===- 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([&](mlir::Operation *op) {
+    if (!mlir::isa<fir::DoLoopOp>(op))
+      return;
+
+    auto loopOp = mlir::dyn_cast<fir::DoLoopOp>(op);
+    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);
+    }
+  });
+}
diff --git a/flang/test/Driver/bbc-mlir-pass-pipeline.f90 b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
index c94b98c7c5805..5520d750e2ce1 100644
--- a/flang/test/Driver/bbc-mlir-pass-pipeline.f90
+++ b/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
diff --git a/flang/test/Driver/mlir-debug-pass-pipeline.f90 b/flang/test/Driver/mlir-debug-pass-pipeline.f90
index 49b1f8c5c3134..6e9846fa422e5 100644
--- a/flang/test/Driver/mlir-debug-pass-pipeline.f90
+++ b/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
diff --git a/flang/test/Driver/mlir-pass-pipeline.f90 b/flang/test/Driver/mlir-pass-pipeline.f90
index 8e1a3d43edd1c..db4551e93fe64 100644
--- a/flang/test/Driver/mlir-pass-pipeline.f90
+++ b/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
diff --git a/flang/test/Fir/basic-program.fir b/flang/test/Fir/basic-program.fir
index dd184d99cb809..7bbfd709b0aaf 100644
--- a/flang/test/Fir/basic-program.fir
+++ b/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
diff --git a/flang/test/Transforms/stack-reclaime.fir b/flang/test/Transforms/stack-reclaime.fir
new file mode 100644
index 0000000000000..b53cc96035751
--- /dev/null
+++ b/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

vzakhari

LGTM. We should improve this for only constant-sized allocas case in future.

clementval · 2024-06-12T21:09:43Z

LGTM. We should improve this for only constant-sized allocas case in future.

Yeah I started by only doing it for alloca with dynamic size but in the example mentioned in the description we are generating alloca of descriptor not really dynamic allocas.

So maybe we should check for dynamic size or descriptor?

vzakhari · 2024-06-12T21:16:01Z

LGTM. We should improve this for only constant-sized allocas case in future.

Yeah I started by only doing it for alloca with dynamic size but in the example mentioned in the description we are generating alloca of descriptor not really dynamic allocas.

So maybe we should check for dynamic size or descriptor?

Yes, we should probably have a utility mayBeDynamicAlloca that will check for the shape/len operands.

tblah

Thank you for running this on all top level operations! One minor comment

flang/lib/Optimizer/Transforms/StackReclaim.cpp

jeanPerier

Thanks for moving this into a generic pass!
LGTM, please take care of @tblah comment.

tblah

Thanks!

[flang] Add stack reclaim pass to reclaim allocas in loop

9451471

clementval requested review from sscalpone, jeanPerier and vzakhari June 12, 2024 20:48

llvmbot added flang:driver flang Flang issues not falling into any other category flang:fir-hlfir labels Jun 12, 2024

clementval mentioned this pull request Jun 12, 2024

[flang] Insert stacksave/stackrestore when alloca are present in loops #95173

Closed

vzakhari approved these changes Jun 12, 2024

View reviewed changes

tblah reviewed Jun 13, 2024

View reviewed changes

flang/lib/Optimizer/Transforms/StackReclaim.cpp Outdated Show resolved Hide resolved

jeanPerier approved these changes Jun 13, 2024

View reviewed changes

Address comment

e639f5b

tblah approved these changes Jun 14, 2024

View reviewed changes

clementval merged commit 29d857f into llvm:main Jun 14, 2024
7 checks passed

clementval deleted the stack_reclaim branch June 14, 2024 17:22

jpienaar mentioned this pull request Jun 17, 2024

[mlirc] Add missing extern C #95829

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

Uh oh!

clementval commented Jun 12, 2024

Uh oh!

llvmbot commented Jun 12, 2024

Uh oh!

llvmbot commented Jun 12, 2024

Uh oh!

vzakhari left a comment

Uh oh!

clementval commented Jun 12, 2024

Uh oh!

vzakhari commented Jun 12, 2024

Uh oh!

tblah left a comment

Uh oh!

Uh oh!

jeanPerier left a comment

Uh oh!

tblah left a comment

Uh oh!

Uh oh!

Uh oh!

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

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

Uh oh!

Conversation

clementval commented Jun 12, 2024

Uh oh!

llvmbot commented Jun 12, 2024

Uh oh!

llvmbot commented Jun 12, 2024

Uh oh!

vzakhari left a comment

Choose a reason for hiding this comment

Uh oh!

clementval commented Jun 12, 2024

Uh oh!

vzakhari commented Jun 12, 2024

Uh oh!

tblah left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jeanPerier left a comment

Choose a reason for hiding this comment

Uh oh!

tblah left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!