LowerBitSets: Align referenced globals.

[pcc]

Summary

This change aligns globals to the next highest power of 2 bytes, up to a
maximum of 128. This makes it more likely that we will be able to compress
bit sets with a greater alignment. In many more cases, we can now take
advantage of a new optimization also introduced in this patch that removes
bit set checks if the bit set is all ones.

The 128 byte maximum was found to provide the best tradeoff between instruction
overhead and data overhead in a recent build of Chromium. It allows us to
remove ~900KB of instructions at the cost of ~250KB of data.

[kcc]

~900KB of instructions

Out of how many?

Depending on the %% savings we may or may not want the extra complexity.

[pcc]

This change reduced the size of the .text section from 6701377 bytes to 5844340 bytes in a binary based on Chrome but with everything except bit set checks removed. So around 15%. I'll also collect figures for Chrome itself overnight.

[kcc]

This also deserves to be mentioned in the design document clang.llvm.org/docs/ControlFlowIntegrityDesign.html
(probably create a section "Optimisations")
Since this is a security feature, we want to make sure that the security researchers who compare
the design document with the actual generated code are not surprised.

[pcc]

Size info for chrome (before, after):

out_cfi_0bf03cb47380bf9bea4e738b7b5f2ded56eb335b/Release/chrome  :
section                       size        addr
.interp                         28         624
.note.ABI-tag                   32         652
.note.gnu.build-id              36         684
.dynsym                      44688         720
.dynstr                      38634       45408
.gnu.hash                      380       84048
.gnu.version                  3724       84428
.gnu.version_r                 960       88152
.rela.dyn                 13011840       89112
.rela.plt                    41232    13100952
.init                           26    13142184
.plt                         27504    13142224
.text                    112043598    13169728
malloc_hook                   4304   125213328
google_malloc                 4769   125217632
.fini                            9   125222404
.rodata                   12601040   125222656
.eh_frame                  7829036   137823696
.eh_frame_hdr              1680676   145652732
.tbss                           16   147341264
.data.rel.ro.local         1927376   147341264
.jcr                             8   149268640
.fini_array                      8   149268648
.init_array                     56   149268656
.data.rel.ro               3470504   149268720
.dynamic                      1152   152739224
.got                          2072   152740376
.got.plt                     13768   152742448
.tm_clone_table                  0   152756224
.data                       153243   152756224
.bss                        444996   152909472
.comment                        92           0
.note.gnu.gold-version          28           0
Total                    153345835


out_cfi_a0ad911453217da8f0d553bc9be5f6162bca9ab8/Release/chrome  :
section                       size        addr
.interp                         28         624
.note.ABI-tag                   32         652
.note.gnu.build-id              36         684
.dynsym                      44688         720
.dynstr                      38634       45408
.gnu.hash                      380       84048
.gnu.version                  3724       84428
.gnu.version_r                 960       88152
.rela.dyn                 13011840       89112
.rela.plt                    41232    13100952
.init                           26    13142184
.plt                         27504    13142224
.text                    109658814    13169728
malloc_hook                   4304   122828544
google_malloc                 4759   122832848
.fini                            9   122837608
.rodata                   12349136   122837760
.eh_frame                  7812260   135186896
.eh_frame_hdr              1680676   142999156
.tbss                           16   144687376
.data.rel.ro.local         2010512   144687376
.jcr                             8   146697888
.fini_array                      8   146697896
.init_array                     56   146697904
.data.rel.ro               3878568   146697968
.dynamic                      1152   150576536
.got                          2072   150577688
.got.plt                     13768   150579760
.tm_clone_table                  0   150593536
.data                       153243   150593536
.bss                        444996   150746784
.comment                        92           0
.note.gnu.gold-version          28           0
Total                    151183561

So over 2MB savings overall, better than I expected!

[kcc]

LGTM

[pcc]

This also deserves to be mentioned in the design document clang.llvm.org/docs/ControlFlowIntegrityDesign.html

r230588.

include/llvm/Transforms/IPO/LowerBitSets.h

struct BitSetInfo {
   // For example, a log2 alignment of 3 means that bits in the bitset
   // represent addresses 8 bytes apart.
   unsigned AlignLog2;
 
