[llvm] Ensure propagated constants in the vtable are aligned

Author: PiJoules

llvm/lib/Transforms/IPO/WholeProgramDevirt.cpp

@@ -226,6 +226,10 @@ struct PatternList {
     return false;
   }
 };
+
+inline uint64_t RoundUp(uint64_t Size, uint64_t Align) {
+  return (Size + (Align - 1)) / Align * Align;
+}
 } // namespace
 
 // Find the minimum offset that we may store a value of size Size bits at. If
@@ -298,7 +302,9 @@ wholeprogramdevirt::findLowestOffset(ArrayRef<VirtualCallTarget> Targets,
           ++Byte;
         }
       }
-      return (MinByte + I) * 8;
+      // Rounding up ensures the constant is always stored at address we
+      // can directly load from without misalignment.
+      return RoundUp((MinByte + I) * 8, Size);
     NextI:;
     }
   }
@@ -1834,9 +1840,19 @@ bool DevirtModule::tryVirtualConstProp(
   if (!RetType)
     return false;
   unsigned BitWidth = RetType->getBitWidth();
+
+  // TODO: Since we can evaluated these constants at compile-time, we can save
+  // some space by calculating the smallest range of values that all these
+  // constants can fit in, then only allocate enough space to fit those values.
+  // At each callsite, we can get the original type by doing a sign/zero
+  // extension. For example, if we would store an i64, but we can see that all
+  // the values fit into an i16, then we can store an i16 before/after the
+  // vtable and at each callsite do a s/zext.
   if (BitWidth > 64)
     return false;
 
+  Align TypeAlignment = M.getDataLayout().getPrefTypeAlign(RetType);
+
   // Make sure that each function is defined, does not access memory, takes at
   // least one argument, does not use its first argument (which we assume is
   // 'this'), and has the same return type.
@@ -1861,6 +1877,18 @@ bool DevirtModule::tryVirtualConstProp(
         Fn->arg_empty() || !Fn->arg_begin()->use_empty() ||
         Fn->getReturnType() != RetType)
       return false;
+
+    // This only works if the integer size is at most the alignment of the
+    // vtable. If the table is underaligned, then we can't guarantee that the
+    // constant will always be aligned to the integer type alignment. For
+    // example, if the table is `align 1`, we can never guarantee that an i32
+    // stored before/after the vtable is 32-bit aligned without changing the
+    // alignment of the new global.
+    GlobalVariable *GV = Target.TM->Bits->GV;
+    Align TableAlignment = M.getDataLayout().getValueOrABITypeAlignment(
+        GV->getAlign(), GV->getValueType());
+    if (TypeAlignment > TableAlignment)
+      return false;
   }
 
   for (auto &&CSByConstantArg : SlotInfo.ConstCSInfo) {
@@ -1880,6 +1908,12 @@ bool DevirtModule::tryVirtualConstProp(
 
     // Find an allocation offset in bits in all vtables associated with the
     // type.
+    // TODO: If a series of i1s would be placed after the vtable, it would
+    // help save some space if they were placed at the very end after all
+    // other larger-size constants. Having these i1s anywhere in the middle
+    // of the allocation would mean extra padding is needed for any subsequent
+    // constants, but having them at the end wouldn't require this padding
+    // at the very end.
     uint64_t AllocBefore =
         findLowestOffset(TargetsForSlot, /*IsAfter=*/false, BitWidth);
     uint64_t AllocAfter =
@@ -1911,6 +1945,14 @@ bool DevirtModule::tryVirtualConstProp(
       setAfterReturnValues(TargetsForSlot, AllocAfter, BitWidth, OffsetByte,
                            OffsetBit);
 
+    // In an earlier check we forbade constant propagation from operating on
+    // tables whose alignment is less than the alignment needed for loading
+    // the constant. Thus, the address we take the offset from will always be
+    // aligned to at least this integer alignment. Now, we need to ensure that
+    // the offset is also aligned to this integer alignment to ensure we always
+    // have an aligned load.
+    assert(OffsetByte % TypeAlignment.value() == 0);
+
     if (RemarksEnabled || AreStatisticsEnabled())
       for (auto &&Target : TargetsForSlot)
         Target.WasDevirt = true;

llvm/test/Transforms/WholeProgramDevirt/virtual-const-prop-begin.ll

@@ -1,34 +1,35 @@
 ; RUN: opt -S -passes=wholeprogramdevirt -whole-program-visibility %s | FileCheck %s
 
 target datalayout = "e-p:64:64"
-target triple = "x86_64-unknown-linux-gnu"
 
-; CHECK: [[VT1DATA:@[^ ]*]] = private constant { [8 x i8], [3 x ptr], [0 x i8] } { [8 x i8] c"\00\00\00\01\00\00\00\02", [3 x ptr] [ptr @vf0i1, ptr @vf1i1, ptr @vf1i32], [0 x i8] zeroinitializer }, section "vt1sec", !type [[T8:![0-9]+]]
+;; Note that i16 is used here such that we can ensure all constants for "typeid"
+;; can come before the vtable.
+; CHECK: [[VT1DATA:@[^ ]*]] = private constant { [8 x i8], [3 x ptr], [0 x i8] } { [8 x i8] c"\00\00\00\00\03\00\00\02", [3 x ptr] [ptr @vf0i1, ptr @vf1i1, ptr @vf1i16], [0 x i8] zeroinitializer }, section "vt1sec", !type [[T8:![0-9]+]]
 @vt1 = constant [3 x ptr] [
 ptr @vf0i1,
 ptr @vf1i1,
-ptr @vf1i32
+ptr @vf1i16
 ], section "vt1sec", !type !0
 
-; CHECK: [[VT2DATA:@[^ ]*]] = private constant { [8 x i8], [3 x ptr], [0 x i8] } { [8 x i8] c"\00\00\00\02\00\00\00\01", [3 x ptr] [ptr @vf1i1, ptr @vf0i1, ptr @vf2i32], [0 x i8] zeroinitializer }, !type [[T8]]
+; CHECK: [[VT2DATA:@[^ ]*]] = private constant { [8 x i8], [3 x ptr], [0 x i8] } { [8 x i8] c"\00\00\00\00\04\00\00\01", [3 x ptr] [ptr @vf1i1, ptr @vf0i1, ptr @vf2i16], [0 x i8] zeroinitializer }, !type [[T8]]
 @vt2 = constant [3 x ptr] [
 ptr @vf1i1,
 ptr @vf0i1,
-ptr @vf2i32
+ptr @vf2i16
 ], !type !0
 
-; CHECK: [[VT3DATA:@[^ ]*]] = private constant { [5 x i8], [3 x ptr], [0 x i8] } { [5 x i8] c"\03\00\00\00\02", [3 x ptr] [ptr @vf0i1, ptr @vf1i1, ptr @vf3i32], [0 x i8] zeroinitializer }, align 1, !type [[T5:![0-9]+]]
+; CHECK: [[VT3DATA:@[^ ]*]] = private constant { [4 x i8], [3 x ptr], [0 x i8] } { [4 x i8] c"\05\00\00\02", [3 x ptr] [ptr @vf0i1, ptr @vf1i1, ptr @vf3i16], [0 x i8] zeroinitializer }, align 2, !type [[T5:![0-9]+]]
 @vt3 = constant [3 x ptr] [
 ptr @vf0i1,
 ptr @vf1i1,
-ptr @vf3i32
-], align 1, !type !0
+ptr @vf3i16
+], align 2, !type !0
 
-; CHECK: [[VT4DATA:@[^ ]*]] = private constant { [16 x i8], [3 x ptr], [0 x i8] } { [16 x i8] c"\00\00\00\00\00\00\00\00\00\00\00\04\00\00\00\01", [3 x ptr] [ptr @vf1i1, ptr @vf0i1, ptr @vf4i32], [0 x i8] zeroinitializer },  align 16, !type [[T16:![0-9]+]]
+; CHECK: [[VT4DATA:@[^ ]*]] = private constant { [16 x i8], [3 x ptr], [0 x i8] } { [16 x i8] c"\00\00\00\00\00\00\00\00\00\00\00\00\06\00\00\01", [3 x ptr] [ptr @vf1i1, ptr @vf0i1, ptr @vf4i16], [0 x i8] zeroinitializer },  align 16, !type [[T16:![0-9]+]]
 @vt4 = constant [3 x ptr] [
 ptr @vf1i1,
 ptr @vf0i1,
-ptr @vf4i32
+ptr @vf4i16
 ], align 16, !type !0
 
 ; CHECK: @vt5 = {{.*}}, !type [[T0:![0-9]+]]
@@ -38,10 +39,35 @@ ptr @__cxa_pure_virtual,
 ptr @__cxa_pure_virtual
 ], !type !0
 
+;; Test relative vtables
+; CHECK:      [[VT6RELDATA:@[^ ]*]] = private constant { [4 x i8], [3 x i32], [0 x i8] } { [4 x i8] c"\00\00\03\00", [3 x i32] [
+; CHECK-SAME:     i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf0i1 to i64), i64 ptrtoint (ptr @vt6_rel to i64)) to i32),
+; CHECK-SAME:     i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf1i1 to i64), i64 ptrtoint (ptr @vt6_rel to i64)) to i32),
+; CHECK-SAME:     i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf1i16 to i64), i64 ptrtoint (ptr @vt6_rel to i64)) to i32)
+; CHECK-SAME: ], [0 x i8] zeroinitializer }, !type [[TREL:![0-9]+]]
+@vt6_rel = constant [3 x i32] [
+i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf0i1 to i64), i64 ptrtoint (ptr @vt6_rel to i64)) to i32),
+i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf1i1 to i64), i64 ptrtoint (ptr @vt6_rel to i64)) to i32),
+i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf1i16 to i64), i64 ptrtoint (ptr @vt6_rel to i64)) to i32)
+], !type !2
+
+; CHECK:      [[VT7RELDATA:@[^ ]*]] = private constant { [4 x i8], [3 x i32], [0 x i8] } { [4 x i8] c"\00\00\04\00", [3 x i32] [
+; CHECK-SAME:     i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf1i1 to i64), i64 ptrtoint (ptr @vt7_rel to i64)) to i32),
+; CHECK-SAME:     i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf0i1 to i64), i64 ptrtoint (ptr @vt7_rel to i64)) to i32),
+; CHECK-SAME:     i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf2i16 to i64), i64 ptrtoint (ptr @vt7_rel to i64)) to i32)
+; CHECK-SAME: ], [0 x i8] zeroinitializer }, !type [[TREL]]
+@vt7_rel = constant [3 x i32] [
+i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf1i1 to i64), i64 ptrtoint (ptr @vt7_rel to i64)) to i32),
+i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf0i1 to i64), i64 ptrtoint (ptr @vt7_rel to i64)) to i32),
+i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @vf2i16 to i64), i64 ptrtoint (ptr @vt7_rel to i64)) to i32)
+], !type !2
+
 ; CHECK: @vt1 = alias [3 x ptr], getelementptr inbounds ({ [8 x i8], [3 x ptr], [0 x i8] }, ptr [[VT1DATA]], i32 0, i32 1)
 ; CHECK: @vt2 = alias [3 x ptr], getelementptr inbounds ({ [8 x i8], [3 x ptr], [0 x i8] }, ptr [[VT2DATA]], i32 0, i32 1)
