[llvm] [BOLT][AArch64] Tweak heuristics for epilogue recognition (PR #169584)

YongKang Zhu via llvm-commits llvm-commits at lists.llvm.org
Tue Nov 25 16:12:19 PST 2025


https://github.com/yozhu created https://github.com/llvm/llvm-project/pull/169584

As commented in the code, if we see a load into LR from stack and there is no instruction saving LR onto stack, we just assume the basic block is epilogue.

This is not meant to accurately recognize epilogue in all possible cases, but to have BOLT be conservative on treating basic block as epilogue and then turning indirect branch with unknown control flow to tail call.

>From 187f547c2752afda617a1c3b2922458701f6c24b Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at fb.com>
Date: Tue, 11 Nov 2025 14:45:12 -0800
Subject: [PATCH] [BOLT][AArch64] Tweak heuristic for epilogue recognition

As commented in the code, if we see a load into LR from stack
and there is no instruction saving LR onto stack, we just assume
the basic block is epilogue.

This is not meant to accurately recognize epilogue in all possible
cases, but to have BOLT be conservative on treating basic block as
epilogue and then turning indirect branch with unknown control flow
to tail call.
---
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 60 ++++++++++---------
 bolt/test/AArch64/epilogue-determination.s    | 50 ++++++++++++++--
 2 files changed, 76 insertions(+), 34 deletions(-)

diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index f1291f676f1b5..a8480d1db1ca3 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -170,46 +170,48 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return isLoadFromStack(Inst);
   }
 
-  // We look for instructions that load from stack or make stack pointer
-  // adjustment, and assume the basic block is an epilogue if and only if
-  // such instructions are present and also immediately precede the branch
-  // instruction that ends the basic block.
+  // We look for instruction that saves LR to or restores LR from stack.
+  //
+  // If we ever see an LR save to stack, we assume this block is not an
+  // epilogue.
+  //
+  // If there is no LR save in the block and we see an LR restore from
+  // stack, we assume it is an epilogue.
+  //
+  // If neither is seen, we assume it is not an epilogue.
+  //
+  // This is not meant to accurately recognize epilogue in all possible
+  // cases, but to have BOLT be conservative on treating basic block as
+  // epilogue and then turning indirect branch with unknown control flow
+  // to tail call.
   bool isEpilogue(const BinaryBasicBlock &BB) const override {
     if (BB.succ_size())
       return false;
 
-    bool SeenLoadFromStack = false;
-    bool SeenStackPointerAdjustment = false;
-    for (const MCInst &Instr : BB) {
+    bool SeenLRRestoreFromStack = false;
+    for (auto It = BB.rbegin(); It != BB.rend(); ++It) {
+      const MCInst &Instr = *It;
       // Skip CFI pseudo instruction.
       if (isCFI(Instr))
         continue;
-
-      bool IsPop = isPop(Instr);
-      // A load from stack instruction could do SP adjustment in pre-index or
-      // post-index form, which we can skip to check for epilogue recognition
-      // purpose.
-      bool IsSPAdj = (isADD(Instr) || isMOVW(Instr)) &&
-                     Instr.getOperand(0).isReg() &&
-                     Instr.getOperand(0).getReg() == AArch64::SP;
-      SeenLoadFromStack |= IsPop;
-      SeenStackPointerAdjustment |= IsSPAdj;
-
-      if (!SeenLoadFromStack && !SeenStackPointerAdjustment)
-        continue;
-      if (IsPop || IsSPAdj || isPAuthOnLR(Instr))
-        continue;
       if (isReturn(Instr))
         return true;
-      if (isBranch(Instr))
-        break;
 
-      // Any previously seen load from stack or stack adjustment instruction
-      // is definitely not part of epilogue code sequence, so reset these two.
-      SeenLoadFromStack = false;
-      SeenStackPointerAdjustment = false;
+      if (isStoreToStack(Instr)) {
+        for (const MCOperand &Operand : useOperands(Instr)) {
+          if (Operand.isReg() && Operand.getReg() == AArch64::LR) {
+            return false;
+          }
+        }
+      } else if (isLoadFromStack(Instr)) {
+        for (const MCOperand &Operand : defOperands(Instr)) {
+          if (Operand.isReg() && Operand.getReg() == AArch64::LR) {
+            SeenLRRestoreFromStack = true;
+          }
+        }
+      }
     }
-    return SeenLoadFromStack || SeenStackPointerAdjustment;
+    return SeenLRRestoreFromStack;
   }
 
   void createCall(MCInst &Inst, const MCSymbol *Target,
diff --git a/bolt/test/AArch64/epilogue-determination.s b/bolt/test/AArch64/epilogue-determination.s
index 437d8149c0d6b..47f9de8fdcd01 100644
--- a/bolt/test/AArch64/epilogue-determination.s
+++ b/bolt/test/AArch64/epilogue-determination.s
@@ -1,5 +1,6 @@
 # Test that we will not incorrectly take the first basic block in function
-# `_foo` as epilogue due to the first load from stack instruction.
+# `_foo` or the second basic block in function `_goo` as epilogue, and will
+# recognize epilogues in the other cases.
 
 # RUN: %clang %cflags %s -o %t.so -Wl,-q
 # RUN: llvm-bolt %t.so -o %t.bolt --print-cfg | FileCheck %s
@@ -13,8 +14,9 @@ _foo:
   ldrsw x9, [x10, x9, lsl #2]
   add x10, x10, x9
   br x10
-# CHECK-NOT: x10 # TAILCALL
-# CHECK: x10 # UNKNOWN CONTROL FLOW
+# CHECK-LABEL: _foo
+# CHECK: br x10
+# CHECK-SAME: # UNKNOWN CONTROL FLOW
   mov x0, 0
   ret
   mov x0, 1
@@ -31,13 +33,51 @@ _bar:
   stp x29, x30, [sp, #-0x10]!
   mov x29, sp
   sub sp, sp, #0x10
+  tbnz x0, #0x3, _L2
   ldr x8, [x29, #0x30]
   blr x8
+_L2:
   add sp, sp, #0x10
   ldp x29, x30, [sp], #0x10
   br x2
-# CHECK-NOT: x2 # UNKNOWN CONTROL FLOW
-# CHECK: x2 # TAILCALL
+# CHECK-LABEL: _bar
+# CHECK: br x2
+# CHECK-SAME: # TAILCALL
+
+  .global _goo
+  .type _goo, %function
+_goo:
+  ldr w8, [sp]
+  adr x10, _jmptbl2
+  ldrsw x9, [x10, x9, lsl #2]
+  add x10, x10, x9
+  str x30, [sp, #-0x10]!
+  bl _bar
+  ldr x30, [sp], #0x10
+  mov x1, x0
+  mov x0, xzr
+  br x10
+# CHECK-LABEL: _goo
+# CHECK: br x10
+# CHECK-SAME: # UNKNOWN CONTROL FLOW
+  mov x0, 0
+  ret
+  mov x0, 1
+  ret
+
+  .global _faz
+  .type _faz, %function
+_faz:
+  str x30, [sp, #-0x10]!
+  tbnz x0, #0x1, _L3
+  bl _bar
+_L3:
+  ldr x30, [sp], #0x10
+  mov x0, x1
+  br x3
+# CHECK-LABEL: _faz
+# CHECK: br x3
+# CHECK-SAME: # TAILCALL
 
   .global _start
   .type _start, %function



More information about the llvm-commits mailing list