[GlobalDCE] Handle relative pointers in VFE (for Swift vtables)

[kubamracek]

Summary

As part of codesize optimization work for Swift, we'd like to add Virtual Function Elimination to Swift, very similarly to how GlobalDCE supports C++ VFE. This PR adds support for Swift vtables which contai "relative pointers" instead of direct pointer references. These are in the form of:

@symbol = ... {
  i32 trunc (i64 sub (i64 ptrtoint (<type> @target to i64), i64 ptrtoint (... @symbol to i64)) to i32)
}

The PR extends GlobalDCE's way of looking up a vtable offset into a dependency to be able to see through this expression and find the target symbol.

See the added test for an example this.

[fhahn]

llvm/lib/Analysis/TypeMetadataUtils.cpp

189

I think it would be good at least to have the layout documented somewhere. Not sure where the right place would be, but I think we should have at least a comment here, explaining what the expected valid layout is.

@pcc, @tejohnson any ideas?

199

The sub case here in particular needs some explanation I think. Is there anything we can verify about the layout? E.g. here the second operand always needs to be the vtable global, right?

llvm/lib/Transforms/IPO/GlobalDCE.cpp

136

Could you also explain *why* those dependencies can be safely ignored?

139

Does this extra code also apply to the non-swift case? If so, could you add a non-swift related test case?

Also, could this part split off from the extension to getPointerAtOffset or are both needed for any transformation to happen in the swift case?

llvm/test/Transforms/GlobalDCE/virtual-functions-swift.ll

1 ↗

(On Diff #364788)

Can this simplify be removed? It is run again below and the output is piped to FileCheck.

[kubamracek]

@fhahn added explaining comment that shows the form of the relative pointer expression, added extra structure checking for "sub", removed the "allow non-vfunc entries in the vtable" part.

[kubamracek]

llvm/lib/Transforms/IPO/GlobalDCE.cpp

136

This is now removed from the patch, but for completeness, there was a typo in the comment that suggested that the change was ignoring more dependencies compared to the existing code -- that was a typo, the change was actually "keeping" some dependencies as "regular" deps, specifically if they don't have an entry in the !type metadata, therefore they are not a vfunc.

[kubamracek]

https://reviews.llvm.org/D108741 [GlobalDCE] Handle non-vfunc entries in vtables during VFE (for Swift vtables)

[fhahn]

llvm/lib/Analysis/TypeMetadataUtils.cpp

129–130

I think this comment would be helpful at the declaration in the header, so it is included in the auto-generated documentation. So it might be good to move this to the header and maybe have an inline code comment mark where the swift-specific handling starts?

179

Could you also add a test case for that?

[kubamracek]

llvm/lib/Analysis/TypeMetadataUtils.cpp

179

That's tough without https://reviews.llvm.org/D108741, because returning "not found" from this function is currently going to make GlobalDCE be *more aggressive* and remove that slot from the vtable. Writing a test for this behavior is... "weird" because that behavior sounds like a bug (to be fixed with https://reviews.llvm.org/D108741).

[fhahn]

LGTM, thanks for splitting this off.

Please wait a couple of days with committing, in case there are additional comments/concerns.

llvm/lib/Analysis/TypeMetadataUtils.cpp

179

It would still provide coverage for that code path.

[kubamracek]

llvm/lib/Analysis/TypeMetadataUtils.cpp

179

Good point. I have added a test for that, and I'll fix the test with the subsequent change to test for the expected behavior.

llvm/include/llvm/Analysis/TypeMetadataUtils.h

 /// Given a call to the intrinsic \@llvm.type.checked.load, find all
 /// devirtualizable call sites based on the call and return them in DevirtCalls.
 void findDevirtualizableCallsForTypeCheckedLoad(
     SmallVectorImpl<DevirtCallSite> &DevirtCalls,
     SmallVectorImpl<Instruction *> &LoadedPtrs,
     SmallVectorImpl<Instruction *> &Preds, bool &HasNonCallUses,
     const CallInst *CI, DominatorTree &DT);
 
-Constant *getPointerAtOffset(Constant *I, uint64_t Offset, Module &M);
+// Processes a Constant recursively looking into elements of arrays, structs and
+// expressions to find a trivial pointer element that is located at the given
+// offset (relative to the beginning of the whole outer Constant).
+//
+// Used for example from GlobalDCE to find an entry in a C++ vtable that matches
+// a vcall offset.
+//
+// To support Swift vtables, getPointerAtOffset can see through "relative
+// pointers", i.e. (sub-)expressions of the form of:
+//
+//   @symbol = ... {
+//     i32 trunc (i64 sub (
+//       i64 ptrtoint (<type> @target to i64), i64 ptrtoint (... @symbol to i64)
+//     ) to i32)
+//   }
+//
+// For such (sub-)expressions, getPointerAtOffset returns the @target pointer.
+Constant *getPointerAtOffset(Constant *I, uint64_t Offset, Module &M,
+                             Constant *TopLevelGlobal = nullptr);
 }
 
 #endif

