Add a type_checked_load_relative to support relative function pointer tables

[aschwaighofer]

Summary

This adds a type_checked_load_relative intrinsic whose semantics it is to
load a relative function pointer.

A relative function pointer is a pointer to a 32bit value that when
added to its address yields the address of the function.

[kubamracek]

Can we also get a test with llvm.type.checked.load.relative where the call *doesn't* get devirtualized?

[aschwaighofer]

Include test case where the llvm.type.checked.load.relative intrinsic is
expanded to compute the pointer underlying the relative pointer value.

[aschwaighofer]

Update to recent ToT

[aschwaighofer]

clang-format

llvm/docs/LangRef.rst

1. The function pointer that would have been loaded from an arbitrarily chosen
      metadata.
 
   2. If the function has a non-void return type, a pointer to a function that
      returns an unspecified value without causing side effects.
 
 If the function's return value's second element is false, the value of the
 first element is undefined.
 
+.. _type.checked.load.relative:
+
+'``llvm.type.checked.load.relative``' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+::
+
+      declare {ptr, i1} @llvm.type.checked.load.relative(ptr %ptr, i32 %offset, metadata %type) argmemonly nounwind readonly
+
+Overview:
+"""""""""
+
+The ``llvm.type.checked.load.relative`` intrinsic loads a relative pointer to a
+function from a virtual table pointer using metadata. Otherwise, its semantic is
+identical to the ``llvm.type.checked.load`` intrinsic.
+
+A relative pointer is a pointer to an offset to the pointed to value. The
+address of the underlying pointer of the relative pointer is obtained by adding
+the offset to the address of the offset value.
 
 '``llvm.arithmetic.fence``' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 Syntax:
 """""""
 
 ::

llvm/include/llvm/IR/Intrinsics.td

 def int_type_test : DefaultAttrsIntrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_metadata_ty],
                               [IntrNoMem, IntrWillReturn, IntrSpeculatable]>;
 
 // Safely loads a function pointer from a virtual table pointer using type metadata.
 def int_type_checked_load : DefaultAttrsIntrinsic<[llvm_ptr_ty, llvm_i1_ty],
                                       [llvm_ptr_ty, llvm_i32_ty, llvm_metadata_ty],
                                       [IntrNoMem, IntrWillReturn]>;
 
+// Safely loads a relative function pointer from a virtual table pointer using type metadata.
+def int_type_checked_load_relative : DefaultAttrsIntrinsic<[llvm_ptr_ty, llvm_i1_ty],
+                                      [llvm_ptr_ty, llvm_i32_ty, llvm_metadata_ty],
+                                      [IntrNoMem, IntrWillReturn]>;
+
 // Test whether a pointer is associated with a type metadata identifier. Used
 // for public visibility classes that may later be refined to private
 // visibility.
 def int_public_type_test : DefaultAttrsIntrinsic<[llvm_i1_ty], [llvm_ptr_ty, llvm_metadata_ty],
                               [IntrNoMem, IntrWillReturn, IntrSpeculatable]>;
 
 // Create a branch funnel that implements an indirect call to a limited set of
 // callees. This needs to be a musttail call.

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

case Intrinsic::public_type_test: {
     findDevirtualizableCallsForTypeTest(DevirtCalls, Assumes, CI, DT);
     for (auto &Call : DevirtCalls)
       addVCallToSet(Call, Guid, TypeTestAssumeVCalls,
                     TypeTestAssumeConstVCalls);
 
     break;
   }
 
+  case Intrinsic::type_checked_load_relative:
   case Intrinsic::type_checked_load: {
     auto *TypeMDVal = cast<MetadataAsValue>(CI->getArgOperand(2));
     auto *TypeId = dyn_cast<MDString>(TypeMDVal->getMetadata());
     if (!TypeId)
       break;
     GlobalValue::GUID Guid = GlobalValue::getGUID(TypeId->getString());
 
     SmallVector<DevirtCallSite, 4> DevirtCalls;

llvm/lib/Analysis/TypeMetadataUtils.cpp

 }
 
 void llvm::findDevirtualizableCallsForTypeCheckedLoad(
     SmallVectorImpl<DevirtCallSite> &DevirtCalls,
     SmallVectorImpl<Instruction *> &LoadedPtrs,
     SmallVectorImpl<Instruction *> &Preds, bool &HasNonCallUses,
     const CallInst *CI, DominatorTree &DT) {
   assert(CI->getCalledFunction()->getIntrinsicID() ==
-         Intrinsic::type_checked_load);
+             Intrinsic::type_checked_load ||
+         CI->getCalledFunction()->getIntrinsicID() ==
+             Intrinsic::type_checked_load_relative);
 
   auto *Offset = dyn_cast<ConstantInt>(CI->getArgOperand(1));
   if (!Offset) {
     HasNonCallUses = true;
     return;
   }
 
   for (const Use &U : CI->uses()) {

llvm/lib/LTO/LTO.cpp

 Error LTO::checkPartiallySplit() {
   if (!ThinLTO.CombinedIndex.partiallySplitLTOUnits())
     return Error::success();
 
   Function *TypeTestFunc = RegularLTO.CombinedModule->getFunction(
       Intrinsic::getName(Intrinsic::type_test));
   Function *TypeCheckedLoadFunc = RegularLTO.CombinedModule->getFunction(
       Intrinsic::getName(Intrinsic::type_checked_load));
+  Function *TypeCheckedLoadRelativeFunc =
+      RegularLTO.CombinedModule->getFunction(
+          Intrinsic::getName(Intrinsic::type_checked_load_relative));
 
   // First check if there are type tests / type checked loads in the
   // merged regular LTO module IR.
   if ((TypeTestFunc && !TypeTestFunc->use_empty()) ||
-      (TypeCheckedLoadFunc && !TypeCheckedLoadFunc->use_empty()))
+      (TypeCheckedLoadFunc && !TypeCheckedLoadFunc->use_empty()) ||
+      (TypeCheckedLoadRelativeFunc &&
+       !TypeCheckedLoadRelativeFunc->use_empty()))
     return make_error<StringError>(
         "inconsistent LTO Unit splitting (recompile with -fsplit-lto-unit)",
         inconvertibleErrorCode());
 
   // Otherwise check if there are any recorded in the combined summary from the
   // ThinLTO modules.
   for (auto &P : ThinLTO.CombinedIndex) {
     for (auto &S : P.second.SummaryList) {

llvm/lib/Transforms/IPO/GlobalDCE.cpp

for (const auto &VTableInfo : TypeIdMap[TypeId]) {
     GVDependencies[Caller].insert(Callee);
   }
 }
 
 void GlobalDCEPass::ScanTypeCheckedLoadIntrinsics(Module &M) {
   LLVM_DEBUG(dbgs() << "Scanning type.checked.load intrinsics\n");
   Function *TypeCheckedLoadFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load));
+  Function *TypeCheckedLoadRelativeFunc =
+      M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load_relative));
 
-  if (!TypeCheckedLoadFunc)
+  auto scan = [&](Function *CheckedLoadFunc) {
+    if (!CheckedLoadFunc)
       return;
 
-  for (auto *U : TypeCheckedLoadFunc->users()) {
+    for (auto *U : CheckedLoadFunc->users()) {
       auto CI = dyn_cast<CallInst>(U);
       if (!CI)
         continue;
 
       auto *Offset = dyn_cast<ConstantInt>(CI->getArgOperand(1));
       Value *TypeIdValue = CI->getArgOperand(2);
       auto *TypeId = cast<MetadataAsValue>(TypeIdValue)->getMetadata();
 
       if (Offset) {
         ScanVTableLoad(CI->getFunction(), TypeId, Offset->getZExtValue());
       } else {
-      // type.checked.load with a non-constant offset, so assume every entry in
-      // every matching vtable is used.
+        // type.checked.load with a non-constant offset, so assume every entry
+        // in every matching vtable is used.
         for (const auto &VTableInfo : TypeIdMap[TypeId]) {
           VFESafeVTables.erase(VTableInfo.first);
         }
       }
     }
+  };
+
+  scan(TypeCheckedLoadFunc);
+  scan(TypeCheckedLoadRelativeFunc);
 }
 
 void GlobalDCEPass::AddVirtualFunctionDependencies(Module &M) {
   if (!ClEnableVFE)
     return;
 
   // If the Virtual Function Elim module flag is present and set to zero, then
   // the vcall_visibility metadata was inserted for another optimization (WPD)

llvm/lib/Transforms/IPO/GlobalSplit.cpp

 static bool splitGlobals(Module &M) {
   // First, see if the module uses either of the llvm.type.test or
   // llvm.type.checked.load intrinsics, which indicates that splitting globals
   // may be beneficial.
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
   Function *TypeCheckedLoadFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load));
+  Function *TypeCheckedLoadRelativeFunc =
+      M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load_relative));
   if ((!TypeTestFunc || TypeTestFunc->use_empty()) &&
-      (!TypeCheckedLoadFunc || TypeCheckedLoadFunc->use_empty()))
+      (!TypeCheckedLoadFunc || TypeCheckedLoadFunc->use_empty()) &&
+      (!TypeCheckedLoadRelativeFunc ||
+       TypeCheckedLoadRelativeFunc->use_empty()))
     return false;
 
   bool Changed = false;
   for (GlobalVariable &GV : llvm::make_early_inc_range(M.globals()))
     Changed |= splitGlobal(GV);
   return Changed;
 }
 
 PreservedAnalyses GlobalSplitPass::run(Module &M, ModuleAnalysisManager &AM) {
   if (!splitGlobals(M))
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }

llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

void promoteTypeIds(Module &M, StringRef ModuleId) {
   if (Function *TypeCheckedLoadFunc =
           M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load))) {
     for (const Use &U : TypeCheckedLoadFunc->uses()) {
       auto CI = cast<CallInst>(U.getUser());
       ExternalizeTypeId(CI, 2);
     }
   }
 
+  if (Function *TypeCheckedLoadRelativeFunc = M.getFunction(
+          Intrinsic::getName(Intrinsic::type_checked_load_relative))) {
+    for (const Use &U : TypeCheckedLoadRelativeFunc->uses()) {
+      auto CI = cast<CallInst>(U.getUser());
+      ExternalizeTypeId(CI, 2);
+    }
+  }
+
   for (GlobalObject &GO : M.global_objects()) {
     SmallVector<MDNode *, 1> MDs;
     GO.getMetadata(LLVMContext::MD_type, MDs);
 
     GO.eraseMetadata(LLVMContext::MD_type);
     for (auto *MD : MDs) {
       auto I = LocalToGlobal.find(MD->getOperand(1));
       if (I == LocalToGlobal.end()) {

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

for (const TypeMemberInfo &TM : TypeMemberInfos) {
 
     // We cannot perform whole program devirtualization analysis on a vtable
     // with public LTO visibility.
     if (TM.Bits->GV->getVCallVisibility() ==
         GlobalObject::VCallVisibilityPublic)
       return false;
 
     Constant *Ptr = getPointerAtOffset(TM.Bits->GV->getInitializer(),
-                                       TM.Offset + ByteOffset, M);
+                                       TM.Offset + ByteOffset, M, TM.Bits->GV);
     if (!Ptr)
       return false;
 
     auto C = Ptr->stripPointerCasts();
     // Make sure this is a function or alias to a function.
     auto Fn = dyn_cast<Function>(C);
     auto A = dyn_cast<GlobalAlias>(C);
     if (!Fn && A)
for (Use &U : llvm::make_early_inc_range(TypeCheckedLoadFunc->uses())) {
     // Start by generating "pessimistic" code that explicitly loads the function
     // pointer from the vtable and performs the type check. If possible, we will
     // eliminate the load and the type check later.
 
     // If possible, only generate the load at the point where it is used.
     // This helps avoid unnecessary spills.
     IRBuilder<> LoadB(
         (LoadedPtrs.size() == 1 && !HasNonCallUses) ? LoadedPtrs[0] : CI);
+
+    Value *LoadedValue = nullptr;
+    if (TypeCheckedLoadFunc->getIntrinsicID() ==
+        Intrinsic::type_checked_load_relative) {
+      Value *GEP = LoadB.CreateGEP(Int8Ty, Ptr, Offset);
+      Value *GEPPtr = LoadB.CreateBitCast(GEP, PointerType::getUnqual(Int32Ty));
+      LoadedValue = LoadB.CreateLoad(Int32Ty, GEPPtr);
+      LoadedValue = LoadB.CreateSExt(LoadedValue, IntPtrTy);
+      GEP = LoadB.CreatePtrToInt(GEP, IntPtrTy);
+      LoadedValue = LoadB.CreateAdd(GEP, LoadedValue);
+      LoadedValue = LoadB.CreateIntToPtr(LoadedValue, Int8PtrTy);
+    } else {
       Value *GEP = LoadB.CreateGEP(Int8Ty, Ptr, Offset);
-    Value *GEPPtr = LoadB.CreateBitCast(GEP, PointerType::getUnqual(Int8PtrTy));
-    Value *LoadedValue = LoadB.CreateLoad(Int8PtrTy, GEPPtr);
+      Value *GEPPtr =
+          LoadB.CreateBitCast(GEP, PointerType::getUnqual(Int8PtrTy));
+      LoadedValue = LoadB.CreateLoad(Int8PtrTy, GEPPtr);
+    }
 
     for (Instruction *LoadedPtr : LoadedPtrs) {
       LoadedPtr->replaceAllUsesWith(LoadedValue);
       LoadedPtr->eraseFromParent();
     }
 
     // Likewise for the type test.
     IRBuilder<> CallB((Preds.size() == 1 && !HasNonCallUses) ? Preds[0] : CI);
bool DevirtModule::run() {
   if ((ExportSummary && ExportSummary->partiallySplitLTOUnits()) ||
       (ImportSummary && ImportSummary->partiallySplitLTOUnits()))
     return false;
 
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));
   Function *TypeCheckedLoadFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load));
+  Function *TypeCheckedLoadRelativeFunc =
+      M.getFunction(Intrinsic::getName(Intrinsic::type_checked_load_relative));
   Function *AssumeFunc = M.getFunction(Intrinsic::getName(Intrinsic::assume));
 
   // Normally if there are no users of the devirtualization intrinsics in the
   // module, this pass has nothing to do. But if we are exporting, we also need
   // to handle any users that appear only in the function summaries.
   if (!ExportSummary &&
       (!TypeTestFunc || TypeTestFunc->use_empty() || !AssumeFunc ||
        AssumeFunc->use_empty()) &&
-      (!TypeCheckedLoadFunc || TypeCheckedLoadFunc->use_empty()))
+      (!TypeCheckedLoadFunc || TypeCheckedLoadFunc->use_empty()) &&
+      (!TypeCheckedLoadRelativeFunc ||
+       TypeCheckedLoadRelativeFunc->use_empty()))
     return false;
 
   // Rebuild type metadata into a map for easy lookup.
   std::vector<VTableBits> Bits;
   DenseMap<Metadata *, std::set<TypeMemberInfo>> TypeIdMap;
   buildTypeIdentifierMap(Bits, TypeIdMap);
 
   if (TypeTestFunc && AssumeFunc)
     scanTypeTestUsers(TypeTestFunc, TypeIdMap);
 
   if (TypeCheckedLoadFunc)
     scanTypeCheckedLoadUsers(TypeCheckedLoadFunc);
 
+  if (TypeCheckedLoadRelativeFunc)
+    scanTypeCheckedLoadUsers(TypeCheckedLoadRelativeFunc);
+
   if (ImportSummary) {
     for (auto &S : CallSlots)
       importResolution(S.first, S.second);
 
     removeRedundantTypeTests();
 
     // We have lowered or deleted the type intrinsics, so we will no longer have
     // enough information to reason about the liveness of virtual function

llvm/test/Transforms/WholeProgramDevirt/devirt-single-impl-check-relative.ll

+; RUN: opt -S -passes=wholeprogramdevirt -whole-program-visibility -pass-remarks=wholeprogramdevirt %s 2>&1 | FileCheck %s
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK: remark: <unknown>:0:0: single-impl: devirtualized a call to vf
+; CHECK: remark: <unknown>:0:0: devirtualized vf
+; CHECK-NOT: devirtualized
+
+; A vtable with "relative pointers", slots don't contain pointers to implementations, but instead have an i32 offset from the vtable itself to the implementation.
+@vtable = internal unnamed_addr constant { [2 x i32] } { [2 x i32] [
+  i32 trunc (i64 sub (i64 ptrtoint (ptr @vfunc1_live to i64), i64 ptrtoint (ptr @vtable to i64)) to i32),
+  i32 trunc (i64 sub (i64 ptrtoint (ptr @vfunc2_dead to i64), i64 ptrtoint (ptr @vtable to i64)) to i32)
+]}, align 8, !type !0, !type !1
+;, !vcall_visibility !{i64 2}
+!0 = !{i64 0, !"vfunc1.type"}
+!1 = !{i64 4, !"vfunc2.type"}
+
+define internal void @vfunc1_live() {
+  ret void
+}
+
+define internal void @vfunc2_dead() {
+  ret void
+}
+
+; CHECK: define void @call
+define void @call(ptr %obj) {
+  %vtable = load ptr, ptr %obj
+  %pair = call {ptr, i1} @llvm.type.checked.load.relative(ptr %vtable, i32 0, metadata !"vfunc1.type")
+  %fptr = extractvalue {ptr, i1} %pair, 0
+  %p = extractvalue {ptr, i1} %pair, 1
+  ; CHECK: br i1 true,
+  br i1 %p, label %cont, label %trap
+
+cont:
+  ; CHECK: call void @vfunc1_live(
+  call void %fptr()
+  ret void
+
+trap:
+  call void @llvm.trap()
+  unreachable
+}
+
+declare {ptr, i1} @llvm.type.checked.load.relative(ptr, i32, metadata)
+declare void @llvm.trap()

llvm/test/Transforms/WholeProgramDevirt/expand-check-relative.ll

+; RUN: opt -S -passes=wholeprogramdevirt %s | FileCheck %s
+
+; Test that we correctly expand the llvm.type.checked.load.relative intrinsic in
+; cases where we cannot devirtualize.
+
+target datalayout = "e-p:64:64"
+target triple = "x86_64-unknown-linux-gnu"
+
+@vt1 = constant { [2 x i32] } { [2 x i32] [
+  i32 trunc (i64 sub (i64 ptrtoint (ptr @vf1 to i64), i64 ptrtoint (ptr @vt1 to i64)) to i32),
+  i32 trunc (i64 sub (i64 ptrtoint (ptr @vf2 to i64), i64 ptrtoint (ptr @vt1 to i64)) to i32)
+]}, align 8, !type !0, !type !1
+
+!0 = !{i64 0, !"vfunc1.type"}
+!1 = !{i64 4, !"vfunc2.type"}
+
+
+define void @vf1(ptr %this) {
+  ret void
+}
+
+define void @vf2(ptr %this) {
+  ret void
+}
+
+; CHECK: define void @call
+; CHECK:  [[TT:%.*]] = call i1 @llvm.type.test(ptr [[VT:%.*]], metadata !"vfunc1.type")
+; CHECK:  br i1 [[TT]]
+
+; Relative pointer computation at the address of the i32 value to the i32 value
+; to get to the pointer value.
+
+; CHECK:  [[T0:%.*]] = getelementptr i8, ptr [[VT]], i32 0
+; CHECK:  [[T1:%.*]] = load i32, ptr [[T0]]
+; CHECK:  [[T2:%.*]] = sext i32 [[T1]] to i64
+; CHECK:  [[T3:%.*]] = ptrtoint ptr [[T0]] to i64
+; CHECK:  [[T4:%.*]] = add i64 [[T3]], [[T2]]
+; CHECK:  [[F:%.*]] = inttoptr i64 [[T4]] to ptr
+; CHECK:  call void [[F]](ptr
+
+define void @call(ptr %obj) {
+  %vtable = load ptr, ptr %obj
+  %pair = call {ptr, i1} @llvm.type.checked.load.relative(ptr %vtable, i32 0, metadata !"vfunc1.type")
+  %p = extractvalue {ptr, i1} %pair, 1
+  br i1 %p, label %cont, label %trap
+
+cont:
+  %fptr = extractvalue {ptr, i1} %pair, 0
+  call void %fptr(ptr %obj)
+  ret void
+
+trap:
+  call void @llvm.trap()
+  unreachable
+}
+
+declare {ptr, i1} @llvm.type.checked.load.relative(ptr, i32, metadata)
+declare void @llvm.trap()