Introduce llvm.speculation_safe_value intrinsic.

[kristof.beyls]

Summary

This is part of implementing a technique to mitigate against Spectre v1,
similar in spirit to what has been proposed by Chandler for X86_64 at
http://lists.llvm.org/pipermail/llvm-dev/2018-March/122085.html.

This patch adds a new llvm intrinsic:

T @llvm.speculation_safe_value(T %val)

where T can be any integer or any pointer type.

• This intrinsic returns either value %val or 0. When this intrinsic is executed on a miss-speculated path, where the miss-speculation is caused by at least one miss-predicted direct conditional branch, the intrinsic guarantees that value 0 is returned.
• This intrinsic is close to what Philip Reames suggested as an alternative on https://reviews.llvm.org/D41761.
• This could be extended to become T @llvm.speculation_safe_value(T %val, T %miss_spec_val) where the intrinsics returns %miss_spec_val instead of 0 when executed on miss-speculated path. However, that is currently not implemented in this patch.

The introduction of this intrinsic has multiple advantages in helping to
mitigate against Spectre v1 attacks:

• It enables users to only mark the specific values they need to be protected under miss-speculation - if they know which values/program locations need protection.
• For automatic protection of all values loaded: this enables relatively easy experimentation with different heuristics of where the intrinsic/protection should get inserted. One such example heuristic is implemented in a follow-on patch. In other words, even when only focussing on a fully automated approach, separating the decisions on where to insert protection vs doing the necessary transforms/lowering on where protection is needed makes the code slightly easier to understand, review and maintain.

The intrinsic (as implemented in this patch) assumes control flow
miss-speculation tracking has been enabled (see previous patch).
A follow-on patch will make this intrinsic also work without needing
the miss-speculation tracking.

docs/LangRef.rst

 None.
 
 Semantics:
 """"""""""
 
 This intrinsic actually does nothing, but optimizers must assume that it
 has externally observable side effects.
 
+
+'``llvm.speculation_safe_value``' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+This is an overloaded intrinsic. You can use llvm.speculation_safe_value on any
+integer type, and any pointer type.  However, not all targets support this
+intrinsic at the moment.
+
+::
+
+      declare T @llvm.speculation_safe_value.T(T %val)
+
+Overview:
+"""""""""
+
+The '``llvm.speculation_safe_value``' intrinsic.
+
+Arguments:
+""""""""""
+
+The first argument is a pointer or integer value.
+
+
+Semantics:
+""""""""""
+
+On a processor that predicts the direction and target of branches, code executes
+speculatively, i.e. before it is known if the code actually should be executed
+according to program logic.
+
+When the processor executes code speculatively that it should not execute
+according to program logic, the code is said to be executing on a
+miss-speculated path. Miss-speculated paths are caused by incorrect prediction
+of the direction or targets of branches.
+
+This intrinsic guarantees that for miss-speculated paths where at least the
+direction of one of the previously executed conditional branches was
+mispredicted, the intrinsic returns 0.
+
+On fully correctly predicted execution paths, it returns %val.
+
+For paths not covered by the above statements, it returns either 0 or %val.
+
+
 Stack Map Intrinsics
 --------------------
 
 LLVM provides experimental intrinsics to support runtime patching
 mechanisms commonly desired in dynamic language JITs. These intrinsics
 are described in :doc:`StackMaps`.
 
 Element Wise Atomic Memory Intrinsics

include/llvm/CodeGen/ISDOpcodes.h

enum NodeType {
     /// These reductions are non-strict, and have a single vector operand.
     VECREDUCE_FADD, VECREDUCE_FMUL,
     VECREDUCE_ADD, VECREDUCE_MUL,
     VECREDUCE_AND, VECREDUCE_OR, VECREDUCE_XOR,
     VECREDUCE_SMAX, VECREDUCE_SMIN, VECREDUCE_UMAX, VECREDUCE_UMIN,
     /// FMIN/FMAX nodes can have flags, for NaN/NoNaN variants.
     VECREDUCE_FMAX, VECREDUCE_FMIN,
 
+    SPECULATION_SAFE_VALUE,
+
     /// BUILTIN_OP_END - This must be the last enum value in this list.
     /// The target-specific pre-isel opcode values start here.
     BUILTIN_OP_END
   };
 
   /// FIRST_TARGET_MEMORY_OPCODE - Target-specific pre-isel operations
   /// which do not reference a specific memory location should be less than
   /// this value. Those that do must not be less than this value, and can

include/llvm/IR/Intrinsics.td

 def int_experimental_vector_reduce_fmin : Intrinsic<[llvm_anyfloat_ty],
                                                     [llvm_anyvector_ty],
                                                     [IntrNoMem]>;
 
 //===----- Intrinsics that are used to provide predicate information -----===//
 
 def int_ssa_copy : Intrinsic<[llvm_any_ty], [LLVMMatchType<0>],
                              [IntrNoMem, Returned<0>]>;
