[LowerTypeTests] Ignore llvm.assume users through phis

[aeubanks]

Summary

When checking for type tests that only have assume users, we should look
through phis since we may have merged two assumes through a phi.

This makes it so we also ignore type tests that don't have any users, so
add some users into some existing tests.

[aeubanks]

as somebody unfamiliar with WPD/CFI, is it necessary that they use the same llvm.type.test mechanism?

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

 
 #include "llvm/Transforms/IPO/LowerTypeTests.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TypeMetadataUtils.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Attributes.h"
void LowerTypeTestsModule::replaceCfiUses(Function *Old, Value *New,
   for (auto *C : Constants)
     C->handleOperandChange(Old, New);
 }
 
 void LowerTypeTestsModule::replaceDirectCalls(Value *Old, Value *New) {
   Old->replaceUsesWithIf(New, isDirectCall);
 }
 
+// Returns if the Value is only used by assumes, permitting indirectly via phis.
+static bool hasOnlyAssumeUsers(Value *V, SmallPtrSet<Value *, 2> &Visited) {
+  return all_of(V->uses(), [&Visited](Use &U) {
+    if (auto *P = dyn_cast<PHINode>(U.getUser()))
+      return !Visited.insert(P).second || hasOnlyAssumeUsers(P, Visited);
+    return isa<AssumeInst>(U.getUser());
+  });
+}
+
 bool LowerTypeTestsModule::lower() {
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
 
   if (DropTypeTests && TypeTestFunc) {
     for (Use &U : llvm::make_early_inc_range(TypeTestFunc->uses())) {
       auto *CI = cast<CallInst>(U.getUser());
       // Find and erase llvm.assume intrinsics for this llvm.type.test call.
if (TypeTestFunc) {
     for (const Use &U : TypeTestFunc->uses()) {
       auto CI = cast<CallInst>(U.getUser());
       // If this type test is only used by llvm.assume instructions, it
       // was used for whole program devirtualization, and is being kept
       // for use by other optimization passes. We do not need or want to
       // lower it here. We also don't want to rewrite any associated globals
       // unnecessarily. These will be removed by a subsequent LTT invocation
       // with the DropTypeTests flag set.
-      bool OnlyAssumeUses = !CI->use_empty();
-      for (const Use &CIU : CI->uses()) {
-        if (isa<AssumeInst>(CIU.getUser()))
-          continue;
-        OnlyAssumeUses = false;
-        break;
-      }
-      if (OnlyAssumeUses)
+      SmallPtrSet<Value *, 2> Visited;
+      if (hasOnlyAssumeUsers(CI, Visited))
         continue;
 
       auto TypeIdMDVal = dyn_cast<MetadataAsValue>(CI->getArgOperand(1));
       if (!TypeIdMDVal)
         report_fatal_error("Second argument of llvm.type.test must be metadata");
       auto TypeId = TypeIdMDVal->getMetadata();
       AddTypeIdUse(TypeId).CallSites.push_back(CI);
     }

llvm/test/Transforms/LowerTypeTests/assume-phi.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=lowertypetests -whole-program-visibility %s 2>&1 | FileCheck %s
+
+; Make sure to skip type tests that feed into assumes perhaps indirectly via phis.
+define void @call(ptr %obj, i1 %i) {
+; CHECK-LABEL: @call(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[I:%.*]], label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[VTABLE:%.*]] = load ptr, ptr [[OBJ:%.*]], align 8
+; CHECK-NEXT:    [[FPTR:%.*]] = load ptr, ptr [[VTABLE]], align 8
+; CHECK-NEXT:    call void [[FPTR]](ptr [[OBJ]])
+; CHECK-NEXT:    [[P:%.*]] = call i1 @llvm.type.test(ptr [[VTABLE]], metadata !"typeid")
+; CHECK-NEXT:    br label [[FIN:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[VTABLE2:%.*]] = load ptr, ptr [[OBJ]], align 8
+; CHECK-NEXT:    [[FPTR2:%.*]] = load ptr, ptr [[VTABLE2]], align 8
+; CHECK-NEXT:    call void [[FPTR2]](ptr [[OBJ]])
+; CHECK-NEXT:    [[P2:%.*]] = call i1 @llvm.type.test(ptr [[VTABLE2]], metadata !"typeid")
+; CHECK-NEXT:    br label [[FIN]]
+; CHECK:       fin:
+; CHECK-NEXT:    [[PN:%.*]] = phi i1 [ [[P]], [[BB1]] ], [ [[P2]], [[BB2]] ]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PN]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %i, label %bb1, label %bb2
+bb1:
+  %vtable = load ptr, ptr %obj
+  %fptr = load ptr, ptr %vtable
+  call void %fptr(ptr %obj)
+  %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid")
+  br label %fin
+bb2:
+  %vtable2 = load ptr, ptr %obj
+  %fptr2 = load ptr, ptr %vtable2
+  call void %fptr2(ptr %obj)
+  %p2 = call i1 @llvm.type.test(ptr %vtable2, metadata !"typeid")
+  br label %fin
+fin:
+  %pn = phi i1 [ %p, %bb1 ], [ %p2, %bb2 ]
+  call void @llvm.assume(i1 %pn)
+  ret void
+}
+
+; Make sure we handle phi cycles.
+define void @call2(ptr %obj, i1 %i) {
+; CHECK-LABEL: @call2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[Z:%.*]] = phi i1 [ false, [[ENTRY:%.*]] ], [ [[P:%.*]], [[BB1]] ], [ [[PN:%.*]], [[BB2:%.*]] ]
+; CHECK-NEXT:    [[VTABLE:%.*]] = load ptr, ptr [[OBJ:%.*]], align 8
+; CHECK-NEXT:    [[FPTR:%.*]] = load ptr, ptr [[VTABLE]], align 8
+; CHECK-NEXT:    call void [[FPTR]](ptr [[OBJ]])
+; CHECK-NEXT:    [[P]] = call i1 @llvm.type.test(ptr [[VTABLE]], metadata !"typeid")
+; CHECK-NEXT:    br i1 [[I:%.*]], label [[BB2]], label [[BB1]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[PN]] = phi i1 [ [[Z]], [[BB1]] ]
+; CHECK-NEXT:    br i1 [[I]], label [[FIN:%.*]], label [[BB1]]
+; CHECK:       fin:
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PN]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %bb1
+bb1:
+  %z = phi i1 [ false, %entry ], [ %p, %bb1 ], [ %pn, %bb2 ]
+  %vtable = load ptr, ptr %obj
+  %fptr = load ptr, ptr %vtable
+  call void %fptr(ptr %obj)
+  %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid")
+  br i1 %i, label %bb2, label %bb1
+bb2:
+  %pn = phi i1 [ %z, %bb1 ]
+  br i1 %i, label %fin, label %bb1
+fin:
+  call void @llvm.assume(i1 %pn)
+  ret void
+}
+
+declare i1 @llvm.type.test(ptr, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i32 0, !"typeid"}

llvm/test/Transforms/LowerTypeTests/blockaddress-2.ll

 @badfileops = internal global %struct.f { void ()* @bad_f, void ()* @bad_f }, align 8
 
 declare i1 @llvm.type.test(i8*, metadata)
 
 define internal void @bad_f() !type !1 {
   ret void
 }
 
+declare void @use(i1)
+
 define internal fastcc void @do_f() unnamed_addr !type !2 {
   %1 = tail call i1 @llvm.type.test(i8* undef, metadata !"_ZTSFiP4fileP3uioP5ucrediP6threadE"), !nosanitize !3
+  call void @use(i1 %1)
   ret void
 }
 
 !1 = !{i64 0, !"_ZTSFiP4fileP3uioP5ucrediP6threadE"}
 !2 = !{i64 0, !"_ZTSFiP6threadiP4fileP3uioliE"}
 !3 = !{}

llvm/test/Transforms/LowerTypeTests/blockaddress.ll

 ; RUN: opt -S %s -lowertypetests | FileCheck %s
 
 
 ; CHECK: define hidden i8* @f2.cfi() !type !0 {
 ; CHECK-NEXT:  br label %b
 ; CHECK: b:
 ; CHECK-NEXT:  ret i8* blockaddress(@f2.cfi, %b)
 ; CHECK-NEXT: }
 
 target triple = "x86_64-unknown-linux"
 
+declare void @use(i1)
+
 define void @f1() {
 entry:
   %0 = call i1 @llvm.type.test(i8* bitcast (i8* ()* @f2 to i8*), metadata !"_ZTSFvP3bioE")
+  call void @use(i1 %0)
   ret void
 }
 
 declare i1 @llvm.type.test(i8*, metadata)
 
 define i8* @f2() !type !5 {
   br label %b
 
 b:
   ret i8* blockaddress(@f2, %b)
 }
 
 !5 = !{i64 0, !"_ZTSFvP3bioE"}

llvm/test/Transforms/LowerTypeTests/cfi-direct-call1.ll

 entry:
   %add = add nsw i32 %i, 44
   ret i32 %add
 }
 
 declare !type !3 !type !4 extern_weak i32 @extern_weak()
 declare !type !3 !type !4 i32 @extern_decl()
 declare i1 @llvm.type.test(i8*, metadata)
