[WPD] Extend checking mode to support fallback to indirect call

[tejohnson]

Summary

Extend -wholeprogramdevirt-check to support both the existing
trapping mode on an incorrect devirtualization, as well as a new
mode to fallback to an indirect call on a mismatch. The new mode is

The new mode is useful in cases where we want to enable
devirtualization but cannot fully guarantee whole program visibility
(e.g in the case where LTO has been disabled for a small set of objects
that could potentially override virtual methods without having a symbol
reference to anything in the base class including the vtable).

[mingmingl]

It occurred to me that the WPD internal doc also mentions --wholeprogramdevirt-check, and current example uses it as a boolean.

Overall changing type of a debugging option makes sense (users would be able to figure out if they see an error, and there are enough time to update g3doc before a release)

llvm/include/llvm/Transforms/Utils/CallPromotionUtils.h

84

nit pick

I'm wondering if it's more conventional to keep examples [1] in cpp or in the header (both are pretty readable IMO)

[1] https://github.com/llvm/llvm-project/blob/fc968bcba4d7e02390f1c1cba76a6a2cad1ae59f/llvm/lib/Transforms/Utils/CallPromotionUtils.cpp#L195-L281

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

1174

llvm::promoteCall [1] clears two types of metadata and then compares function type.

Function types are the same (for virtual functions); I'm wondering if it's needed to clear two types of metadata here (for my understanding)

[1] https://github.com/llvm/llvm-project/blob/fc968bcba4d7e02390f1c1cba76a6a2cad1ae59f/llvm/lib/Transforms/Utils/CallPromotionUtils.cpp#L456

[tejohnson]

Address comments

[tejohnson]

In D121419#3374225, @luna wrote:

It occurred to me that the WPD internal doc also mentions --wholeprogramdevirt-check, and current example uses it as a boolean.

Yep, I'll be updating the internal doc - afaik no one is using it (was previously just used by me for debugging).

llvm/include/llvm/Transforms/Utils/CallPromotionUtils.h

84

Not sure of the convention, but I wanted to keep the comment block in the header reasonable and similar in size to the existing ones. The examples comment block in the source file is really large for this function. I'll keep it there for now.

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

1174

Good catch. Looks like leaving this metadata on direct calls doesn't cause any issues, since it hasn't been stripped for the normal WPD (non-fallback) case before. But it causes a real issue when left on the fallback indirect case in the new handling, because ICP will come along and promote again. This isn't a correctness issue, but could bloat codesize - although in the case I tested the optimizer smartly merged the redundant promoted direct calls. But I have fixed it in all cases (fallback and for the direct calls under all options) and added tests for all of these cases.

[mingmingl]

In D121419#3375351, @tejohnson wrote:

In D121419#3374225, @luna wrote:

It occurred to me that the WPD internal doc also mentions --wholeprogramdevirt-check, and current example uses it as a boolean.

Yep, I'll be updating the internal doc - afaik no one is using it (was previously just used by me for debugging).

Thanks for making the change!

llvm/include/llvm/Transforms/Utils/CallPromotionUtils.h

84

Not sure of the convention, but I wanted to keep the comment block in the header reasonable and similar in size to the existing ones. The examples comment block in the source file is really large for this function. I'll keep it there for now.

Yeah the current way is reasonable and readable, so sounds good to keep it as is.

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

1174

thanks for figuring it out and making the fix!

