[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

Author: samitolvanen

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

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -2080,6 +2080,18 @@ class TargetLoweringBase {
     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

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&);

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();

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

llvm/lib/Target/AArch64/AArch64KCFI.cpp renamed to llvm/lib/CodeGen/KCFI.cpp

@@ -1,87 +1,78 @@
-//===---- AArch64KCFI.cpp - Implements KCFI -------------------------------===//
+//===---- 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 file implements KCFI indirect call checking.
+// This pass implements KCFI indirect call check lowering.
 //
 //===----------------------------------------------------------------------===//
 
-#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"
+#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 "aarch64-kcfi"
-#define AARCH64_KCFI_PASS_NAME "Insert KCFI indirect call checks"
+#define DEBUG_TYPE "kcfi"
+#define KCFI_PASS_NAME "Insert KCFI indirect call checks"
 
 STATISTIC(NumKCFIChecksAdded, "Number of indirect call checks added");
 
 namespace {
-class AArch64KCFI : public MachineFunctionPass {
+class KCFI : public MachineFunctionPass {
 public:
   static char ID;
 
-  AArch64KCFI() : MachineFunctionPass(ID) {}
+  KCFI() : MachineFunctionPass(ID) {}
 
-  StringRef getPassName() const override { return AARCH64_KCFI_PASS_NAME; }
+  StringRef getPassName() const override { return KCFI_PASS_NAME; }
   bool runOnMachineFunction(MachineFunction &MF) override;
 
 private:
   /// Machine instruction info used throughout the class.
-  const AArch64InstrInfo *TII = nullptr;
+  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 AArch64KCFI::ID = 0;
+char KCFI::ID = 0;
 } // end anonymous namespace
 
-INITIALIZE_PASS(AArch64KCFI, DEBUG_TYPE, AARCH64_KCFI_PASS_NAME, false, false)
+INITIALIZE_PASS(KCFI, DEBUG_TYPE, KCFI_PASS_NAME, false, false)
 
-FunctionPass *llvm::createAArch64KCFIPass() { return new AArch64KCFI(); }
+FunctionPass *llvm::createKCFIPass() { return new KCFI(); }
 
-bool AArch64KCFI::emitCheck(MachineBasicBlock &MBB,
-                            MachineBasicBlock::instr_iterator MBBI) const {
+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");
 
-  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);
+  // 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);
 
-  MachineInstr *Check =
-      BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(AArch64::KCFI_CHECK))
-          .addReg(Target.getReg())
-          .addImm(MBBI->getCFIType())
-          .getInstr();
+  // 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
@@ -93,16 +84,18 @@ bool AArch64KCFI::emitCheck(MachineBasicBlock &MBB,
   return true;
 }
 
-bool AArch64KCFI::runOnMachineFunction(MachineFunction &MF) {
+bool KCFI::runOnMachineFunction(MachineFunction &MF) {
   const Module *M = MF.getMMI().getModule();
   if (!M->getModuleFlag("kcfi"))
     return false;
 
-  const auto &SubTarget = MF.getSubtarget<AArch64Subtarget>();
+  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) {

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 &);

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();
 }

llvm/lib/Target/AArch64/AArch64ISelLowering.h

@@ -862,6 +862,10 @@ class AArch64TargetLowering : public TargetLowering {
 
   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;
 

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.

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

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

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 &);

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();

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;

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;