[IR][GlobalDCE] Extend !vcall_visibility to be able to specify an offset range that it applies to

kubamracek · kubamracek · commit d4a606be3e2b · 2021-10-24T07:20:22.000-07:00
This PR adds support to GlobalDCE for vtables containing references that are not
virtual functions and should not participate in VFE. This is achieved by
extending !vcall_visibility to be able to specify an offset range that it
applies to.
diff --git a/llvm/docs/TypeMetadata.rst b/llvm/docs/TypeMetadata.rst
@@ -288,3 +288,11 @@ In addition, all function pointer loads from a vtable marked with the
 calls sites can be correlated with the vtables which they might load from.
 Other parts of the vtable (RTTI, offset-to-top, ...) can still be accessed with
 normal loads.
+
+Alternatively, the ``!vcall_visibility`` metadata attachment can have an
+extended format of a tuple with two additional integer values representing the
+begin and end offset (a half-open range, the begin offset is included, the end
+offset is excluded) within the vtable that the visibility applies to. When the
+range is missing, the meaning is the same as a range covering the entire vtable.
+Any part of the vtable that is not covered by the specified range is not
+eligible for elimination of virtual functions.
diff --git a/llvm/include/llvm/IR/GlobalObject.h b/llvm/include/llvm/IR/GlobalObject.h
@@ -138,6 +138,7 @@ class GlobalObject : public GlobalValue {
   void addTypeMetadata(unsigned Offset, Metadata *TypeID);
   void setVCallVisibilityMetadata(VCallVisibility Visibility);
   VCallVisibility getVCallVisibility() const;
+  std::pair<uint64_t, uint64_t> getVTableOffsetRange() const;
 
   /// Returns true if the alignment of the value can be unilaterally
   /// increased.
diff --git a/llvm/include/llvm/Transforms/IPO/GlobalDCE.h b/llvm/include/llvm/Transforms/IPO/GlobalDCE.h
@@ -47,9 +47,11 @@ class GlobalDCEPass : public PassInfoMixin<GlobalDCEPass> {
   DenseMap<Metadata *, SmallSet<std::pair<GlobalVariable *, uint64_t>, 4>>
       TypeIdMap;
 
-  // Global variables which are vtables, and which we have enough information
-  // about to safely do dead virtual function elimination.
-  SmallPtrSet<GlobalValue *, 32> VFESafeVTables;
+  /// VTable -> set of vfuncs. This only contains vtables for which we have
+  /// enough information to safely do dead virtual function elimination, and
+  /// only contains vfuncs that are within the range specified in
+  /// !vcall_visibility).
+  DenseMap<GlobalValue *, SmallPtrSet<GlobalValue *, 8>> VFESafeVTablesAndFns;
 
   void UpdateGVDependencies(GlobalValue &GV);
   void MarkLive(GlobalValue &GV,
diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
@@ -1525,6 +1525,23 @@ GlobalObject::VCallVisibility GlobalObject::getVCallVisibility() const {
   return VCallVisibility::VCallVisibilityPublic;
 }
 
+std::pair<uint64_t, uint64_t> GlobalObject::getVTableOffsetRange() const {
+  if (MDNode *MD = getMetadata(LLVMContext::MD_vcall_visibility)) {
+    if (MD->getNumOperands() >= 3) {
+      uint64_t RangeStart =
+          cast<ConstantInt>(
+              cast<ConstantAsMetadata>(MD->getOperand(1))->getValue())
+              ->getZExtValue();
+      uint64_t RangeEnd =
+          cast<ConstantInt>(
+              cast<ConstantAsMetadata>(MD->getOperand(2))->getValue())
+              ->getZExtValue();
+      return std::make_pair(RangeStart, RangeEnd);
+    }
+  }
+  return std::make_pair(0, std::numeric_limits<uint64_t>::max());
+}
+
 void Function::setSubprogram(DISubprogram *SP) {
   setMetadata(LLVMContext::MD_dbg, SP);
 }
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
@@ -749,6 +749,40 @@ void Verifier::visitGlobalVariable(const GlobalVariable &GV) {
                       "DIGlobalVariableExpression");
   }
 
+  MDs.clear();
+  GV.getMetadata(LLVMContext::MD_vcall_visibility, MDs);
+  for (auto *MD : MDs) {
+    Assert(MD->getNumOperands() >= 1, "bad !vcall_visibility attachment");
+    Assert(isa<ConstantAsMetadata>(MD->getOperand(0)),
+           "bad !vcall_visibility attachment");
+    auto *Op0Val = cast<ConstantAsMetadata>(MD->getOperand(0))->getValue();
+    Assert(isa<ConstantInt>(Op0Val), "bad !vcall_visibility attachment");
+    auto Op0Int = cast<ConstantInt>(Op0Val)->getValue();
+    Assert(Op0Int.uge(0) && Op0Int.ult(std::numeric_limits<uint64_t>::max()),
+           "bad !vcall_visibility attachment");
+    if (MD->getNumOperands() == 3) {
+      Assert(isa<ConstantAsMetadata>(MD->getOperand(1)),
+             "bad !vcall_visibility attachment");
+      auto *Op1Val = cast<ConstantAsMetadata>(MD->getOperand(1))->getValue();
+      Assert(isa<ConstantInt>(Op1Val), "bad !vcall_visibility attachment");
+      auto Op1Int = cast<ConstantInt>(Op1Val)->getValue();
+      Assert(Op1Int.uge(0) && Op1Int.ult(std::numeric_limits<uint64_t>::max()),
+             "bad !vcall_visibility attachment");
+
+      Assert(isa<ConstantAsMetadata>(MD->getOperand(2)),
+             "bad !vcall_visibility attachment");
+      auto *Op2Val = cast<ConstantAsMetadata>(MD->getOperand(2))->getValue();
+      Assert(isa<ConstantInt>(Op2Val), "bad !vcall_visibility attachment");
+      auto Op2Int = cast<ConstantInt>(Op2Val)->getValue();
+      Assert(Op2Int.uge(0) && Op2Int.ult(std::numeric_limits<uint64_t>::max()),
+             "bad !vcall_visibility attachment");
+
+      Assert(Op1Int.ule(Op2Int), "bad !vcall_visibility attachment");
+    } else {
+      Assert(MD->getNumOperands() == 1, "bad !vcall_visibility attachment");
+    }
+  }
+
   // Scalable vectors cannot be global variables, since we don't know
   // the runtime size. If the global is an array containing scalable vectors,
   // that will be caught by the isValidElementType methods in StructType or
diff --git a/llvm/lib/Transforms/IPO/GlobalDCE.cpp b/llvm/lib/Transforms/IPO/GlobalDCE.cpp
@@ -127,11 +127,12 @@ void GlobalDCEPass::UpdateGVDependencies(GlobalValue &GV) {
     ComputeDependencies(User, Deps);
   Deps.erase(&GV); // Remove self-reference.
   for (GlobalValue *GVU : Deps) {
-    // If this is a dep from a vtable to a virtual function, and we have
-    // complete information about all virtual call sites which could call
-    // though this vtable, then skip it, because the call site information will
-    // be more precise.
-    if (VFESafeVTables.count(GVU) && isa<Function>(&GV)) {
+    // If this is a dep from a vtable to a virtual function, and it's within the
+    // range specified in !vcall_visibility, and we have complete information
+    // about all virtual call sites which could call though this vtable, then
+    // skip it, because the call site information will be more precise.
+    if (isa<Function>(&GV) && VFESafeVTablesAndFns.count(GVU) &&
+        VFESafeVTablesAndFns[GVU].contains(&GV)) {
       LLVM_DEBUG(dbgs() << "Ignoring dep " << GVU->getName() << " -> "
                         << GV.getName() << "\n");
       continue;
@@ -157,6 +158,47 @@ void GlobalDCEPass::MarkLive(GlobalValue &GV,
   }
 }
 
+/// Recursively iterate over the (sub-)constants in the vtable and look for
+/// vptrs, if their offset is within [RangeStart..RangeEnd), add them to VFuncs.
+static void FindVirtualFunctionsInVTable(Module &M, Constant *C,
+                                         uint64_t RangeStart, uint64_t RangeEnd,
+                                         SmallPtrSet<GlobalValue *, 8> *VFuncs,
+                                         uint64_t BaseOffset = 0) {
+  if (auto *GV = dyn_cast<GlobalValue>(C)) {
+    if (auto *F = dyn_cast<Function>(GV))
+      if (RangeStart <= BaseOffset && BaseOffset < RangeEnd)
+        VFuncs->insert(F);
+
+    // Do not recurse outside of the current global.
+    return;
+  }
+
+  if (auto *S = dyn_cast<ConstantStruct>(C)) {
+    StructType *STy = dyn_cast<StructType>(S->getType());
+    const StructLayout *SL = M.getDataLayout().getStructLayout(STy);
+    for (auto EI : llvm::enumerate(STy->elements())) {
+      auto Offset = SL->getElementOffset(EI.index());
+      unsigned Op = SL->getElementContainingOffset(Offset);
+      FindVirtualFunctionsInVTable(M, cast<Constant>(S->getOperand(Op)),
+                                   RangeStart, RangeEnd, VFuncs,
+                                   BaseOffset + Offset);
+    }
+  } else if (auto *A = dyn_cast<ConstantArray>(C)) {
+    ArrayType *ATy = A->getType();
+    auto EltSize = M.getDataLayout().getTypeAllocSize(ATy->getElementType());
+    for (unsigned i = 0, e = ATy->getNumElements(); i != e; ++i) {
+      FindVirtualFunctionsInVTable(M, cast<Constant>(A->getOperand(i)),
+                                   RangeStart, RangeEnd, VFuncs,
+                                   BaseOffset + EltSize * i);
+    }
+  } else {
+    for (auto &Op : C->operands()) {
+      FindVirtualFunctionsInVTable(M, cast<Constant>(Op), RangeStart, RangeEnd,
+                                   VFuncs, BaseOffset);
+    }
+  }
+}
+
 void GlobalDCEPass::ScanVTables(Module &M) {
   SmallVector<MDNode *, 2> Types;
   LLVM_DEBUG(dbgs() << "Building type info -> vtable map\n");
@@ -196,7 +238,14 @@ void GlobalDCEPass::ScanVTables(Module &M) {
           (LTOPostLink &&
            TypeVis == GlobalObject::VCallVisibilityLinkageUnit)) {
         LLVM_DEBUG(dbgs() << GV.getName() << " is safe for VFE\n");
-        VFESafeVTables.insert(&GV);
+
+        // Find and record all the vfunctions that are within the offset range
+        // specified in the !vcall_visibility attribute.
+        auto Range = GO->getVTableOffsetRange();
+        SmallPtrSet<GlobalValue *, 8> VFuncs;
+        FindVirtualFunctionsInVTable(M, GV.getInitializer(), std::get<0>(Range),
+                                     std::get<1>(Range), &VFuncs);
+        VFESafeVTablesAndFns[&GV] = VFuncs;
       }
     }
   }
@@ -213,14 +262,14 @@ void GlobalDCEPass::ScanVTableLoad(Function *Caller, Metadata *TypeId,
                            *Caller->getParent(), VTable);
     if (!Ptr) {
       LLVM_DEBUG(dbgs() << "can't find pointer in vtable!\n");
-      VFESafeVTables.erase(VTable);
+      VFESafeVTablesAndFns.erase(VTable);
       return;
     }
 
     auto Callee = dyn_cast<Function>(Ptr->stripPointerCasts());
     if (!Callee) {
       LLVM_DEBUG(dbgs() << "vtable entry is not function pointer!\n");
-      VFESafeVTables.erase(VTable);
+      VFESafeVTablesAndFns.erase(VTable);
       return;
     }
 
@@ -253,7 +302,7 @@ void GlobalDCEPass::ScanTypeCheckedLoadIntrinsics(Module &M) {
       // type.checked.load with a non-constant offset, so assume every entry in
       // every matching vtable is used.
       for (auto &VTableInfo : TypeIdMap[TypeId]) {
-        VFESafeVTables.erase(VTableInfo.first);
+        VFESafeVTablesAndFns.erase(VTableInfo.first);
       }
     }
   }
@@ -274,16 +323,15 @@ void GlobalDCEPass::AddVirtualFunctionDependencies(Module &M) {
 
   ScanVTables(M);
 
-  if (VFESafeVTables.empty())
+  if (VFESafeVTablesAndFns.empty())
     return;
 
   ScanTypeCheckedLoadIntrinsics(M);
 
-  LLVM_DEBUG(
-    dbgs() << "VFE safe vtables:\n";
-    for (auto *VTable : VFESafeVTables)
-      dbgs() << "  " << VTable->getName() << "\n";
-  );
+  LLVM_DEBUG(dbgs() << "VFE safe vtables:\n";
+             for (auto &Entry
+                  : VFESafeVTablesAndFns) dbgs()
+             << "  " << Entry.first->getName() << "\n";);
 }
 
 PreservedAnalyses GlobalDCEPass::run(Module &M, ModuleAnalysisManager &MAM) {
@@ -449,7 +497,7 @@ PreservedAnalyses GlobalDCEPass::run(Module &M, ModuleAnalysisManager &MAM) {
   GVDependencies.clear();
   ComdatMembers.clear();
   TypeIdMap.clear();
-  VFESafeVTables.clear();
+  VFESafeVTablesAndFns.clear();
 
   if (Changed)
     return PreservedAnalyses::none();
diff --git a/llvm/test/Transforms/GlobalDCE/virtual-functions-non-vfunc-entries.ll b/llvm/test/Transforms/GlobalDCE/virtual-functions-non-vfunc-entries.ll
@@ -0,0 +1,116 @@
+; RUN: opt < %s -globaldce -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+declare { i8*, i1 } @llvm.type.checked.load(i8*, i32, metadata)
+
+; A vtable that contains a non-nfunc entry, @regular_non_virtual_funcA, but
+; without a range specific in !vcall_visibility, which means *all* function
+; pointers are eligible for VFE, so GlobalDCE will treat the
+; @regular_non_virtual_funcA slot as eligible for VFE, and remove it.
+@vtableA = internal unnamed_addr constant { [3 x i8*] } { [3 x i8*] [
+  i8* bitcast (void ()* @vfunc1_live to i8*),
+  i8* bitcast (void ()* @vfunc2_dead to i8*),
+  i8* bitcast (void ()* @regular_non_virtual_funcA to i8*)
+]}, align 8, !type !{i64 0, !"vfunc1.type"}, !type !{i64 8, !"vfunc2.type"}, !vcall_visibility !{i64 2}
+
+; CHECK:      @vtableA = internal unnamed_addr constant { [3 x i8*] } { [3 x i8*] [
+; CHECK-SAME:   i8* bitcast (void ()* @vfunc1_live to i8*),
+; CHECK-SAME:   i8* null,
+; CHECK-SAME:   i8* null
+; CHECK-SAME: ] }, align 8
+
+
+; A vtable that contains a non-nfunc entry, @regular_non_virtual_funcB, with a
+; range of [0,16) which means only the first two entries are eligible for VFE.
+; GlobalDCE should keep @regular_non_virtual_funcB in the vtable.
+@vtableB = internal unnamed_addr constant { [3 x i8*] } { [3 x i8*] [
+  i8* bitcast (void ()* @vfunc1_live to i8*),
+  i8* bitcast (void ()* @vfunc2_dead to i8*),
+  i8* bitcast (void ()* @regular_non_virtual_funcB to i8*)
+]}, align 8, !type !{i64 0, !"vfunc1.type"}, !type !{i64 8, !"vfunc2.type"}, !vcall_visibility !{i64 2, i64 0, i64 16}
+
+; CHECK:      @vtableB = internal unnamed_addr constant { [3 x i8*] } { [3 x i8*] [
+; CHECK-SAME:   i8* bitcast (void ()* @vfunc1_live to i8*),
+; CHECK-SAME:   i8* null,
+; CHECK-SAME:   i8* bitcast (void ()* @regular_non_virtual_funcB to i8*)
+; CHECK-SAME: ] }, align 8
+
+; A vtable that contains a non-nfunc entry, @regular_non_virtual_funcB, with a
+; range of [0,16) which means only the first two entries are eligible for VFE.
+; GlobalDCE should keep @regular_non_virtual_funcB in the vtable.
+@vtableC = internal unnamed_addr constant { [3 x i8*] } { [3 x i8*] [
+  i8* bitcast (void ()* @regular_non_virtual_funcC to i8*),
+  i8* bitcast (void ()* @vfunc1_live to i8*),
+  i8* bitcast (void ()* @vfunc2_dead to i8*)
+]}, align 8, !type !{i64 8, !"vfunc1.type"}, !type !{i64 16, !"vfunc2.type"}, !vcall_visibility !{i64 2, i64 8, i64 24}
+
+; CHECK:      @vtableC = internal unnamed_addr constant { [3 x i8*] } { [3 x i8*] [
+; CHECK-SAME:   i8* bitcast (void ()* @regular_non_virtual_funcC to i8*),
+; CHECK-SAME:   i8* bitcast (void ()* @vfunc1_live to i8*),
+; CHECK-SAME:   i8* null
+; CHECK-SAME: ] }, align 8
+
+; A vtable with "relative pointers"
+@vtableD = internal unnamed_addr constant { i32, i32, i8* } {
+  i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc1_live              to i64), i64 ptrtoint ({ i32, i32, i8* }* @vtableD to i64)) to i32),
+  i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc2_dead              to i64), i64 ptrtoint ({ i32, i32, i8* }* @vtableD to i64)) to i32),
+  i8* bitcast (void ()* @regular_non_virtual_funcD to i8*)
+}, align 8, !type !{i64 0, !"vfunc1.type"}, !type !{i64 4, !"vfunc2.type"}, !vcall_visibility !{i64 2, i64 0, i64 8}
+
+; CHECK:      @vtableD = internal unnamed_addr constant { i32, i32, i8* } {
+; CHECK-SAME:   i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc1_live              to i64), i64 ptrtoint ({ i32, i32, i8* }* @vtableD to i64)) to i32),
+; CHECK-SAME:   i32 0,
+; CHECK-SAME:   i8* bitcast (void ()* @regular_non_virtual_funcD to i8*)
+; CHECK-SAME: }, align 8
+
+; (1) vfunc1_live is referenced from @main, stays alive
+define internal void @vfunc1_live() {
+  ; CHECK: define internal void @vfunc1_live(
+  ret void
+}
+
+; (2) vfunc2_dead is never referenced, gets removed and vtable slot is null'd
+define internal void @vfunc2_dead() {
+  ; CHECK-NOT: define internal void @vfunc2_dead(
+  ret void
+}
+
+; (3) not using a range in !vcall_visibility, global gets removed
+define internal void @regular_non_virtual_funcA() {
+  ; CHECK-NOT: define internal void @regular_non_virtual_funcA(
+  ret void
+}
+
+; (4) using a range in !vcall_visibility, pointer is outside of range, so should
+; stay alive
+define internal void @regular_non_virtual_funcB() {
+  ; CHECK: define internal void @regular_non_virtual_funcB(
+  ret void
+}
+
+; (5) using a range in !vcall_visibility, pointer is outside of range, so should
+; stay alive
+define internal void @regular_non_virtual_funcC() {
+  ; CHECK: define internal void @regular_non_virtual_funcC(
+  ret void
+}
+
+; (6) using a range in !vcall_visibility, pointer is outside of range, so should
+; stay alive
+define internal void @regular_non_virtual_funcD() {
+  ; CHECK: define internal void @regular_non_virtual_funcD(
+  ret void
+}
+
+define void @main() {
+  %1 = ptrtoint { [3 x i8*] }* @vtableA to i64 ; to keep @vtableA alive
+  %2 = ptrtoint { [3 x i8*] }* @vtableB to i64 ; to keep @vtableB alive
+  %3 = ptrtoint { [3 x i8*] }* @vtableC to i64 ; to keep @vtableC alive
+  %4 = ptrtoint { i32, i32, i8* }* @vtableD to i64 ; to keep @vtableD alive
+  %5 = tail call { i8*, i1 } @llvm.type.checked.load(i8* null, i32 0, metadata !"vfunc1.type")
+  ret void
+}
+
+!999 = !{i32 1, !"Virtual Function Elim", i32 1}
+!llvm.module.flags = !{!999}
diff --git a/llvm/test/Verifier/vcall-visibility.ll b/llvm/test/Verifier/vcall-visibility.ll
@@ -0,0 +1,19 @@
+; RUN: not opt -S -verify < %s 2>&1 | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; vcall_visibility must have either 1 or 3 operands
+@vtableA = internal unnamed_addr constant { [1 x i8*] } { [1 x i8*] [i8* null]}, align 8, !vcall_visibility !{i64 2, i64 42}
+; CHECK: bad !vcall_visibility attachment
+
+; range start cannot be greater than range end
+@vtableB = internal unnamed_addr constant { [1 x i8*] } { [1 x i8*] [i8* null]}, align 8, !vcall_visibility !{i64 2, i64 10, i64 8}
+; CHECK: bad !vcall_visibility attachment
+
+; vcall_visibility range cannot be over 64 bits
+@vtableC = internal unnamed_addr constant { [1 x i8*] } { [1 x i8*] [i8* null]}, align 8, !vcall_visibility !{i64 2, i128 0, i128 9223372036854775808000}
+; CHECK: bad !vcall_visibility attachment
+
+; range must be two integers
+@vtableD = internal unnamed_addr constant { [1 x i8*] } { [1 x i8*] [i8* null]}, align 8, !vcall_visibility !{i64 2, i64 0, !"string"}
+; CHECK: bad !vcall_visibility attachment
