[ARM] Implement harden-sls-retbr for Thumb mode

[kristof.beyls]

Summary

The only non-trivial consideration in this patch is that the formation
of TBB/TBH instructions, which is done in the constant island pass, does
not understand the speculation barriers inserted by the SLSHardening
pass. As such, when harden-sls-retbr is enabled for a function, the
formation of TBB/TBH instructions in the constant island pass is
disabled.

[ostannard]

LGTM

llvm/lib/Target/ARM/ARMAsmPrinter.cpp

case ARM::SpeculationBarrierISBDSBEndBB: {
     TmpInstDSB.addOperand(MCOperand::createImm(0xf));
     EmitToStreamer(*OutStreamer, TmpInstDSB);
     MCInst TmpInstISB;
     TmpInstISB.setOpcode(ARM::ISB);
     TmpInstISB.addOperand(MCOperand::createImm(0xf));
     EmitToStreamer(*OutStreamer, TmpInstISB);
     return;
   }
+  case ARM::t2SpeculationBarrierISBDSBEndBB: {
+    // Print DSB SYS + ISB
+    MCInst TmpInstDSB;
+    TmpInstDSB.setOpcode(ARM::t2DSB);
+    TmpInstDSB.addOperand(MCOperand::createImm(0xf));
+    TmpInstDSB.addOperand(MCOperand::createImm(ARMCC::AL));
+    TmpInstDSB.addOperand(MCOperand::createReg(0));
+    EmitToStreamer(*OutStreamer, TmpInstDSB);
+    MCInst TmpInstISB;
+    TmpInstISB.setOpcode(ARM::t2ISB);
+    TmpInstISB.addOperand(MCOperand::createImm(0xf));
+    TmpInstISB.addOperand(MCOperand::createImm(ARMCC::AL));
+    TmpInstISB.addOperand(MCOperand::createReg(0));
+    EmitToStreamer(*OutStreamer, TmpInstISB);
+    return;
+  }
   case ARM::SpeculationBarrierSBEndBB: {
     // Print SB
     MCInst TmpInstSB;
     TmpInstSB.setOpcode(ARM::SB);
     EmitToStreamer(*OutStreamer, TmpInstSB);
     return;
   }
+  case ARM::t2SpeculationBarrierSBEndBB: {
+    // Print SB
+    MCInst TmpInstSB;
+    TmpInstSB.setOpcode(ARM::t2SB);
+    EmitToStreamer(*OutStreamer, TmpInstSB);
+    return;
+  }
   }
 
   MCInst TmpInst;
   LowerARMMachineInstrToMCInst(MI, TmpInst, *this);
 
   EmitToStreamer(*OutStreamer, TmpInst);
 }
 

llvm/lib/Target/ARM/ARMBaseInstrInfo.h

 static inline bool isIndirectControlFlowNotComingBack(const MachineInstr &MI) {
   int opc = MI.getOpcode();
   return MI.isReturn() || isIndirectBranchOpcode(MI.getOpcode()) ||
          isJumpTableBranchOpcode(opc);
 }
 
 static inline bool isSpeculationBarrierEndBBOpcode(int Opc) {
   return Opc == ARM::SpeculationBarrierISBDSBEndBB ||
-         Opc == ARM::SpeculationBarrierSBEndBB;
+         Opc == ARM::SpeculationBarrierSBEndBB ||
+         Opc == ARM::t2SpeculationBarrierISBDSBEndBB ||
+         Opc == ARM::t2SpeculationBarrierSBEndBB;
 }
 
 static inline bool isPopOpcode(int Opc) {
   return Opc == ARM::tPOP_RET || Opc == ARM::LDMIA_RET ||
          Opc == ARM::t2LDMIA_RET || Opc == ARM::tPOP || Opc == ARM::LDMIA_UPD ||
          Opc == ARM::t2LDMIA_UPD || Opc == ARM::VLDMDIA_UPD;
 }
 

