[LTT] Handle merged llvm.assume when dropping type tests

[tejohnson]

Summary

When the lower type test pass is invoked a second time with
DropTypeTests set to true, it expects that all remaining type tests feed
assume instructions, which are removed along with the type tests.

In some cases the llvm.assume might have been merged with another one,
i.e. from a builtin_assume instruction, in which case the type test
would actually feed a phi that in turn feeds the merged assume
instruction. In this case we can simply replace that operand of the phi
with "true" before removing the type test.

[wmi]

LGTM.

[MaskRay]

llvm/test/Transforms/LowerTypeTests/drop_type_test_phi.ll

38

I am a bit puzzled by the comment "If this type test is only used by llvm.assume instructions": this test has exactly one assume. Where did the merge happen?

[tejohnson]

llvm/test/Transforms/LowerTypeTests/drop_type_test_phi.ll

38

The comment you are referring to is from LowerTypeTestsModule::lower() the first time we invoke LTT, which happens before the optimization passes in the LTO backends (right after importing and WPD). There is a second invocation of LTT which simply drops the remaining type test sequences later on, and that is what is being fixed by this patch. It is after a number of standard optimization passes. It has been awhile since I fixed this but so I don't recall exactly which pass was combining the llvm.assume instructions, but it could have been GVN or InstCombine? There are no special restrictions on being able to combine llvm.assume in general.

However, you raise a good point, because the code you referred to will I believe cause the type tests to be removed earlier than desired on the first pass if it already feeds a phi (e.g. if the assume merging happened during the pre-LTO link optimization passes), so that should be loosened up a bit. IIRC the merging of assumes in my case happened after ThinLTO function importing, so didn't hit that case. Let me look at adding a patch to that code to handle the phi case. BTW this is not a correctness issue, it just means they wouldn't be around for subsequent use for optimization.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

static cl::opt<std::string> ClReadSummary(
     cl::desc("Read summary from given YAML file before running pass"),
     cl::Hidden);
 
 static cl::opt<std::string> ClWriteSummary(
     "lowertypetests-write-summary",
     cl::desc("Write summary to given YAML file after running pass"),
     cl::Hidden);
 