   bool isSingleOffset() const {
     return Bits.size() == 1 && Bits[0] == 1;
   }
 
+  bool isAllOnes() const {
+    for (unsigned I = 0; I != Bits.size() - 1; ++I)
+      if (Bits[I] != 0xFF)
+        return false;
+
+    if (BitSize % 8 == 0)
+      return Bits[Bits.size() - 1] == 0xFF;
+
+    return Bits[Bits.size() - 1] == (1 << (BitSize % 8)) - 1;
+  }
+
   bool containsGlobalOffset(uint64_t Offset) const;
 
   bool containsValue(const DataLayout *DL,
                      const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout,
                      Value *V, uint64_t COffset = 0) const;
 
 };
 

lib/Transforms/IPO/LowerBitSets.cpp

 struct LowerBitSets : public ModulePass {
   static char ID;
   LowerBitSets() : ModulePass(ID) {
     initializeLowerBitSetsPass(*PassRegistry::getPassRegistry());
   }
 
   const DataLayout *DL;
   IntegerType *Int1Ty;
+  IntegerType *Int8Ty;
   IntegerType *Int32Ty;
   Type *Int32PtrTy;
   IntegerType *Int64Ty;
   Type *IntPtrTy;
 
   // The llvm.bitsets named metadata.
   NamedMDNode *BitSetNM;
 
   // Mapping from bitset mdstrings to the call sites that test them.
   DenseMap<MDString *, std::vector<CallInst *>> BitSetTestCallSites;
 
   BitSetInfo
   buildBitSet(MDString *BitSet,
               const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout);
   Value *createBitSetTest(IRBuilder<> &B, const BitSetInfo &BSI,
                           GlobalVariable *BitSetGlobal, Value *BitOffset);
-  void
+  Value *
   lowerBitSetCall(CallInst *CI, const BitSetInfo &BSI,
                   GlobalVariable *BitSetGlobal, GlobalVariable *CombinedGlobal,
                   const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout);
   void buildBitSetsFromGlobals(Module &M,
                                const std::vector<MDString *> &BitSets,
                                const std::vector<GlobalVariable *> &Globals);
   bool buildBitSets(Module &M);
   bool eraseBitSetMetadata(Module &M);

 ModulePass *llvm::createLowerBitSetsPass() { return new LowerBitSets; }
 
 bool LowerBitSets::doInitialization(Module &M) {
   DL = M.getDataLayout();
   if (!DL)
     report_fatal_error("Data layout required");
 
   Int1Ty = Type::getInt1Ty(M.getContext());
+  Int8Ty = Type::getInt8Ty(M.getContext());
   Int32Ty = Type::getInt32Ty(M.getContext());
   Int32PtrTy = PointerType::getUnqual(Int32Ty);
   Int64Ty = Type::getInt64Ty(M.getContext());
   IntPtrTy = DL->getIntPtrType(M.getContext(), 0);
 
   BitSetNM = M.getNamedMetadata("llvm.bitsets");
 
   BitSetTestCallSites.clear();
if (BSI.Bits.size() <= 8) {
     Value *BitSetEntryAddr = B.CreateGEP(
         ConstantExpr::getBitCast(BitSetGlobal, Int32PtrTy), BitSetGlobalOffset);
     Value *BitSetEntry = B.CreateLoad(BitSetEntryAddr);
 
     return createMaskedBitTest(B, BitSetEntry, BitOffset);
   }
 }
 
-/// Lower a llvm.bitset.test call to its implementation.
-void LowerBitSets::lowerBitSetCall(
+/// Lower a llvm.bitset.test call to its implementation. Returns the value to
+/// replace the call with.
+Value *LowerBitSets::lowerBitSetCall(
     CallInst *CI, const BitSetInfo &BSI, GlobalVariable *BitSetGlobal,
     GlobalVariable *CombinedGlobal,
     const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout) {
   Value *Ptr = CI->getArgOperand(0);
 
-  if (BSI.containsValue(DL, GlobalLayout, Ptr)) {
-    CI->replaceAllUsesWith(
-        ConstantInt::getTrue(BitSetGlobal->getParent()->getContext()));
-    CI->eraseFromParent();
-    return;
-  }
+  if (BSI.containsValue(DL, GlobalLayout, Ptr))
+    return ConstantInt::getTrue(BitSetGlobal->getParent()->getContext());
 
   Constant *GlobalAsInt = ConstantExpr::getPtrToInt(CombinedGlobal, IntPtrTy);
   Constant *OffsetedGlobalAsInt = ConstantExpr::getAdd(
       GlobalAsInt, ConstantInt::get(IntPtrTy, BSI.ByteOffset));
 
   BasicBlock *InitialBB = CI->getParent();
 
   IRBuilder<> B(CI);
 
   Value *PtrAsInt = B.CreatePtrToInt(Ptr, IntPtrTy);
 
-  if (BSI.isSingleOffset()) {
-    Value *Eq = B.CreateICmpEQ(PtrAsInt, OffsetedGlobalAsInt);
-    CI->replaceAllUsesWith(Eq);
-    CI->eraseFromParent();
-    return;
-  }
+  if (BSI.isSingleOffset())
+    return B.CreateICmpEQ(PtrAsInt, OffsetedGlobalAsInt);
 
   Value *PtrOffset = B.CreateSub(PtrAsInt, OffsetedGlobalAsInt);
 
   Value *BitOffset;
   if (BSI.AlignLog2 == 0) {
     BitOffset = PtrOffset;
   } else {
     // We need to check that the offset both falls within our range and is
Value *OffsetSHL = B.CreateShl(
         PtrOffset, ConstantInt::get(IntPtrTy, DL->getPointerSizeInBits(0) -
                                                   BSI.AlignLog2));
     BitOffset = B.CreateOr(OffsetSHR, OffsetSHL);
   }
 
   Constant *BitSizeConst = ConstantInt::get(IntPtrTy, BSI.BitSize);
   Value *OffsetInRange = B.CreateICmpULT(BitOffset, BitSizeConst);
 
+  // If the bit set is all ones, testing against it is unnecessary.
+  if (BSI.isAllOnes())
+    return OffsetInRange;
+
   TerminatorInst *Term = SplitBlockAndInsertIfThen(OffsetInRange, CI, false);
   IRBuilder<> ThenB(Term);
 
   // Now that we know that the offset is in range and aligned, load the
   // appropriate bit from the bitset.
   Value *Bit = createBitSetTest(ThenB, BSI, BitSetGlobal, BitOffset);
 
   // The value we want is 0 if we came directly from the initial block
   // (having failed the range or alignment checks), or the loaded bit if
   // we came from the block in which we loaded it.
   B.SetInsertPoint(CI);
   PHINode *P = B.CreatePHI(Int1Ty, 2);
   P->addIncoming(ConstantInt::get(Int1Ty, 0), InitialBB);
   P->addIncoming(Bit, ThenB.GetInsertBlock());
-
-  CI->replaceAllUsesWith(P);
-  CI->eraseFromParent();
+  return P;
 }
 
 /// Given a disjoint set of bitsets and globals, layout the globals, build the
 /// bit sets and lower the llvm.bitset.test calls.
 void LowerBitSets::buildBitSetsFromGlobals(
     Module &M,
     const std::vector<MDString *> &BitSets,
     const std::vector<GlobalVariable *> &Globals) {
   // Build a new global with the combined contents of the referenced globals.
   std::vector<Constant *> GlobalInits;
-  for (GlobalVariable *G : Globals)
+  for (GlobalVariable *G : Globals) {
     GlobalInits.push_back(G->getInitializer());
+    uint64_t InitSize = DL->getTypeAllocSize(G->getInitializer()->getType());
+
+    // Compute the amount of padding required to align the next element to the
+    // next power of 2.
+    uint64_t Padding = NextPowerOf2(InitSize - 1) - InitSize;
+
+    // Cap at 128 was found experimentally to have a good data/instruction
+    // overhead tradeoff.
+    if (Padding > 128)
+      Padding = RoundUpToAlignment(InitSize, 128) - InitSize;
+
+    GlobalInits.push_back(
+        ConstantAggregateZero::get(ArrayType::get(Int8Ty, Padding)));
+  }
+  if (!GlobalInits.empty())
+    GlobalInits.pop_back();
   Constant *NewInit = ConstantStruct::getAnon(M.getContext(), GlobalInits);
   auto CombinedGlobal =
       new GlobalVariable(M, NewInit->getType(), /*isConstant=*/true,
                          GlobalValue::PrivateLinkage, NewInit);
 
   const StructLayout *CombinedGlobalLayout =
       DL->getStructLayout(cast<StructType>(NewInit->getType()));
 
   // Compute the offsets of the original globals within the new global.
   DenseMap<GlobalVariable *, uint64_t> GlobalLayout;
   for (unsigned I = 0; I != Globals.size(); ++I)
-    GlobalLayout[Globals[I]] = CombinedGlobalLayout->getElementOffset(I);
+    // Multiply by 2 to account for padding elements.
+    GlobalLayout[Globals[I]] = CombinedGlobalLayout->getElementOffset(I * 2);
 
   // For each bitset in this disjoint set...
   for (MDString *BS : BitSets) {
     // Build the bitset.
     BitSetInfo BSI = buildBitSet(BS, GlobalLayout);
 
     // Create a global in which to store it.
     ++NumBitSetsCreated;
     Constant *BitsConst = ConstantDataArray::get(M.getContext(), BSI.Bits);
     auto BitSetGlobal = new GlobalVariable(
         M, BitsConst->getType(), /*isConstant=*/true,
         GlobalValue::PrivateLinkage, BitsConst, BS->getString() + ".bits");
 
     // Lower each call to llvm.bitset.test for this bitset.
     for (CallInst *CI : BitSetTestCallSites[BS]) {
       ++NumBitSetCallsLowered;
+      Value *Lowered =
           lowerBitSetCall(CI, BSI, BitSetGlobal, CombinedGlobal, GlobalLayout);
+      CI->replaceAllUsesWith(Lowered);
+      CI->eraseFromParent();
     }
   }
 
   // Build aliases pointing to offsets into the combined global for each
   // global from which we built the combined global, and replace references
   // to the original globals with references to the aliases.
   for (unsigned I = 0; I != Globals.size(); ++I) {
+    // Multiply by 2 to account for padding elements.
     Constant *CombinedGlobalIdxs[] = {ConstantInt::get(Int32Ty, 0),
-                                      ConstantInt::get(Int32Ty, I)};
+                                      ConstantInt::get(Int32Ty, I * 2)};
     Constant *CombinedGlobalElemPtr =
         ConstantExpr::getGetElementPtr(CombinedGlobal, CombinedGlobalIdxs);
     GlobalAlias *GAlias = GlobalAlias::create(
         Globals[I]->getType()->getElementType(),
         Globals[I]->getType()->getAddressSpace(), Globals[I]->getLinkage(),
         "", CombinedGlobalElemPtr, &M);
     GAlias->takeName(Globals[I]);
     Globals[I]->replaceAllUsesWith(GAlias);

test/Transforms/LowerBitSets/layout.ll

 ; RUN: opt -S -lowerbitsets < %s | FileCheck %s
 
 target datalayout = "e-p:32:32"
 
 ; Tests that this set of globals is laid out according to our layout algorithm
 ; (see GlobalLayoutBuilder in include/llvm/Transforms/IPO/LowerBitSets.h).
 ; The chosen layout in this case is a, e, b, d, c.
 
-; CHECK: private constant { i32, i32, i32, i32, i32 } { i32 1, i32 5, i32 2, i32 4, i32 3 }
+; CHECK: private constant { i32, [0 x i8], i32, [0 x i8], i32, [0 x i8], i32, [0 x i8], i32 } { i32 1, [0 x i8] zeroinitializer, i32 5, [0 x i8] zeroinitializer, i32 2, [0 x i8] zeroinitializer, i32 4, [0 x i8] zeroinitializer, i32 3 }
 @a = constant i32 1
 @b = constant i32 2
 @c = constant i32 3
 @d = constant i32 4
 @e = constant i32 5
 
 !0 = !{!"bitset1", i32* @a, i32 0}
 !1 = !{!"bitset1", i32* @b, i32 0}

test/Transforms/LowerBitSets/simple.ll

 ; RUN: opt -S -lowerbitsets < %s | FileCheck %s
 ; RUN: opt -S -O3 < %s | FileCheck -check-prefix=CHECK-NODISCARD %s
 
 target datalayout = "e-p:32:32"
 
-; CHECK: [[G:@[^ ]*]] = private constant { i32, [63 x i32], i32, [2 x i32] } { i32 1, [63 x i32] zeroinitializer, i32 3, [2 x i32] [i32 4, i32 5] }
+; CHECK: [[G:@[^ ]*]] = private constant { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] } { i32 1, [0 x i8] zeroinitializer, [63 x i32] zeroinitializer, [4 x i8] zeroinitializer, i32 3, [0 x i8] zeroinitializer, [2 x i32] [i32 4, i32 5] }
 @a = constant i32 1
 @b = constant [63 x i32] zeroinitializer
 @c = constant i32 3
 @d = constant [2 x i32] [i32 4, i32 5]
 
 ; Offset 0, 4 byte alignment
-; CHECK: @bitset1.bits = private constant [9 x i8] c"\03\00\00\00\00\00\00\00\04"
+; CHECK: @bitset1.bits = private constant [9 x i8] c"\03\00\00\00\00\00\00\00\08"
 !0 = !{!"bitset1", i32* @a, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset1", i32* @a, i32 0}
 !1 = !{!"bitset1", [63 x i32]* @b, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset1", [63 x i32]* @b, i32 0}
 !2 = !{!"bitset1", [2 x i32]* @d, i32 4}
 ; CHECK-NODISCARD-DAG: !{!"bitset1", [2 x i32]* @d, i32 4}
 
-; Offset 4, 4 byte alignment
-; CHECK: @bitset2.bits = private constant [8 x i8] c"\01\00\00\00\00\00\00\80"
+; Offset 4, 256 byte alignment
+; CHECK: @bitset2.bits = private constant [1 x i8] c"\03"
 !3 = !{!"bitset2", [63 x i32]* @b, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset2", [63 x i32]* @b, i32 0}
 !4 = !{!"bitset2", i32* @c, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset2", i32* @c, i32 0}
 
-; Offset 0, 256 byte alignment
-; CHECK: @bitset3.bits = private constant [1 x i8] c"\03"
+; Offset 0, 4 byte alignment
+; CHECK: @bitset3.bits = private constant [9 x i8] c"\01\00\00\00\00\00\00\00\02"
 !5 = !{!"bitset3", i32* @a, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset3", i32* @a, i32 0}
 !6 = !{!"bitset3", i32* @c, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset3", i32* @c, i32 0}
 
 ; Entries whose second operand is null (the result of a global being DCE'd)
 ; should be ignored.
 !7 = !{!"bitset2", null, i32 0}
 
 !llvm.bitsets = !{ !0, !1, !2, !3, !4, !5, !6, !7 }
 
-; CHECK: @a = alias getelementptr inbounds ({ i32, [63 x i32], i32, [2 x i32] }* [[G]], i32 0, i32 0)
-; CHECK: @b = alias getelementptr inbounds ({ i32, [63 x i32], i32, [2 x i32] }* [[G]], i32 0, i32 1)
-; CHECK: @c = alias getelementptr inbounds ({ i32, [63 x i32], i32, [2 x i32] }* [[G]], i32 0, i32 2)
-; CHECK: @d = alias getelementptr inbounds ({ i32, [63 x i32], i32, [2 x i32] }* [[G]], i32 0, i32 3)
+; CHECK: @a = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 0)
+; CHECK: @b = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 2)
+; CHECK: @c = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 4)
+; CHECK: @d = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 6)
 
 declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
 
 ; CHECK: @foo(i32* [[A0:%[^ ]*]])
 define i1 @foo(i32* %p) {
   ; CHECK-NOT: llvm.bitset.test
 
   ; CHECK: [[R0:%[^ ]*]] = bitcast i32* [[A0]] to i8*
   %pi8 = bitcast i32* %p to i8*
   ; CHECK: [[R1:%[^ ]*]] = ptrtoint i8* [[R0]] to i32
-  ; CHECK: [[R2:%[^ ]*]] = sub i32 [[R1]], ptrtoint ({ i32, [63 x i32], i32, [2 x i32] }* [[G]] to i32)
+  ; CHECK: [[R2:%[^ ]*]] = sub i32 [[R1]], ptrtoint ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]] to i32)
   ; CHECK: [[R3:%[^ ]*]] = lshr i32 [[R2]], 2
   ; CHECK: [[R4:%[^ ]*]] = shl i32 [[R2]], 30
   ; CHECK: [[R5:%[^ ]*]] = or i32 [[R3]], [[R4]]
-  ; CHECK: [[R6:%[^ ]*]] = icmp ult i32 [[R5]], 67
+  ; CHECK: [[R6:%[^ ]*]] = icmp ult i32 [[R5]], 68
   ; CHECK: br i1 [[R6]]
 
   ; CHECK: [[R8:%[^ ]*]] = lshr i32 [[R5]], 5
   ; CHECK: [[R9:%[^ ]*]] = getelementptr i32* bitcast ([9 x i8]* @bitset1.bits to i32*), i32 [[R8]]
   ; CHECK: [[R10:%[^ ]*]] = load i32* [[R9]]
   ; CHECK: [[R11:%[^ ]*]] = and i32 [[R5]], 31
   ; CHECK: [[R12:%[^ ]*]] = shl i32 1, [[R11]]
   ; CHECK: [[R13:%[^ ]*]] = and i32 [[R10]], [[R12]]
define i1 @foo(i32* %p) {
   ret i1 %x
 }
 
 ; CHECK: @bar(i32* [[B0:%[^ ]*]])
 define i1 @bar(i32* %p) {
   ; CHECK: [[S0:%[^ ]*]] = bitcast i32* [[B0]] to i8*
   %pi8 = bitcast i32* %p to i8*
   ; CHECK: [[S1:%[^ ]*]] = ptrtoint i8* [[S0]] to i32
-  ; CHECK: [[S2:%[^ ]*]] = sub i32 [[S1]], add (i32 ptrtoint ({ i32, [63 x i32], i32, [2 x i32] }* [[G]] to i32), i32 4)
-  ; CHECK: [[S3:%[^ ]*]] = lshr i32 [[S2]], 2
-  ; CHECK: [[S4:%[^ ]*]] = shl i32 [[S2]], 30
+  ; CHECK: [[S2:%[^ ]*]] = sub i32 [[S1]], add (i32 ptrtoint ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]] to i32), i32 4)
+  ; CHECK: [[S3:%[^ ]*]] = lshr i32 [[S2]], 8
+  ; CHECK: [[S4:%[^ ]*]] = shl i32 [[S2]], 24
   ; CHECK: [[S5:%[^ ]*]] = or i32 [[S3]], [[S4]]
-  ; CHECK: [[S6:%[^ ]*]] = icmp ult i32 [[S5]], 64
-  ; CHECK: br i1 [[S6]]
-
-  ; CHECK: [[S8:%[^ ]*]] = zext i32 [[S5]] to i64
-  ; CHECK: [[S9:%[^ ]*]] = and i64 [[S8]], 63
-  ; CHECK: [[S10:%[^ ]*]] = shl i64 1, [[S9]]
-  ; CHECK: [[S11:%[^ ]*]] = and i64 -9223372036854775807, [[S10]]
-  ; CHECK: [[S12:%[^ ]*]] = icmp ne i64 [[S11]], 0
-
-  ; CHECK: [[S16:%[^ ]*]] = phi i1 [ false, {{%[^ ]*}} ], [ [[S12]], {{%[^ ]*}} ]
+  ; CHECK: [[S6:%[^ ]*]] = icmp ult i32 [[S5]], 2
   %x = call i1 @llvm.bitset.test(i8* %pi8, metadata !"bitset2")
-  ; CHECK: ret i1 [[S16]]
+
+  ; CHECK: ret i1 [[S6]]
   ret i1 %x
 }
 
 ; CHECK: @baz(i32* [[C0:%[^ ]*]])
 define i1 @baz(i32* %p) {
   ; CHECK: [[T0:%[^ ]*]] = bitcast i32* [[C0]] to i8*
   %pi8 = bitcast i32* %p to i8*
   ; CHECK: [[T1:%[^ ]*]] = ptrtoint i8* [[T0]] to i32
-  ; CHECK: [[T2:%[^ ]*]] = sub i32 [[T1]], ptrtoint ({ i32, [63 x i32], i32, [2 x i32] }* [[G]] to i32)
-  ; CHECK: [[T3:%[^ ]*]] = lshr i32 [[T2]], 8
-  ; CHECK: [[T4:%[^ ]*]] = shl i32 [[T2]], 24
+  ; CHECK: [[T2:%[^ ]*]] = sub i32 [[T1]], ptrtoint ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]] to i32)
+  ; CHECK: [[T3:%[^ ]*]] = lshr i32 [[T2]], 2
+  ; CHECK: [[T4:%[^ ]*]] = shl i32 [[T2]], 30
   ; CHECK: [[T5:%[^ ]*]] = or i32 [[T3]], [[T4]]