llvm/lib/Target/ARM/ARMBaseInstrInfo.cpp

unsigned ARMBaseInstrInfo::getInstSizeInBytes(const MachineInstr &MI) const {
   case ARM::INLINEASM_BR: {
     // If this machine instr is an inline asm, measure it.
     unsigned Size = getInlineAsmLength(MI.getOperand(0).getSymbolName(), *MAI);
     if (!MF->getInfo<ARMFunctionInfo>()->isThumbFunction())
       Size = alignTo(Size, 4);
     return Size;
   }
   case ARM::SpeculationBarrierISBDSBEndBB:
+  case ARM::t2SpeculationBarrierISBDSBEndBB:
     // This gets lowered to 2 4-byte instructions.
     return 8;
   case ARM::SpeculationBarrierSBEndBB:
+  case ARM::t2SpeculationBarrierSBEndBB:
     // This gets lowered to 1 4-byte instructions.
     return 4;
   }
 }
 
 unsigned ARMBaseInstrInfo::getInstBundleLength(const MachineInstr &MI) const {
   unsigned Size = 0;
   MachineBasicBlock::const_instr_iterator I = MI.getIterator();

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

bool ARMConstantIslands::runOnMachineFunction(MachineFunction &mf) {
   AFI = MF->getInfo<ARMFunctionInfo>();
   DT = &getAnalysis<MachineDominatorTree>();
 
   isThumb = AFI->isThumbFunction();
   isThumb1 = AFI->isThumb1OnlyFunction();
   isThumb2 = AFI->isThumb2Function();
 
   bool GenerateTBB = isThumb2 || (isThumb1 && SynthesizeThumb1TBB);
+  // TBB generation code in this constant island pass has not been adapted to
+  // deal with speculation barriers.
+  if (STI->hardenSlsRetBr())
+    GenerateTBB = false;
 
   // Renumber all of the machine basic blocks in the function, guaranteeing that
   // the numbers agree with the position of the block in the function.
   MF->RenumberBlocks();
 
   // Try to reorder and otherwise adjust the block layout to make good use
   // of the TB[BH] instructions.
   bool MadeChange = false;

llvm/lib/Target/ARM/ARMInstrThumb2.td

 // Non-predicable aliases of a predicable DSB: the predicate is (14, 0) where
 // 14 = AL (always execute) and 0 = "instruction doesn't read the CPSR".
 def : InstAlias<"ssbb", (t2DSB 0x0, 14, 0), 1>, Requires<[HasDB, IsThumb2]>;
 def : InstAlias<"pssbb", (t2DSB 0x4, 14, 0), 1>, Requires<[HasDB, IsThumb2]>;
 
 // Armv8-R 'Data Full Barrier'
 def : InstAlias<"dfb${p}", (t2DSB 0xc, pred:$p), 1>, Requires<[HasDFB]>;
 
+// SpeculationBarrierEndBB must only be used after an unconditional control
+// flow, i.e. after a terminator for which isBarrier is True.
+let hasSideEffects = 1, isCodeGenOnly = 1, isTerminator = 1, isBarrier = 1 in {
+  def t2SpeculationBarrierISBDSBEndBB
+      : PseudoInst<(outs), (ins), NoItinerary, []>, Sched<[]>;
+  def t2SpeculationBarrierSBEndBB
+      : PseudoInst<(outs), (ins), NoItinerary, []>, Sched<[]>;
+}
+
 // Alias for LDR, LDRB, LDRH, LDRSB, and LDRSH without the ".w" optional
 // width specifier.
 def : t2InstAlias<"ldr${p} $Rt, $addr",
                   (t2LDRi12 GPR:$Rt, t2addrmode_imm12:$addr, pred:$p)>;
 def : t2InstAlias<"ldrb${p} $Rt, $addr",
                   (t2LDRBi12 rGPR:$Rt, t2addrmode_imm12:$addr, pred:$p)>;
 def : t2InstAlias<"ldrh${p} $Rt, $addr",
                   (t2LDRHi12 rGPR:$Rt, t2addrmode_imm12:$addr, pred:$p)>;

llvm/lib/Target/ARM/ARMSLSHardening.cpp

static void insertSpeculationBarrier(const ARMSubtarget *ST,
   assert(MBBI != MBB.begin() &&
          "Must not insert SpeculationBarrierEndBB as only instruction in MBB.");
   assert(std::prev(MBBI)->isBarrier() &&
          "SpeculationBarrierEndBB must only follow unconditional control flow "
          "instructions.");
   assert(std::prev(MBBI)->isTerminator() &&
          "SpeculationBarrierEndBB must only follow terminators.");
   const TargetInstrInfo *TII = ST->getInstrInfo();
-  unsigned BarrierOpc = ST->hasSB() && !AlwaysUseISBDSB
-                            ? ARM::SpeculationBarrierSBEndBB
-                            : ARM::SpeculationBarrierISBDSBEndBB;
+  assert(ST->hasDataBarrier() || ST->hasSB());
+  bool ProduceSB = ST->hasSB() && !AlwaysUseISBDSB;
+  unsigned BarrierOpc =
+      ProduceSB ? (ST->isThumb() ? ARM::t2SpeculationBarrierSBEndBB
+                                 : ARM::SpeculationBarrierSBEndBB)
+                : (ST->isThumb() ? ARM::t2SpeculationBarrierISBDSBEndBB
+                                 : ARM::SpeculationBarrierISBDSBEndBB);
   if (MBBI == MBB.end() || !isSpeculationBarrierEndBBOpcode(MBBI->getOpcode()))
     BuildMI(MBB, MBBI, DL, TII->get(BarrierOpc));
 }
 
 bool ARMSLSHardening::runOnMachineFunction(MachineFunction &MF) {
   ST = &MF.getSubtarget<ARMSubtarget>();
   TII = MF.getSubtarget().getInstrInfo();
 
   bool Modified = false;
   for (auto &MBB : MF)
     Modified |= hardenReturnsAndBRs(MBB);
 
   return Modified;
 }
 
 bool ARMSLSHardening::hardenReturnsAndBRs(MachineBasicBlock &MBB) const {
   if (!ST->hardenSlsRetBr())
     return false;
+  assert(!ST->isThumb1Only());
   bool Modified = false;
   MachineBasicBlock::iterator MBBI = MBB.getFirstTerminator(), E = MBB.end();
   MachineBasicBlock::iterator NextMBBI;
   for (; MBBI != E; MBBI = NextMBBI) {
     MachineInstr &MI = *MBBI;
     NextMBBI = std::next(MBBI);
     if (isIndirectControlFlowNotComingBack(MI)) {
       assert(MI.isTerminator());

llvm/test/CodeGen/ARM/speculation-hardening-sls.ll

 ; RUN: llc -mattr=harden-sls-retbr -verify-machineinstrs -mtriple=armv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,ARM,HARDEN,ISBDSB,ISBDSBDAGISEL -dump-input-context=100
+; RUN: llc -mattr=harden-sls-retbr -verify-machineinstrs -mtriple=thumbv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,THUMB,HARDENTHUMB,HARDEN,ISBDSB,ISBDSBDAGISEL -dump-input-context=100
 ; RUN: llc -mattr=harden-sls-retbr -mattr=+sb -verify-machineinstrs -mtriple=armv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,ARM,HARDEN,SB,SBDAGISEL -dump-input-context=100
+; RUN: llc -mattr=harden-sls-retbr -mattr=+sb -verify-machineinstrs -mtriple=thumbv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,THUMB,HARDENTHUMB,HARDEN,SB,SBDAGISEL -dump-input-context=100
 ; RUN: llc -verify-machineinstrs -mtriple=armv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,ARM,NOHARDEN,NOHARDENARM -dump-input-context=100
+; RUN: llc -verify-machineinstrs -mtriple=thumbv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,THUMB,NOHARDEN,NOHARDENTHUMB
 ; RUN: llc -global-isel -global-isel-abort=0 -mattr=harden-sls-retbr -verify-machineinstrs -mtriple=armv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,ARM,HARDEN,ISBDSB
+; RUN: llc -global-isel -global-isel-abort=0 -mattr=harden-sls-retbr -verify-machineinstrs -mtriple=thumbv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,THUMB,HARDENTHUMB,HARDEN,ISBDSB
 ; RUN: llc -global-isel -global-isel-abort=0 -mattr=harden-sls-retbr -mattr=+sb -verify-machineinstrs -mtriple=armv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,ARM,HARDEN,SB
+; RUN: llc -global-isel -global-isel-abort=0 -mattr=harden-sls-retbr -mattr=+sb -verify-machineinstrs -mtriple=thumbv8-linux-gnueabi < %s | FileCheck %s --check-prefixes=CHECK,THUMB,HARDENTHUMB,HARDEN,SB
 
 ; Function Attrs: norecurse nounwind readnone
 define dso_local i32 @double_return(i32 %a, i32 %b) local_unnamed_addr {
 entry:
   %cmp = icmp sgt i32 %a, 0
   br i1 %cmp, label %if.then, label %if.else
 
 if.then:                                          ; preds = %entry
   ; Make a very easy, very likely to predicate return (BX LR), to test that
   ; it will not get predicated when sls-hardening is enabled.
   %mul = mul i32 %b, %a
   ret i32 %mul
 ; CHECK-LABEL: double_return:
 ; HARDEN:          {{bx lr$}}
 ; NOHARDENARM:     {{bxge lr$}}
+; NOHARDENTHUMB:   {{bx lr$}}
 ; ISBDSB-NEXT: dsb sy
 ; ISBDSB-NEXT: isb
 ; SB-NEXT:     {{ sb$}}
 
 if.else:                                          ; preds = %entry
   %div3 = sdiv i32 %a, %b
   %div2 = sdiv i32 %a, %div3
   %div1 = sdiv i32 %a, %div2

 define dso_local i32 @indirect_branch(i32 %a, i32 %b, i32 %i) {
 ; CHECK-LABEL: indirect_branch:
 entry:
   %idxprom = sext i32 %i to i64
   %arrayidx = getelementptr inbounds [2 x i8*], [2 x i8*]* @__const.indirect_branch.ptr, i64 0, i64 %idxprom
   %0 = load i8*, i8** %arrayidx, align 8
   indirectbr i8* %0, [label %return, label %l2]
 ; ARM:       bx r0
+; THUMB:     mov pc, r0
 ; ISBDSB-NEXT: dsb sy
 ; ISBDSB-NEXT: isb
 ; SB-NEXT:     {{ sb$}}
 
 l2:                                               ; preds = %entry
   br label %return
 ; CHECK:       {{bx lr$}}
 ; ISBDSB-NEXT: dsb sy

 entry:
   switch i32 %b, label %sw.epilog [
     i32 0, label %sw.bb
     i32 1, label %sw.bb1
     i32 3, label %sw.bb3
     i32 4, label %sw.bb5
   ]
 ; ARM:             ldr pc, [{{r[0-9]}}, {{r[0-9]}}, lsl #2]
+; NOHARDENTHUMB:   tbb [pc, {{r[0-9]}}]
+; HARDENTHUMB:     mov pc, {{r[0-9]}}
 ; ISBDSB-NEXT:     dsb sy
 ; ISBDSB-NEXT:     isb
 ; SB-NEXT:         {{ sb$}}
 
 
 sw.bb:                                            ; preds = %entry
   %add = shl nsw i32 %a, 1
   br label %sw.bb1