split KCFI_CHECK into subarch instructions to have Size defined with the PseudoInst

Author: kees

llvm/lib/Target/ARM/ARMAsmPrinter.cpp

@@ -1885,16 +1885,19 @@ void ARMAsmPrinter::LowerKCFI_CHECK(const MachineInstr &MI) {
       .getValueAsString()
       .getAsInteger(10, PrefixNops);
 
-  // Emit ARM32 or Thumb (Thumb1/Thumb2) instruction sequence.
-  const ARMSubtarget &STI = MI.getMF()->getSubtarget<ARMSubtarget>();
-  if (STI.isThumb()) {
-    if (STI.isThumb2()) {
-      EmitKCFI_CHECK_Thumb2(AddrReg, Type, Call, PrefixNops);
-    } else {
-      EmitKCFI_CHECK_Thumb1(AddrReg, Type, Call, PrefixNops);
-    }
-  } else {
+  // Emit the appropriate instruction sequence based on the opcode variant.
+  switch (MI.getOpcode()) {
+  case ARM::KCFI_CHECK_ARM:
     EmitKCFI_CHECK_ARM32(AddrReg, Type, Call, PrefixNops);
+    break;
+  case ARM::KCFI_CHECK_Thumb2:
+    EmitKCFI_CHECK_Thumb2(AddrReg, Type, Call, PrefixNops);
+    break;
+  case ARM::KCFI_CHECK_Thumb1:
+    EmitKCFI_CHECK_Thumb1(AddrReg, Type, Call, PrefixNops);
+    break;
+  default:
+    llvm_unreachable("Unexpected KCFI_CHECK opcode");
   }
 }
 
@@ -1931,7 +1934,9 @@ void ARMAsmPrinter::emitInstruction(const MachineInstr *MI) {
   switch (Opc) {
   case ARM::t2MOVi32imm: llvm_unreachable("Should be lowered by thumb2it pass");
   case ARM::DBG_VALUE: llvm_unreachable("Should be handled by generic printing");
-  case ARM::KCFI_CHECK:
+  case ARM::KCFI_CHECK_ARM:
+  case ARM::KCFI_CHECK_Thumb2:
+  case ARM::KCFI_CHECK_Thumb1:
     LowerKCFI_CHECK(*MI);
     return;
   case ARM::LEApcrel:

llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp

@@ -616,35 +616,6 @@ unsigned ARMBaseInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
     // contrast to AArch64 instructions which have a default size of 4 bytes for
     // example.
     return MCID.getSize();
-  case ARM::KCFI_CHECK: {
-    // KCFI_CHECK is a pseudo-instruction that expands to a sequence of
-    // instructions during AsmPrinter. We need to return the size of the
-    // expanded sequence so that branch distance calculations are correct.
-    //
-    // The expansion depends on the target architecture:
-    // - ARM32: 7 instructions = 28 bytes
-    //   (bic, ldr, 4x eor, beq, udf)
-    // - Thumb2: 7-9 instructions = 28-32 bytes
-    //   (optional push, bic, ldr, 4x eor, optional pop, beq.w, udf)
-    // - Thumb1: 22-25 instructions = 44-50 bytes
-    //   (pushes, bic, movs, lsls, adds, cmp, pops)
-    //
-    // We return a conservative estimate to ensure branch distance calculations
-    // don't underestimate the size.
-    const ARMSubtarget &STI = MF->getSubtarget<ARMSubtarget>();
-    if (STI.isThumb()) {
-      if (STI.isThumb2()) {
-        // Thumb2 (worst case)
-        return 32;
-      } else {
-        // Thumb1 (worst case)
-        return 50;
-      }
-    } else {
-      // ARM32
-      return 28;
-    }
-  }
   case TargetOpcode::BUNDLE:
     return getInstBundleLength(MI);
   case ARM::CONSTPOOL_ENTRY:

llvm/lib/Target/ARM/ARMISelLowering.cpp

@@ -12058,7 +12058,19 @@ ARMTargetLowering::EmitKCFICheck(MachineBasicBlock &MBB,
   assert(TargetOp && TargetOp->isReg() && "Invalid target operand");
   TargetOp->setIsRenamable(false);
 
-  return BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(ARM::KCFI_CHECK))
+  // Select the appropriate KCFI_CHECK variant based on the instruction set
+  unsigned KCFICheckOpcode;
+  if (Subtarget->isThumb()) {
+    if (Subtarget->isThumb2()) {
+      KCFICheckOpcode = ARM::KCFI_CHECK_Thumb2;
+    } else {
+      KCFICheckOpcode = ARM::KCFI_CHECK_Thumb1;
+    }
+  } else {
+    KCFICheckOpcode = ARM::KCFI_CHECK_ARM;
+  }
+
+  return BuildMI(MBB, MBBI, MBBI->getDebugLoc(), TII->get(KCFICheckOpcode))
       .addReg(TargetOp->getReg())
       .addImm(MBBI->getCFIType())
       .getInstr();

llvm/lib/Target/ARM/ARMInstrInfo.td

@@ -6538,9 +6538,27 @@ def : Pat<(atomic_fence (timm), 0), (MEMBARRIER)>;
 //===----------------------------------------------------------------------===//
 // KCFI check pseudo-instruction.
 //===----------------------------------------------------------------------===//
-def KCFI_CHECK
+// KCFI_CHECK pseudo-instruction for Kernel Control-Flow Integrity.
+// Expands to a sequence that verifies the function pointer's type hash.
+// Different sizes for different architectures due to different expansions.
+
+def KCFI_CHECK_ARM
+    : PseudoInst<(outs), (ins GPR:$ptr, i32imm:$type), NoItinerary, []>,
+      Sched<[]>, Requires<[IsARM]> {
+  let Size = 28;  // 7 instructions (bic, ldr, 4x eor, beq, udf)
+}
+
+def KCFI_CHECK_Thumb2
     : PseudoInst<(outs), (ins GPR:$ptr, i32imm:$type), NoItinerary, []>,
-      Sched<[]>;
+      Sched<[]>, Requires<[IsThumb2]> {
+  let Size = 32;  // worst-case 9 instructions (push, bic, ldr, 4x eor, pop, beq.w, udf)
+}
+
+def KCFI_CHECK_Thumb1
+    : PseudoInst<(outs), (ins GPR:$ptr, i32imm:$type), NoItinerary, []>,
+      Sched<[]>, Requires<[IsThumb1Only]> {
+  let Size = 50;  // worst-case 25 instructions (pushes, bic helper, type building, cmp, pops)
+}
 
 //===----------------------------------------------------------------------===//
 // Instructions used for emitting unwind opcodes on Windows.

llvm/test/CodeGen/ARM/kcfi-arm.ll

@@ -11,7 +11,7 @@
 ; ISEL:     BLX %0, csr_aapcs,{{.*}} cfi-type 12345678
 
 ; KCFI:       BUNDLE{{.*}} {
-; KCFI-NEXT:    KCFI_CHECK $r0, 12345678
+; KCFI-NEXT:    KCFI_CHECK_ARM $r0, 12345678
 ; KCFI-NEXT:    BLX killed $r0, csr_aapcs,{{.*}}
 ; KCFI-NEXT:  }
 
@@ -21,7 +21,7 @@
 ; ISEL:     TCRETURNri %0, 0, csr_aapcs, implicit $sp, cfi-type 12345678
 
 ; KCFI:       BUNDLE{{.*}} {
-; KCFI-NEXT:    KCFI_CHECK $r0, 12345678
+; KCFI-NEXT:    KCFI_CHECK_ARM $r0, 12345678
 ; KCFI-NEXT:    TAILJMPr killed $r0, csr_aapcs, implicit $sp, implicit $sp
 ; KCFI-NEXT:  }
 

llvm/test/CodeGen/ARM/kcfi-thumb2.ll

@@ -11,7 +11,7 @@
 ; ISEL:     tBLXr 14 /* CC::al */, $noreg, %0, csr_aapcs,{{.*}} cfi-type 12345678
 
 ; KCFI:       BUNDLE{{.*}} {
-; KCFI-NEXT:    KCFI_CHECK $r0, 12345678
+; KCFI-NEXT:    KCFI_CHECK_Thumb2 $r0, 12345678
 ; KCFI-NEXT:    tBLXr 14 /* CC::al */, $noreg, {{(killed )?}}$r0, csr_aapcs,{{.*}}
 ; KCFI-NEXT:  }
 
@@ -21,7 +21,7 @@
 ; ISEL:     TCRETURNri %0, 0, csr_aapcs, implicit $sp, cfi-type 12345678
 
 ; KCFI:       BUNDLE{{.*}} {
-; KCFI-NEXT:    KCFI_CHECK $r0, 12345678
+; KCFI-NEXT:    KCFI_CHECK_Thumb2 $r0, 12345678
 ; KCFI-NEXT:    tTAILJMPr {{(killed )?}}$r0, csr_aapcs, implicit $sp, implicit $sp
 ; KCFI-NEXT:  }