+declare void @use(i1)
 
 define hidden i32 @main(i32 %argc) {
 entry:
   %cmp.i = icmp sgt i32 %argc, 1
   %fptr1 = select i1 %cmp.i, i32 ()* @local_func1, i32 ()* @local_func2
   %fptr2 = select i1 %cmp.i, i32 ()* @extern_weak, i32 ()* @extern_decl
   %0 = bitcast i32 ()* %fptr1 to i8*
   %1 = tail call i1 @llvm.type.test(i8* nonnull %0, metadata !"_ZTSFivE")
+  call void @use(i1 %1)
 
   %call2 = tail call i32 %fptr1()
   %2 = bitcast i32 ()* %fptr2 to i8*
   %3 = tail call i1 @llvm.type.test(i8* %2, metadata !"_ZTSFivE")
+  call void @use(i1 %3)
 
   %call4 = tail call i32 %fptr2()
   %call5 = tail call i32 @extern_decl()
   %call7 = tail call i32 @extern_weak()
   %call6 = tail call i32 @local_func1()
   %call8 = tail call i32 @local_func3(i32 4)
   ret i32 12
 }

llvm/test/Transforms/LowerTypeTests/layout.ll

 @d = constant i32 4, !type !1
 @e = constant i32 5, !type !2
 
 !0 = !{i32 0, !"typeid1"}
 !1 = !{i32 0, !"typeid2"}
 !2 = !{i32 0, !"typeid3"}
 
 declare i1 @llvm.type.test(i8* %ptr, metadata %bitset) nounwind readnone
