[WPD] Don't optimize calls more than once

[aeubanks]

Summary

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.

[tejohnson]

WPD currently assumes that each vtable load corresponds to one virtual call. However, with -fstrict-vtable-pointers this may not be true.

In the test case there is still a one to one correlation between the vtable load and virtual call. Do you mean that it can correspond to more than one type test assume sequence?

[aeubanks]

This probably requires more test coverage, but I'm not super familiar

In D104798#2836614, @tejohnson wrote:

WPD currently assumes that each vtable load corresponds to one virtual call. However, with -fstrict-vtable-pointers this may not be true.

In the test case there is still a one to one correlation between the vtable load and virtual call. Do you mean that it can correspond to more than one type test assume sequence?

Originally there are two vtable loads, each with a type test assume sequence. The vtable loads become coalesced, leaving two type test assume sequences for one vtable load.

[tejohnson]

In D104798#2836775, @aeubanks wrote:

This probably requires more test coverage, but I'm not super familiar

In D104798#2836614, @tejohnson wrote:

WPD currently assumes that each vtable load corresponds to one virtual call. However, with -fstrict-vtable-pointers this may not be true.

In the test case there is still a one to one correlation between the vtable load and virtual call. Do you mean that it can correspond to more than one type test assume sequence?

Originally there are two vtable loads, each with a type test assume sequence. The vtable loads become coalesced, leaving two type test assume sequences for one vtable load.

I assume -fstrict-vtable-pointers was what allowed the vtable loads to be coalesced? The description seems wrong though, since it is talking about one vtable load to one virtual call, and here that is still the case. I'm a little confused - in the test case why would there originally have been 2 vtable loads for that single virtual call?

It looks like the issue is that we no longer have a one to one correspondence between type test assume sequences and virtual calls - is that correct?

Also, do we need to do something similar for singleImplDevirt, or why does this issue not occur there?

[aeubanks]

also handle singleImplDevirt

[aeubanks]

In D104798#2837187, @tejohnson wrote:

In D104798#2836775, @aeubanks wrote:

This probably requires more test coverage, but I'm not super familiar

In D104798#2836614, @tejohnson wrote:

WPD currently assumes that each vtable load corresponds to one virtual call. However, with -fstrict-vtable-pointers this may not be true.

In the test case there is still a one to one correlation between the vtable load and virtual call. Do you mean that it can correspond to more than one type test assume sequence?

Originally there are two vtable loads, each with a type test assume sequence. The vtable loads become coalesced, leaving two type test assume sequences for one vtable load.

I assume -fstrict-vtable-pointers was what allowed the vtable loads to be coalesced? The description seems wrong though, since it is talking about one vtable load to one virtual call, and here that is still the case. I'm a little confused - in the test case why would there originally have been 2 vtable loads for that single virtual call?

Oh my test case is confusing. I omitted the second virtual call in the test case and the crash still reproduced.

It looks like the issue is that we no longer have a one to one correspondence between type test assume sequences and virtual calls - is that correct?

Yeah I think your statement makes more sense.

Also, do we need to do something similar for singleImplDevirt, or why does this issue not occur there?

Ah I missed that. Probably it just happened that it didn't trigger in Chrome. Added to that part.

I'll try to come up with tests for all of the cases.

[aeubanks]

add tests

[aeubanks]

add back accidentally removed assert in applySingleImplDevirt()

[tejohnson]

lgtm

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

struct DevirtModule {
   /// when multiple unique return values occur in the same vtable.
   ArrayType *Int8Arr0Ty;
 
   bool RemarksEnabled;
   function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter;
 
   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
   // 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.

 void DevirtModule::applySingleImplDevirt(VTableSlotInfo &SlotInfo,
                                          Constant *TheFn, bool &IsExported) {
   // Don't devirtualize function if we're told to skip it
   // in -wholeprogramdevirt-skip.
   if (FunctionsToSkip.match(TheFn->stripPointerCasts()->getName()))
     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;
       IRBuilder<> Builder(&CB);
       Value *Callee =
           Builder.CreateBitCast(TheFn, CB.getCalledOperand()->getType());
 
if (!Eval.EvaluateFunction(Target.Fn, RetVal, EvalArgs) ||
       return false;
     Target.RetVal = cast<ConstantInt>(RetVal)->getZExtValue();
   }
   return true;
 }
 
 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();
 }
 
 bool DevirtModule::tryUniformRetValOpt(
     MutableArrayRef<VirtualCallTarget> TargetsForSlot, CallSiteInfo &CSInfo,
     WholeProgramDevirtResolution::ByArg *Res) {
   // Uniform return value optimization. If all functions return the same
   // constant, replace all calls with that constant.
else
     SetAbsRange(0, 1ull << AbsWidth);
   return C;
 }
 
 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,
                      B.CreateBitCast(UniqueMemberAddr, Call.VTable->getType()));
     Cmp = B.CreateZExt(Cmp, Call.CB.getType());
     Call.replaceAndErase("unique-ret-val", FnName, RemarksEnabled, OREGetter,
                          Cmp);
   }
if (tryUniqueRetValOptFor(false))
       return true;
   }
   return false;
 }
 
 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 =
         B.CreateGEP(Int8Ty, B.CreateBitCast(Call.VTable, Int8PtrTy), Byte);
     if (RetType->getBitWidth() == 1) {
       Value *Bits = B.CreateLoad(Int8Ty, Addr);
       Value *BitsAndBit = B.CreateAnd(Bits, Bit);
       auto IsBitSet = B.CreateICmpNE(BitsAndBit, ConstantInt::get(Int8Ty, 0));

llvm/test/Transforms/WholeProgramDevirt/devirt-single-impl-multiple-assumes.ll

+; 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"}

llvm/test/Transforms/WholeProgramDevirt/uniform-retval-multiple-assumes.ll

+; 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"}

llvm/test/Transforms/WholeProgramDevirt/unique-retval-multiple-assumes.ll

+; 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"}

llvm/test/Transforms/WholeProgramDevirt/virtual-const-prop-multiple-assumes.ll

+; 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"}
 No newline at end of file