-  ; CHECK: [[T6:%[^ ]*]] = icmp ult i32 [[T5]], 2
+  ; CHECK: [[T6:%[^ ]*]] = icmp ult i32 [[T5]], 66
   ; CHECK: br i1 [[T6]]
 
-  ; CHECK: [[T8:%[^ ]*]] = and i32 [[T5]], 31
-  ; CHECK: [[T9:%[^ ]*]] = shl i32 1, [[T8]]
-  ; CHECK: [[T10:%[^ ]*]] = and i32 3, [[T9]]
-  ; CHECK: [[T11:%[^ ]*]] = icmp ne i32 [[T10]], 0
+  ; CHECK: [[T8:%[^ ]*]] = lshr i32 [[T5]], 5
+  ; CHECK: [[T9:%[^ ]*]] = getelementptr i32* bitcast ([9 x i8]* @bitset3.bits to i32*), i32 [[T8]]
+  ; CHECK: [[T10:%[^ ]*]] = load i32* [[T9]]
+  ; CHECK: [[T11:%[^ ]*]] = and i32 [[T5]], 31
+  ; CHECK: [[T12:%[^ ]*]] = shl i32 1, [[T11]]
+  ; CHECK: [[T13:%[^ ]*]] = and i32 [[T10]], [[T12]]
+  ; CHECK: [[T14:%[^ ]*]] = icmp ne i32 [[T13]], 0
 
