Skip to content

[CodeGen][RISCV] Inline stack probes immediately after allocateStack in eliminateCallFramePseudoInstr#195456

Merged
lenary merged 2 commits into
llvm:mainfrom
CSharperMantle:fix-pei-probe-all-bb
May 8, 2026
Merged

[CodeGen][RISCV] Inline stack probes immediately after allocateStack in eliminateCallFramePseudoInstr#195456
lenary merged 2 commits into
llvm:mainfrom
CSharperMantle:fix-pei-probe-all-bb

Conversation

@CSharperMantle

@CSharperMantle CSharperMantle commented May 2, 2026

Copy link
Copy Markdown
Contributor

This PR adds a call to inlineStackProbe immediately after allocateStack in eliminateCallFramePseudoInstr. This allows code generation for stack probe pseudoinstructions in non-entry BBs.

Fixes #195454.

@github-actions

github-actions Bot commented May 2, 2026

Copy link
Copy Markdown

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

@CSharperMantle CSharperMantle marked this pull request as ready for review May 2, 2026 15:33
@llvmorg-github-actions

Copy link
Copy Markdown

@llvm/pr-subscribers-backend-risc-v

Author: Rong "Mantle" Bao (CSharperMantle)

Changes

This PR extends the search range for inline stack probe insertion to all BBs, instead of the current "SaveBlocks only". This allows code generation for stack probe pseudoinstructions in non-entry BBs.

Fixes #195454.


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

2 Files Affected:

  • (modified) llvm/lib/CodeGen/PrologEpilogInserter.cpp (+8-2)
  • (added) llvm/test/CodeGen/RISCV/stack-probing-dynamic-nonentry.ll (+115)
