Add IR and codegen support for deactivation symbols.

llvm/docs/LangRef.rst

@@ -3234,6 +3234,22 @@ A "convergencectrl" operand bundle is only valid on a ``convergent`` operation.
 When present, the operand bundle must contain exactly one value of token type.
 See the :doc:`ConvergentOperations` document for details.
 
+Deactivation Symbol Operand Bundles
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+A ``"deactivation-symbol"`` operand bundle is valid on the following
+instructions (AArch64 only):
+
+- Call to a normal function with ``notail`` attribute and a first argument and
+  return value of type ``ptr``.
+- Call to ``llvm.ptrauth.sign`` or ``llvm.ptrauth.auth`` intrinsics.
+
+This operand bundle specifies that if the deactivation symbol is defined
+to a valid value for the target, the marked instruction will return the
+value of its first argument instead of calling the specified function
+or intrinsic. This is achieved with ``PATCHINST`` relocations on the
+target instructions (see the AArch64 psABI for details).
+
 .. _moduleasm:
 
 Module-Level Inline Assembly

llvm/include/llvm/CodeGen/GlobalISel/CallLowering.h

@@ -159,6 +159,8 @@ class LLVM_ABI CallLowering {
 
     /// True if this call results in convergent operations.
     bool IsConvergent = true;
+
+    GlobalValue *DeactivationSymbol = nullptr;
   };
 
   /// Argument handling is mostly uniform between the four places that

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