+static cl::opt<bool>
+    ClDropTypeTests("lowertypetests-drop-type-tests",
+                    cl::desc("Simply drop type test assume sequences"),
+                    cl::Hidden, cl::init(false));
+
 bool BitSetInfo::containsGlobalOffset(uint64_t Offset) const {
   if (Offset < ByteOffset)
     return false;
 
   if ((Offset - ByteOffset) % (uint64_t(1) << AlignLog2) != 0)
     return false;
 
   uint64_t BitOffset = (Offset - ByteOffset) >> AlignLog2;
struct LowerTypeTests : public ModulePass {
 
   LowerTypeTests() : ModulePass(ID), UseCommandLine(true) {
     initializeLowerTypeTestsPass(*PassRegistry::getPassRegistry());
   }
 
   LowerTypeTests(ModuleSummaryIndex *ExportSummary,
                  const ModuleSummaryIndex *ImportSummary, bool DropTypeTests)
       : ModulePass(ID), ExportSummary(ExportSummary),
-        ImportSummary(ImportSummary), DropTypeTests(DropTypeTests) {
+        ImportSummary(ImportSummary),
+        DropTypeTests(DropTypeTests || ClDropTypeTests) {
     initializeLowerTypeTestsPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnModule(Module &M) override {
     if (UseCommandLine)
       return LowerTypeTestsModule::runForTesting(M);
     return LowerTypeTestsModule(M, ExportSummary, ImportSummary, DropTypeTests)
         .lower();
else
     buildBitSetsFromFunctions(TypeIds, OrderedGTMs);
 }
 
 /// Lower all type tests in this module.
 LowerTypeTestsModule::LowerTypeTestsModule(
     Module &M, ModuleSummaryIndex *ExportSummary,
     const ModuleSummaryIndex *ImportSummary, bool DropTypeTests)
     : M(M), ExportSummary(ExportSummary), ImportSummary(ImportSummary),
-      DropTypeTests(DropTypeTests) {
+      DropTypeTests(DropTypeTests || ClDropTypeTests) {
   assert(!(ExportSummary && ImportSummary));
   Triple TargetTriple(M.getTargetTriple());
   Arch = TargetTriple.getArch();
   OS = TargetTriple.getOS();
   ObjectFormat = TargetTriple.getObjectFormat();
 }
 
 bool LowerTypeTestsModule::runForTesting(Module &M) {
bool LowerTypeTestsModule::lower() {
   if (DropTypeTests && TypeTestFunc) {
     for (auto UI = TypeTestFunc->use_begin(), UE = TypeTestFunc->use_end();
          UI != UE;) {
       auto *CI = cast<CallInst>((*UI++).getUser());
       // Find and erase llvm.assume intrinsics for this llvm.type.test call.
       for (auto CIU = CI->use_begin(), CIUE = CI->use_end(); CIU != CIUE;)
         if (auto *Assume = dyn_cast<AssumeInst>((*CIU++).getUser()))
           Assume->eraseFromParent();
+      // If the assume was merged with another assume, we might have a use on a
+      // phi (which will feed the assume). Simply replace the use on the phi
+      // with "true" and leave the merged assume.
+      if (!CI->use_empty()) {
+        assert(all_of(CI->users(),
+                      [](User *U) -> bool { return isa<PHINode>(U); }));
+        CI->replaceAllUsesWith(ConstantInt::getTrue(M.getContext()));
+      }
       CI->eraseFromParent();
     }
 
     // We have deleted the type intrinsics, so we 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);

llvm/test/Transforms/LowerTypeTests/drop_type_test_phi.ll

+; Test to ensure dropping of type tests can handle a phi feeding the assume.
+; RUN: opt -S -lowertypetests -lowertypetests-drop-type-tests -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+%struct.A = type { i32 (...)** }
+%struct.B = type { %struct.A }
+%struct.C = type { %struct.A }
+
+@_ZTV1B = constant { [4 x i8*] } { [4 x i8*] [i8* null, i8* undef, i8* bitcast (i32 (%struct.B*, i32)* @_ZN1B1fEi to i8*), i8* bitcast (i32 (%struct.A*, i32)* @_ZN1A1nEi to i8*)] }, !type !0, !type !1
+@_ZTV1C = constant { [4 x i8*] } { [4 x i8*] [i8* null, i8* undef, i8* bitcast (i32 (%struct.C*, i32)* @_ZN1C1fEi to i8*), i8* bitcast (i32 (%struct.A*, i32)* @_ZN1A1nEi to i8*)] }, !type !0, !type !2
+
+; CHECK-LABEL: define i32 @test
+define i32 @test(%struct.A* %obj, i32 %a, i32 %b) {
+entry:
+  %tobool.not = icmp eq i32 %a, 0
+  br i1 %tobool.not, label %if.else, label %if.then
+
+if.then:
+  %0 = bitcast %struct.A* %obj to i8***
+  %vtable = load i8**, i8*** %0
+  %1 = bitcast i8** %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %1, metadata !"_ZTS1A")
+; CHECK-NOT: @llvm.type.test
+  %fptrptr = getelementptr i8*, i8** %vtable, i32 1
+  %2 = bitcast i8** %fptrptr to i32 (%struct.A*, i32)**
+  %fptr1 = load i32 (%struct.A*, i32)*, i32 (%struct.A*, i32)** %2, align 8
+  %call = tail call i32 %fptr1(%struct.A* nonnull %obj, i32 %a)
+  br label %if.end
+
+if.else:
+  %3 = icmp ne i32 %b, 0
+  br label %if.end
+
+if.end:
+  %4 = phi i1 [ %3, %if.else ], [ %p, %if.then ]
+  call void @llvm.assume(i1 %4)
+; Still have the assume, but the type test target replaced with true.
+; CHECK: %4 = phi i1 [ %3, %if.else ], [ true, %if.then ]
+; CHECK: call void @llvm.assume(i1 %4)
+
+  ret i32 0
+}
+; CHECK-LABEL: ret i32
+; CHECK-LABEL: }
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+define i32 @_ZN1B1fEi(%struct.B* %this, i32 %a) #0 {
+   ret i32 0;
+}
+
+define i32 @_ZN1A1nEi(%struct.A* %this, i32 %a) #0 {
+   ret i32 0;
+}
+
+define i32 @_ZN1C1fEi(%struct.C* %this, i32 %a) #0 {
+   ret i32 0;
+}
+
+attributes #0 = { noinline optnone }
+
+!0 = !{i64 16, !"_ZTS1A"}
+!1 = !{i64 16, !"_ZTS1B"}
+!2 = !{i64 16, !"_ZTS1C"}