[CodeGen][KCFI] Move cfi-type lowering to TargetLowering

KCFI machine function passes transform indirect calls with a
cfi-type attribute into architecture-specific type checks bundled
together with the calls. Instead of having a separate pass for each
architecture, add a generic machine function pass for KCFI and
move the architecture-specific code that emits the actual check to
TargetLowering. This avoids unnecessary duplication and makes it
easier to add KCFI support to other architectures.

Reviewed By: nickdesaulniers

Differential Revision: https://reviews.llvm.org/D149915

diff --git a/llvm/include/llvm/CodeGen/Passes.h b/llvm/include/llvm/CodeGen/Passes.h
index 69dc805fcf58a5db13c8813e230a9cb83dd72a20..11bc1d48a93d7c9b8f2d30454f4ba4d128957767 100644 (file)
--- a/llvm/include/llvm/CodeGen/Passes.h
+++ b/llvm/include/llvm/CodeGen/Passes.h
@@ -598,6 +598,9 @@ namespace llvm {
   FunctionPass *createSelectOptimizePass();
 
   FunctionPass *createCallBrPass();
+
+  /// Lowers KCFI operand bundles for indirect calls.
+  FunctionPass *createKCFIPass();
 } // End llvm namespace
 
 #endif

diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index c458c820c8a184fe0f79832dc205475a7cb55364..074a485dbd9368aa21c56d4106b6e01051cc5449 100644 (file)
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -2080,6 +2080,18 @@ public:
     llvm_unreachable("Masked cmpxchg expansion unimplemented on this target");
   }
 
+  //===--------------------------------------------------------------------===//
+  /// \name KCFI check lowering.
+  /// @{
+
+  virtual MachineInstr *EmitKCFICheck(MachineBasicBlock &MBB,
+                                      MachineBasicBlock::instr_iterator &MBBI,
+                                      const TargetInstrInfo *TII) const {
+    llvm_unreachable("KCFI is not supported on this target");
+  }
+
+  /// @}
+
   /// Inserts in the IR a target-specific intrinsic specifying a fence.
   /// It is called by AtomicExpandPass before expanding an
   ///   AtomicRMW/AtomicCmpXchg/AtomicStore/AtomicLoad

diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 9e9aea5841710ed938e640cffd0705527eed87ac..a65bd4ef025b9f020210e4fe1488bf3d8518a447 100644 (file)
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -165,6 +165,7 @@ void initializeInterleavedLoadCombinePass(PassRegistry &);
 void initializeIntervalPartitionPass(PassRegistry&);
 void initializeJMCInstrumenterPass(PassRegistry&);
 void initializeJumpThreadingPass(PassRegistry&);
+void initializeKCFIPass(PassRegistry &);
 void initializeLCSSAVerificationPassPass(PassRegistry&);
 void initializeLCSSAWrapperPassPass(PassRegistry&);
 void initializeLazyBlockFrequencyInfoPassPass(PassRegistry&);