+
+//===----- Intrinsics to mitigate against miss-speculation exploits -------===//
+
+def int_speculationsafevalue : Intrinsic<[llvm_any_ty], [LLVMMatchType<0>], []>;
+
 //===----------------------------------------------------------------------===//
 // Target-specific intrinsics
 //===----------------------------------------------------------------------===//
 
 include "llvm/IR/IntrinsicsPowerPC.td"
 include "llvm/IR/IntrinsicsX86.td"
 include "llvm/IR/IntrinsicsARM.td"
 include "llvm/IR/IntrinsicsAArch64.td"
 include "llvm/IR/IntrinsicsXCore.td"
 include "llvm/IR/IntrinsicsHexagon.td"
 include "llvm/IR/IntrinsicsNVVM.td"
 include "llvm/IR/IntrinsicsMips.td"
 include "llvm/IR/IntrinsicsAMDGPU.td"
 include "llvm/IR/IntrinsicsBPF.td"
 include "llvm/IR/IntrinsicsSystemZ.td"
 include "llvm/IR/IntrinsicsWebAssembly.td"

include/llvm/Target/TargetSelectionDAG.td

 def SDTAtomicLoad : SDTypeProfile<1, 1, [
   SDTCisInt<0>, SDTCisPtrTy<1>
 ]>;
 
 def SDTConvertOp : SDTypeProfile<1, 5, [ //cvtss, su, us, uu, ff, fs, fu, sf, su
   SDTCisVT<2, OtherVT>, SDTCisVT<3, OtherVT>, SDTCisPtrTy<4>, SDTCisPtrTy<5>
 ]>;
 
+def SDTSpeculationSafe: SDTypeProfile<1, 1, [
+  SDTCisInt<1>, SDTCisSameAs<1, 0>
+]>;
+
 class SDCallSeqStart<list<SDTypeConstraint> constraints> :
         SDTypeProfile<0, 2, constraints>;
 class SDCallSeqEnd<list<SDTypeConstraint> constraints> :
         SDTypeProfile<0, 2, constraints>;
 
 //===----------------------------------------------------------------------===//
 // Selection DAG Node definitions.
 //

 def intrinsic_wo_chain : SDNode<"ISD::INTRINSIC_WO_CHAIN",
                                 SDTypeProfile<1, -1, [SDTCisPtrTy<1>]>, []>;
 
 def SDT_assertext : SDTypeProfile<1, 1,
   [SDTCisInt<0>, SDTCisInt<1>, SDTCisSameAs<1, 0>]>;
 def assertsext : SDNode<"ISD::AssertSext", SDT_assertext>;
 def assertzext : SDNode<"ISD::AssertZext", SDT_assertext>;
 
+def speculationsafevalue : SDNode<"ISD::SPECULATION_SAFE_VALUE",
+                                  SDTSpeculationSafe, []>;
 
 //===----------------------------------------------------------------------===//
 // Selection DAG Condition Codes
 
 class CondCode; // ISD::CondCode enums
 def SETOEQ : CondCode; def SETOGT : CondCode;
 def SETOGE : CondCode; def SETOLT : CondCode; def SETOLE : CondCode;
 def SETONE : CondCode; def SETO   : CondCode; def SETUO  : CondCode;

lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

#endif
   case ISD::ATOMIC_LOAD_UMAX:
   case ISD::ATOMIC_SWAP:
     Res = PromoteIntRes_Atomic1(cast<AtomicSDNode>(N)); break;
 
   case ISD::ATOMIC_CMP_SWAP:
   case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS:
     Res = PromoteIntRes_AtomicCmpSwap(cast<AtomicSDNode>(N), ResNo);
     break;
+  case ISD::SPECULATION_SAFE_VALUE:
+    Res = PromoteIntRes_SpeculationSafeValue(N);
+    break;
   }
 
   // If the result is null then the sub-method took care of registering it.
   if (Res.getNode())
     SetPromotedInteger(SDValue(N, ResNo), Res);
 }
 
 SDValue DAGTypeLegalizer::PromoteIntRes_MERGE_VALUES(SDNode *N,
SDValue DAGTypeLegalizer::PromoteIntRes_SimpleIntBinOp(SDNode *N) {
   // these operations don't care.  They may have weird bits going out, but
   // that too is okay if they are integer operations.
   SDValue LHS = GetPromotedInteger(N->getOperand(0));
   SDValue RHS = GetPromotedInteger(N->getOperand(1));
   return DAG.getNode(N->getOpcode(), SDLoc(N),
                      LHS.getValueType(), LHS, RHS);
 }
 
+SDValue DAGTypeLegalizer::PromoteIntRes_SpeculationSafeValue(SDNode *N) {
+  // Propagate size promotion through the intrinsic.
+  SDValue Op = GetPromotedInteger(N->getOperand(0));
+  return DAG.getNode(N->getOpcode(), SDLoc(N),
+                     Op.getValueType(), Op);
+}
+
 SDValue DAGTypeLegalizer::PromoteIntRes_SExtIntBinOp(SDNode *N) {
   // Sign extend the input.
   SDValue LHS = SExtPromotedInteger(N->getOperand(0));
   SDValue RHS = SExtPromotedInteger(N->getOperand(1));
   return DAG.getNode(N->getOpcode(), SDLoc(N),
                      LHS.getValueType(), LHS, RHS);
 }
 
#endif
   case ISD::SRL: ExpandIntRes_Shift(N, Lo, Hi); break;
 
   case ISD::SADDO:
   case ISD::SSUBO: ExpandIntRes_SADDSUBO(N, Lo, Hi); break;
   case ISD::UADDO:
   case ISD::USUBO: ExpandIntRes_UADDSUBO(N, Lo, Hi); break;
   case ISD::UMULO:
   case ISD::SMULO: ExpandIntRes_XMULO(N, Lo, Hi); break;
+
+  case ISD::SPECULATION_SAFE_VALUE:
+                   ExpandIntRes_SPECULATION_SAFE_VALUE(N, Lo, Hi); break;
   }
 
   // If Lo/Hi is null, the sub-method took care of registering results etc.
   if (Lo.getNode())
     SetExpandedInteger(SDValue(N, ResNo), Lo, Hi);
 }
 
 /// Lower an atomic node to the appropriate builtin call.