+declare void @use(i1)
 
 define void @foo() {
   %x = call i1 @llvm.type.test(i8* undef, metadata !"typeid1")
+  call void @use(i1 %x)
   %y = call i1 @llvm.type.test(i8* undef, metadata !"typeid2")
+  call void @use(i1 %y)
   %z = call i1 @llvm.type.test(i8* undef, metadata !"typeid3")
+  call void @use(i1 %z)
   ret void
 }

llvm/test/Transforms/LowerTypeTests/simple.ll

 ; CHECK: @b = hidden alias [63 x i32], getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 2)
 ; CHECK: @c = protected alias i32, getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 4)
 ; CHECK: @d = alias [2 x i32], getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 6)
 
 ; CHECK: @bits{{[0-9]*}} = private alias i8, getelementptr inbounds ([68 x i8], [68 x i8]* [[BA]], i32 0, i32 0)
 ; CHECK: @bits.{{[0-9]*}} = private alias i8, getelementptr inbounds ([68 x i8], [68 x i8]* [[BA]], i32 0, i32 0)
 
 declare i1 @llvm.type.test(i8* %ptr, metadata %bitset) nounwind readnone
+declare void @use(i1)
 
 ; CHECK: @foo(i32* [[A0:%[^ ]*]])
 define i1 @foo(i32* %p) {
   ; CHECK-NOT: llvm.type.test
 
   ; CHECK: [[R0:%[^ ]*]] = bitcast i32* [[A0]] to i8*
   %pi8 = bitcast i32* %p to i8*
   ; CHECK: [[R1:%[^ ]*]] = ptrtoint i8* [[R0]] to i32
   ; CHECK: [[R2:%[^ ]*]] = sub i32 [[R1]], ptrtoint ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]] to i32)
   ; CHECK: [[R3:%[^ ]*]] = lshr i32 [[R2]], 2
   ; CHECK: [[R4:%[^ ]*]] = shl i32 [[R2]], 30
   ; CHECK: [[R5:%[^ ]*]] = or i32 [[R3]], [[R4]]
   ; CHECK: [[R6:%[^ ]*]] = icmp ule i32 [[R5]], 67
   ; CHECK: br i1 [[R6]]
 
   ; CHECK: [[R8:%[^ ]*]] = getelementptr i8, i8* @bits_use.{{[0-9]*}}, i32 [[R5]]
   ; CHECK: [[R9:%[^ ]*]] = load i8, i8* [[R8]]
   ; CHECK: [[R10:%[^ ]*]] = and i8 [[R9]], 1
   ; CHECK: [[R11:%[^ ]*]] = icmp ne i8 [[R10]], 0
 
   ; CHECK: [[R16:%[^ ]*]] = phi i1 [ false, {{%[^ ]*}} ], [ [[R11]], {{%[^ ]*}} ]
   %x = call i1 @llvm.type.test(i8* %pi8, metadata !"typeid1")
