[llvm] [PrologEpilogInserter] Scan all blocks for inline stack probe insertion (PR #195456)

[Rong Mantle Bao via llvm-commits]

https://github.com/CSharperMantle created https://github.com/llvm/llvm-project/pull/195456

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.

>From b3467c85c83849b6d57f5d7531bd36ba23df40ff Mon Sep 17 00:00:00 2001
From: Rong Bao <rong.bao at csmantle.top>
Date: Sat, 2 May 2026 22:45:53 +0800
Subject: [PATCH] [PrologEpilogInserter] Scan all blocks for inline stack probe
 insertion

This prevents incomplete stack probe generation in non-entry blocks, for
example in this snippet targeting RISC-V:

        target triple = "riscv64-unknown-linux-gnu"

        define void @f(i64 %n) #0 {
        entry:
          %v = alloca i32, i64 %n, align 4
          call void @g(ptr %v, [3000 x i64] poison)
          ret void
        }

        declare void @g(ptr, [3000 x i64])

        attributes #0 = { uwtable "frame-pointer"="none" "probe-stack"="inline-asm" }
---
 llvm/lib/CodeGen/PrologEpilogInserter.cpp | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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