Add llvm.protected.field.ptr intrinsic and pre-ISel lowering.

llvm/docs/LangRef.rst

@@ -31721,3 +31721,57 @@ Semantics:
 
 The '``llvm.preserve.struct.access.index``' intrinsic produces the same result
 as a getelementptr with base ``base`` and access operands ``{0, gep_index}``.
+
+'``llvm.protected.field.ptr``' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+::
+
+      declare ptr @llvm.protected.field.ptr(ptr ptr, i64 disc, i1 use_hw_encoding)
+
+Overview:
+"""""""""
+
+The '``llvm.protected.field.ptr``' intrinsic returns a pointer to the
+storage location of a pointer that has special properties as described
+below.
+
+Arguments:
+""""""""""
+
+The first argument is the pointer specifying the location to store the
+pointer. The second argument is the discriminator, which is used as an
+input for the pointer encoding. The third argument specifies whether to
+use a target-specific mechanism to encode the pointer.
+
+Semantics:
+""""""""""
+
+This intrinsic returns a pointer which may be used to store a
+pointer at the specified address that is encoded using the specified
+discriminator. Stores via the pointer will cause the stored pointer to be
+blended with the second argument before being stored. The blend operation
+shall be either a weak but cheap and target-independent operation (if
+the third argument is 0) or a stronger target-specific operation (if the
+third argument is 1). When loading from the pointer, the inverse operation
+is done on the loaded pointer after it is loaded. Specifically, when the
+third argument is 1, the pointer is signed (using pointer authentication
+instructions or emulated PAC if not supported by the hardware) using
+the discriminator before being stored, and authenticated after being
+loaded. Note that it is currently unsupported to have the third argument
+be 1 on targets other than AArch64. When the third argument is 0, it is
+rotated left by 16 bits and the discriminator is subtracted before being
+stored, and the discriminator is added and the pointer is rotated right
+by 16 bits after being loaded.
+
+If the pointer is used other than for loading or storing (e.g. its
+address escapes), that will disable all blending operations using
+the deactivation symbol specified in the intrinsic's operand bundle.
+The deactivation symbol operand bundle is copied onto any sign and auth
+intrinsics that this intrinsic is lowered into. The intent is that the
+deactivation symbol represents a field identifier.
+
+This intrinsic is used to implement structure protection.

llvm/include/llvm/IR/Intrinsics.td

@@ -2925,6 +2925,13 @@ def int_experimental_convergence_anchor
 def int_experimental_convergence_loop
   : DefaultAttrsIntrinsic<[llvm_token_ty], [], [IntrNoMem, IntrConvergent]>;
 
+//===----------------- Structure Protection Intrinsics --------------------===//
+
+def int_protected_field_ptr :
+  DefaultAttrsIntrinsic<[llvm_anyptr_ty],
+                        [LLVMMatchType<0>, llvm_i64_ty, llvm_i1_ty],
+                        [IntrNoMem, ImmArg<ArgIndex<2>>]>;
+
 //===----------------------------------------------------------------------===//
 // Target-specific intrinsics
 //===----------------------------------------------------------------------===//

llvm/lib/CodeGen/PreISelIntrinsicLowering.cpp

@@ -21,9 +21,11 @@
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/RuntimeLibcalls.h"
 #include "llvm/IR/Type.h"
@@ -461,6 +463,159 @@ bool PreISelIntrinsicLowering::expandMemIntrinsicUses(
   return Changed;
 }
 
+static bool expandProtectedFieldPtr(Function &Intr) {
+  Module &M = *Intr.getParent();
+
+  SmallPtrSet<GlobalValue *, 2> DSsToDeactivate;
+
+  Type *Int8Ty = Type::getInt8Ty(M.getContext());
+  Type *Int64Ty = Type::getInt64Ty(M.getContext());
+  PointerType *PtrTy = PointerType::get(M.getContext(), 0);
+
+  Function *SignIntr =
+      Intrinsic::getOrInsertDeclaration(&M, Intrinsic::ptrauth_sign, {});
+  Function *AuthIntr =
+      Intrinsic::getOrInsertDeclaration(&M, Intrinsic::ptrauth_auth, {});
+
+  auto *EmuFnTy = FunctionType::get(Int64Ty, {Int64Ty, Int64Ty}, false);
+  FunctionCallee EmuSignIntr = M.getOrInsertFunction("__emupac_pacda", EmuFnTy);
+  FunctionCallee EmuAuthIntr = M.getOrInsertFunction("__emupac_autda", EmuFnTy);
+
+  auto CreateSign = [&](IRBuilder<> &B, Value *Val, Value *Disc,
+                        OperandBundleDef DSBundle) {
+    Function *F = B.GetInsertBlock()->getParent();
+    Attribute FSAttr = F->getFnAttribute("target-features");
+    if (FSAttr.isValid() && FSAttr.getValueAsString().contains("+pauth"))
+      return B.CreateCall(SignIntr, {Val, B.getInt32(2), Disc}, DSBundle);
+    return B.CreateCall(EmuSignIntr, {Val, Disc}, DSBundle);
+  };
+
+  auto CreateAuth = [&](IRBuilder<> &B, Value *Val, Value *Disc,
+                        OperandBundleDef DSBundle) {
+    Function *F = B.GetInsertBlock()->getParent();
+    Attribute FSAttr = F->getFnAttribute("target-features");
+    if (FSAttr.isValid() && FSAttr.getValueAsString().contains("+pauth"))
+      return B.CreateCall(AuthIntr, {Val, B.getInt32(2), Disc}, DSBundle);
+    return B.CreateCall(EmuAuthIntr, {Val, Disc}, DSBundle);
+  };
+
+  auto GetDeactivationSymbol = [&](CallInst *Call) -> GlobalValue * {
+    if (auto Bundle =
+            Call->getOperandBundle(LLVMContext::OB_deactivation_symbol))
+      return cast<GlobalValue>(Bundle->Inputs[0]);
+    return nullptr;
+  };
+
+  for (User *U : llvm::make_early_inc_range(Intr.users())) {
+    auto *Call = cast<CallInst>(U);
+
+    auto *Pointer = Call->getArgOperand(0);
+    auto *Disc = Call->getArgOperand(1);
+    bool UseHWEncoding =
+        cast<ConstantInt>(Call->getArgOperand(2))->getZExtValue();
+
+    auto *DS = GetDeactivationSymbol(Call);
+    OperandBundleDef DSBundle("deactivation-symbol", DS);
+
+    for (Use &U : llvm::make_early_inc_range(Call->uses())) {
+      // Insert code to encode each pointer stored to the pointer returned by
+      // the intrinsic.
+      if (auto *SI = dyn_cast<StoreInst>(U.getUser())) {
+        if (U.getOperandNo() == 1 &&
+            isa<PointerType>(SI->getValueOperand()->getType())) {
+          IRBuilder<> B(SI);
+          auto *SIValInt =
+              B.CreatePtrToInt(SI->getValueOperand(), B.getInt64Ty());
+          Value *Sign;
+          if (UseHWEncoding) {
+            Sign = CreateSign(B, SIValInt, Disc, DSBundle);
+          } else {
+            // FIXME: These don't have deactivation symbol attachments, we'll
+            // need to figure out how to add them.
+            Sign =
+                B.CreateIntrinsic(SIValInt->getType(), Intrinsic::fshl,
+                                  {SIValInt, SIValInt,
+                                   ConstantInt::get(SIValInt->getType(), 16)});
+            Sign = B.CreateSub(Sign, Disc);
+          }
+          SI->setOperand(0, B.CreateIntToPtr(Sign, B.getPtrTy()));
+          SI->setOperand(1, Pointer);
+          continue;
+        }
+      }
+
+      // Insert code to decode each pointer loaded from the pointer returned by
+      // the intrinsic. This is the inverse of the encode operation implemented
+      // above.
+      if (auto *LI = dyn_cast<LoadInst>(U.getUser())) {
+        if (isa<PointerType>(LI->getType())) {
+          IRBuilder<> B(LI);
+          auto *NewLI = cast<LoadInst>(LI->clone());
+          NewLI->setOperand(0, Pointer);
+          B.Insert(NewLI);
+          auto *LIInt = B.CreatePtrToInt(NewLI, B.getInt64Ty());
+          Value *Auth;
+          if (UseHWEncoding) {
+            Auth = CreateAuth(B, LIInt, Disc, DSBundle);
+          } else {
+            // FIXME: These don't have deactivation symbol attachments, we'll
+            // need to figure out how to add them.
+            Auth = B.CreateAdd(LIInt, Disc);
+            Auth = B.CreateIntrinsic(
+                Auth->getType(), Intrinsic::fshr,
+                {Auth, Auth, ConstantInt::get(Auth->getType(), 16)});
+          }
+          LI->replaceAllUsesWith(B.CreateIntToPtr(Auth, B.getPtrTy()));
+          LI->eraseFromParent();
+          continue;
+        }
+      }
+      // Comparisons against null cannot be used to recover the original
+      // pointer so we replace them with comparisons against the original
+      // pointer.
+      if (auto *CI = dyn_cast<ICmpInst>(U.getUser())) {
+        if (auto *Op = dyn_cast<Constant>(CI->getOperand(0))) {
+          if (Op->isNullValue()) {
+            CI->setOperand(1, Pointer);
+            continue;
+          }
+        }
+        if (auto *Op = dyn_cast<Constant>(CI->getOperand(1))) {
+          if (Op->isNullValue()) {
+            CI->setOperand(0, Pointer);
+            continue;
+          }
+        }
+      }
+
+      // We couldn't rewrite away this use of the intrinsic. Replace it with the
+      // pointer operand, and arrange to define a deactivation symbol.
+      U.set(Pointer);
+      if (DS)
+        DSsToDeactivate.insert(DS);
+    }
+
+    Call->eraseFromParent();
+  }
+
+  if (!DSsToDeactivate.empty()) {
+    // This is an AArch64 NOP instruction. When the deactivation symbol support
+    // is expanded to more architectures, there will likely need to be an API
+    // for retrieving this constant.
+    Constant *Nop =
+        ConstantExpr::getIntToPtr(ConstantInt::get(Int64Ty, 0xd503201f), PtrTy);
+    for (GlobalValue *OldDS : DSsToDeactivate) {
+      GlobalValue *DS = GlobalAlias::create(
+          Int8Ty, 0, GlobalValue::ExternalLinkage, OldDS->getName(), Nop, &M);
+      DS->setVisibility(GlobalValue::HiddenVisibility);
+      DS->takeName(OldDS);
+      OldDS->replaceAllUsesWith(DS);
+      OldDS->eraseFromParent();
+    }
+  }
+  return true;
+}
+
 bool PreISelIntrinsicLowering::lowerIntrinsics(Module &M) const {
   // Map unique constants to globals.
   DenseMap<Constant *, GlobalVariable *> CMap;
@@ -603,6 +758,9 @@ bool PreISelIntrinsicLowering::lowerIntrinsics(Module &M) const {
         return lowerUnaryVectorIntrinsicAsLoop(M, CI);
       });
       break;
+    case Intrinsic::protected_field_ptr:
+      Changed |= expandProtectedFieldPtr(F);
+      break;
     }
   }
   return Changed;