@@ -56,6 +56,7 @@ struct MachineIRBuilderState {
   MDNode *PCSections = nullptr;
   /// MMRA Metadata to be set on any instruction we create.
   MDNode *MMRA = nullptr;
+  Value *DS = nullptr;
 
   /// \name Fields describing the insertion point.
   /// @{
@@ -369,6 +370,7 @@ class LLVM_ABI MachineIRBuilder {
     State.II = MI.getIterator();
     setPCSections(MI.getPCSections());
     setMMRAMetadata(MI.getMMRAMetadata());
+    setDeactivationSymbol(MI.getDeactivationSymbol());
   }
   /// @}
 
@@ -405,6 +407,9 @@ class LLVM_ABI MachineIRBuilder {
   /// Set the PC sections metadata to \p MD for all the next build instructions.
   void setMMRAMetadata(MDNode *MMRA) { State.MMRA = MMRA; }
 
+  Value *getDeactivationSymbol() { return State.DS; }
+  void setDeactivationSymbol(Value *DS) { State.DS = DS; }
+
   /// Get the current instruction's MMRA metadata.
   MDNode *getMMRAMetadata() { return State.MMRA; }
 

llvm/include/llvm/CodeGen/ISDOpcodes.h

@@ -1579,6 +1579,10 @@ enum NodeType {
   // Outputs: Output Chain
   CLEAR_CACHE,
 
+  // Untyped node storing deactivation symbol reference
+  // (DeactivationSymbolSDNode).
+  DEACTIVATION_SYMBOL,
+
   /// 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

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -1207,7 +1207,7 @@ class LLVM_ABI MachineFunction {
       ArrayRef<MachineMemOperand *> MMOs, MCSymbol *PreInstrSymbol = nullptr,
       MCSymbol *PostInstrSymbol = nullptr, MDNode *HeapAllocMarker = nullptr,
       MDNode *PCSections = nullptr, uint32_t CFIType = 0,
-      MDNode *MMRAs = nullptr);
+      MDNode *MMRAs = nullptr, Value *DS = nullptr);
 
   /// Allocate a string and populate it with the given external symbol name.
   const char *createExternalSymbolName(StringRef Name);

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -160,29 +160,33 @@ class MachineInstr
   ///
   /// This has to be defined eagerly due to the implementation constraints of
   /// `PointerSumType` where it is used.
-  class ExtraInfo final : TrailingObjects<ExtraInfo, MachineMemOperand *,
-                                          MCSymbol *, MDNode *, uint32_t> {
+  class ExtraInfo final
+      : TrailingObjects<ExtraInfo, MachineMemOperand *, MCSymbol *, MDNode *,
+                        uint32_t, Value *> {
   public:
     static ExtraInfo *create(BumpPtrAllocator &Allocator,
                              ArrayRef<MachineMemOperand *> MMOs,
                              MCSymbol *PreInstrSymbol = nullptr,
                              MCSymbol *PostInstrSymbol = nullptr,
                              MDNode *HeapAllocMarker = nullptr,
                              MDNode *PCSections = nullptr, uint32_t CFIType = 0,
-                             MDNode *MMRAs = nullptr) {
+                             MDNode *MMRAs = nullptr, Value *DS = nullptr) {
       bool HasPreInstrSymbol = PreInstrSymbol != nullptr;
       bool HasPostInstrSymbol = PostInstrSymbol != nullptr;
       bool HasHeapAllocMarker = HeapAllocMarker != nullptr;
       bool HasMMRAs = MMRAs != nullptr;
       bool HasCFIType = CFIType != 0;
       bool HasPCSections = PCSections != nullptr;
+      bool HasDS = DS != nullptr;
       auto *Result = new (Allocator.Allocate(
-          totalSizeToAlloc<MachineMemOperand *, MCSymbol *, MDNode *, uint32_t>(
+          totalSizeToAlloc<MachineMemOperand *, MCSymbol *, MDNode *, uint32_t,
+                           Value *>(
               MMOs.size(), HasPreInstrSymbol + HasPostInstrSymbol,
-              HasHeapAllocMarker + HasPCSections + HasMMRAs, HasCFIType),
+              HasHeapAllocMarker + HasPCSections + HasMMRAs, HasCFIType, HasDS),
           alignof(ExtraInfo)))
           ExtraInfo(MMOs.size(), HasPreInstrSymbol, HasPostInstrSymbol,
-                    HasHeapAllocMarker, HasPCSections, HasCFIType, HasMMRAs);
+                    HasHeapAllocMarker, HasPCSections, HasCFIType, HasMMRAs,
+                    HasDS);
 
       // Copy the actual data into the trailing objects.
       llvm::copy(MMOs, Result->getTrailingObjects<MachineMemOperand *>());
@@ -202,6 +206,8 @@ class MachineInstr
         Result->getTrailingObjects<uint32_t>()[0] = CFIType;
       if (HasMMRAs)
         Result->getTrailingObjects<MDNode *>()[MDNodeIdx++] = MMRAs;
+      if (HasDS)
+        Result->getTrailingObjects<Value *>()[0] = DS;
 
       return Result;
     }
@@ -240,6 +246,10 @@ class MachineInstr
                       : nullptr;
     }
 
+    Value *getDeactivationSymbol() const {
+      return HasDS ? getTrailingObjects<Value *>()[0] : 0;
+    }
+
   private:
     friend TrailingObjects;
 
@@ -255,6 +265,7 @@ class MachineInstr
     const bool HasPCSections;
     const bool HasCFIType;
     const bool HasMMRAs;
+    const bool HasDS;
 
     // Implement the `TrailingObjects` internal API.
     size_t numTrailingObjects(OverloadToken<MachineMemOperand *>) const {
@@ -269,16 +280,17 @@ class MachineInstr
     size_t numTrailingObjects(OverloadToken<uint32_t>) const {
       return HasCFIType;
     }
+    size_t numTrailingObjects(OverloadToken<Value *>) const { return HasDS; }
 
     // Just a boring constructor to allow us to initialize the sizes. Always use
     // the `create` routine above.
     ExtraInfo(int NumMMOs, bool HasPreInstrSymbol, bool HasPostInstrSymbol,
               bool HasHeapAllocMarker, bool HasPCSections, bool HasCFIType,
-              bool HasMMRAs)
+              bool HasMMRAs, bool HasDS)
         : NumMMOs(NumMMOs), HasPreInstrSymbol(HasPreInstrSymbol),
           HasPostInstrSymbol(HasPostInstrSymbol),
           HasHeapAllocMarker(HasHeapAllocMarker), HasPCSections(HasPCSections),
-          HasCFIType(HasCFIType), HasMMRAs(HasMMRAs) {}
+          HasCFIType(HasCFIType), HasMMRAs(HasMMRAs), HasDS(HasDS) {}
   };
 
   /// Enumeration of the kinds of inline extra info available. It is important
@@ -867,6 +879,14 @@ class MachineInstr
     return nullptr;
   }
 
+  Value *getDeactivationSymbol() const {
+    if (!Info)
+      return nullptr;
+    if (ExtraInfo *EI = Info.get<EIIK_OutOfLine>())
+      return EI->getDeactivationSymbol();
+    return nullptr;
+  }
+
   /// Helper to extract a CFI type hash if one has been added.
   uint32_t getCFIType() const {
     if (!Info)
@@ -1969,6 +1989,8 @@ class MachineInstr
   /// Set the CFI type for the instruction.
   LLVM_ABI void setCFIType(MachineFunction &MF, uint32_t Type);
 
+  LLVM_ABI void setDeactivationSymbol(MachineFunction &MF, Value *DS);
+
   /// Return the MIFlags which represent both MachineInstrs. This
   /// should be used when merging two MachineInstrs into one. This routine does
   /// not modify the MIFlags of this MachineInstr.
@@ -2088,7 +2110,7 @@ class MachineInstr
   void setExtraInfo(MachineFunction &MF, ArrayRef<MachineMemOperand *> MMOs,
                     MCSymbol *PreInstrSymbol, MCSymbol *PostInstrSymbol,
                     MDNode *HeapAllocMarker, MDNode *PCSections,
-                    uint32_t CFIType, MDNode *MMRAs);
+                    uint32_t CFIType, MDNode *MMRAs, Value *DS);
 };
 
 /// Special DenseMapInfo traits to compare MachineInstr* by *value* of the

