[BOLT][AArch64] Add more heuristics on epilogue determination

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -784,6 +784,11 @@ class MCPlusBuilder {
 
   virtual bool isPop(const MCInst &Inst) const { return false; }
 
+  /// Determine if a basic block looks like an epilogue. For now it is only
+  /// called at the final stage of building CFG to check basic block ending
+  /// with an indirect call that has unknown control flow attribute.
+  virtual bool isEpilogue(const BinaryBasicBlock &BB) const { return false; }
+
   /// Return true if the instruction is used to terminate an indirect branch.
   virtual bool isTerminateBranch(const MCInst &Inst) const {
     llvm_unreachable("not implemented");

bolt/lib/Core/BinaryFunction.cpp

@@ -2177,13 +2177,10 @@ bool BinaryFunction::postProcessIndirectBranches(
         continue;
       }
 
-      // If this block contains an epilogue code and has an indirect branch,
-      // then most likely it's a tail call. Otherwise, we cannot tell for sure
-      // what it is and conservatively reject the function's CFG.
-      bool IsEpilogue = llvm::any_of(BB, [&](const MCInst &Instr) {
-        return BC.MIB->isLeave(Instr) || BC.MIB->isPop(Instr);
-      });
-      if (IsEpilogue) {
+      // If this block contains epilogue code and has an indirect branch,
+      // then most likely it's a tail call. Otherwise, we cannot tell for
+      // sure what it is and conservatively reject the function's CFG.
+      if (BC.MIB->isEpilogue(BB)) {
         BC.MIB->convertJmpToTailCall(Instr);
         BB.removeAllSuccessors();
         continue;

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -164,11 +164,53 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
 
   bool isPush(const MCInst &Inst) const override {
     return isStoreToStack(Inst);
-  };
+  }
 
   bool isPop(const MCInst &Inst) const override {
     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.
+  bool isEpilogue(const BinaryBasicBlock &BB) const override {
+    if (BB.succ_size())
+      return false;
+
+    bool SeenLoadFromStack = false;
+    bool SeenStackPointerAdjustment = false;
+    for (const MCInst &Instr : BB) {
+      // 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) || isPAuthAndRet(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;
+    }
+    return SeenLoadFromStack || SeenStackPointerAdjustment;
+  }
 
   void createCall(MCInst &Inst, const MCSymbol *Target,
                   MCContext *Ctx) override {

bolt/lib/Target/X86/X86MCPlusBuilder.cpp

@@ -219,6 +219,12 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     return getPopSize(Inst) == 0 ? false : true;
   }
 
+  bool isEpilogue(const BinaryBasicBlock &BB) const override {
+    return ::llvm::any_of(BB, [&](const MCInst &Instr) {
+      return isLeave(Instr) || isPop(Instr);
+    });
+  }
+
   bool isTerminateBranch(const MCInst &Inst) const override {
     return Inst.getOpcode() == X86::ENDBR32 || Inst.getOpcode() == X86::ENDBR64;
   }

bolt/test/AArch64/epilogue-determination.s

@@ -0,0 +1,48 @@
+# Test that we will not incorrectly take the first basic block in function
+# `_foo` as epilogue due to the first load from stack instruction.
+
+# RUN: %clang %cflags %s -o %t.so -Wl,-q
+# RUN: llvm-bolt %t.so -o %t.bolt --print-cfg | FileCheck %s
+
+  .text
+  .global _foo
+  .type _foo, %function
+_foo:
+  ldr w8, [sp]
+  adr x10, _jmptbl
+  ldrsw x9, [x10, x9, lsl #2]
+  add x10, x10, x9
+  br x10
+# CHECK-NOT: x10 # TAILCALL
+# CHECK: x10 # UNKNOWN CONTROL FLOW
+  mov x0, 0
+  ret
+  mov x0, 1
+  ret
+
+  .balign 4
+_jmptbl:
+  .long -16
+  .long -8
+
+  .global _bar
+  .type _bar, %function
+_bar:
+  stp x29, x30, [sp, #-0x10]!
+  mov x29, sp
+  sub sp, sp, #0x10
+  ldr x8, [x29, #0x30]
+  blr x8
+  add sp, sp, #0x10
+  ldp x29, x30, [sp], #0x10
+  br x2
+# CHECK-NOT: x2 # UNKNOWN CONTROL FLOW
+# CHECK: x2 # TAILCALL
+
+  .global _start
+  .type _start, %function
+_start:
+  ret
+
+  # Dummy relocation to force relocation mode
+  .reloc 0, R_AARCH64_NONE