[llvm] [BOLT] Don't terminate on trap instruction for Linux kernel (PR #87021)

[Maksim Panchenko via llvm-commits]

https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/87021

Under normal circumstances, we terminate basic blocks on a trap instruction. However, Linux kernel may resume execution after hitting a trap (ud2 on x86). Thus, we introduce "--terminal-trap" option that will specify if the trap instruction should terminate the control flow. The option is on by default except for the Linux kernel mode when it's off.

>From ab3ff5f98112ac8a571cb7415c96790042e28797 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Thu, 28 Mar 2024 14:59:05 -0700
Subject: [PATCH] [BOLT] Don't terminate on trap instruction for Linux kernel

Under normal circumstances, we terminate basic blocks on a trap
instruction. However, Linux kernel may resume execution after hitting a
trap (ud2 on x86). Thus, we introduce "--terminal-trap" option that will
specify if the trap instruction should terminate the control flow. The
option is on by default except for the Linux kernel mode when it's off.
---
 bolt/include/bolt/Core/MCPlusBuilder.h   |  4 +---
 bolt/lib/Core/MCPlusBuilder.cpp          | 14 ++++++++++++++
 bolt/lib/Rewrite/RewriteInstance.cpp     | 11 +++++++++--
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp |  7 -------
 bolt/test/X86/linux-bug-table.s          |  4 ++++
 5 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 198a8d8bf48f82..3818e3d66c994a 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -533,9 +533,7 @@ class MCPlusBuilder {
     return Analysis->isReturn(Inst);
   }
 
-  virtual bool isTerminator(const MCInst &Inst) const {
-    return Analysis->isTerminator(Inst);
-  }
+  virtual bool isTerminator(const MCInst &Inst) const;
 
   virtual bool isNoop(const MCInst &Inst) const {
     llvm_unreachable("not implemented");
diff --git a/bolt/lib/Core/MCPlusBuilder.cpp b/bolt/lib/Core/MCPlusBuilder.cpp
index 5b14ad5cdb880f..f22b35a90bc5d6 100644
--- a/bolt/lib/Core/MCPlusBuilder.cpp
+++ b/bolt/lib/Core/MCPlusBuilder.cpp
@@ -12,12 +12,14 @@
 
 #include "bolt/Core/MCPlusBuilder.h"
 #include "bolt/Core/MCPlus.h"
+#include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstrAnalysis.h"
 #include "llvm/MC/MCInstrDesc.h"
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include <cstdint>
 #include <queue>
@@ -28,6 +30,13 @@ using namespace llvm;
 using namespace bolt;
 using namespace MCPlus;
 
+namespace opts {
+cl::opt<bool>
+    TerminalTrap("terminal-trap",
+                 cl::desc("Assume that execution stops at trap instruction"),
+                 cl::init(true), cl::Hidden, cl::cat(BoltCategory));
+}
+
 bool MCPlusBuilder::equals(const MCInst &A, const MCInst &B,
                            CompFuncTy Comp) const {
   if (A.getOpcode() != B.getOpcode())
@@ -121,6 +130,11 @@ bool MCPlusBuilder::equals(const MCTargetExpr &A, const MCTargetExpr &B,
   llvm_unreachable("target-specific expressions are unsupported");
 }
 
+bool MCPlusBuilder::isTerminator(const MCInst &Inst) const {
+  return Analysis->isTerminator(Inst) ||
+         (opts::TerminalTrap && Info->get(Inst.getOpcode()).isTrap());
+}
+
 void MCPlusBuilder::setTailCall(MCInst &Inst) const {
   assert(!hasAnnotation(Inst, MCAnnotation::kTailCall));
   setAnnotationOpValue(Inst, MCAnnotation::kTailCall, true);
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 2ead51ff6a1286..5ca5594117a62f 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -84,6 +84,7 @@ extern cl::opt<JumpTableSupportLevel> JumpTables;
 extern cl::opt<bool> KeepNops;
 extern cl::list<std::string> ReorderData;
 extern cl::opt<bolt::ReorderFunctions::ReorderType> ReorderFunctions;
+extern cl::opt<bool> TerminalTrap;
 extern cl::opt<bool> TimeBuild;
 
 cl::opt<bool> AllowStripped("allow-stripped",
@@ -2033,8 +2034,14 @@ void RewriteInstance::adjustCommandLineOptions() {
   if (opts::Lite)
     BC->outs() << "BOLT-INFO: enabling lite mode\n";
 
-  if (BC->IsLinuxKernel && !opts::KeepNops.getNumOccurrences())
-    opts::KeepNops = true;
+  if (BC->IsLinuxKernel) {
+    if (!opts::KeepNops.getNumOccurrences())
+      opts::KeepNops = true;
+
+    // Linux kernel may resume execution after a trap instruction in some cases.
+    if (!opts::TerminalTrap.getNumOccurrences())
+      opts::TerminalTrap = false;
+  }
 }
 
 namespace {
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 15f95f82177765..8b1894953f3757 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -211,13 +211,6 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     return false;
   }
 
-  // FIXME: For compatibility with old LLVM only!
-  bool isTerminator(const MCInst &Inst) const override {
-    unsigned Opcode = Inst.getOpcode();
-    return Info->get(Opcode).isTerminator() || X86::isUD1(Opcode) ||
-           X86::isUD2(Opcode);
-  }
-
   bool isIndirectCall(const MCInst &Inst) const override {
     return isCall(Inst) &&
            ((getMemoryOperandNo(Inst) != -1) || Inst.getOperand(0).isReg());
diff --git a/bolt/test/X86/linux-bug-table.s b/bolt/test/X86/linux-bug-table.s
index f688a60c977191..63f70a0b35d9fe 100644
--- a/bolt/test/X86/linux-bug-table.s
+++ b/bolt/test/X86/linux-bug-table.s
@@ -40,6 +40,10 @@ _start:
 # CHECK-REOPT-SAME: BugEntry: 2
 
   ret
+## The return instruction is reachable only via preceding ud2. Test that it is
+## treated as a reachable instruction in the Linux kernel mode.
+
+# CHECK-REOPT-NEXT: ret
   .size _start, .-_start