-; CHECK: @vt3 = alias [3 x ptr], getelementptr inbounds ({ [5 x i8], [3 x ptr], [0 x i8] }, ptr [[VT3DATA]], i32 0, i32 1)
+; CHECK: @vt3 = alias [3 x ptr], getelementptr inbounds ({ [4 x i8], [3 x ptr], [0 x i8] }, ptr [[VT3DATA]], i32 0, i32 1)
 ; CHECK: @vt4 = alias [3 x ptr], getelementptr inbounds ({ [16 x i8], [3 x ptr], [0 x i8] }, ptr [[VT4DATA]], i32 0, i32 1)
+; CHECK: @vt6_rel = alias [3 x i32], getelementptr inbounds ({ [4 x i8], [3 x i32], [0 x i8] }, ptr [[VT6RELDATA]], i32 0, i32 1)
+; CHECK: @vt7_rel = alias [3 x i32], getelementptr inbounds ({ [4 x i8], [3 x i32], [0 x i8] }, ptr [[VT7RELDATA]], i32 0, i32 1)
 
 define i1 @vf0i1(ptr %this) readnone {
   ret i1 0
@@ -51,20 +77,20 @@ define i1 @vf1i1(ptr %this) readnone {
   ret i1 1
 }
 
-define i32 @vf1i32(ptr %this) readnone {
-  ret i32 1
+define i16 @vf1i16(ptr %this) readnone {
+  ret i16 3
 }
 
-define i32 @vf2i32(ptr %this) readnone {
-  ret i32 2
+define i16 @vf2i16(ptr %this) readnone {
+  ret i16 4
 }
 
-define i32 @vf3i32(ptr %this) readnone {
-  ret i32 3
+define i16 @vf3i16(ptr %this) readnone {
+  ret i16 5
 }
 
-define i32 @vf4i32(ptr %this) readnone {
-  ret i32 4
+define i16 @vf4i16(ptr %this) readnone {
+  ret i16 6
 }
 
 ; CHECK: define i1 @call1(
@@ -87,7 +113,7 @@ define i1 @call2(ptr %obj) {
   %vtable = load ptr, ptr %obj
   %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid")
   call void @llvm.assume(i1 %p)
-  %fptrptr = getelementptr [3 x ptr], ptr %vtable, i32 0, i32 1
+  %fptrptr = getelementptr [3 x ptr], ptr %vtable, i16 0, i16 1
   %fptr = load ptr, ptr %fptrptr
   ; CHECK: [[VTGEP2:%[^ ]*]] = getelementptr i8, ptr %vtable, i32 -1
   ; CHECK: [[VTLOAD2:%[^ ]*]] = load i8, ptr [[VTGEP2]]
@@ -98,27 +124,68 @@ define i1 @call2(ptr %obj) {
   ret i1 %result
 }
 
-; CHECK: define i32 @call3(
-define i32 @call3(ptr %obj) {
+; CHECK: define i16 @call3(
+define i16 @call3(ptr %obj) {
   %vtable = load ptr, ptr %obj
   %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid")
   call void @llvm.assume(i1 %p)
-  %fptrptr = getelementptr [3 x ptr], ptr %vtable, i32 0, i32 2
+  %fptrptr = getelementptr [3 x ptr], ptr %vtable, i16 0, i16 2
   %fptr = load ptr, ptr %fptrptr
-  ; CHECK: [[VTGEP3:%[^ ]*]] = getelementptr i8, ptr %vtable, i32 -5
-  ; CHECK: [[VTLOAD3:%[^ ]*]] = load i32, ptr [[VTGEP3]]
-  %result = call i32 %fptr(ptr %obj)
-  ; CHECK: ret i32 [[VTLOAD3]]
-  ret i32 %result
+  ; CHECK: [[VTGEP3:%[^ ]*]] = getelementptr i8, ptr %vtable, i32 -4
+  ; CHECK: [[VTLOAD3:%[^ ]*]] = load i16, ptr [[VTGEP3]]
+  %result = call i16 %fptr(ptr %obj)
+  ; CHECK: ret i16 [[VTLOAD3]]
+  ret i16 %result
+}
+
+; CHECK: define i1 @call1_rel(
+define i1 @call1_rel(ptr %obj) {
+  %vtable = load ptr, ptr %obj
+  %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid3")
+  call void @llvm.assume(i1 %p)
+  %fptr = call ptr @llvm.load.relative.i32(ptr %vtable, i32 0)
+  %result = call i1 %fptr(ptr %obj)
+  ret i1 %result
+  ; CHECK: [[RES:%[^ ]*]] = icmp eq ptr %vtable, @vt7_rel
+  ; CHECK: ret i1 [[RES]]
+}
+
+; CHECK: define i1 @call2_rel(
+define i1 @call2_rel(ptr %obj) {
+  %vtable = load ptr, ptr %obj
+  %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid3")
+  call void @llvm.assume(i1 %p)
+  %fptr = call ptr @llvm.load.relative.i32(ptr %vtable, i32 4)
+  %result = call i1 %fptr(ptr %obj)
+  ret i1 %result
+  ; CHECK: [[RES:%[^ ]*]] = icmp eq ptr %vtable, @vt6_rel
+  ; CHECK: ret i1 [[RES]]
+}
+
+; CHECK: define i16 @call3_rel(
+define i16 @call3_rel(ptr %obj) {
+  %vtable = load ptr, ptr %obj
+  %p = call i1 @llvm.type.test(ptr %vtable, metadata !"typeid3")
+  call void @llvm.assume(i1 %p)
+  %fptr = call ptr @llvm.load.relative.i32(ptr %vtable, i32 8)
+  ; CHECK: [[VTGEP3:%[^ ]*]] = getelementptr i8, ptr %vtable, i32 -2
+  ; CHECK: [[VTLOAD3:%[^ ]*]] = load i16, ptr [[VTGEP3]]
+  %result = call i16 %fptr(ptr %obj)
+  ; CHECK: ret i16 [[VTLOAD3]]
+  ret i16 %result
 }
 
 declare i1 @llvm.type.test(ptr, metadata)
 declare void @llvm.assume(i1)
 declare void @__cxa_pure_virtual()
+declare ptr @llvm.load.relative.i32(ptr, i32)
 
 ; CHECK: [[T8]] = !{i32 8, !"typeid"}
-; CHECK: [[T5]] = !{i32 5, !"typeid"}
+; CHECK: [[T5]] = !{i32 4, !"typeid"}
 ; CHECK: [[T16]] = !{i32 16, !"typeid"}
 ; CHECK: [[T0]] = !{i32 0, !"typeid"}
+; CHECK: [[TREL]] = !{i32 4, !"typeid3"}
 
 !0 = !{i32 0, !"typeid"}
+!1 = !{i32 0, !"typeid2"}
+!2 = !{i32 0, !"typeid3"}