llvm/include/llvm/CodeGen/MachineInstrBuilder.h

@@ -70,29 +70,44 @@ enum {
 } // end namespace RegState
 
 /// Set of metadata that should be preserved when using BuildMI(). This provides
-/// a more convenient way of preserving DebugLoc, PCSections and MMRA.
+/// a more convenient way of preserving certain data from the original
+/// instruction.
 class MIMetadata {
 public:
   MIMetadata() = default;
-  MIMetadata(DebugLoc DL, MDNode *PCSections = nullptr, MDNode *MMRA = nullptr)
-      : DL(std::move(DL)), PCSections(PCSections), MMRA(MMRA) {}
+  MIMetadata(DebugLoc DL, MDNode *PCSections = nullptr, MDNode *MMRA = nullptr,
+             Value *DeactivationSymbol = nullptr)
+      : DL(std::move(DL)), PCSections(PCSections), MMRA(MMRA),
+        DeactivationSymbol(DeactivationSymbol) {}
   MIMetadata(const DILocation *DI, MDNode *PCSections = nullptr,
              MDNode *MMRA = nullptr)
       : DL(DI), PCSections(PCSections), MMRA(MMRA) {}
   explicit MIMetadata(const Instruction &From)
       : DL(From.getDebugLoc()),
-        PCSections(From.getMetadata(LLVMContext::MD_pcsections)) {}
+        PCSections(From.getMetadata(LLVMContext::MD_pcsections)),
+        DeactivationSymbol(getDeactivationSymbol(&From)) {}
   explicit MIMetadata(const MachineInstr &From)
-      : DL(From.getDebugLoc()), PCSections(From.getPCSections()) {}
+      : DL(From.getDebugLoc()), PCSections(From.getPCSections()),
+        DeactivationSymbol(From.getDeactivationSymbol()) {}
 
   const DebugLoc &getDL() const { return DL; }
   MDNode *getPCSections() const { return PCSections; }
   MDNode *getMMRAMetadata() const { return MMRA; }
+  Value *getDeactivationSymbol() const { return DeactivationSymbol; }
 
 private:
   DebugLoc DL;
   MDNode *PCSections = nullptr;
   MDNode *MMRA = nullptr;
+  Value *DeactivationSymbol = nullptr;
+
+  static inline Value *getDeactivationSymbol(const Instruction *I) {
+    if (auto *CB = dyn_cast<CallBase>(I))
+      if (auto Bundle =
+              CB->getOperandBundle(llvm::LLVMContext::OB_deactivation_symbol))
+        return Bundle->Inputs[0].get();
+    return nullptr;
+  }
 };
 
 class MachineInstrBuilder {
@@ -348,6 +363,8 @@ class MachineInstrBuilder {
       MI->setPCSections(*MF, MIMD.getPCSections());
     if (MIMD.getMMRAMetadata())
       MI->setMMRAMetadata(*MF, MIMD.getMMRAMetadata());
+    if (MIMD.getDeactivationSymbol())
+      MI->setDeactivationSymbol(*MF, MIMD.getDeactivationSymbol());
     return *this;
   }
 

llvm/include/llvm/CodeGen/SelectionDAG.h

@@ -759,6 +759,7 @@ class SelectionDAG {
                                  int64_t offset = 0, unsigned TargetFlags = 0) {
     return getGlobalAddress(GV, DL, VT, offset, true, TargetFlags);
   }
+  LLVM_ABI SDValue getDeactivationSymbol(const GlobalValue *GV);
   LLVM_ABI SDValue getFrameIndex(int FI, EVT VT, bool isTarget = false);
   SDValue getTargetFrameIndex(int FI, EVT VT) {
     return getFrameIndex(FI, VT, true);

llvm/include/llvm/CodeGen/SelectionDAGISel.h

@@ -151,6 +151,7 @@ class SelectionDAGISel {
     OPC_RecordChild7,
     OPC_RecordMemRef,
     OPC_CaptureGlueInput,
+    OPC_CaptureDeactivationSymbol,
     OPC_MoveChild,
     OPC_MoveChild0,
     OPC_MoveChild1,

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

@@ -2005,6 +2005,22 @@ class GlobalAddressSDNode : public SDNode {
   }
 };
 
+class DeactivationSymbolSDNode : public SDNode {
+  friend class SelectionDAG;
+
+  const GlobalValue *TheGlobal;
+
+  DeactivationSymbolSDNode(const GlobalValue *GV, SDVTList VTs)
+      : SDNode(ISD::DEACTIVATION_SYMBOL, 0, DebugLoc(), VTs), TheGlobal(GV) {}
+
+public:
+  const GlobalValue *getGlobal() const { return TheGlobal; }
+
+  static bool classof(const SDNode *N) {
+    return N->getOpcode() == ISD::DEACTIVATION_SYMBOL;
+  }
+};
+
 class FrameIndexSDNode : public SDNode {
   friend class SelectionDAG;
 

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -4765,6 +4765,7 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
     SmallVector<SDValue, 4> InVals;
     const ConstantInt *CFIType = nullptr;
     SDValue ConvergenceControlToken;
+    GlobalValue *DeactivationSymbol = nullptr;
 
     std::optional<PtrAuthInfo> PAI;
 
@@ -4918,6 +4919,11 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
       return *this;
     }
 
+    CallLoweringInfo &setDeactivationSymbol(GlobalValue *Sym) {
+      DeactivationSymbol = Sym;
+      return *this;
+    }
+
     ArgListTy &getArgs() {
       return Args;
     }

llvm/include/llvm/IR/LLVMContext.h

@@ -98,7 +98,8 @@ class LLVMContext {
     OB_kcfi = 8,                   // "kcfi"
     OB_convergencectrl = 9,        // "convergencectrl"
     OB_align = 10,                 // "align"
-    OB_LastBundleID = OB_align     // Marker for last bundle ID
+    OB_deactivation_symbol = 11,   // "deactivation-symbol"
+    OB_LastBundleID = OB_deactivation_symbol
   };
 
   /// getMDKindID - Return a unique non-zero ID for the specified metadata kind.

llvm/include/llvm/Target/Target.td

@@ -694,6 +694,7 @@ class Instruction : InstructionEncoding {
                                 // If so, make sure to override
                                 // TargetInstrInfo::getInsertSubregLikeInputs.
   bit variadicOpsAreDefs = false; // Are variadic operands definitions?
+  bit supportsDeactivationSymbol = false;
 
   // Does the instruction have side effects that are not captured by any
   // operands of the instruction or other flags?

llvm/lib/CodeGen/GlobalISel/CallLowering.cpp

@@ -196,6 +196,10 @@ bool CallLowering::lowerCall(MachineIRBuilder &MIRBuilder, const CallBase &CB,
     assert(Info.CFIType->getType()->isIntegerTy(32) && "Invalid CFI type");
   }
 
+  if (auto Bundle = CB.getOperandBundle(LLVMContext::OB_deactivation_symbol)) {
+    Info.DeactivationSymbol = cast<GlobalValue>(Bundle->Inputs[0]);
+  }
+
   Info.CB = &CB;
   Info.KnownCallees = CB.getMetadata(LLVMContext::MD_callees);
   Info.CallConv = CallConv;

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -2917,6 +2917,9 @@ bool IRTranslator::translateIntrinsic(
     }
   }
 
+  if (auto Bundle = CB.getOperandBundle(LLVMContext::OB_deactivation_symbol))
+    MIB->setDeactivationSymbol(*MF, Bundle->Inputs[0].get());
+
   return true;
 }
 

llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

@@ -38,8 +38,10 @@ void MachineIRBuilder::setMF(MachineFunction &MF) {
 //------------------------------------------------------------------------------
 
 MachineInstrBuilder MachineIRBuilder::buildInstrNoInsert(unsigned Opcode) {
-  return BuildMI(getMF(), {getDL(), getPCSections(), getMMRAMetadata()},
-                 getTII().get(Opcode));
+  return BuildMI(
+      getMF(),
+      {getDL(), getPCSections(), getMMRAMetadata(), getDeactivationSymbol()},
+      getTII().get(Opcode));
 }
 
 MachineInstrBuilder MachineIRBuilder::insertInstr(MachineInstrBuilder MIB) {

llvm/lib/CodeGen/MIRParser/MILexer.cpp

@@ -281,6 +281,7 @@ static MIToken::TokenKind getIdentifierKind(StringRef Identifier) {
       .Case("heap-alloc-marker", MIToken::kw_heap_alloc_marker)
       .Case("pcsections", MIToken::kw_pcsections)
       .Case("cfi-type", MIToken::kw_cfi_type)
+      .Case("deactivation-symbol", MIToken::kw_deactivation_symbol)
       .Case("bbsections", MIToken::kw_bbsections)
       .Case("bb_id", MIToken::kw_bb_id)
       .Case("unknown-size", MIToken::kw_unknown_size)

llvm/lib/CodeGen/MIRParser/MILexer.h

@@ -136,6 +136,7 @@ struct MIToken {
     kw_heap_alloc_marker,
     kw_pcsections,
     kw_cfi_type,
+    kw_deactivation_symbol,
     kw_bbsections,
     kw_bb_id,
     kw_unknown_size,