hwasan: Use llvm.read_register intrinsic to read the PC on aarch64 instead of taking the function's address.

[pcc]

Summary

This shaves an instruction (and a GOT entry in PIC code) off prologues of
functions with stack variables.

Depends on D63469

Depends on D63471

[pcc]

• Add test

[eugenis]

LGTM nice!

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

public:
   bool isInterestingAlloca(const AllocaInst &AI);
   bool tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
   bool instrumentStack(
       SmallVectorImpl<AllocaInst *> &Allocas,
       DenseMap<AllocaInst *, std::vector<DbgDeclareInst *>> &AllocaDeclareMap,
       SmallVectorImpl<Instruction *> &RetVec, Value *StackTag);
+  Value *readRegister(IRBuilder<> &IRB, StringRef Name);
   bool instrumentLandingPads(SmallVectorImpl<Instruction *> &RetVec);
   Value *getNextTagWithCall(IRBuilder<> &IRB);
   Value *getStackBaseTag(IRBuilder<> &IRB);
   Value *getAllocaTag(IRBuilder<> &IRB, Value *StackTag, AllocaInst *AI,
                      unsigned AllocaNo);
   Value *getUARTag(IRBuilder<> &IRB, Value *StackTag);
 
   Value *getHwasanThreadSlotPtr(IRBuilder<> &IRB, Type *Ty);
void HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord) {
   // Extract the address field from ThreadLong. Unnecessary on AArch64 with TBI.
   Value *ThreadLongMaybeUntagged =
       TargetTriple.isAArch64() ? ThreadLong : untagPointer(IRB, ThreadLong);
 
   if (WithFrameRecord) {
     StackBaseTag = IRB.CreateAShr(ThreadLong, 3);
 
     // Prepare ring buffer data.
-    auto PC = IRB.CreatePtrToInt(F, IntptrTy);
+    Value *PC;
+    if (TargetTriple.getArch() == Triple::aarch64)
+      PC = readRegister(IRB, "pc");
+    else
+      PC = IRB.CreatePtrToInt(F, IntptrTy);
     auto GetStackPointerFn =
         Intrinsic::getDeclaration(F->getParent(), Intrinsic::frameaddress);
     Value *SP = IRB.CreatePtrToInt(
         IRB.CreateCall(GetStackPointerFn,
                        {Constant::getNullValue(IRB.getInt32Ty())}),
         IntptrTy);
     // Mix SP and PC.
     // Assumptions:
void HWAddressSanitizer::emitPrologue(IRBuilder<> &IRB, bool WithFrameRecord) {
   LocalDynamicShadow = IRB.CreateAdd(
       IRB.CreateOr(
           ThreadLongMaybeUntagged,
           ConstantInt::get(IntptrTy, (1ULL << kShadowBaseAlignment) - 1)),
       ConstantInt::get(IntptrTy, 1), "hwasan.shadow");
   LocalDynamicShadow = IRB.CreateIntToPtr(LocalDynamicShadow, Int8PtrTy);
 }
 
-bool HWAddressSanitizer::instrumentLandingPads(
-    SmallVectorImpl<Instruction *> &LandingPadVec) {
-  Module *M = LandingPadVec[0]->getModule();
+Value *HWAddressSanitizer::readRegister(IRBuilder<> &IRB, StringRef Name) {
+  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
   Function *ReadRegister =
       Intrinsic::getDeclaration(M, Intrinsic::read_register, IntptrTy);
-  const char *RegName =
-      (TargetTriple.getArch() == Triple::x86_64) ? "rsp" : "sp";
-  MDNode *MD = MDNode::get(*C, {MDString::get(*C, RegName)});
+  MDNode *MD = MDNode::get(*C, {MDString::get(*C, Name)});
   Value *Args[] = {MetadataAsValue::get(*C, MD)};
+  return IRB.CreateCall(ReadRegister, Args);
+}
 
+bool HWAddressSanitizer::instrumentLandingPads(
+    SmallVectorImpl<Instruction *> &LandingPadVec) {
   for (auto *LP : LandingPadVec) {
     IRBuilder<> IRB(LP->getNextNode());
-    IRB.CreateCall(HWAsanHandleVfork, {IRB.CreateCall(ReadRegister, Args)});
+    IRB.CreateCall(
+        HWAsanHandleVfork,
+        {readRegister(IRB, (TargetTriple.getArch() == Triple::x86_64) ? "rsp"
+                                                                      : "sp")});
   }
   return true;
 }
 
 bool HWAddressSanitizer::instrumentStack(
     SmallVectorImpl<AllocaInst *> &Allocas,
     DenseMap<AllocaInst *, std::vector<DbgDeclareInst *>> &AllocaDeclareMap,
     SmallVectorImpl<Instruction *> &RetVec, Value *StackTag) {

llvm/test/Instrumentation/HWAddressSanitizer/prologue.ll

 ; CHECK-TLS:   %[[A:[^ ]*]] = call i8* @llvm.thread.pointer()
 ; CHECK-TLS:   %[[B:[^ ]*]] = getelementptr i8, i8* %[[A]], i32 48
 ; CHECK-TLS:   %[[C:[^ ]*]] = bitcast i8* %[[B]] to i64*
 ; CHECK-TLS:   %[[D:[^ ]*]] = load i64, i64* %[[C]]
 ; CHECK-TLS:   %[[E:[^ ]*]] = ashr i64 %[[D]], 3
 
 ; CHECK-NOHISTORY-NOT: store i64
 
+; CHECK-HISTORY: call i64 @llvm.read_register.i64(metadata [[MD:![0-9]*]])
 ; CHECK-HISTORY: %[[PTR:[^ ]*]] = inttoptr i64 %[[D]] to i64*
 ; CHECK-HISTORY: store i64 %{{.*}}, i64* %[[PTR]]
 ; CHECK-HISTORY: %[[D1:[^ ]*]] = ashr i64 %[[D]], 56
 ; CHECK-HISTORY: %[[D2:[^ ]*]] = shl nuw nsw i64 %[[D1]], 12
 ; CHECK-HISTORY: %[[D3:[^ ]*]] = xor i64 %[[D2]], -1
 ; CHECK-HISTORY: %[[D4:[^ ]*]] = add i64 %[[D]], 8
 ; CHECK-HISTORY: %[[D5:[^ ]*]] = and i64 %[[D4]], %[[D3]]
 ; CHECK-HISTORY: store i64 %[[D5]], i64* %[[C]]
 
 ; CHECK-TLS:   %[[F:[^ ]*]] = or i64 %[[D]], 4294967295
 ; CHECK-TLS:   = add i64 %[[F]], 1
 
 ; CHECK-HISTORY: = xor i64 %[[E]], 0
 
 ; CHECK-NOHISTORY-NOT: store i64
 
 
 entry:
   %x = alloca i32, align 4
   call void @use(i32* %x)
   ret void
 }
+
+; CHECK-HISTORY: [[MD]] = !{!"pc"}