llvm/lib/Analysis/TypeMetadataUtils.cpp

for (const Use &U : CI->uses()) {
     HasNonCallUses = true;
   }
 
   for (Value *LoadedPtr : LoadedPtrs)
     findCallsAtConstantOffset(DevirtCalls, &HasNonCallUses, LoadedPtr,
                               Offset->getZExtValue(), CI, DT);
 }
 
-Constant *llvm::getPointerAtOffset(Constant *I, uint64_t Offset, Module &M) {
+Constant *llvm::getPointerAtOffset(Constant *I, uint64_t Offset, Module &M,
+                                   Constant *TopLevelGlobal) {
   if (I->getType()->isPointerTy()) {
     if (Offset == 0)
       return I;
     return nullptr;
   }
 
   const DataLayout &DL = M.getDataLayout();
 
   if (auto *C = dyn_cast<ConstantStruct>(I)) {
     const StructLayout *SL = DL.getStructLayout(C->getType());
     if (Offset >= SL->getSizeInBytes())
       return nullptr;
 
     unsigned Op = SL->getElementContainingOffset(Offset);
     return getPointerAtOffset(cast<Constant>(I->getOperand(Op)),
-                              Offset - SL->getElementOffset(Op), M);
+                              Offset - SL->getElementOffset(Op), M,
+                              TopLevelGlobal);
   }
   if (auto *C = dyn_cast<ConstantArray>(I)) {
     ArrayType *VTableTy = C->getType();
     uint64_t ElemSize = DL.getTypeAllocSize(VTableTy->getElementType());
 
     unsigned Op = Offset / ElemSize;
     if (Op >= C->getNumOperands())
       return nullptr;
 
     return getPointerAtOffset(cast<Constant>(I->getOperand(Op)),
-                              Offset % ElemSize, M);
+                              Offset % ElemSize, M, TopLevelGlobal);
+  }
+
+  // (Swift-specific) relative-pointer support starts here.
+  if (auto *CI = dyn_cast<ConstantInt>(I)) {
+    if (Offset == 0 && CI->getZExtValue() == 0) {
+      return I;
+    }
+  }
+  if (auto *C = dyn_cast<ConstantExpr>(I)) {
+    switch (C->getOpcode()) {
+    case Instruction::Trunc:
+    case Instruction::PtrToInt:
+      return getPointerAtOffset(cast<Constant>(C->getOperand(0)), Offset, M,
+                                TopLevelGlobal);
+    case Instruction::Sub: {
+      auto *Operand0 = cast<Constant>(C->getOperand(0));
+      auto *Operand1 = cast<Constant>(C->getOperand(1));
+      auto *Operand1TargetGlobal = getPointerAtOffset(Operand1, 0, M);
+
+      // Check that in the "sub (@a, @b)" expression, @b points back to the top
+      // level global that we're processing. Otherwise bail.
+      if (Operand1TargetGlobal != TopLevelGlobal)
+        return nullptr;
+
+      return getPointerAtOffset(Operand0, Offset, M, TopLevelGlobal);
+    }
+    default:
+      return nullptr;
+    }
   }
   return nullptr;
 }

llvm/lib/Transforms/IPO/GlobalDCE.cpp