diff --git a/llvm/include/llvm/LinkAllPasses.h b/llvm/include/llvm/LinkAllPasses.h
index 2b844595765b44b2fb175ecf718e37bb87afa5cb..9fc7d817c53c30cb01fe55f841bf509f75643606 100644 (file)
--- a/llvm/include/llvm/LinkAllPasses.h
+++ b/llvm/include/llvm/LinkAllPasses.h
@@ -93,6 +93,7 @@ namespace {
       (void) llvm::createInstSimplifyLegacyPass();
       (void) llvm::createInstructionCombiningPass();
       (void) llvm::createJMCInstrumenterPass();
+      (void) llvm::createKCFIPass();
       (void) llvm::createLCSSAPass();
       (void) llvm::createLICMPass();
       (void) llvm::createLoopSinkPass();

diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 4c4a0ea24cde8497f82e805d6f4c50aacde6a8b1..106571b9061bebd5aaa8cd446005d5fbf4e86648 100644 (file)
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -93,6 +93,7 @@ add_llvm_component_library(LLVMCodeGen
   InterleavedLoadCombinePass.cpp
   IntrinsicLowering.cpp
   JMCInstrumenter.cpp
+  KCFI.cpp
   LatencyPriorityQueue.cpp
   LazyMachineBlockFrequencyInfo.cpp
   LexicalScopes.cpp

diff --git a/llvm/lib/CodeGen/KCFI.cpp b/llvm/lib/CodeGen/KCFI.cpp
new file mode 100644 (file)
index 0000000..6f3779f
--- /dev/null
+++ b/llvm/lib/CodeGen/KCFI.cpp
@@ -0,0 +1,108 @@
+//===---- KCFI.cpp - Implements KCFI --------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass implements KCFI indirect call check lowering.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineInstrBundle.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/InitializePasses.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "kcfi"
+#define KCFI_PASS_NAME "Insert KCFI indirect call checks"
+
+STATISTIC(NumKCFIChecksAdded, "Number of indirect call checks added");
+
+namespace {
+class KCFI : public MachineFunctionPass {
+public:
+  static char ID;
+
+  KCFI() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override { return KCFI_PASS_NAME; }
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+private:
+  /// Machine instruction info used throughout the class.
+  const TargetInstrInfo *TII = nullptr;
+
+  /// Target lowering for arch-specific parts.
+  const TargetLowering *TLI = nullptr;
+
+  /// Emits a KCFI check before an indirect call.
+  /// \returns true if the check was added and false otherwise.
+  bool emitCheck(MachineBasicBlock &MBB,
+                 MachineBasicBlock::instr_iterator I) const;
+};
+
+char KCFI::ID = 0;
+} // end anonymous namespace
+
+INITIALIZE_PASS(KCFI, DEBUG_TYPE, KCFI_PASS_NAME, false, false)
+
+FunctionPass *llvm::createKCFIPass() { return new KCFI(); }
+
+bool KCFI::emitCheck(MachineBasicBlock &MBB,
+                     MachineBasicBlock::instr_iterator MBBI) const {
+  assert(TII && "Target instruction info was not initialized");
+  assert(TLI && "Target lowering was not initialized");
+
+  // If the call instruction is bundled, we can only emit a check safely if
+  // it's the first instruction in the bundle.
+  if (MBBI->isBundled() && !std::prev(MBBI)->isBundle())
+    report_fatal_error("Cannot emit a KCFI check for a bundled call");
+
+  // Emit a KCFI check for the call instruction at MBBI. The implementation
+  // must unfold memory operands if applicable.
+  MachineInstr *Check = TLI->EmitKCFICheck(MBB, MBBI, TII);
+
+  // Clear the original call's CFI type.
+  assert(MBBI->isCall() && "Unexpected instruction type");
+  MBBI->setCFIType(*MBB.getParent(), 0);
+
+  // If not already bundled, bundle the check and the call to prevent
+  // further changes.
+  if (!MBBI->isBundled())
+    finalizeBundle(MBB, Check->getIterator(), std::next(MBBI->getIterator()));
+
+  ++NumKCFIChecksAdded;
+  return true;
+}
+
+bool KCFI::runOnMachineFunction(MachineFunction &MF) {
+  const Module *M = MF.getMMI().getModule();
+  if (!M->getModuleFlag("kcfi"))
+    return false;
+
+  const auto &SubTarget = MF.getSubtarget();
+  TII = SubTarget.getInstrInfo();
+  TLI = SubTarget.getTargetLowering();
+
+  bool Changed = false;
+  for (MachineBasicBlock &MBB : MF) {
+    // Use instr_iterator because we don't want to skip bundles.
+    for (MachineBasicBlock::instr_iterator MII = MBB.instr_begin(),
+                                           MIE = MBB.instr_end();
+         MII != MIE; ++MII) {
+      if (MII->isCall() && MII->getCFIType())
+        Changed |= emitCheck(MBB, MII);
+    }
+  }
+
+  return Changed;
+}

diff --git a/llvm/lib/Target/AArch64/AArch64.h b/llvm/lib/Target/AArch64/AArch64.h
index fad620169eaefa950f7f5d9e8a9340ee17a781f5..76f55666e7435a7944a5e6ed2cb602585502a5a9 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64.h
+++ b/llvm/lib/Target/AArch64/AArch64.h
@@ -42,7 +42,6 @@ FunctionPass *createAArch64ExpandPseudoPass();
 FunctionPass *createAArch64SLSHardeningPass();
 FunctionPass *createAArch64IndirectThunks();
 FunctionPass *createAArch64SpeculationHardeningPass();
-FunctionPass *createAArch64KCFIPass();
 FunctionPass *createAArch64LoadStoreOptimizationPass();
 ModulePass *createAArch64LowerHomogeneousPrologEpilogPass();
 FunctionPass *createAArch64SIMDInstrOptPass();
@@ -86,7 +85,6 @@ void initializeAArch64DAGToDAGISelPass(PassRegistry &);
 void initializeAArch64DeadRegisterDefinitionsPass(PassRegistry&);
 void initializeAArch64ExpandPseudoPass(PassRegistry &);
 void initializeAArch64GlobalsTaggingPass(PassRegistry &);
-void initializeAArch64KCFIPass(PassRegistry &);
 void initializeAArch64LoadStoreOptPass(PassRegistry&);
 void initializeAArch64LowerHomogeneousPrologEpilogPass(PassRegistry &);
 void initializeAArch64MIPeepholeOptPass(PassRegistry &);

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 670781b771fb49c1221770dc38bd91afa2ab6258..0fb8e4245d48e15532307bddb4a111fd1811e776 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -23856,6 +23856,33 @@ bool AArch64TargetLowering::shouldConvertFpToSat(unsigned Op, EVT FPVT,
   return TargetLowering::shouldConvertFpToSat(Op, FPVT, VT);
 }
 
+MachineInstr *
+AArch64TargetLowering::EmitKCFICheck(MachineBasicBlock &MBB,
+                                     MachineBasicBlock::instr_iterator &MBBI,
+                                     const TargetInstrInfo *TII) const {
+  assert(MBBI->isCall() && MBBI->getCFIType() &&
+         "Invalid call instruction for a KCFI check");
+
+  switch (MBBI->getOpcode()) {
+  case AArch64::BLR:
+  case AArch64::BLRNoIP:
+  case AArch64::TCRETURNri:
+  case AArch64::TCRETURNriBTI:
+    break;
+  default:
+    llvm_unreachable("Unexpected CFI call opcode");
+  }
+
+  MachineOperand &Target = MBBI->getOperand(0);
+  assert(Target.isReg() && "Invalid target operand for an indirect call");
+  Target.setIsRenamable(false);
+
+  return BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(AArch64::KCFI_CHECK))
+      .addReg(Target.getReg())
+      .addImm(MBBI->getCFIType())
+      .getInstr();
+}
+
 bool AArch64TargetLowering::enableAggressiveFMAFusion(EVT VT) const {
   return Subtarget->hasAggressiveFMA() && VT.isFloatingPoint();
 }

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index 2df9fc869eb05aa6a297f701c8c7f7cf7182ebe0..2430ad923013801fb0cbcc2aeac718ce11aa80dd 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -862,6 +862,10 @@ public:
 
   bool supportKCFIBundles() const override { return true; }
 
+  MachineInstr *EmitKCFICheck(MachineBasicBlock &MBB,
+                              MachineBasicBlock::instr_iterator &MBBI,
+                              const TargetInstrInfo *TII) const override;
+
   /// Enable aggressive FMA fusion on targets that want it.
   bool enableAggressiveFMAFusion(EVT VT) const override;
 

diff --git a/llvm/lib/Target/AArch64/AArch64KCFI.cpp b/llvm/lib/Target/AArch64/AArch64KCFI.cpp
deleted file mode 100644 (file)
index 271001c..0000000
--- a/llvm/lib/Target/AArch64/AArch64KCFI.cpp
+++ /dev/null
@@ -1,115 +0,0 @@
-//===---- AArch64KCFI.cpp - Implements KCFI -------------------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements KCFI indirect call checking.
-//
-//===----------------------------------------------------------------------===//
-
-#include "AArch64.h"
-#include "AArch64InstrInfo.h"
-#include "AArch64Subtarget.h"
-#include "AArch64TargetMachine.h"
-#include "llvm/ADT/Statistic.h"
-#include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineInstrBuilder.h"
-#include "llvm/CodeGen/MachineInstrBundle.h"
-#include "llvm/CodeGen/MachineModuleInfo.h"
-
-using namespace llvm;
-
-#define DEBUG_TYPE "aarch64-kcfi"
-#define AARCH64_KCFI_PASS_NAME "Insert KCFI indirect call checks"
-
-STATISTIC(NumKCFIChecksAdded, "Number of indirect call checks added");
-
-namespace {
-class AArch64KCFI : public MachineFunctionPass {
-public:
-  static char ID;
-
-  AArch64KCFI() : MachineFunctionPass(ID) {}
-
-  StringRef getPassName() const override { return AARCH64_KCFI_PASS_NAME; }
-  bool runOnMachineFunction(MachineFunction &MF) override;
-
-private:
-  /// Machine instruction info used throughout the class.
-  const AArch64InstrInfo *TII = nullptr;
-
-  /// Emits a KCFI check before an indirect call.
-  /// \returns true if the check was added and false otherwise.
-  bool emitCheck(MachineBasicBlock &MBB,
-                 MachineBasicBlock::instr_iterator I) const;
-};
-
-char AArch64KCFI::ID = 0;
-} // end anonymous namespace
-
-INITIALIZE_PASS(AArch64KCFI, DEBUG_TYPE, AARCH64_KCFI_PASS_NAME, false, false)
-
-FunctionPass *llvm::createAArch64KCFIPass() { return new AArch64KCFI(); }
-
-bool AArch64KCFI::emitCheck(MachineBasicBlock &MBB,
-                            MachineBasicBlock::instr_iterator MBBI) const {
-  assert(TII && "Target instruction info was not initialized");
-
-  // If the call instruction is bundled, we can only emit a check safely if
-  // it's the first instruction in the bundle.
-  if (MBBI->isBundled() && !std::prev(MBBI)->isBundle())
-    report_fatal_error("Cannot emit a KCFI check for a bundled call");
-
-  switch (MBBI->getOpcode()) {
-  case AArch64::BLR:
-  case AArch64::BLRNoIP:
-  case AArch64::TCRETURNri:
-  case AArch64::TCRETURNriBTI:
-    break;
-  default:
-    llvm_unreachable("Unexpected CFI call opcode");
-  }
-
-  MachineOperand &Target = MBBI->getOperand(0);
-  assert(Target.isReg() && "Invalid target operand for an indirect call");
-  Target.setIsRenamable(false);
-
-  MachineInstr *Check =
-      BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(AArch64::KCFI_CHECK))
-          .addReg(Target.getReg())
-          .addImm(MBBI->getCFIType())
-          .getInstr();
-  MBBI->setCFIType(*MBB.getParent(), 0);
-
-  // If not already bundled, bundle the check and the call to prevent
-  // further changes.
-  if (!MBBI->isBundled())
-    finalizeBundle(MBB, Check->getIterator(), std::next(MBBI->getIterator()));
-
-  ++NumKCFIChecksAdded;
-  return true;
-}
-
-bool AArch64KCFI::runOnMachineFunction(MachineFunction &MF) {
-  const Module *M = MF.getMMI().getModule();
-  if (!M->getModuleFlag("kcfi"))
-    return false;
-
-  const auto &SubTarget = MF.getSubtarget<AArch64Subtarget>();
-  TII = SubTarget.getInstrInfo();
-
-  bool Changed = false;
-  for (MachineBasicBlock &MBB : MF) {
-    for (MachineBasicBlock::instr_iterator MII = MBB.instr_begin(),
-                                           MIE = MBB.instr_end();
-         MII != MIE; ++MII) {
-      if (MII->isCall() && MII->getCFIType())
-        Changed |= emitCheck(MBB, MII);
-    }
-  }
-
-  return Changed;
-}