void DAGTypeLegalizer::ExpandIntRes_CTPOP(SDNode *N,
   // ctpop(HiLo) -> ctpop(Hi)+ctpop(Lo)
   GetExpandedInteger(N->getOperand(0), Lo, Hi);
   EVT NVT = Lo.getValueType();
   Lo = DAG.getNode(ISD::ADD, dl, NVT, DAG.getNode(ISD::CTPOP, dl, NVT, Lo),
                    DAG.getNode(ISD::CTPOP, dl, NVT, Hi));
   Hi = DAG.getConstant(0, dl, NVT);
 }
 
+void DAGTypeLegalizer::ExpandIntRes_SPECULATION_SAFE_VALUE(SDNode *N,
+                                                           SDValue &Lo,
+                                                           SDValue &Hi) {
+  SDLoc dl(N);
+  GetExpandedInteger(N->getOperand(0), Lo, Hi);
+  Lo = DAG.getNode(N->getOpcode(), dl, Lo.getValueType(), Lo);
+  Hi = DAG.getNode(N->getOpcode(), dl, Hi.getValueType(), Hi);
+}
+
 void DAGTypeLegalizer::ExpandIntRes_CTTZ(SDNode *N,
                                          SDValue &Lo, SDValue &Hi) {
   SDLoc dl(N);
   // cttz (HiLo) -> Lo != 0 ? cttz(Lo) : (cttz(Hi)+32)
   GetExpandedInteger(N->getOperand(0), Lo, Hi);
   EVT NVT = Lo.getValueType();
 
   SDValue LoNotZero = DAG.getSetCC(dl, getSetCCResultType(NVT), Lo,

lib/CodeGen/SelectionDAG/LegalizeTypes.h

private:
   SDValue PromoteIntRes_SRA(SDNode *N);
   SDValue PromoteIntRes_SRL(SDNode *N);
   SDValue PromoteIntRes_TRUNCATE(SDNode *N);
   SDValue PromoteIntRes_UADDSUBO(SDNode *N, unsigned ResNo);
   SDValue PromoteIntRes_ADDSUBCARRY(SDNode *N, unsigned ResNo);
   SDValue PromoteIntRes_UNDEF(SDNode *N);
   SDValue PromoteIntRes_VAARG(SDNode *N);
   SDValue PromoteIntRes_XMULO(SDNode *N, unsigned ResNo);
+  SDValue PromoteIntRes_SpeculationSafeValue(SDNode *N);
 
   // Integer Operand Promotion.
   bool PromoteIntegerOperand(SDNode *N, unsigned OperandNo);
   SDValue PromoteIntOp_ANY_EXTEND(SDNode *N);
   SDValue PromoteIntOp_ATOMIC_STORE(AtomicSDNode *N);
   SDValue PromoteIntOp_BITCAST(SDNode *N);
   SDValue PromoteIntOp_BUILD_PAIR(SDNode *N);
   SDValue PromoteIntOp_BR_CC(SDNode *N, unsigned OpNo);
private:
   void ExpandIntRes_UREM              (SDNode *N, SDValue &Lo, SDValue &Hi);
   void ExpandIntRes_Shift             (SDNode *N, SDValue &Lo, SDValue &Hi);
 
   void ExpandIntRes_MINMAX            (SDNode *N, SDValue &Lo, SDValue &Hi);
 
   void ExpandIntRes_SADDSUBO          (SDNode *N, SDValue &Lo, SDValue &Hi);
   void ExpandIntRes_UADDSUBO          (SDNode *N, SDValue &Lo, SDValue &Hi);
   void ExpandIntRes_XMULO             (SDNode *N, SDValue &Lo, SDValue &Hi);
+  void ExpandIntRes_SPECULATION_SAFE_VALUE(SDNode *N, SDValue &Lo, SDValue &Hi);
 
   void ExpandIntRes_ATOMIC_LOAD       (SDNode *N, SDValue &Lo, SDValue &Hi);
 
   void ExpandShiftByConstant(SDNode *N, const APInt &Amt,
                              SDValue &Lo, SDValue &Hi);
   bool ExpandShiftWithKnownAmountBit(SDNode *N, SDValue &Lo, SDValue &Hi);
   bool ExpandShiftWithUnknownAmountBit(SDNode *N, SDValue &Lo, SDValue &Hi);
 

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

case Intrinsic::ctlz: {
     return nullptr;
   }
   case Intrinsic::ctpop: {
     SDValue Arg = getValue(I.getArgOperand(0));
     EVT Ty = Arg.getValueType();
     setValue(&I, DAG.getNode(ISD::CTPOP, sdl, Ty, Arg));
     return nullptr;
   }
+  case Intrinsic::speculationsafevalue: {
+    SDValue Arg = getValue(I.getArgOperand(0));
+    EVT Ty = Arg.getValueType();
+    setValue(&I, DAG.getNode(ISD::SPECULATION_SAFE_VALUE, sdl, Ty, Arg));
+    return nullptr;
+  }
   case Intrinsic::stacksave: {
     SDValue Op = getRoot();
     Res = DAG.getNode(
         ISD::STACKSAVE, sdl,
         DAG.getVTList(TLI.getPointerTy(DAG.getDataLayout()), MVT::Other), Op);
     setValue(&I, Res);
     DAG.setRoot(Res.getValue(1));
     return nullptr;

lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

if (G) {
       if (Name) return Name;
       return "<<Unknown Target Node #" + utostr(getOpcode()) + ">>";
     }
     return "<<Unknown Node #" + utostr(getOpcode()) + ">>";
 
 #ifndef NDEBUG
   case ISD::DELETED_NODE:               return "<<Deleted Node!>>";
 #endif
+  case ISD::SPECULATION_SAFE_VALUE:     return "SpeculationSafeValue";
   case ISD::PREFETCH:                   return "Prefetch";
   case ISD::ATOMIC_FENCE:               return "AtomicFence";
   case ISD::ATOMIC_CMP_SWAP:            return "AtomicCmpSwap";
   case ISD::ATOMIC_CMP_SWAP_WITH_SUCCESS: return "AtomicCmpSwapWithSuccess";
   case ISD::ATOMIC_SWAP:                return "AtomicSwap";
   case ISD::ATOMIC_LOAD_ADD:            return "AtomicLoadAdd";
   case ISD::ATOMIC_LOAD_SUB:            return "AtomicLoadSub";
   case ISD::ATOMIC_LOAD_AND:            return "AtomicLoadAnd";

lib/Target/AArch64/AArch64ISelLowering.h

enum NodeType : unsigned {
 
   SMULL,
   UMULL,
 
   // Reciprocal estimates and steps.
   FRECPE, FRECPS,
   FRSQRTE, FRSQRTS,
 
+  // control flow miss-speculation protection
+  SpeculationSafeValue,
+
   // NEON Load/Store with post-increment base updates
   LD2post = ISD::FIRST_TARGET_MEMORY_OPCODE,
   LD3post,
   LD4post,
   ST2post,
   ST3post,
   ST4post,
   LD1x2post,
private:
   SDValue LowerFSINCOS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECREDUCE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerATOMIC_LOAD_SUB(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerATOMIC_LOAD_AND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerWindowsDYNAMIC_STACKALLOC(SDValue Op, SDValue Chain,
                                          SDValue &Size,
                                          SelectionDAG &DAG) const;
+  SDValue LowerSPECULATION_SAFE_VALUE(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue BuildSDIVPow2(SDNode *N, const APInt &Divisor, SelectionDAG &DAG,
                         std::vector<SDNode *> *Created) const override;
   SDValue getSqrtEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                           int &ExtraSteps, bool &UseOneConst,
                           bool Reciprocal) const override;
   SDValue getRecipEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                            int &ExtraSteps) const override;

lib/Target/AArch64/AArch64ISelLowering.cpp

for (unsigned im = (unsigned)ISD::PRE_INC;
     setIndexedStoreAction(im, MVT::f64, Legal);
     setIndexedStoreAction(im, MVT::f32, Legal);
     setIndexedStoreAction(im, MVT::f16, Legal);
   }
 
   // Trap.
   setOperationAction(ISD::TRAP, MVT::Other, Legal);
 
+  setOperationAction(ISD::SPECULATION_SAFE_VALUE, MVT::i32, Custom);
+  setOperationAction(ISD::SPECULATION_SAFE_VALUE, MVT::i64, Custom);
+
   // We combine OR nodes for bitfield operations.
   setTargetDAGCombine(ISD::OR);
 
   // Vector add and sub nodes may conceal a high-half opportunity.
   // Also, try to fold ADD into CSINC/CSINV..
   setTargetDAGCombine(ISD::ADD);
   setTargetDAGCombine(ISD::SUB);
   setTargetDAGCombine(ISD::SRL);
const char *AArch64TargetLowering::getTargetNodeName(unsigned Opcode) const {
   case AArch64ISD::ST3LANEpost:       return "AArch64ISD::ST3LANEpost";
   case AArch64ISD::ST4LANEpost:       return "AArch64ISD::ST4LANEpost";
   case AArch64ISD::SMULL:             return "AArch64ISD::SMULL";
   case AArch64ISD::UMULL:             return "AArch64ISD::UMULL";
   case AArch64ISD::FRECPE:            return "AArch64ISD::FRECPE";
   case AArch64ISD::FRECPS:            return "AArch64ISD::FRECPS";
   case AArch64ISD::FRSQRTE:           return "AArch64ISD::FRSQRTE";
   case AArch64ISD::FRSQRTS:           return "AArch64ISD::FRSQRTS";
+  case AArch64ISD::SpeculationSafeValue:
+                                      return "AArch64ISD::SpeculationSafeValue";
   }
   return nullptr;
 }
 
 MachineBasicBlock *
 AArch64TargetLowering::EmitF128CSEL(MachineInstr &MI,
                                     MachineBasicBlock *MBB) const {
   // We materialise the F128CSEL pseudo-instruction as some control flow and a
static SDValue LowerPREFETCH(SDValue Op, SelectionDAG &DAG) {
   unsigned PrfOp = (IsWrite << 4) |     // Load/Store bit
                    (!IsData << 3) |     // IsDataCache bit
                    (Locality << 1) |    // Cache level bits
                    (unsigned)IsStream;  // Stream bit
   return DAG.getNode(AArch64ISD::PREFETCH, DL, MVT::Other, Op.getOperand(0),
                      DAG.getConstant(PrfOp, DL, MVT::i32), Op.getOperand(1));
 }
 
+SDValue
+AArch64TargetLowering::LowerSPECULATION_SAFE_VALUE(SDValue Op,
+                                                   SelectionDAG &DAG) const {
+  assert((Op.getValueType() == MVT::i32 || Op.getValueType() == MVT::i64) &&
+         "Unexpected lowering");
+
+  SDLoc DL(Op);
+  return DAG.getNode(AArch64ISD::SpeculationSafeValue, DL, Op.getValueType(),
+                     Op.getOperand(0));
+}
+
 SDValue AArch64TargetLowering::LowerFP_EXTEND(SDValue Op,
                                               SelectionDAG &DAG) const {
   assert(Op.getValueType() == MVT::f128 && "Unexpected lowering");
 
   RTLIB::Libcall LC;
   LC = RTLIB::getFPEXT(Op.getOperand(0).getValueType(), Op.getValueType());
 
   return LowerF128Call(Op, DAG, LC);
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::VECREDUCE_FMIN:
     return LowerVECREDUCE(Op, DAG);
   case ISD::ATOMIC_LOAD_SUB:
     return LowerATOMIC_LOAD_SUB(Op, DAG);
   case ISD::ATOMIC_LOAD_AND:
     return LowerATOMIC_LOAD_AND(Op, DAG);
   case ISD::DYNAMIC_STACKALLOC:
     return LowerDYNAMIC_STACKALLOC(Op, DAG);
+  case ISD::SPECULATION_SAFE_VALUE:
+    return LowerSPECULATION_SAFE_VALUE(Op, DAG);
   }
 }
 
 //===----------------------------------------------------------------------===//
 //                      Calling Convention Implementation
 //===----------------------------------------------------------------------===//
 
 #include "AArch64GenCallingConv.inc"

lib/Target/AArch64/AArch64InstrInfo.h

public:
   bool isExynosResetFast(const MachineInstr &MI) const;
   /// Returns true if the instruction has a shift left that can be executed
   /// more efficiently.
   bool isExynosShiftLeftFast(const MachineInstr &MI) const;
   /// Returns true if the instruction has a shift by immediate that can be
   /// executed in one cycle less.
   bool isFalkorShiftExtFast(const MachineInstr &MI) const;
 
+  bool isSchedulingBoundary(const MachineInstr &MI,
+                            const MachineBasicBlock *MBB,
+                            const MachineFunction &MF) const override;
+
 private:
   /// Sets the offsets on outlined instructions in \p MBB which use SP
   /// so that they will be valid post-outlining.
   ///
   /// \param MBB A \p MachineBasicBlock in an outlined function.
   void fixupPostOutline(MachineBasicBlock &MBB) const;
 
   void instantiateCondBranch(MachineBasicBlock &MBB, const DebugLoc &DL,

lib/Target/AArch64/AArch64InstrInfo.cpp

MachineInstr *LDRXpost = BuildMI(MF, DebugLoc(), get(AArch64::LDRXpost))
                                .addReg(AArch64::SP, RegState::Define)
                                .addReg(AArch64::LR, RegState::Define)
                                .addReg(AArch64::SP)
                                .addImm(16);
   It = MBB.insert(It, LDRXpost);
 
   return CallPt;
 }
+
+bool AArch64InstrInfo::isSchedulingBoundary(const MachineInstr &MI,
+                                            const MachineBasicBlock *MBB,
+                                            const MachineFunction &MF) const {
+  if (AArch64GenInstrInfo::isSchedulingBoundary(MI, MBB, MF))
+    return true;
+
+  switch (MI.getOpcode()) {
+  case AArch64::HINT:
+    // CSDB hints are scheduling barriers.
+    if (MI.getOperand(0).getImm() == 0x14)
+      return true;
+    break;
+  default:
+    ;
+  }
+
+  return false;
+}

lib/Target/AArch64/AArch64InstrInfo.td

 def AArch64frsqrts  : SDNode<"AArch64ISD::FRSQRTS", SDTFPBinOp>;
 
 def AArch64saddv    : SDNode<"AArch64ISD::SADDV", SDT_AArch64UnaryVec>;
 def AArch64uaddv    : SDNode<"AArch64ISD::UADDV", SDT_AArch64UnaryVec>;
 def AArch64sminv    : SDNode<"AArch64ISD::SMINV", SDT_AArch64UnaryVec>;
 def AArch64uminv    : SDNode<"AArch64ISD::UMINV", SDT_AArch64UnaryVec>;
 def AArch64smaxv    : SDNode<"AArch64ISD::SMAXV", SDT_AArch64UnaryVec>;
 def AArch64umaxv    : SDNode<"AArch64ISD::UMAXV", SDT_AArch64UnaryVec>;
+def AArch64SpeculationSafeValue : SDNode<"AArch64ISD::SpeculationSafeValue",
+                                            SDTIntUnaryOp>;
 
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 
 // AArch64 Instruction Predicate Definitions.
 // We could compute these on a per-module basis but doing so requires accessing
 // the Function object through the <Target>Subtarget and objections were raised
let isCall = 1, Defs = [LR, X0, X1], hasSideEffects = 1,
     isCodeGenOnly = 1 in
 def TLSDESC_CALLSEQ
     : Pseudo<(outs), (ins i64imm:$sym),
              [(AArch64tlsdesc_callseq tglobaltlsaddr:$sym)]>,
       Sched<[WriteI, WriteLD, WriteI, WriteBrReg]>;
 def : Pat<(AArch64tlsdesc_callseq texternalsym:$sym),
           (TLSDESC_CALLSEQ texternalsym:$sym)>;
 
+let hasSideEffects = 1, isCodeGenOnly = 1 in {
+  def SpeculationSafeValueX
+      : Pseudo<(outs GPR64:$dst), (ins GPR64:$src),
+               [(set GPR64:$dst, (AArch64SpeculationSafeValue GPR64:$src))]>,
+        Sched<[]>;
+  def SpeculationSafeValueW
+      : Pseudo<(outs GPR32:$dst), (ins GPR32:$src),
+               [(set GPR32:$dst, (AArch64SpeculationSafeValue GPR32:$src))]>,
+        Sched<[]>;
+}
+
 //===----------------------------------------------------------------------===//
 // Conditional branch (immediate) instruction.
 //===----------------------------------------------------------------------===//
 def Bcc : BranchCond;
 
 //===----------------------------------------------------------------------===//
 // Compare-and-branch instructions.
 //===----------------------------------------------------------------------===//

lib/Target/AArch64/AArch64SpeculationHardening.cpp

 
 #define AARCH64_SPECULATION_HARDENING_NAME "AArch64 speculation hardening pass"
 
 namespace {
 
 class AArch64SpeculationHardening : public MachineFunctionPass {
 public:
   const AArch64InstrInfo *TII;
+  MachineRegisterInfo *MRI;
+  const TargetRegisterInfo *TRI;
 
   static char ID;
 
   AArch64SpeculationHardening() : MachineFunctionPass(ID) {
     initializeAArch64SpeculationHardeningPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &Fn) override;
 
   StringRef getPassName() const override {
     return AARCH64_SPECULATION_HARDENING_NAME;
   }
 
 private:
   unsigned MisspeculatingTaintVR;
   unsigned MisspeculatingTaintVR32Bit;
+  BitVector RegsNeedingBarrierBeforeUse;
 
   bool instrumentCFST(MachineBasicBlock &MBB);
+  bool lowerSpeculationSafeValuePseudo(MachineBasicBlock &MBB);
+  bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
+                                  MachineBasicBlock::iterator MBBI);
+  bool insertSpeculationBarrier(MachineBasicBlock &MBB,
+                                MachineBasicBlock::iterator MBBI, DebugLoc DL);
   bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                      MachineBasicBlock *&FBB, AArch64CC::CondCode &Cond) const;
   bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                                MachineBasicBlock *&FBB,
                                AArch64CC::CondCode &CondCode) const;
   void insertSPToRegTaintPropagation(MachineBasicBlock *MBB,
                                      MachineBasicBlock::iterator MBBI) const;
   void insertRegToSPTaintPropagation(MachineBasicBlock *MBB,
void AArch64SpeculationHardening::insertRegToSPTaintPropagation(
   // mov   SP, Xtmp === ADD SP, Xtmp, #0
   BuildMI(*MBB, MBBI, DebugLoc(), TII->get(AArch64::ADDXri))
       .addDef(AArch64::SP)
       .addUse(TmpReg, RegState::Kill)
       .addImm(0)
       .addImm(0); // no shift
 }
 
+/// \brief If MBBI references a pseudo instruction that should be expanded here,
+/// do the expansion and return true.  Otherwise return false.
+bool AArch64SpeculationHardening::expandSpeculationSafeValue(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
+  MachineInstr &MI = *MBBI;
+  unsigned Opcode = MI.getOpcode();
+  bool Is64Bit = true;
+  switch (Opcode) {
+  default:
+    break;
+  case AArch64::SpeculationSafeValueW:
+    Is64Bit = false;
+    LLVM_FALLTHROUGH;
+  case AArch64::SpeculationSafeValueX:
+    unsigned DstReg = MI.getOperand(0).getReg();
+    unsigned SrcReg = MI.getOperand(1).getReg();
+    for (MachineOperand Op : MI.defs()) {
+      assert(Op.isReg());
+      // Mark this register and all its super-registers as needing to be
+      // value speculation protected before its next use, by using a CSDB
+      // barrier instruction.
+      TRI->markSuperRegs(RegsNeedingBarrierBeforeUse, Op.getReg());
+    }
+
+    assert(EnableSpeculationTracking);
+    // mask off with CFST state
+    BuildMI(MBB, MBBI, MI.getDebugLoc(),
+            Is64Bit ? TII->get(AArch64::ANDXrs) : TII->get(AArch64::ANDWrs))
+        .addDef(DstReg)
+        .addUse(SrcReg, RegState::Kill)
+        .addUse(Is64Bit ? MisspeculatingTaintVR : MisspeculatingTaintVR32Bit)
+        .addImm(0);
+    MI.eraseFromParent();
+    return true;
+  }
+  return false;
+}
+
+bool AArch64SpeculationHardening::insertSpeculationBarrier(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI, DebugLoc DL) {
+  assert(EnableSpeculationTracking && "Speculation tracking must be enabled "
+                                      "when inserting speculation barriers");
+  // insert data value speculation barrier (CSDB)
+  BuildMI(MBB, MBBI, DL, TII->get(AArch64::HINT)).addImm(0x14);
+  RegsNeedingBarrierBeforeUse.reset();
+  return true;
+}
+
+bool AArch64SpeculationHardening::lowerSpeculationSafeValuePseudo(
+    MachineBasicBlock &MBB) {
+  bool Modified = false;
+
+  RegsNeedingBarrierBeforeUse.reset();
+
+  // The following loop iterates over all instructions in the basic block, and
+  // performs 2 operations:
+  // 1. Insert a CSDB at this location if needed.
+  // 2. Insert a masking instruction (AND) with the CFST to lower the
+  //    ValueSpeculationSafe pseudo instruction.
+  //
+  // The insertion of the CSDB is done as late as possible (i.e. just before
+  // the use of a masked register), in the hope that that will reduce the total
+  // number of CSDBs in a block when there are multiple masked registers in the
+  // block.
+  MachineBasicBlock::iterator MBBI = MBB.begin(), E = MBB.end();
+  DebugLoc DL;
+  while (MBBI != E) {
+    MachineInstr &MI = *MBBI;
+    DL = MI.getDebugLoc();
+    MachineBasicBlock::iterator NMBBI = std::next(MBBI);
+
+    // First check if a barrier needs to be inserted due to earlier registers
+    // that were masked and that are used by the next instruction.
+    // Also emit the barrier on any potential control flow changes.
+    bool NeedToEmitBarrier = false;
+    if (RegsNeedingBarrierBeforeUse.any() && (MI.isCall() || MI.isTerminator()))
+      NeedToEmitBarrier = true;
+    if (!NeedToEmitBarrier)
+      for (MachineOperand Op : MI.uses()) {
+        if (Op.isReg() && RegsNeedingBarrierBeforeUse[Op.getReg()]) {
+          NeedToEmitBarrier = true;
+          break;
+        }
+      }
+
+    if (NeedToEmitBarrier)
+      Modified |= insertSpeculationBarrier(MBB, MBBI, DL);
+
+    Modified |= expandSpeculationSafeValue(MBB, MBBI);
+
+    MBBI = NMBBI;
+  }
+
+  if (RegsNeedingBarrierBeforeUse.any())
+    Modified |= insertSpeculationBarrier(MBB, MBBI, DL);
+
+  return Modified;
+}
+
 bool AArch64SpeculationHardening::runOnMachineFunction(MachineFunction &MF) {
   TII = static_cast<const AArch64InstrInfo *>(MF.getSubtarget().getInstrInfo());
+  MRI = &MF.getRegInfo();
+  TRI = MF.getSubtarget().getRegisterInfo();
+
+  RegsNeedingBarrierBeforeUse.resize(TRI->getNumRegs());
 
   bool Modified = false;
 
   MisspeculatingTaintVR = AArch64::X16;
   MisspeculatingTaintVR32Bit = AArch64::W16;
 
-  // Instrument control flow speculation tracking, if requested.
+  // Step 2: instrument control flow speculation tracking, if requested.
   LLVM_DEBUG(
       dbgs()
       << "***** AArch64SpeculationHardening - track control flow *****\n");
 
   if (EnableSpeculationTracking) {
     // 2.1. Add instrumentation code to function entry and exits.
     SmallVector<MachineBasicBlock *, 2> EntryBlocks;
     EntryBlocks.push_back(&MF.front());
     for (const LandingPadInfo &LPI : MF.getLandingPads())
       EntryBlocks.push_back(LPI.LandingPadBlock);
     for (auto Entry : EntryBlocks)
       insertSPToRegTaintPropagation(
           Entry, Entry->SkipPHIsLabelsAndDebug(Entry->begin()));
 
     // 2.2. Add instrumentation code to every basic block.
     for (auto &MBB : MF)
       Modified |= instrumentCFST(MBB);
   }
 
+  LLVM_DEBUG(
+      dbgs()
+      << "***** AArch64SpeculationHardening - Lowering SpeculationSafeValue "
+         "Pseudos *****\n");
+  // Step 3: Lower SpeculationSafeValue pseudo instructions.
+  for (auto &MBB : MF)
+    Modified |= lowerSpeculationSafeValuePseudo(MBB);
+
   return Modified;
 }
 
 /// \brief Returns an instance of the pseudo instruction expansion pass.
 FunctionPass *llvm::createAArch64SpeculationHardeningPass() {
   return new AArch64SpeculationHardening();
 }

test/CodeGen/AArch64/speculation-hardening.ll

 ; RUN: llc < %s -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -aarch64-track-speculation | FileCheck %s --check-prefix=CFST
 
+declare i8 @llvm.speculationsafevalue.i8(i8)
+declare i96 @llvm.speculationsafevalue.i96(i96)
+declare i8* @llvm.speculationsafevalue.p0(i8*)
+
 declare void @function_without_arguments()
 
 define i32 @f(i8* nocapture readonly %p, i32 %i, i32 %N) local_unnamed_addr {
 ; CFST-LABEL: f
 entry:
 ; CFST:  cmp sp, #0
 ; CFST:  csetm x16, ne
 
 ; CFST:  mov x17, sp
 ; CFST:  and x17, x17, x16
 ; CFST:  mov sp, x17
   %call = tail call i32 @tail_callee(i32 %i)
 ; CFST:  cmp sp, #0
 ; CFST:  csetm x16, ne
   %cmp = icmp slt i32 %call, %N
   br i1 %cmp, label %if.then, label %return
 ; CFST: b.ge
 
 if.then:                                          ; preds = %entry
 ; CFST: csel x16, x16, xzr, lt
   %idxprom = sext i32 %i to i64
   %arrayidx = getelementptr inbounds i8, i8* %p, i64 %idxprom
   %0 = load i8, i8* %arrayidx, align 1
+  %1 = tail call i8 @llvm.speculationsafevalue.i8(i8 %0)
 ; CFST:      ldrb [[LOADED:w[0-9]+]],
-  %conv = zext i8 %0 to i32
+; CFST:      and  [[DST:w[0-9]+]], [[LOADED]], w16
+; CFST-NEXT: csdb
+  %conv = zext i8 %1 to i32
   br label %return
 
 ; CFST: csel x16, x16, xzr, ge
 return:                                           ; preds = %entry, %if.then
   %retval.0 = phi i32 [ %conv, %if.then ], [ 0, %entry ]
 ; CFST:  mov x17, sp
 ; CFST:  and x17, x17, x16
 ; CFST:  mov sp, x17
   ret i32 %retval.0
 }
 
+define i64 @nonpow2_narrowing(i128* %p) {
+; CFST-LABEL:  nonpow2_narrowing:
+  %a = load i128, i128* %p
+; CFST: ldp [[SRC1:x[0-9]+]], [[SRC2:x[0-9]+]], [x0]
+  %b = trunc i128 %a to i96
+  %c = tail call i96 @llvm.speculationsafevalue.i96(i96 %b)
+; CFST: and [[SRC1]], [[SRC1]], x16
+; CFST-NOT: csdb
+; CFST: and [[SRC2]], [[SRC2]], x16
+; CFST-NEXT: csdb
+; CFST: add x0, 
+  %lo = trunc i96 %c to i64
+  %hi96 = lshr i96 %c, 64
+  %hi = trunc i96 %hi96 to i64
+  %d = add i64 %lo, %hi
+  ret i64 %d
+}
+
+define i8* @ptr(i8* %p) {
+; CFST-LABEL:  ptr:
+  %a = tail call i8* @llvm.speculationsafevalue.p0(i8* %p)
+; CFST: and x0, x0, x16
+; CFST: csdb
+; CFST: ret
+  ret i8* %a
+}
+
+define i8* @csdb_before_call(i8* %p) {
+; CFST-LABEL:  csdb_before_call:
+  %a = tail call i8* @llvm.speculationsafevalue.p0(i8* %p)
+  call void @function_without_arguments()
+  ret i8* %a
+; CFST: and [[DST:x[0-9]+]], x0, x16
+; CFST: csdb
+; CFST: bl function_without_arguments
+; CFST: ret
+}
+
+define i8* @csdb_on_unused_in_bb_with_terminator(i8* %p, i32 %v1, i32 %v2) {
+; CFST-LABEL:  csdb_on_unused_in_bb_with_terminator:
+entry:
+  %a = tail call i8* @llvm.speculationsafevalue.p0(i8* %p)
+  %cmp = icmp slt i32 %v1, %v2
+  br i1 %cmp, label %if.then, label %return
+; CFST: and [[DST:x[0-9]+]], x0, x16
+; CFST: csdb
+; CFST: b.ge
+; CFST-NOT: csdb
+if.then:
+  %tmp1 = sdiv i32 %v1, %v2
+  %b = getelementptr i8, i8* %a, i32 %tmp1
+  br label %return
+
+return:
+  %c = phi i8* [ %a, %entry ] , [ %b, %if.then ]
+  ret i8* %c
+; CFST: ret
+}
+
+define i8* @csdb_on_unused_in_bb_without_terminator(i8* %p, i32 %v1, i32 %v2) {
+; CFST-LABEL:  csdb_on_unused_in_bb_without_terminator:
+entry:
+  %cmp = icmp slt i32 %v1, %v2
+  br i1 %cmp, label %if.then, label %return
+; CFST: b.ge
+if.then:
+  %tmp1 = sdiv i32 %v1, %v2
+  %b1 = getelementptr i8, i8* %p, i32 %tmp1
+  %b = tail call i8* @llvm.speculationsafevalue.p0(i8* %b1)
+  br label %return
+; CFST: and [[DST:x[0-9]+]], [[SRC:x[0-9]+]], x16
+; CFST: csdb
+; CFST-NOT: csdb
+; CFST: ret
+ 
+return:
+  %c = phi i8* [ %p, %entry ] , [ %b, %if.then ]
+  ret i8* %c
+}
+
 ; Make sure that for a tail call, taint doesn't get put into SP twice.
 define i32 @tail_caller(i32 %a) local_unnamed_addr {
 ; CFST-LABEL: tail_caller:
 ; CFST:     mov     x17, sp
 ; CFST:     and     x17, x17, x16
 ; CFST:     mov     sp, x17
 ; CFST-NOT: mov     x17, sp
   %call = tail call i32 @tail_callee(i32 %a)
   ret i32 %call
 }
 
 declare i32 @tail_callee(i32) local_unnamed_addr
+declare i32 @g(i32) local_unnamed_addr