+  call void @use(i1 %x)
 
   ; CHECK-NOT: llvm.type.test
   %y = call i1 @llvm.type.test(i8* %pi8, metadata !"typeid1")
+  call void @use(i1 %y)
 
   ; CHECK: ret i1 [[R16]]
   ret i1 %x
 }
 
 ; CHECK: @bar(i32* [[B0:%[^ ]*]])
 define i1 @bar(i32* %p) {
   ; CHECK: [[S0:%[^ ]*]] = bitcast i32* [[B0]] to i8*
   %pi8 = bitcast i32* %p to i8*
   ; CHECK: [[S1:%[^ ]*]] = ptrtoint i8* [[S0]] to i32
   ; CHECK: [[S2:%[^ ]*]] = sub i32 [[S1]], ptrtoint (i8* getelementptr (i8, i8* bitcast ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]] to i8*), i32 4) to i32)
   ; CHECK: [[S3:%[^ ]*]] = lshr i32 [[S2]], 8
   ; CHECK: [[S4:%[^ ]*]] = shl i32 [[S2]], 24
   ; CHECK: [[S5:%[^ ]*]] = or i32 [[S3]], [[S4]]
   ; CHECK: [[S6:%[^ ]*]] = icmp ule i32 [[S5]], 1
   %x = call i1 @llvm.type.test(i8* %pi8, metadata !"typeid2")
+  call void @use(i1 %x)
 
   ; CHECK: ret i1 [[S6]]
   ret i1 %x
 }
 
 ; CHECK: @baz(i32* [[C0:%[^ ]*]])
 define i1 @baz(i32* %p) {
   ; CHECK: [[T0:%[^ ]*]] = bitcast i32* [[C0]] to i8*