diff --git a/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/llvm/lib/CodeGen/PrologEpilogInserter.cpp
index 68fd54cf00146..5c7dc6f48e274 100644
--- a/llvm/lib/CodeGen/PrologEpilogInserter.cpp
+++ b/llvm/lib/CodeGen/PrologEpilogInserter.cpp
@@ -1181,8 +1181,14 @@ void PEIImpl::insertPrologEpilogCode(MachineFunction &MF) {
   // Zero call used registers before restoring callee-saved registers.
   insertZeroCallUsedRegs(MF);
 
-  for (MachineBasicBlock *SaveBlock : SaveBlocks)
-    TFI.inlineStackProbe(MF, *SaveBlock);
+  {
+    SmallVector<MachineBasicBlock *, 8> Blocks;
+    for (MachineBasicBlock &MBB : MF)
+      Blocks.push_back(&MBB);
+
+    for (MachineBasicBlock *MBB : Blocks)
+      TFI.inlineStackProbe(MF, *MBB);
+  }
 
   // Emit additional code that is required to support segmented stacks, if
   // we've been asked for it.  This, when linked with a runtime with support
diff --git a/llvm/test/CodeGen/RISCV/stack-probing-dynamic-nonentry.ll b/llvm/test/CodeGen/RISCV/stack-probing-dynamic-nonentry.ll
new file mode 100644
index 0000000000000..411168d445a95
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/stack-probing-dynamic-nonentry.ll
@@ -0,0 +1,115 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=riscv64 -mattr=+m -O2 < %s | FileCheck %s -check-prefix=RV64I
+; RUN: llc -mtriple=riscv32 -mattr=+m -O2 < %s | FileCheck %s -check-prefix=RV32I
+
+; Test that very large outgoing call frames in functions with variable-sized
+; objects get proper stack probing. The outgoing args are large enough to force
+; the PROBED_STACKALLOC path, which must be expanded in a non-entry block.
+
+define void @f(i64 %n) #0 {
+; RV64I-LABEL: f:
+; RV64I:       # %bb.0: # %entry
+; RV64I-NEXT:    addi sp, sp, -16
+; RV64I-NEXT:    .cfi_def_cfa_offset 16
+; RV64I-NEXT:    sd zero, 0(sp)
+; RV64I-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
+; RV64I-NEXT:    sd s0, 0(sp) # 8-byte Folded Spill
+; RV64I-NEXT:    .cfi_offset ra, -8
+; RV64I-NEXT:    .cfi_offset s0, -16
+; RV64I-NEXT:    addi s0, sp, 16
+; RV64I-NEXT:    .cfi_def_cfa s0, 0
+; RV64I-NEXT:    slli a0, a0, 2
+; RV64I-NEXT:    addi a0, a0, 15
+; RV64I-NEXT:    andi a0, a0, -16
+; RV64I-NEXT:    sub a0, sp, a0
+; RV64I-NEXT:    lui a1, 1
+; RV64I-NEXT:  .LBB0_1: # %entry
+; RV64I-NEXT:    # =>This Inner Loop Header: Depth=1
+; RV64I-NEXT:    sub sp, sp, a1
+; RV64I-NEXT:    sd zero, 0(sp)
+; RV64I-NEXT:    bltu a0, sp, .LBB0_1
+; RV64I-NEXT:  # %bb.2: # %entry
+; RV64I-NEXT:    mv sp, a0
+; RV64I-NEXT:    lui a1, 5
+; RV64I-NEXT:    sub t1, sp, a1
+; RV64I-NEXT:    lui t2, 1
+; RV64I-NEXT:  .LBB0_3: # %entry
+; RV64I-NEXT:    # =>This Inner Loop Header: Depth=1
+; RV64I-NEXT:    sub sp, sp, t2
+; RV64I-NEXT:    sd zero, 0(sp)
+; RV64I-NEXT:    bne sp, t1, .LBB0_3
+; RV64I-NEXT:  # %bb.4: # %entry
+; RV64I-NEXT:    addi sp, sp, -2048
+; RV64I-NEXT:    addi sp, sp, -1424
+; RV64I-NEXT:    sd zero, 0(sp)
+; RV64I-NEXT:    call g
+; RV64I-NEXT:    lui a0, 6
+; RV64I-NEXT:    addi a0, a0, -624
+; RV64I-NEXT:    add sp, sp, a0
+; RV64I-NEXT:    addi sp, s0, -16
+; RV64I-NEXT:    .cfi_def_cfa sp, 16
+; RV64I-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
+; RV64I-NEXT:    ld s0, 0(sp) # 8-byte Folded Reload
+; RV64I-NEXT:    .cfi_restore ra
+; RV64I-NEXT:    .cfi_restore s0
+; RV64I-NEXT:    addi sp, sp, 16
+; RV64I-NEXT:    .cfi_def_cfa_offset 0
+; RV64I-NEXT:    ret
+;
+; RV32I-LABEL: f:
+; RV32I:       # %bb.0: # %entry
+; RV32I-NEXT:    addi sp, sp, -16
+; RV32I-NEXT:    .cfi_def_cfa_offset 16
+; RV32I-NEXT:    sw zero, 0(sp)
+; RV32I-NEXT:    sw ra, 12(sp) # 4-byte Folded Spill
+; RV32I-NEXT:    sw s0, 8(sp) # 4-byte Folded Spill
+; RV32I-NEXT:    .cfi_offset ra, -4
+; RV32I-NEXT:    .cfi_offset s0, -8
+; RV32I-NEXT:    addi s0, sp, 16
+; RV32I-NEXT:    .cfi_def_cfa s0, 0
+; RV32I-NEXT:    slli a0, a0, 2
+; RV32I-NEXT:    addi a0, a0, 15
+; RV32I-NEXT:    andi a0, a0, -16
+; RV32I-NEXT:    sub a0, sp, a0
+; RV32I-NEXT:    lui a1, 1
+; RV32I-NEXT:  .LBB0_1: # %entry
+; RV32I-NEXT:    # =>This Inner Loop Header: Depth=1
+; RV32I-NEXT:    sub sp, sp, a1
+; RV32I-NEXT:    sw zero, 0(sp)
+; RV32I-NEXT:    bltu a0, sp, .LBB0_1
+; RV32I-NEXT:  # %bb.2: # %entry
+; RV32I-NEXT:    mv sp, a0
+; RV32I-NEXT:    lui a1, 5
+; RV32I-NEXT:    sub t1, sp, a1
+; RV32I-NEXT:    lui t2, 1
+; RV32I-NEXT:  .LBB0_3: # %entry
+; RV32I-NEXT:    # =>This Inner Loop Header: Depth=1
+; RV32I-NEXT:    sub sp, sp, t2
+; RV32I-NEXT:    sw zero, 0(sp)
+; RV32I-NEXT:    bne sp, t1, .LBB0_3
+; RV32I-NEXT:  # %bb.4: # %entry
+; RV32I-NEXT:    addi sp, sp, -2048
+; RV32I-NEXT:    addi sp, sp, -1456
+; RV32I-NEXT:    sw zero, 0(sp)
+; RV32I-NEXT:    call g
+; RV32I-NEXT:    lui a0, 6
+; RV32I-NEXT:    addi a0, a0, -592
+; RV32I-NEXT:    add sp, sp, a0
+; RV32I-NEXT:    addi sp, s0, -16
+; RV32I-NEXT:    .cfi_def_cfa sp, 16
+; RV32I-NEXT:    lw ra, 12(sp) # 4-byte Folded Reload
+; RV32I-NEXT:    lw s0, 8(sp) # 4-byte Folded Reload
+; RV32I-NEXT:    .cfi_restore ra
+; RV32I-NEXT:    .cfi_restore s0
+; RV32I-NEXT:    addi sp, sp, 16
+; RV32I-NEXT:    .cfi_def_cfa_offset 0
+; RV32I-NEXT:    ret
+entry:
+  %v = alloca i32, i64 %n
+  call void @g(ptr %v, [3000 x i64] poison)
+  ret void
+}
+
+declare void @g(ptr, [3000 x i64])
+
+attributes #0 = { "probe-stack"="inline-asm" }

