[llvm] Allow always dropping all llvm.type.test sequences

[ilovepi]

Currently, the DropTypeTests parameter only fully works with phi nodes
and llvm.assume instructions. However, we'd like CFI to work in
conjunction with FatLTO, in so far as the bitcode section should be able
to contain the CFI instrumentation, while any incompatible bits are
dropped when compiling the object code.

To do that, we need to drop the llvm.type.test instructions everywhere,
and not just their uses in phi nodes. This patch updates the
LowerTypeTest pass so that uses are removed, and replaced with true in
all cases, and not just in phi nodes.

Addressing this will allow us to fix #112053 by modifying the FatLTO
pipeline.

[llvmbot]

@llvm/pr-subscribers-clang-codegen
@llvm/pr-subscribers-clang

@llvm/pr-subscribers-llvm-transforms

Author: Paul Kirth (ilovepi)

Changes

Currently, the DropTypeTests parameter only fully works with phi nodes
and llvm.assume instructions. However, we'd like CFI to work in
conjunction with FatLTO, in so far as the bitcode section should be able
to contain the CFI instrumentation, while any incompatible bits are
dropped when compiling the object code.

To do that, we need to drop the llvm.type.test instructions everywhere,
and not just their uses in phi nodes. This patch updates the
LowerTypeTest pass so that uses are removed, and replaced with true in
all cases, and not just in phi nodes.

Addressing this will allow us to fix #112053 by modifying the FatLTO
pipeline.

---

Full diff: https://github.com/llvm/llvm-project/pull/112787.diff

3 Files Affected:

• (modified) llvm/include/llvm/Transforms/IPO/LowerTypeTests.h (+3-2)
• (modified) llvm/lib/Transforms/IPO/LowerTypeTests.cpp (+20-12)
• (added) llvm/test/Transforms/LowerTypeTests/drop_type_test.ll (+19)

diff --git a/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h b/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h
index eb682c437b94bc..3f5cced844062c 100644
--- a/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h
+++ b/llvm/include/llvm/Transforms/IPO/LowerTypeTests.h
@@ -203,14 +203,15 @@ class LowerTypeTestsPass : public PassInfoMixin<LowerTypeTestsPass> {
   ModuleSummaryIndex *ExportSummary = nullptr;
   const ModuleSummaryIndex *ImportSummary = nullptr;
   bool DropTypeTests = true;
+  bool AlwaysDropTests = false;
 
 public:
   LowerTypeTestsPass() : UseCommandLine(true) {}
   LowerTypeTestsPass(ModuleSummaryIndex *ExportSummary,
                      const ModuleSummaryIndex *ImportSummary,
-                     bool DropTypeTests = false)
+                     bool DropTypeTests = false, bool AlwaysDropTests = false)
       : ExportSummary(ExportSummary), ImportSummary(ImportSummary),
-        DropTypeTests(DropTypeTests) {}
+        DropTypeTests(DropTypeTests), AlwaysDropTests(AlwaysDropTests) {}
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
diff --git a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
index 519a4e9314a26b..864c8cc29533a6 100644
--- a/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
+++ b/llvm/lib/Transforms/IPO/LowerTypeTests.cpp
@@ -122,6 +122,10 @@ static cl::opt<bool>
     ClDropTypeTests("lowertypetests-drop-type-tests",
                     cl::desc("Simply drop type test assume sequences"),
                     cl::Hidden, cl::init(false));