llvm/include/llvm/Transforms/Utils/CallPromotionUtils.h

 #ifndef LLVM_TRANSFORMS_UTILS_CALLPROMOTIONUTILS_H
 #define LLVM_TRANSFORMS_UTILS_CALLPROMOTIONUTILS_H
 
 namespace llvm {
 class CallBase;
 class CastInst;
 class Function;
 class MDNode;
+class Value;
 
 /// Return true if the given indirect call site can be made to call \p Callee.
 ///
 /// This function ensures that the number and type of the call site's arguments
 /// and return value match those of the given function. If the types do not
 /// match exactly, they must at least be bitcast compatible. If \p FailureReason
 /// is non-null and the indirect call cannot be promoted, the failure reason
 /// will be stored in it.

 ///
 /// @_ZTV4Impl = linkonce_odr dso_local unnamed_addr constant { [3 x i8*] }
 ///     { [3 x i8*]
 ///     [i8* null, i8* bitcast ({ i8*, i8*, i8* }* @_ZTI4Impl to i8*),
 ///     i8* bitcast (void (%class.Impl*)* @_ZN4Impl3RunEv to i8*)] }
 ///
 bool tryPromoteCall(CallBase &CB);
 
+/// Predicate and clone the given call site.
+///
+/// This function creates an if-then-else structure at the location of the call
+/// site. The "if" condition compares the call site's called value to the given
+/// callee. The original call site is moved into the "else" block, and a clone
+/// of the call site is placed in the "then" block. The cloned instruction is
+/// returned.
+CallBase &versionCallSite(CallBase &CB, Value *Callee, MDNode *BranchWeights);
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_CALLPROMOTIONUTILS_H

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ModuleSummaryIndexYAML.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/PassRegistry.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/GlobPattern.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/FunctionAttrs.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/CallPromotionUtils.h"
 #include "llvm/Transforms/Utils/Evaluator.h"
 #include <algorithm>
 #include <cstddef>
 #include <map>
 #include <set>
 #include <string>
 
 using namespace llvm;
static cl::opt<bool> DisableWholeProgramVisibility(
     cl::desc("Disable whole program visibility (overrides enabling options)"));
 
 /// Provide way to prevent certain function from being devirtualized
 static cl::list<std::string>
     SkipFunctionNames("wholeprogramdevirt-skip",
                       cl::desc("Prevent function(s) from being devirtualized"),
                       cl::Hidden, cl::ZeroOrMore, cl::CommaSeparated);
 
-/// Mechanism to add runtime checking of devirtualization decisions, trapping on
-/// any that are not correct. Useful for debugging undefined behavior leading to
-/// failures with WPD.
-static cl::opt<bool>
-    CheckDevirt("wholeprogramdevirt-check", cl::init(false), cl::Hidden,
-                cl::ZeroOrMore,
-                cl::desc("Add code to trap on incorrect devirtualizations"));
+/// Mechanism to add runtime checking of devirtualization decisions, optionally
+/// trapping or falling back to indirect call on any that are not correct.
+/// Trapping mode is useful for debugging undefined behavior leading to failures
+/// with WPD. Fallback mode is useful for ensuring safety when whole program
+/// visibility may be compromised.
+enum WPDCheckMode { None, Trap, Fallback };
+static cl::opt<WPDCheckMode> DevirtCheckMode(
+    "wholeprogramdevirt-check", cl::Hidden, cl::ZeroOrMore,
+    cl::desc("Type of checking for incorrect devirtualizations"),
+    cl::values(clEnumValN(WPDCheckMode::None, "none", "No checking"),
+               clEnumValN(WPDCheckMode::Trap, "trap", "Trap when incorrect"),
+               clEnumValN(WPDCheckMode::Fallback, "fallback",
+                          "Fallback to indirect when incorrect")));
 
 namespace {
 struct PatternList {
   std::vector<GlobPattern> Patterns;
   template <class T> void init(const T &StringList) {
     for (const auto &S : StringList)
       if (Expected<GlobPattern> Pat = GlobPattern::create(S))
         Patterns.push_back(std::move(*Pat));
for (auto &&VCallSite : CSInfo.CallSites) {
         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());
 
-      // If checking is enabled, add support to compare the virtual function
-      // pointer to the devirtualized target. In case of a mismatch, perform a
-      // debug trap.
-      if (CheckDevirt) {
+      // If trap checking is enabled, add support to compare the virtual
+      // function pointer to the devirtualized target. In case of a mismatch,
+      // perform a debug trap.
+      if (DevirtCheckMode == WPDCheckMode::Trap) {
         auto *Cond = Builder.CreateICmpNE(CB.getCalledOperand(), Callee);
         Instruction *ThenTerm =
             SplitBlockAndInsertIfThen(Cond, &CB, /*Unreachable=*/false);
         Builder.SetInsertPoint(ThenTerm);
         Function *TrapFn = Intrinsic::getDeclaration(&M, Intrinsic::debugtrap);
         auto *CallTrap = Builder.CreateCall(TrapFn);
         CallTrap->setDebugLoc(CB.getDebugLoc());
       }
 
-      // Devirtualize.
+      // If fallback checking is enabled, add support to compare the virtual
+      // function pointer to the devirtualized target. In case of a mismatch,
+      // fall back to indirect call.
+      if (DevirtCheckMode == WPDCheckMode::Fallback) {
+        MDNode *Weights =
+            MDBuilder(M.getContext()).createBranchWeights((1U << 20) - 1, 1);
+        // Version the indirect call site. If the called value is equal to the
+        // given callee, 'NewInst' will be executed, otherwise the original call
+        // site will be executed.
+        CallBase &NewInst = versionCallSite(CB, Callee, Weights);
+        NewInst.setCalledOperand(Callee);
+        // Since the new call site is direct, we must clear metadata that
+        // is only appropriate for indirect calls. This includes !prof and
+        // !callees metadata.
+        NewInst.setMetadata(LLVMContext::MD_prof, nullptr);
+        NewInst.setMetadata(LLVMContext::MD_callees, nullptr);
+        // Additionally, we should remove them from the fallback indirect call,
+        // so that we don't attempt to perform indirect call promotion later.
+        CB.setMetadata(LLVMContext::MD_prof, nullptr);
+        CB.setMetadata(LLVMContext::MD_callees, nullptr);
+      }
+
+      // In either trapping or non-checking mode, devirtualize original call.
+      else {
+        // Devirtualize unconditionally.
         CB.setCalledOperand(Callee);
+        // Since the call site is now direct, we must clear metadata that
+        // is only appropriate for indirect calls. This includes !prof and
+        // !callees metadata.
+        CB.setMetadata(LLVMContext::MD_prof, nullptr);
+        CB.setMetadata(LLVMContext::MD_callees, nullptr);
+      }
 
       // This use is no longer unsafe.
       if (VCallSite.NumUnsafeUses)
         --*VCallSite.NumUnsafeUses;
     }
     if (CSInfo.isExported())
       IsExported = true;
     CSInfo.markDevirt();

llvm/lib/Transforms/Utils/CallPromotionUtils.cpp

 ///     ; block. It is not yet promoted.
 ///     %t1 = musttail call i32 %ptr()
 ///     ret %t1
 ///
 ///   orig_bb:
 ///     ; The original call instruction stays in its original block.
 ///     %t0 = musttail call i32 %ptr()
 ///     ret %t0
-static CallBase &versionCallSite(CallBase &CB, Value *Callee,
-                                 MDNode *BranchWeights) {
+CallBase &llvm::versionCallSite(CallBase &CB, Value *Callee,
 
   IRBuilder<> Builder(&CB);
   CallBase *OrigInst = &CB;
   BasicBlock *OrigBlock = OrigInst->getParent();
 
   // Create the compare. The called value and callee must have the same type to
   // be compared.
   if (CB.getCalledOperand()->getType() != Callee->getType())

llvm/test/ThinLTO/X86/devirt.ll

entry:
   %p = call i1 @llvm.type.test(i8* %1, metadata !"_ZTS1A")
   call void @llvm.assume(i1 %p)
   %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
 
   ; Check that the call was devirtualized.
   ; CHECK-IR: %call = tail call i32 @_ZN1A1nEi
-  %call = tail call i32 %fptr1(%struct.A* nonnull %obj, i32 %a)
+  ; Ensure !prof and !callees metadata for indirect call promotion removed.
+  ; CHECK-IR-NOT: prof
+  ; CHECK-IR-NOT: callees
+  %call = tail call i32 %fptr1(%struct.A* nonnull %obj, i32 %a), !prof !5, !callees !6
 
   %3 = bitcast i8** %vtable to i32 (%struct.A*, i32)**
   %fptr22 = load i32 (%struct.A*, i32)*, i32 (%struct.A*, i32)** %3, align 8
 
   ; We still have to call it as virtual.
   ; CHECK-IR: %call3 = tail call i32 %fptr22
   %call3 = tail call i32 %fptr22(%struct.A* nonnull %obj, i32 %call)
 

 ; Make sure we don't inline or otherwise optimize out the direct calls.
 attributes #0 = { noinline optnone }
 
 !0 = !{i64 16, !"_ZTS1A"}
 !1 = !{i64 16, !"_ZTS1B"}
 !2 = !{i64 16, !"_ZTS1C"}
 !3 = !{i64 16, !4}
 !4 = distinct !{}
+!5 = !{!"VP", i32 0, i64 1, i64 1621563287929432257, i64 1}
+!6 = !{i32 (%struct.A*, i32)* @_ZN1A1nEi}

llvm/test/ThinLTO/X86/devirt_check.ll

 ; REQUIRES: x86-registered-target
 
 ; Test that devirtualization option -wholeprogramdevirt-check adds code to check
-; that the devirtualization decision was correct and trap if not.
+; that the devirtualization decision was correct and trap or fallback if not.
 
 ; The vtables have vcall_visibility metadata with hidden visibility, to enable
 ; devirtualization.
 
 ; Generate unsplit module with summary for ThinLTO index-based WPD.
 ; RUN: opt -thinlto-bc -o %t2.o %s
+
+; Check first in trapping mode.
 ; RUN: llvm-lto2 run %t2.o -save-temps -use-new-pm -pass-remarks=. \
-; RUN:   -wholeprogramdevirt-check \
+; RUN:   -wholeprogramdevirt-check=trap \
 ; RUN:   -o %t3 \
 ; RUN:   -r=%t2.o,test,px \
 ; RUN:   -r=%t2.o,_ZN1A1nEi,p \
 ; RUN:   -r=%t2.o,_ZN1B1fEi,p \
 ; RUN:   -r=%t2.o,_ZTV1B,px 2>&1 | FileCheck %s --check-prefix=REMARK
-; RUN: llvm-dis %t3.1.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK-IR
+; RUN: llvm-dis %t3.1.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK --check-prefix=TRAP
+
+; Check next in fallback mode.
+; RUN: llvm-lto2 run %t2.o -save-temps -use-new-pm -pass-remarks=. \
+; RUN:   -wholeprogramdevirt-check=fallback \
+; RUN:   -o %t3 \
+; RUN:   -r=%t2.o,test,px \
+; RUN:   -r=%t2.o,_ZN1A1nEi,p \
+; RUN:   -r=%t2.o,_ZN1B1fEi,p \
+; RUN:   -r=%t2.o,_ZTV1B,px 2>&1 | FileCheck %s --check-prefix=REMARK
+; RUN: llvm-dis %t3.1.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK --check-prefix=FALLBACK
 
 ; REMARK-DAG: single-impl: devirtualized a call to _ZN1A1nEi
 
 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 }
 
 @_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, !vcall_visibility !5
 
 
-; CHECK-IR-LABEL: define i32 @test
+; CHECK-LABEL: define i32 @test
 define i32 @test(%struct.A* %obj, i32 %a) {
 entry:
   %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")
   call void @llvm.assume(i1 %p)
   %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
 
   ; Check that the call was devirtualized, but preceeded by a check guarding
   ; a trap if the function pointer doesn't match.
-  ; CHECK-IR:   %.not = icmp eq i32 (%struct.A*, i32)* %fptr1, @_ZN1A1nEi
-  ; CHECK-IR:   br i1 %.not, label %3, label %2
-  ; CHECK-IR: 2:
-  ; CHECK-IR:   tail call void @llvm.debugtrap()
-  ; CHECK-IR:   br label %3
-  ; CHECK-IR: 3:
-  ; CHECK-IR:   tail call i32 @_ZN1A1nEi
-  %call = tail call i32 %fptr1(%struct.A* nonnull %obj, i32 %a)
+  ; TRAP:   %.not = icmp eq i32 (%struct.A*, i32)* %fptr1, @_ZN1A1nEi
+  ; Ensure !prof and !callees metadata for indirect call promotion removed.
+  ; TRAP-NOT: prof
+  ; TRAP-NOT: callees
+  ; TRAP:   br i1 %.not, label %3, label %2
+  ; TRAP: 2:
+  ; TRAP:   tail call void @llvm.debugtrap()
+  ; TRAP:   br label %3
+  ; TRAP: 3:
+  ; TRAP:   tail call i32 @_ZN1A1nEi
+  ; Check that the call was devirtualized, but preceeded by a check guarding
+  ; a fallback if the function pointer doesn't match.
+  ; FALLBACK:   %2 = icmp eq i32 (%struct.A*, i32)* %fptr1, @_ZN1A1nEi
+  ; FALLBACK:   br i1 %2, label %if.true.direct_targ, label %if.false.orig_indirect
+  ; FALLBACK: if.true.direct_targ:
+  ; FALLBACK:   tail call i32 @_ZN1A1nEi
+  ; Ensure !prof and !callees metadata for indirect call promotion removed.
+  ; FALLBACK-NOT: prof
+  ; FALLBACK-NOT: callees
+  ; FALLBACK:   br label %if.end.icp
+  ; FALLBACK: if.false.orig_indirect:
+  ; FALLBACK:   tail call i32 %fptr1
+  ; Ensure !prof and !callees metadata for indirect call promotion removed.
+  ; In particular, if left on the fallback indirect call ICP may perform an
+  ; additional round of promotion.
+  ; FALLBACK-NOT: prof
+  ; FALLBACK-NOT: callees
+  ; FALLBACK:   br label %if.end.icp
+  %call = tail call i32 %fptr1(%struct.A* nonnull %obj, i32 %a), !prof !6, !callees !7
 
   ret i32 %call
 }
-; CHECK-IR-LABEL:   ret i32
-; CHECK-IR-LABEL: }
+; 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;
 }
 
 ; Make sure we don't inline or otherwise optimize out the direct calls.
 attributes #0 = { noinline optnone }
 
 !0 = !{i64 16, !"_ZTS1A"}
 !1 = !{i64 16, !"_ZTS1B"}
 !3 = !{i64 16, !4}
 !4 = distinct !{}
 !5 = !{i64 1}
+!6 = !{!"VP", i32 0, i64 1, i64 1621563287929432257, i64 1}
+!7 = !{i32 (%struct.A*, i32)* @_ZN1A1nEi}