[flang-commits] [flang] 29d857f - [flang] Add stack reclaim pass to reclaim allocas in loop (#95309)

via flang-commits flang-commits at lists.llvm.org
Fri Jun 14 10:22:01 PDT 2024


Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-06-14T10:21:58-07:00
New Revision: 29d857f183fe6159c3265d6cee8c87419eb615ad

URL: https://github.com/llvm/llvm-project/commit/29d857f183fe6159c3265d6cee8c87419eb615ad
DIFF: https://github.com/llvm/llvm-project/commit/29d857f183fe6159c3265d6cee8c87419eb615ad.diff

LOG: [flang] Add stack reclaim pass to reclaim allocas in loop (#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 #95173

Added: 
    flang/lib/Optimizer/Transforms/StackReclaim.cpp
    flang/test/Transforms/stack-reclaime.fir

Modified: 
    flang/include/flang/Optimizer/Transforms/Passes.h
    flang/include/flang/Optimizer/Transforms/Passes.td
    flang/include/flang/Tools/CLOptions.inc
    flang/lib/Optimizer/Transforms/CMakeLists.txt
    flang/test/Driver/bbc-mlir-pass-pipeline.f90
    flang/test/Driver/mlir-debug-pass-pipeline.f90
    flang/test/Driver/mlir-pass-pipeline.f90
    flang/test/Fir/basic-program.fir

Removed: 
    


################################################################################
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..e5e0e4eab8298
--- /dev/null
+++ b/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);
+    }
+  });
+}

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


        


More information about the flang-commits mailing list