void GlobalDCEPass::UpdateGVDependencies(GlobalValue &GV) {
   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)) {
       LLVM_DEBUG(dbgs() << "Ignoring dep " << GVU->getName() << " -> "
                         << GV.getName() << "\n");
       continue;
     }
     GVDependencies[GVU].insert(&GV);
   }
 }
 
 /// Mark Global value as Live
 void GlobalDCEPass::MarkLive(GlobalValue &GV,
                              SmallVectorImpl<GlobalValue *> *Updates) {
   auto const Ret = AliveGlobals.insert(&GV);
   if (!Ret.second)

 void GlobalDCEPass::ScanVTableLoad(Function *Caller, Metadata *TypeId,
                                    uint64_t CallOffset) {
   for (auto &VTableInfo : TypeIdMap[TypeId]) {
     GlobalVariable *VTable = VTableInfo.first;
     uint64_t VTableOffset = VTableInfo.second;
 
     Constant *Ptr =
         getPointerAtOffset(VTable->getInitializer(), VTableOffset + CallOffset,
-                           *Caller->getParent());
+                           *Caller->getParent(), VTable);
     if (!Ptr) {
       LLVM_DEBUG(dbgs() << "can't find pointer in vtable!\n");
       VFESafeVTables.erase(VTable);
       return;
     }
 
     auto Callee = dyn_cast<Function>(Ptr->stripPointerCasts());
     if (!Callee) {

llvm/test/Transforms/GlobalDCE/virtual-functions-relative-pointers-bad.ll

+; 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)
+
+@vtable = internal unnamed_addr constant { [3 x i32] } { [3 x i32] [
+  i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc1              to i64), i64 ptrtoint ({ [3 x i32] }* @vtable to i64)) to i32),
+  i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc2              to i64), i64 ptrtoint ({ [3 x i32] }* @vtable to i64)) to i32),
+
+  ; a "bad" relative pointer because it's base is not the @vtable symbol
+  i32 trunc (i64 sub (i64 ptrtoint (void ()* @weird_ref_1         to i64), i64 ptrtoint (void ()* @weird_ref_2 to i64)) to i32)
+]}, align 8, !type !0, !type !1, !vcall_visibility !{i64 2}
+!0 = !{i64 0, !"vfunc1.type"}
+!1 = !{i64 4, !"vfunc2.type"}
+
+; CHECK:      @vtable = internal unnamed_addr constant { [3 x i32] } { [3 x i32] [
+; CHECK-SAME:   i32 trunc (i64 sub (i64 0, i64 ptrtoint ({ [3 x i32] }* @vtable to i64)) to i32),
+; CHECK-SAME:   i32 trunc (i64 sub (i64 0, i64 ptrtoint ({ [3 x i32] }* @vtable to i64)) to i32),
+; CHECK-SAME:   i32 trunc (i64 sub (i64 0, i64 ptrtoint (void ()* @weird_ref_2 to i64)) to i32)
+; CHECK-SAME: ] }, align 8, !type !0, !type !1, !vcall_visibility !2
+
+define internal void @vfunc1() { ret void }
+define internal void @vfunc2() { ret void }
+define internal void @weird_ref_1() { ret void }
+define internal void @weird_ref_2() { ret void }
+
+define void @main() {
+  %1 = ptrtoint { [3 x i32] }* @vtable to i64 ; to keep @vtable alive
+  call void @weird_ref_2()
+  ret void
+}
+
+!999 = !{i32 1, !"Virtual Function Elim", i32 1}
+!llvm.module.flags = !{!999}

llvm/test/Transforms/GlobalDCE/virtual-functions-relative-pointers.ll

+; 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 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 (void ()* @vfunc1_live              to i64), i64 ptrtoint ({ [2 x i32] }* @vtable to i64)) to i32),
+  i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc2_dead              to i64), i64 ptrtoint ({ [2 x i32] }* @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"}
+
+; CHECK:      @vtable = internal unnamed_addr constant { [2 x i32] } { [2 x i32] [
+; CHECK-SAME:   i32 trunc (i64 sub (i64 ptrtoint (void ()* @vfunc1_live              to i64), i64 ptrtoint ({ [2 x i32] }* @vtable to i64)) to i32),
+; CHECK-SAME:   i32 trunc (i64 sub (i64 0,                                                    i64 ptrtoint ({ [2 x i32] }* @vtable to i64)) to i32)
+; CHECK-SAME: ] }, align 8, !type !0, !type !1, !vcall_visibility !2
+
+; (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
+}
+
+define void @main() {
+  %1 = ptrtoint { [2 x i32] }* @vtable to i64 ; to keep @vtable alive
+  %2 = 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}