-  ; CHECK: [[T16:%[^ ]*]] = phi i1 [ false, {{%[^ ]*}} ], [ [[T11]], {{%[^ ]*}} ]
+  ; CHECK: [[T16:%[^ ]*]] = phi i1 [ false, {{%[^ ]*}} ], [ [[T14]], {{%[^ ]*}} ]
   %x = call i1 @llvm.bitset.test(i8* %pi8, metadata !"bitset3")
   ; CHECK: ret i1 [[T16]]
   ret i1 %x
 }
 
 ; CHECK-NOT: !llvm.bitsets

test/Transforms/LowerBitSets/single-offset.ll

 ; RUN: opt -S -lowerbitsets < %s | FileCheck %s
 
 target datalayout = "e-p:32:32"
 
-; CHECK: [[G:@[^ ]*]] = private constant { i32, i32 }
+; CHECK: [[G:@[^ ]*]] = private constant { i32, [0 x i8], i32 }
 @a = constant i32 1
 @b = constant i32 2
 
 !0 = !{!"bitset1", i32* @a, i32 0}
 !1 = !{!"bitset1", i32* @b, i32 0}
 !2 = !{!"bitset2", i32* @a, i32 0}
 !3 = !{!"bitset3", i32* @b, i32 0}
 
 !llvm.bitsets = !{ !0, !1, !2, !3 }
 
 declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
 
 ; CHECK: @foo(i8* [[A0:%[^ ]*]])
 define i1 @foo(i8* %p) {
   ; CHECK: [[R0:%[^ ]*]] = ptrtoint i8* [[A0]] to i32
-  ; CHECK: [[R1:%[^ ]*]] = icmp eq i32 [[R0]], ptrtoint ({ i32, i32 }* [[G]] to i32)
+  ; CHECK: [[R1:%[^ ]*]] = icmp eq i32 [[R0]], ptrtoint ({ i32, [0 x i8], i32 }* [[G]] to i32)
   %x = call i1 @llvm.bitset.test(i8* %p, metadata !"bitset2")
   ; CHECK: ret i1 [[R1]]
   ret i1 %x
 }
 
 ; CHECK: @bar(i8* [[B0:%[^ ]*]])
 define i1 @bar(i8* %p) {
   ; CHECK: [[S0:%[^ ]*]] = ptrtoint i8* [[B0]] to i32
-  ; CHECK: [[S1:%[^ ]*]] = icmp eq i32 [[S0]], add (i32 ptrtoint ({ i32, i32 }* [[G]] to i32), i32 4)
+  ; CHECK: [[S1:%[^ ]*]] = icmp eq i32 [[S0]], add (i32 ptrtoint ({ i32, [0 x i8], i32 }* [[G]] to i32), i32 4)
   %x = call i1 @llvm.bitset.test(i8* %p, metadata !"bitset3")
   ; CHECK: ret i1 [[S1]]
   ret i1 %x
 }
 
 ; CHECK: @x(
 define i1 @x(i8* %p) {
   %x = call i1 @llvm.bitset.test(i8* %p, metadata !"bitset1")
   ret i1 %x
 }

unittests/Transforms/IPO/LowerBitSets.cpp

 TEST(LowerBitSets, BitSetBuilder) {
   struct {
     std::vector<uint64_t> Offsets;
     std::vector<uint8_t> Bits;
     uint64_t ByteOffset;
     uint64_t BitSize;
     unsigned AlignLog2;
     bool IsSingleOffset;
+    bool IsAllOnes;
   } BSBTests[] = {
-      {{}, {0}, 0, 1, 0, false},
-      {{0}, {1}, 0, 1, 0, true},
-      {{4}, {1}, 4, 1, 0, true},
-      {{37}, {1}, 37, 1, 0, true},
-      {{0, 1}, {3}, 0, 2, 0, false},
-      {{0, 4}, {3}, 0, 2, 2, false},
-      {{0, uint64_t(1) << 33}, {3}, 0, 2, 33, false},
-      {{3, 7}, {3}, 3, 2, 2, false},
-      {{0, 1, 7}, {131}, 0, 8, 0, false},
-      {{0, 2, 14}, {131}, 0, 8, 1, false},
-      {{0, 1, 8}, {3, 1}, 0, 9, 0, false},
-      {{0, 2, 16}, {3, 1}, 0, 9, 1, false},
+      {{}, {0}, 0, 1, 0, false, false},
+      {{0}, {1}, 0, 1, 0, true, true},
+      {{4}, {1}, 4, 1, 0, true, true},
+      {{37}, {1}, 37, 1, 0, true, true},
+      {{0, 1}, {3}, 0, 2, 0, false, true},
+      {{0, 4}, {3}, 0, 2, 2, false, true},
+      {{0, uint64_t(1) << 33}, {3}, 0, 2, 33, false, true},
+      {{3, 7}, {3}, 3, 2, 2, false, true},
+      {{0, 1, 7}, {131}, 0, 8, 0, false, false},
+      {{0, 2, 14}, {131}, 0, 8, 1, false, false},
+      {{0, 1, 8}, {3, 1}, 0, 9, 0, false, false},
+      {{0, 2, 16}, {3, 1}, 0, 9, 1, false, false},
+      {{0, 1, 2, 3, 4, 5, 6, 7}, {255}, 0, 8, 0, false, true},
+      {{0, 1, 2, 3, 4, 5, 6, 7, 8}, {255, 1}, 0, 9, 0, false, true},
   };
 
   for (auto &&T : BSBTests) {
     BitSetBuilder BSB;
     for (auto Offset : T.Offsets)
       BSB.addOffset(Offset);
 
     BitSetInfo BSI = BSB.build();
 
     EXPECT_EQ(T.Bits, BSI.Bits);
     EXPECT_EQ(T.ByteOffset, BSI.ByteOffset);
     EXPECT_EQ(T.BitSize, BSI.BitSize);
     EXPECT_EQ(T.AlignLog2, BSI.AlignLog2);
     EXPECT_EQ(T.IsSingleOffset, BSI.isSingleOffset());
+    EXPECT_EQ(T.IsAllOnes, BSI.isAllOnes());
 
     for (auto Offset : T.Offsets)
       EXPECT_TRUE(BSI.containsGlobalOffset(Offset));
 
     auto I = T.Offsets.begin();
     for (uint64_t NonOffset = 0; NonOffset != 256; ++NonOffset) {
       if (I != T.Offsets.end() && *I == NonOffset) {
         ++I;