+static cl::opt<bool>
+    ClForceDropTypeTests("lowertypetests-force-drop-type-tests",
+                         cl::desc("Always drop all type test sequences"),
+                         cl::Hidden, cl::init(false));
 
 bool BitSetInfo::containsGlobalOffset(uint64_t Offset) const {
   if (Offset < ByteOffset)
@@ -400,6 +404,7 @@ class LowerTypeTestsModule {
   // Set when the client has invoked this to simply drop all type test assume
   // sequences.
   bool DropTypeTests;
+  bool AlwaysDropTests;
 
   Triple::ArchType Arch;
   Triple::OSType OS;
@@ -542,7 +547,7 @@ class LowerTypeTestsModule {
   LowerTypeTestsModule(Module &M, ModuleAnalysisManager &AM,
                        ModuleSummaryIndex *ExportSummary,
                        const ModuleSummaryIndex *ImportSummary,
-                       bool DropTypeTests);
+                       bool DropTypeTests, bool AlwaysDropTests);
 
   bool lower();
 
@@ -1828,9 +1833,11 @@ void LowerTypeTestsModule::buildBitSetsFromDisjointSet(
 /// Lower all type tests in this module.
 LowerTypeTestsModule::LowerTypeTestsModule(
     Module &M, ModuleAnalysisManager &AM, ModuleSummaryIndex *ExportSummary,
-    const ModuleSummaryIndex *ImportSummary, bool DropTypeTests)
+    const ModuleSummaryIndex *ImportSummary, bool DropTypeTests,
+    bool AlwaysDropTests)
     : M(M), ExportSummary(ExportSummary), ImportSummary(ImportSummary),
-      DropTypeTests(DropTypeTests || ClDropTypeTests) {
+      DropTypeTests(DropTypeTests || ClDropTypeTests),
+      AlwaysDropTests(AlwaysDropTests || ClForceDropTypeTests) {
   assert(!(ExportSummary && ImportSummary));
   Triple TargetTriple(M.getTargetTriple());
   Arch = TargetTriple.getArch();
@@ -1882,7 +1889,7 @@ bool LowerTypeTestsModule::runForTesting(Module &M, ModuleAnalysisManager &AM) {
           M, AM,
           ClSummaryAction == PassSummaryAction::Export ? &Summary : nullptr,
           ClSummaryAction == PassSummaryAction::Import ? &Summary : nullptr,
-          /*DropTypeTests*/ false)
+          /*DropTypeTests*/ false, /*AlwaysDropTests=*/false)
           .lower();
 
   if (!ClWriteSummary.empty()) {
@@ -1949,7 +1956,7 @@ void LowerTypeTestsModule::replaceDirectCalls(Value *Old, Value *New) {
   Old->replaceUsesWithIf(New, isDirectCall);
 }
 
-static void dropTypeTests(Module &M, Function &TypeTestFunc) {
+static void dropTypeTests(Module &M, Function &TypeTestFunc, bool AlwaysDrop) {
   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.
@@ -1960,8 +1967,9 @@ static void dropTypeTests(Module &M, Function &TypeTestFunc) {
     // 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); }));
+      assert(AlwaysDrop || all_of(CI->users(), [](User *U) -> bool {
+               return isa<PHINode>(U);
+             }));
       CI->replaceAllUsesWith(ConstantInt::getTrue(M.getContext()));
     }
     CI->eraseFromParent();
@@ -1974,14 +1982,14 @@ bool LowerTypeTestsModule::lower() {
 
   if (DropTypeTests) {
     if (TypeTestFunc)
-      dropTypeTests(M, *TypeTestFunc);
+      dropTypeTests(M, *TypeTestFunc, AlwaysDropTests);
     // Normally we'd have already removed all @llvm.public.type.test calls,
     // except for in the case where we originally were performing ThinLTO but
     // decided not to in the backend.
     Function *PublicTypeTestFunc =
         M.getFunction(Intrinsic::getName(Intrinsic::public_type_test));
     if (PublicTypeTestFunc)
-      dropTypeTests(M, *PublicTypeTestFunc);
+      dropTypeTests(M, *PublicTypeTestFunc, AlwaysDropTests);
     if (TypeTestFunc || PublicTypeTestFunc) {
       // We have deleted the type intrinsics, so we no longer have enough
       // information to reason about the liveness of virtual function pointers
@@ -2449,9 +2457,9 @@ PreservedAnalyses LowerTypeTestsPass::run(Module &M,
   if (UseCommandLine)
     Changed = LowerTypeTestsModule::runForTesting(M, AM);
   else
-    Changed =
-        LowerTypeTestsModule(M, AM, ExportSummary, ImportSummary, DropTypeTests)
-            .lower();
+    Changed = LowerTypeTestsModule(M, AM, ExportSummary, ImportSummary,
+                                   DropTypeTests, AlwaysDropTests)
+                  .lower();
   if (!Changed)
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
diff --git a/llvm/test/Transforms/LowerTypeTests/drop_type_test.ll b/llvm/test/Transforms/LowerTypeTests/drop_type_test.ll
new file mode 100644
index 00000000000000..23fa0b97790cf0
--- /dev/null
+++ b/llvm/test/Transforms/LowerTypeTests/drop_type_test.ll
@@ -0,0 +1,19 @@
+; RUN: opt -S -passes=lowertypetests -lowertypetests-force-drop-type-tests -lowertypetests-drop-type-tests < %s | FileCheck %s
+
+define void @func() {
+entry:
+  %0 = tail call i1 @llvm.type.test(ptr null, metadata !"foo")
+  br i1 %0, label %exit, label %trap
+  ;      CHECK: entry:
+  ; CHECK-NEXT: br i1 true, label %exit, label %trap
+  ; CHECK-NOT: @llvm.type.test
+
+trap:
+  unreachable
+
+exit:
+  ret void
+}
+
+declare i1 @llvm.type.test(ptr, metadata) #0
+attributes #0 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }

[ilovepi]

Happy to drop the assert instead of passing through all the extra booleans, but I wasn't sure if that was somehow load bearing outside of the FatLTO use case.

[ilovepi]

@pcc I think I've addressed your comments from #112788