@CSharperMantle CSharperMantle force-pushed the fix-pei-probe-all-bb branch 3 times, most recently from c35749c to d8fa841 Compare May 2, 2026 16:06
@CSharperMantle

Copy link
Copy Markdown
Contributor Author

CC: @topperc, @rzinsly

Comment on lines +1184 to +1191
{
SmallVector<MachineBasicBlock *> Blocks;
for (MachineBasicBlock &MBB : MF)
Blocks.push_back(&MBB);

for (MachineBasicBlock *MBB : Blocks)
TFI.inlineStackProbe(MF, *MBB);
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this is the best place to fix it, this inlineStackProbe call is intended to expand the PROBED_STACKALLOC added by emitPrologue which is also called only for the SaveBlocks.
Also, this is run by other targets like PowerPC which doesn't need all blocks to be searched here.

I would prefer to add a inlineStackProbe call after allocateStack in RISCVFrameLowering::eliminateCallFramePseudoInstr.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this is the best place to fix it, this inlineStackProbe call is intended to expand the PROBED_STACKALLOC added by emitPrologue which is also called only for the SaveBlocks. Also, this is run by other targets like PowerPC which doesn't need all blocks to be searched here.

Ooh, sorry, I didn't know about that. This was the only call site to TargetFrameLowering::inlineStackProbe(), so I thought it was the place to fix...

I would prefer to add a inlineStackProbe call after allocateStack in RISCVFrameLowering::eliminateCallFramePseudoInstr.

Thanks, I'll fix this in v2.

@CSharperMantle CSharperMantle force-pushed the fix-pei-probe-all-bb branch from d8fa841 to 1862e39 Compare May 6, 2026 14:32
@CSharperMantle CSharperMantle changed the title [PrologEpilogInserter] Scan all blocks for inline stack probe insertion [CodeGen][RISCV] Inline stack probes immediately after allocateStack in eliminateCallFramePseudoInstr May 6, 2026
@CSharperMantle CSharperMantle requested a review from rzinsly May 6, 2026 14:33
@rzinsly

rzinsly commented May 6, 2026

Copy link
Copy Markdown
Contributor

The code looks good to me now, but it's in an out-of-date tree.
You need to rebase it with the latest main so we can merge it.

@CSharperMantle CSharperMantle force-pushed the fix-pei-probe-all-bb branch from 1862e39 to c9b0f97 Compare May 7, 2026 01:01
@CSharperMantle

Copy link
Copy Markdown
Contributor Author

@rzinsly: Sure.

…in eliminateCallFramePseudoInstr

This prevents the PROBED_STACKALLOC pseudoinstr being leaked to and
discarded in later stages, which could lead to underadjusted sp.
@CSharperMantle CSharperMantle force-pushed the fix-pei-probe-all-bb branch from 2619f42 to c8fc1c9 Compare May 7, 2026 15:05

@rzinsly rzinsly left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM
Thank you for finding and fixing this!

@lenary lenary left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

@CSharperMantle

Copy link
Copy Markdown
Contributor Author

@rzinsly, @lenary: Would you mind helping me merge this? Thanks.

@lenary lenary merged commit 589faed into llvm:main May 8, 2026
10 checks passed
@github-actions

github-actions Bot commented May 8, 2026

Copy link
Copy Markdown

@CSharperMantle Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

@CSharperMantle CSharperMantle deleted the fix-pei-probe-all-bb branch May 8, 2026 00:17
@llvm-ci

llvm-ci commented May 8, 2026

Copy link
Copy Markdown

LLVM Buildbot has detected a new failure on builder lldb-aarch64-windows running on linaro-armv8-windows-msvc-05 while building llvm at step 6 "test".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/141/builds/18334

Here is the relevant piece of the build log for the reference
Step 6 (test) failure: build (failure)
...
PASS: lldb-api :: functionalities/breakpoint/breakpoint_locations/TestBreakpointLocations.py (341 of 2524)
PASS: lldb-api :: functionalities/breakpoint/breakpoint_on_lambda_capture/TestBreakOnLambdaCapture.py (342 of 2524)
PASS: lldb-api :: functionalities/breakpoint/breakpoint_names/TestBreakpointNames.py (343 of 2524)
PASS: lldb-api :: functionalities/breakpoint/breakpoint_on_overload/TestBreakOnOverload.py (344 of 2524)
PASS: lldb-api :: functionalities/breakpoint/breakpoint_options/TestBreakpointOptions.py (345 of 2524)
PASS: lldb-api :: functionalities/breakpoint/breakpoint_reset_upon_run/TestBreakpointResetUponRun.py (346 of 2524)
UNSUPPORTED: lldb-api :: functionalities/breakpoint/breakpoint_with_realpath_and_source_map/TestBreakpoint.py (347 of 2524)
PASS: lldb-api :: functionalities/breakpoint/callback_deletes_breakpoints/TestCallbackDeletesBreakpoints.py (348 of 2524)
UNSUPPORTED: lldb-api :: functionalities/breakpoint/comp_dir_symlink/TestCompDirSymLink.py (349 of 2524)
PASS: lldb-api :: functionalities/breakpoint/breakpoint_set_restart/TestBreakpointSetRestart.py (350 of 2524)
FAIL: lldb-api :: functionalities/breakpoint/consecutive_breakpoints/TestConsecutiveBreakpoints.py (351 of 2524)
******************** TEST 'lldb-api :: functionalities/breakpoint/consecutive_breakpoints/TestConsecutiveBreakpoints.py' FAILED ********************
Script:
--
C:/Users/tcwg/scoop/apps/python/current/python.exe C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/llvm-project/lldb\test\API\dotest.py -u CXXFLAGS -u CFLAGS --env LLVM_LIBS_DIR=C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./lib --env LLVM_INCLUDE_DIR=C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/include --env LLVM_TOOLS_DIR=C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./bin --triple aarch64-pc-windows-msvc --build-dir C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/lldb-test-build.noindex --lldb-module-cache-dir C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/lldb-test-build.noindex/module-cache-lldb\lldb-api --clang-module-cache-dir C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/lldb-test-build.noindex/module-cache-clang\lldb-api --executable C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./bin/lldb.exe --compiler C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./bin/clang.exe --dsymutil C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./bin/dsymutil.exe --make C:/Users/tcwg/scoop/shims/make.exe --llvm-tools-dir C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./bin --lldb-obj-root C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/tools/lldb --lldb-libs-dir C:/Users/tcwg/llvm-worker/lldb-aarch64-windows/build/./lib --cmake-build-type Release --skip-category=watchpoint --env LLDB_LAUNCH_FLAG_USE_PIPES=1 C:\Users\tcwg\llvm-worker\lldb-aarch64-windows\llvm-project\lldb\test\API\functionalities\breakpoint\consecutive_breakpoints -p TestConsecutiveBreakpoints.py
--
Exit Code: 1

Command Output (stdout):
--
lldb version 23.0.0git (https://github.com/llvm/llvm-project.git revision 589faedadf141e5e63f7a1e92a0327fc9bdc9b09)
  clang revision 589faedadf141e5e63f7a1e92a0327fc9bdc9b09
  llvm revision 589faedadf141e5e63f7a1e92a0327fc9bdc9b09
Skipping the following test categories: watchpoint, libc++, libstdcxx, dwo, dsym, gmodules, debugserver, objc, fork, pexpect


--
Command Output (stderr):
--
PASS: LLDB (C:\Users\tcwg\llvm-worker\lldb-aarch64-windows\build\bin\clang.exe-aarch64) :: test_continue (TestConsecutiveBreakpoints.ConsecutiveBreakpointsTestCase.test_continue)

PASS: LLDB (C:\Users\tcwg\llvm-worker\lldb-aarch64-windows\build\bin\clang.exe-aarch64) :: test_single_step (TestConsecutiveBreakpoints.ConsecutiveBreakpointsTestCase.test_single_step)

FAIL: LLDB (C:\Users\tcwg\llvm-worker\lldb-aarch64-windows\build\bin\clang.exe-aarch64) :: test_single_step_thread_specific (TestConsecutiveBreakpoints.ConsecutiveBreakpointsTestCase.test_single_step_thread_specific)

Log Files:

 - C:\Users\tcwg\llvm-worker\lldb-aarch64-windows\build\lldb-test-build.noindex\functionalities\breakpoint\consecutive_breakpoints\TestConsecutiveBreakpoints\Failure_test_single_step_thread_specific.log

======================================================================

FAIL: test_single_step_thread_specific (TestConsecutiveBreakpoints.ConsecutiveBreakpointsTestCase.test_single_step_thread_specific)

   Test that single step stops, even though the second breakpoint is not valid.

----------------------------------------------------------------------

Traceback (most recent call last):


@lenary

lenary commented May 8, 2026

Copy link
Copy Markdown
Member

That failure looks irrelevant.

moar55 pushed a commit to moar55/llvm-project that referenced this pull request May 12, 2026
…` in `eliminateCallFramePseudoInstr` (llvm#195456)

This PR adds a call to `inlineStackProbe` immediately after
`allocateStack` in `eliminateCallFramePseudoInstr`. This allows code
generation for stack probe pseudoinstructions in non-entry BBs.

Fixes llvm#195454.
EuphoricThinking pushed a commit to EuphoricThinking/llvm-project that referenced this pull request May 14, 2026
…` in `eliminateCallFramePseudoInstr` (llvm#195456)

This PR adds a call to `inlineStackProbe` immediately after
`allocateStack` in `eliminateCallFramePseudoInstr`. This allows code
generation for stack probe pseudoinstructions in non-entry BBs.

Fixes llvm#195454.
pedroMVicente pushed a commit to pedroMVicente/llvm-project that referenced this pull request May 19, 2026
…` in `eliminateCallFramePseudoInstr` (llvm#195456)

This PR adds a call to `inlineStackProbe` immediately after
`allocateStack` in `eliminateCallFramePseudoInstr`. This allows code
generation for stack probe pseudoinstructions in non-entry BBs.

Fixes llvm#195454.
dyung pushed a commit to CSharperMantle/llvm-project that referenced this pull request Jun 12, 2026
…llFramePseudoInstr` (llvm#195456)

[ Upstream commit 589faed ]

Revert `bltu` in probing loops to `blt` because commit
f162be2 isn't applied on release/22.x
yet.

Link: llvm#192485 ("[RISCV] Use
 unsigned comparison for stack clash probing loop")

---

This PR adds a call to `inlineStackProbe` immediately after
`allocateStack` in `eliminateCallFramePseudoInstr`. This allows code
generation for stack probe pseudoinstructions in non-entry BBs.

Fixes llvm#195454.
daunabomba pushed a commit to daunabomba/llvm-project that referenced this pull request Jun 17, 2026
…llFramePseudoInstr` (llvm#195456)

[ Upstream commit 589faed ]

Revert `bltu` in probing loops to `blt` because commit
f162be2 isn't applied on release/22.x
yet.

Link: llvm#192485 ("[RISCV] Use
 unsigned comparison for stack clash probing loop")

---

This PR adds a call to `inlineStackProbe` immediately after
`allocateStack` in `eliminateCallFramePseudoInstr`. This allows code
generation for stack probe pseudoinstructions in non-entry BBs.

Fixes llvm#195454.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[RISCV] -fstack-clash-protection generate incomplete probes for large outgoing frames in non-entry blocks

4 participants