diff --git a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
index 5ccba29f5d430b92c24c05dce0503f2eda26e5eb..c7a6bb563ce6684f6291ec77846a0b80619fb0ca 100644 (file)
--- a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
@@ -215,7 +215,6 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAArch64Target() {
   initializeAArch64ConditionOptimizerPass(*PR);
   initializeAArch64DeadRegisterDefinitionsPass(*PR);
   initializeAArch64ExpandPseudoPass(*PR);
-  initializeAArch64KCFIPass(*PR);
   initializeAArch64LoadStoreOptPass(*PR);
   initializeAArch64MIPeepholeOptPass(*PR);
   initializeAArch64SIMDInstrOptPass(*PR);
@@ -230,6 +229,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAArch64Target() {
   initializeFalkorHWPFFixPass(*PR);
   initializeFalkorMarkStridedAccessesLegacyPass(*PR);
   initializeLDTLSCleanupPass(*PR);
+  initializeKCFIPass(*PR);
   initializeSMEABIPass(*PR);
   initializeSVEIntrinsicOptsPass(*PR);
   initializeAArch64SpeculationHardeningPass(*PR);
@@ -772,7 +772,7 @@ void AArch64PassConfig::addPreSched2() {
       addPass(createAArch64LoadStoreOptimizationPass());
   }
   // Emit KCFI checks for indirect calls.
-  addPass(createAArch64KCFIPass());
+  addPass(createKCFIPass());
 
   // The AArch64SpeculationHardeningPass destroys dominator tree and natural
   // loop info, which is needed for the FalkorHWPFFixPass and also later on.

diff --git a/llvm/lib/Target/AArch64/CMakeLists.txt b/llvm/lib/Target/AArch64/CMakeLists.txt
index 0dfcc080878a117242a429dd581f24a0b243e07a..54157d2be2c74ffabcd4f7c4923fb0af195b2b04 100644 (file)
--- a/llvm/lib/Target/AArch64/CMakeLists.txt
+++ b/llvm/lib/Target/AArch64/CMakeLists.txt
@@ -63,7 +63,6 @@ add_llvm_target(AArch64CodeGen
   AArch64ISelDAGToDAG.cpp
   AArch64ISelLowering.cpp
   AArch64InstrInfo.cpp
-  AArch64KCFI.cpp
   AArch64LoadStoreOptimizer.cpp
   AArch64LowerHomogeneousPrologEpilog.cpp
   AArch64MachineFunctionInfo.cpp

diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index 8388ddd567c308ac893c459639d90aa61271ad01..9aa5fce16cb29cbbbde2e9c1aaa5f91b69bde080 100644 (file)
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/llvm/lib/Target/X86/CMakeLists.txt
@@ -63,7 +63,6 @@ set(sources
   X86InstrFoldTables.cpp
   X86InstrInfo.cpp
   X86EvexToVex.cpp
-  X86KCFI.cpp
   X86LegalizerInfo.cpp
   X86LoadValueInjectionLoadHardening.cpp
   X86LoadValueInjectionRetHardening.cpp

diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 5ff9d7facc878eb5d19f3a2d11de37ddc03a3762..e15f71828da162e2023d1277855432c78e685430 100644 (file)
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -51,9 +51,6 @@ FunctionPass *createX86IssueVZeroUpperPass();
 /// destinations as part of CET IBT mechanism.
 FunctionPass *createX86IndirectBranchTrackingPass();
 
-/// This pass inserts KCFI checks before indirect calls.
-FunctionPass *createX86KCFIPass();
-
 /// Return a pass that pads short functions with NOOPs.
 /// This will prevent a stall when returning on the Atom.
 FunctionPass *createX86PadShortFunctions();
@@ -187,7 +184,6 @@ void initializeX86FastPreTileConfigPass(PassRegistry &);
 void initializeX86FastTileConfigPass(PassRegistry &);
 void initializeX86FixupSetCCPassPass(PassRegistry &);
 void initializeX86FlagsCopyLoweringPassPass(PassRegistry &);
-void initializeX86KCFIPass(PassRegistry &);
 void initializeX86LoadValueInjectionLoadHardeningPassPass(PassRegistry &);
 void initializeX86LoadValueInjectionRetHardeningPassPass(PassRegistry &);
 void initializeX86LowerAMXIntrinsicsLegacyPassPass(PassRegistry &);

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index b09655b78586a3adef8aadb29030eb0729e15b55..b4790508790649d64258cb99a7e2068ccbc3bc76 100644 (file)
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -59144,6 +59144,71 @@ bool X86TargetLowering::supportSwiftError() const {
   return Subtarget.is64Bit();
 }
 
+MachineInstr *
+X86TargetLowering::EmitKCFICheck(MachineBasicBlock &MBB,
+                                 MachineBasicBlock::instr_iterator &MBBI,
+                                 const TargetInstrInfo *TII) const {
+  assert(MBBI->isCall() && MBBI->getCFIType() &&
+         "Invalid call instruction for a KCFI check");
+
+  MachineFunction &MF = *MBB.getParent();
+  // If the call target is a memory operand, unfold it and use R11 for the
+  // call, so KCFI_CHECK won't have to recompute the address.
+  switch (MBBI->getOpcode()) {
+  case X86::CALL64m:
+  case X86::CALL64m_NT:
+  case X86::TAILJMPm64:
+  case X86::TAILJMPm64_REX: {
+    MachineBasicBlock::instr_iterator OrigCall = MBBI;
+    SmallVector<MachineInstr *, 2> NewMIs;
+    if (!TII->unfoldMemoryOperand(MF, *OrigCall, X86::R11, /*UnfoldLoad=*/true,
+                                  /*UnfoldStore=*/false, NewMIs))
+      report_fatal_error("Failed to unfold memory operand for a KCFI check");
+    for (auto *NewMI : NewMIs)
+      MBBI = MBB.insert(OrigCall, NewMI);
+    assert(MBBI->isCall() &&
+           "Unexpected instruction after memory operand unfolding");
+    if (OrigCall->shouldUpdateCallSiteInfo())
+      MF.moveCallSiteInfo(&*OrigCall, &*MBBI);
+    MBBI->setCFIType(MF, OrigCall->getCFIType());
+    OrigCall->eraseFromParent();
+    break;
+  }
+  default:
+    break;
+  }
+
+  MachineOperand &Target = MBBI->getOperand(0);
+  Register TargetReg;
+  switch (MBBI->getOpcode()) {
+  case X86::CALL64r:
+  case X86::CALL64r_NT:
+  case X86::TAILJMPr64:
+  case X86::TAILJMPr64_REX:
+    assert(Target.isReg() && "Unexpected target operand for an indirect call");
+    Target.setIsRenamable(false);
+    TargetReg = Target.getReg();
+    break;
+  case X86::CALL64pcrel32:
+  case X86::TAILJMPd64:
+    assert(Target.isSymbol() && "Unexpected target operand for a direct call");
+    // X86TargetLowering::EmitLoweredIndirectThunk always uses r11 for
+    // 64-bit indirect thunk calls.
+    assert(StringRef(Target.getSymbolName()).endswith("_r11") &&
+           "Unexpected register for an indirect thunk call");
+    TargetReg = X86::R11;
+    break;
+  default:
+    llvm_unreachable("Unexpected CFI call opcode");
+    break;
+  }
+
+  return BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(X86::KCFI_CHECK))
+      .addReg(TargetReg)
+      .addImm(MBBI->getCFIType())
+      .getInstr();
+}
+
 /// Returns true if stack probing through a function call is requested.
 bool X86TargetLowering::hasStackProbeSymbol(const MachineFunction &MF) const {
   return !getStackProbeSymbolName(MF).empty();

diff --git a/llvm/lib/Target/X86/X86ISelLowering.h b/llvm/lib/Target/X86/X86ISelLowering.h
index d431bd6fd270498bf2074fda3ab5862773508735..4eb079f7cadd7d96312b9158f31261e053e80a4b 100644 (file)
--- a/llvm/lib/Target/X86/X86ISelLowering.h
+++ b/llvm/lib/Target/X86/X86ISelLowering.h
@@ -1521,6 +1521,10 @@ namespace llvm {
 
     bool supportKCFIBundles() const override { return true; }
 
+    MachineInstr *EmitKCFICheck(MachineBasicBlock &MBB,
+                                MachineBasicBlock::instr_iterator &MBBI,
+                                const TargetInstrInfo *TII) const override;
+
     bool hasStackProbeSymbol(const MachineFunction &MF) const override;
     bool hasInlineStackProbe(const MachineFunction &MF) const override;
     StringRef getStackProbeSymbolName(const MachineFunction &MF) const override;

diff --git a/llvm/lib/Target/X86/X86InstrInfo.cpp b/llvm/lib/Target/X86/X86InstrInfo.cpp
index 2f1c197e6b56160c16358d410324469760e3896a..859079496710a9a34ff8f38defa3f385652e942f 100644 (file)
--- a/llvm/lib/Target/X86/X86InstrInfo.cpp
+++ b/llvm/lib/Target/X86/X86InstrInfo.cpp
@@ -6169,7 +6169,7 @@ MachineInstr *X86InstrInfo::foldMemoryOperandImpl(
     return nullptr;
 
   // Don't fold loads into indirect calls that need a KCFI check as we'll
-  // have to unfold these in X86KCFIPass anyway.
+  // have to unfold these in X86TargetLowering::EmitKCFICheck anyway.
   if (MI.isCall() && MI.getCFIType())
     return nullptr;
 

diff --git a/llvm/lib/Target/X86/X86KCFI.cpp b/llvm/lib/Target/X86/X86KCFI.cpp
deleted file mode 100644 (file)
index 4086f28..0000000
--- a/llvm/lib/Target/X86/X86KCFI.cpp
+++ /dev/null
@@ -1,150 +0,0 @@
-//===---- X86KCFI.cpp - Implements KCFI -----------------------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements KCFI indirect call checking.
-//
-//===----------------------------------------------------------------------===//
-
-#include "X86.h"
-#include "X86InstrInfo.h"
-#include "X86Subtarget.h"
-#include "X86TargetMachine.h"
-#include "llvm/ADT/Statistic.h"
-#include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineInstrBuilder.h"
-#include "llvm/CodeGen/MachineInstrBundle.h"
-#include "llvm/CodeGen/MachineModuleInfo.h"
-
-using namespace llvm;
-
-#define DEBUG_TYPE "x86-kcfi"
-#define X86_KCFI_PASS_NAME "Insert KCFI indirect call checks"
-
-STATISTIC(NumKCFIChecksAdded, "Number of indirect call checks added");
-
-namespace {
-class X86KCFI : public MachineFunctionPass {
-public:
-  static char ID;
-
-  X86KCFI() : MachineFunctionPass(ID) {}
-
-  StringRef getPassName() const override { return X86_KCFI_PASS_NAME; }
-  bool runOnMachineFunction(MachineFunction &MF) override;
-
-private:
-  /// Machine instruction info used throughout the class.
-  const X86InstrInfo *TII = nullptr;
-
-  /// Emits a KCFI check before an indirect call.
-  /// \returns true if the check was added and false otherwise.
-  bool emitCheck(MachineBasicBlock &MBB,
-                 MachineBasicBlock::instr_iterator I) const;
-};
-
-char X86KCFI::ID = 0;
-} // end anonymous namespace
-
-INITIALIZE_PASS(X86KCFI, DEBUG_TYPE, X86_KCFI_PASS_NAME, false, false)
-
-FunctionPass *llvm::createX86KCFIPass() { return new X86KCFI(); }
-
-bool X86KCFI::emitCheck(MachineBasicBlock &MBB,
-                        MachineBasicBlock::instr_iterator MBBI) const {
-  assert(TII && "Target instruction info was not initialized");
-
-  // If the call instruction is bundled, we can only emit a check safely if
-  // it's the first instruction in the bundle.
-  if (MBBI->isBundled() && !std::prev(MBBI)->isBundle())
-    report_fatal_error("Cannot emit a KCFI check for a bundled call");
-
-  MachineFunction &MF = *MBB.getParent();
-  // If the call target is a memory operand, unfold it and use R11 for the
-  // call, so KCFI_CHECK won't have to recompute the address.
-  switch (MBBI->getOpcode()) {
-  case X86::CALL64m:
-  case X86::CALL64m_NT:
-  case X86::TAILJMPm64:
-  case X86::TAILJMPm64_REX: {
-    MachineBasicBlock::instr_iterator OrigCall = MBBI;
-    SmallVector<MachineInstr *, 2> NewMIs;
-    if (!TII->unfoldMemoryOperand(MF, *OrigCall, X86::R11, /*UnfoldLoad=*/true,
-                                  /*UnfoldStore=*/false, NewMIs))
-      report_fatal_error("Failed to unfold memory operand for a KCFI check");
-    for (auto *NewMI : NewMIs)
-      MBBI = MBB.insert(OrigCall, NewMI);
-    assert(MBBI->isCall() &&
-           "Unexpected instruction after memory operand unfolding");
-    if (OrigCall->shouldUpdateCallSiteInfo())
-      MF.moveCallSiteInfo(&*OrigCall, &*MBBI);
-    MBBI->setCFIType(MF, OrigCall->getCFIType());
-    OrigCall->eraseFromParent();
-    break;
-  }
-  default:
-    break;
-  }
-
-  MachineInstr *Check =
-      BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(X86::KCFI_CHECK))
-          .getInstr();
-  MachineOperand &Target = MBBI->getOperand(0);
-  switch (MBBI->getOpcode()) {
-  case X86::CALL64r:
-  case X86::CALL64r_NT:
-  case X86::TAILJMPr64:
-  case X86::TAILJMPr64_REX:
-    assert(Target.isReg() && "Unexpected target operand for an indirect call");
-    Check->addOperand(MachineOperand::CreateReg(Target.getReg(), false));
-    Target.setIsRenamable(false);
-    break;
-  case X86::CALL64pcrel32:
-  case X86::TAILJMPd64:
-    assert(Target.isSymbol() && "Unexpected target operand for a direct call");
-    // X86TargetLowering::EmitLoweredIndirectThunk always uses r11 for
-    // 64-bit indirect thunk calls.
-    assert(StringRef(Target.getSymbolName()).endswith("_r11") &&
-           "Unexpected register for an indirect thunk call");
-    Check->addOperand(MachineOperand::CreateReg(X86::R11, false));
-    break;
-  default:
-    llvm_unreachable("Unexpected CFI call opcode");
-  }
-
-  Check->addOperand(MachineOperand::CreateImm(MBBI->getCFIType()));
-  MBBI->setCFIType(MF, 0);
-
-  // If not already bundled, bundle the check and the call to prevent
-  // further changes.
-  if (!MBBI->isBundled())
-    finalizeBundle(MBB, Check->getIterator(), std::next(MBBI->getIterator()));
-
-  ++NumKCFIChecksAdded;
-  return true;
-}
-
-bool X86KCFI::runOnMachineFunction(MachineFunction &MF) {
-  const Module *M = MF.getMMI().getModule();
-  if (!M->getModuleFlag("kcfi"))
-    return false;
-
-  const auto &SubTarget = MF.getSubtarget<X86Subtarget>();
-  TII = SubTarget.getInstrInfo();
-
-  bool Changed = false;
-  for (MachineBasicBlock &MBB : MF) {
-    for (MachineBasicBlock::instr_iterator MII = MBB.instr_begin(),
-                                           MIE = MBB.instr_end();
-         MII != MIE; ++MII) {
-      if (MII->isCall() && MII->getCFIType())
-        Changed |= emitCheck(MBB, MII);
-    }
-  }
-
-  return Changed;
-}

diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index a8895453c14ab4a52c12437f5fd8596b5ead1bba..1afc596345322a2d25936e72638c05b0e8efe0aa 100644 (file)
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -87,7 +87,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeX86Target() {
   initializeX86TileConfigPass(PR);
   initializeX86FastPreTileConfigPass(PR);
   initializeX86FastTileConfigPass(PR);
-  initializeX86KCFIPass(PR);
+  initializeKCFIPass(PR);
   initializeX86LowerTileCopyPass(PR);
   initializeX86ExpandPseudoPass(PR);
   initializeX86ExecutionDomainFixPass(PR);
@@ -556,7 +556,7 @@ void X86PassConfig::addPostRegAlloc() {
 
 void X86PassConfig::addPreSched2() {
   addPass(createX86ExpandPseudoPass());
-  addPass(createX86KCFIPass());
+  addPass(createKCFIPass());
 }
 
 void X86PassConfig::addPreEmitPass() {

diff --git a/llvm/test/CodeGen/AArch64/kcfi-bti.ll b/llvm/test/CodeGen/AArch64/kcfi-bti.ll
index 62e2fdcebee35c7479e96e9de64956b2debe4034..0e8dbad1f7c7595173ed5dfb01c371b93fc18a03 100644 (file)
--- a/llvm/test/CodeGen/AArch64/kcfi-bti.ll
+++ b/llvm/test/CodeGen/AArch64/kcfi-bti.ll
@@ -1,6 +1,6 @@
 ; RUN: llc -mtriple=aarch64-- -verify-machineinstrs < %s | FileCheck %s --check-prefix=ASM
 ; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -stop-after=finalize-isel < %s | FileCheck %s --check-prefixes=MIR,ISEL
-; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -stop-after=aarch64-kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
+; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -stop-after=kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
 
 ; ASM:       .word 12345678
 define void @f1(ptr noundef %x) !kcfi_type !2 {

diff --git a/llvm/test/CodeGen/AArch64/kcfi.ll b/llvm/test/CodeGen/AArch64/kcfi.ll
index 088c93645cc09ad2622ad39fc7afa85c16510bec..7c824c19ee69adba9982e751cbcfc2f2333a70b5 100644 (file)
--- a/llvm/test/CodeGen/AArch64/kcfi.ll
+++ b/llvm/test/CodeGen/AArch64/kcfi.ll
@@ -7,8 +7,8 @@
 ; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -mattr=harden-sls-blr -stop-after=finalize-isel < %s | FileCheck %s --check-prefixes=MIR,ISEL-SLS
 ; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -mattr=harden-sls-blr -stop-after=finalize-isel -global-isel < %s | FileCheck %s --check-prefixes=MIR,ISEL-SLS
 
-; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -stop-after=aarch64-kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
-; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -mattr=harden-sls-blr -stop-after=aarch64-kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI-SLS
+; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -stop-after=kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
+; RUN: llc -mtriple=aarch64-- -verify-machineinstrs -mattr=harden-sls-blr -stop-after=kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI-SLS
 
 ; ASM:       .word 12345678
 define void @f1(ptr noundef %x) !kcfi_type !1 {

diff --git a/llvm/test/CodeGen/X86/kcfi.ll b/llvm/test/CodeGen/X86/kcfi.ll
index 4b93fdddb82d7c5db45533453a7c56d11349a233..566a88b76c4fc9b9f8a2653ff75a3a93dae446d4 100644 (file)
--- a/llvm/test/CodeGen/X86/kcfi.ll
+++ b/llvm/test/CodeGen/X86/kcfi.ll
@@ -1,6 +1,6 @@
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs < %s | FileCheck %s --check-prefix=ASM
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs -stop-after=finalize-isel < %s | FileCheck %s --check-prefixes=MIR,ISEL
-; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs -stop-after=x86-kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs -stop-after=kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
 
 ; ASM:       .p2align 4, 0x90
 ; ASM:       .type __cfi_f1,@function