[WPD] Don't optimize calls more than once

aeubanks · aeubanks · commit 7110510ecacf · 2021-06-24T13:28:09.000-07:00
WPD currently assumes that there is a one to one correspondence between type test assume sequences and virtual calls. However, with -fstrict-vtable-pointers this may not be true. This ends up causing crashes when we try to optimize a virtual call more than once ( applyUniformRetValOpt()/applyUniqueRetValOpt()/applyVirtualConstProp()/applySingleImplDevirt()). applySingleImplDevirt() actually didn't previous crash because it would replace the devirtualized call with the same direct call. Adding an assert that the call is indirect causes the corresponding test to crash with the rest of the patch. This makes Chrome successfully build with -fstrict-vtable-pointers + WPD. Reviewed By: tejohnson Differential Revision: https://reviews.llvm.org/D104798
diff --git a/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp b/llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp
@@ -518,6 +518,11 @@ 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;
+
   // 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
@@ -1064,10 +1069,14 @@ void DevirtModule::applySingleImplDevirt(VTableSlotInfo &SlotInfo,
     return;
   auto Apply = [&](CallSiteInfo &CSInfo) {
     for (auto &&VCallSite : CSInfo.CallSites) {
+      if (!OptimizedCalls.insert(&VCallSite.CB).second)
+        continue;
+
       if (RemarksEnabled)
         VCallSite.emitRemark("single-impl",
                              TheFn->stripPointerCasts()->getName(), OREGetter);
       auto &CB = VCallSite.CB;
+      assert(!CB.getCalledFunction() && "devirtualizing direct call?");
       IRBuilder<> Builder(&CB);
       Value *Callee =
           Builder.CreateBitCast(TheFn, CB.getCalledOperand()->getType());
@@ -1406,10 +1415,13 @@ bool DevirtModule::tryEvaluateFunctionsWithArgs(
 
 void DevirtModule::applyUniformRetValOpt(CallSiteInfo &CSInfo, StringRef FnName,
                                          uint64_t TheRetVal) {
-  for (auto Call : CSInfo.CallSites)
+  for (auto Call : CSInfo.CallSites) {
+    if (!OptimizedCalls.insert(&Call.CB).second)
+      continue;
     Call.replaceAndErase(
         "uniform-ret-val", FnName, RemarksEnabled, OREGetter,
         ConstantInt::get(cast<IntegerType>(Call.CB.getType()), TheRetVal));
+  }
   CSInfo.markDevirt();
 }
 
@@ -1515,6 +1527,8 @@ void DevirtModule::applyUniqueRetValOpt(CallSiteInfo &CSInfo, StringRef FnName,
                                         bool IsOne,
                                         Constant *UniqueMemberAddr) {
   for (auto &&Call : CSInfo.CallSites) {
+    if (!OptimizedCalls.insert(&Call.CB).second)
+      continue;
     IRBuilder<> B(&Call.CB);
     Value *Cmp =
         B.CreateICmp(IsOne ? ICmpInst::ICMP_EQ : ICmpInst::ICMP_NE, Call.VTable,
@@ -1583,6 +1597,8 @@ bool DevirtModule::tryUniqueRetValOpt(
 void DevirtModule::applyVirtualConstProp(CallSiteInfo &CSInfo, StringRef FnName,
                                          Constant *Byte, Constant *Bit) {
   for (auto Call : CSInfo.CallSites) {
+    if (!OptimizedCalls.insert(&Call.CB).second)
+      continue;
     auto *RetType = cast<IntegerType>(Call.CB.getType());
     IRBuilder<> B(&Call.CB);
     Value *Addr =
diff --git a/llvm/test/Transforms/WholeProgramDevirt/devirt-single-impl-multiple-assumes.ll b/llvm/test/Transforms/WholeProgramDevirt/devirt-single-impl-multiple-assumes.ll
@@ -0,0 +1,33 @@
+; RUN: opt -S -wholeprogramdevirt -whole-program-visibility %s | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+@vt1 = constant [1 x i8*] [i8* bitcast (void (i8*)* @vf to i8*)], !type !0
+@vt2 = constant [1 x i8*] [i8* bitcast (void (i8*)* @vf to i8*)], !type !0
+
+define void @vf(i8* %this) {
+  ret void
+}
+
+; CHECK: define void @call
+define void @call(i8* %obj) {
+  %vtableptr = bitcast i8* %obj to [1 x i8*]**
+  %vtable = load [1 x i8*]*, [1 x i8*]** %vtableptr
+  %vtablei8 = bitcast [1 x i8*]* %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p)
+  %p2 = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p2)
+  %fptrptr = getelementptr [1 x i8*], [1 x i8*]* %vtable, i32 0, i32 0
+  %fptr = load i8*, i8** %fptrptr
+  %fptr_casted = bitcast i8* %fptr to void (i8*)*
+  ; CHECK: call void @vf(
+  call void %fptr_casted(i8* %obj)
+  ret void
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i32 0, !"typeid"}
diff --git a/llvm/test/Transforms/WholeProgramDevirt/uniform-retval-multiple-assumes.ll b/llvm/test/Transforms/WholeProgramDevirt/uniform-retval-multiple-assumes.ll
@@ -0,0 +1,38 @@
+; RUN: opt -S -wholeprogramdevirt -whole-program-visibility %s | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+@vt1 = constant [1 x i8*] [i8* bitcast (i32 (i8*)* @vf1 to i8*)], !type !0
+@vt2 = constant [1 x i8*] [i8* bitcast (i32 (i8*)* @vf2 to i8*)], !type !0
+
+define i32 @vf1(i8* %this) readnone {
+  ret i32 123
+}
+
+define i32 @vf2(i8* %this) readnone {
+  ret i32 123
+}
+
+; CHECK: define i32 @call
+define i32 @call(i8* %obj) {
+  %vtableptr = bitcast i8* %obj to [1 x i8*]**
+  %vtable = load [1 x i8*]*, [1 x i8*]** %vtableptr
+  %vtablei8 = bitcast [1 x i8*]* %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p)
+  %p2 = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p2)
+  %fptrptr = getelementptr [1 x i8*], [1 x i8*]* %vtable, i32 0, i32 0
+  %fptr = load i8*, i8** %fptrptr
+  %fptr_casted = bitcast i8* %fptr to i32 (i8*)*
+  %result = call i32 %fptr_casted(i8* %obj)
+  ; CHECK-NOT: call i32 %
+  ; CHECK: ret i32 123
+  ret i32 %result
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i32 0, !"typeid"}
diff --git a/llvm/test/Transforms/WholeProgramDevirt/unique-retval-multiple-assumes.ll b/llvm/test/Transforms/WholeProgramDevirt/unique-retval-multiple-assumes.ll
@@ -0,0 +1,41 @@
+; RUN: opt -S -wholeprogramdevirt -whole-program-visibility %s | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+@vt1 = constant [1 x i8*] [i8* bitcast (i1 (i8*)* @vf0 to i8*)], !type !0
+@vt2 = constant [1 x i8*] [i8* bitcast (i1 (i8*)* @vf0 to i8*)], !type !0, !type !1
+@vt3 = constant [1 x i8*] [i8* bitcast (i1 (i8*)* @vf1 to i8*)], !type !0, !type !1
+@vt4 = constant [1 x i8*] [i8* bitcast (i1 (i8*)* @vf1 to i8*)], !type !1
+
+define i1 @vf0(i8* %this) readnone {
+  ret i1 0
+}
+
+define i1 @vf1(i8* %this) readnone {
+  ret i1 1
+}
+
+; CHECK: define i1 @call
+define i1 @call(i8* %obj) {
+  %vtableptr = bitcast i8* %obj to [1 x i8*]**
+  %vtable = load [1 x i8*]*, [1 x i8*]** %vtableptr
+  %vtablei8 = bitcast [1 x i8*]* %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid1")
+  call void @llvm.assume(i1 %p)
+  %p2 = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid1")
+  call void @llvm.assume(i1 %p2)
+  %fptrptr = getelementptr [1 x i8*], [1 x i8*]* %vtable, i32 0, i32 0
+  %fptr = load i8*, i8** %fptrptr
+  %fptr_casted = bitcast i8* %fptr to i1 (i8*)*
+  ; CHECK: [[RES1:%[^ ]*]] = icmp eq [1 x i8*]* %vtable, @vt3
+  %result = call i1 %fptr_casted(i8* %obj)
+  ; CHECK: ret i1 [[RES1]]
+  ret i1 %result
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i32 0, !"typeid1"}
+!1 = !{i32 0, !"typeid2"}
diff --git a/llvm/test/Transforms/WholeProgramDevirt/virtual-const-prop-multiple-assumes.ll b/llvm/test/Transforms/WholeProgramDevirt/virtual-const-prop-multiple-assumes.ll
@@ -0,0 +1,43 @@
+; RUN: opt -S -wholeprogramdevirt -whole-program-visibility %s | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+@vt2 = constant [3 x i8*] [
+i8* bitcast (i1 (i8*)* @vf1i1 to i8*),
+i8* bitcast (i1 (i8*)* @vf0i1 to i8*),
+i8* bitcast (i32 (i8*)* @vf2i32 to i8*)
+], !type !0
+
+define i1 @vf0i1(i8* %this) readnone {
+  ret i1 0
+}
+
+define i1 @vf1i1(i8* %this) readnone {
+  ret i1 1
+}
+
+define i32 @vf2i32(i8* %this) readnone {
+  ret i32 2
+}
+
+; CHECK: define i1 @call1(
+define i1 @call1(i8* %obj) {
+  %vtableptr = bitcast i8* %obj to [3 x i8*]**
+  %vtable = load [3 x i8*]*, [3 x i8*]** %vtableptr
+  %vtablei8 = bitcast [3 x i8*]* %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p)
+  %p2 = call i1 @llvm.type.test(i8* %vtablei8, metadata !"typeid")
+  call void @llvm.assume(i1 %p2)
+  %fptrptr = getelementptr [3 x i8*], [3 x i8*]* %vtable, i32 0, i32 0
+  %fptr = load i8*, i8** %fptrptr
+  %fptr_casted = bitcast i8* %fptr to i1 (i8*)*
+  %result = call i1 %fptr_casted(i8* %obj)
+  ret i1 %result
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i32 0, !"typeid"}
