Add IR and codegen support for deactivation symbols.

Deactivation symbols are a mechanism for allowing object files to disable
specific instructions in other object files at link time. The initial use
case is for pointer field protection.

For more information, see the RFC:
https://discourse.llvm.org/t/rfc-deactivation-symbols/85556

Reviewers: ojhunt, nikic, fmayer, arsenm, ahmedbougacha

Reviewed By: fmayer

Pull Request: llvm#133536

Author: pcc

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;