WholeProgramDevirt: Fix call target propagation for ptrauth architectures

[aschwaighofer]

Summary

We can't have a call with a constant target with a ptrauth bundle. Remove the
ptrauth bundle operand in such a case

rdar://105696396

[aschwaighofer]

Run clang-format.

[kubamracek]

AFAICT, this looks reasonable. If we devirtualize, we need to drop the ptrauth bundle. @ab ?

[aschwaighofer]

Rebase

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

struct DevirtModule {
 
   MapVector<VTableSlot, VTableSlotInfo> CallSlots;
 
   // Calls that have already been optimized. We may add a call to multiple
   // VTableSlotInfos if vtable loads are coalesced and need to make sure not to
   // optimize a call more than once.
   SmallPtrSet<CallBase *, 8> OptimizedCalls;
 
+  // Store calls that had their ptrauth bundle removed. They are to be deleted
+  // at the end of the optimization.
+  SmallVector<CallBase *, 8> CallsWithPtrAuthBundleRemoved;
+
   // This map keeps track of the number of "unsafe" uses of a loaded function
   // pointer. The key is the associated llvm.type.test intrinsic call generated
   // by this pass. An unsafe use is one that calls the loaded function pointer
   // directly. Every time we eliminate an unsafe use (for example, by
   // devirtualizing it or by applying virtual constant propagation), we
   // decrement the value stored in this map. If a value reaches zero, we can
   // eliminate the type check by RAUWing the associated llvm.type.test call with
   // true.
for (auto &&VCallSite : CSInfo.CallSites) {
       else {
         // Devirtualize unconditionally.
         CB.setCalledOperand(Callee);
         // Since the call site is now direct, we must clear metadata that
         // is only appropriate for indirect calls. This includes !prof and
         // !callees metadata.
         CB.setMetadata(LLVMContext::MD_prof, nullptr);
         CB.setMetadata(LLVMContext::MD_callees, nullptr);
+        if (CB.getCalledOperand() &&
+            CB.getOperandBundle(LLVMContext::OB_ptrauth)) {
+          auto *NewCS = CallBase::removeOperandBundle(&CB,
+                                                      LLVMContext::OB_ptrauth,
+                                                      &CB);
+          CB.replaceAllUsesWith(NewCS);
+          // Schedule for deletion at the end of pass run.
+          CallsWithPtrAuthBundleRemoved.push_back(&CB);
+        }
       }
 
       // This use is no longer unsafe.
       if (VCallSite.NumUnsafeUses)
         --*VCallSite.NumUnsafeUses;
     }
     if (CSInfo.isExported())
       IsExported = true;
for (VTableBits &B : Bits)
       rebuildGlobal(B);
 
   // We have lowered or deleted the type intrinsics, so we will no longer have
   // enough information to reason about the liveness of virtual function
   // pointers in GlobalDCE.
   for (GlobalVariable &GV : M.globals())
     GV.eraseMetadata(LLVMContext::MD_vcall_visibility);
 
+  for (auto *CI : CallsWithPtrAuthBundleRemoved)
+    CI->eraseFromParent();
+
   return true;
 }
 
 void DevirtIndex::run() {
   if (ExportSummary.typeIdCompatibleVtableMap().empty())
     return;
 
   DenseMap<GlobalValue::GUID, std::vector<StringRef>> NameByGUID;

llvm/test/Transforms/WholeProgramDevirt/devirt-single-impl-check-ptrauth.ll

+; RUN: opt -S -passes=wholeprogramdevirt,verify -whole-program-visibility -pass-remarks=wholeprogramdevirt %s 2>&1 | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK: remark: <unknown>:0:0: single-impl: devirtualized a call to vf
+; CHECK: remark: <unknown>:0:0: devirtualized vf
+; CHECK-NOT: devirtualized
+
+@vt1 = constant [1 x ptr] [ptr @vf], !type !0
+@vt2 = constant [1 x ptr] [ptr @vf], !type !0
+
+define void @vf(ptr %this) {
+  ret void
+}
+
+; CHECK: define void @call
+define void @call(ptr %obj) {
+  %vtable = load ptr, ptr %obj
+  %pair = call {ptr, i1} @llvm.type.checked.load(ptr %vtable, i32 0, metadata !"typeid")
+  %fptr = extractvalue {ptr, i1} %pair, 0
+  %p = extractvalue {ptr, i1} %pair, 1
+  ; CHECK: br i1 true,
+  br i1 %p, label %cont, label %trap
+
+cont:
+  ; CHECK: call void @vf(
+  call void %fptr(ptr %obj) [ "ptrauth"(i32 5, i64 120) ]
+  ret void
+
+trap:
+  call void @llvm.trap()
+  unreachable
+}
+
+declare {ptr, i1} @llvm.type.checked.load(ptr, i32, metadata)
+declare void @llvm.trap()
+
+!0 = !{i32 0, !"typeid"}