[TypePromotion] Use vectors instead of sets

[samparker]

Summary

Remove the chance of non-deterministic insertion of zexts of the sources by using a SetVector, instead of SmallPtrSet. Do the same for sinks for consistency and to negate the issue from possibly happening. The SafeWrap instructions are now also stored in a SmallVector. The IRPromoter members of these structures have been changed to references.

[hans]

Looks good to me. Thanks!

llvm/lib/CodeGen/TypePromotion.cpp

 // that we can mutate these instruction and perform type promotion when it is
 // legal to do so.
 
 namespace {
 class IRPromoter {
   LLVMContext &Ctx;
   IntegerType *OrigTy = nullptr;
   unsigned PromotedWidth = 0;
+  SetVector<Value*> &Visited;
+  SetVector<Value*> &Sources;
+  SetVector<Instruction*> &Sinks;
+  SmallVectorImpl<Instruction*> &SafeWrap;
   IntegerType *ExtTy = nullptr;
   SmallPtrSet<Value*, 8> NewInsts;
   SmallPtrSet<Instruction*, 4> InstsToRemove;
   DenseMap<Value*, SmallVector<Type*, 4>> TruncTysMap;
   SmallPtrSet<Value*, 8> Promoted;
-  SetVector<Value*> *Visited;
-  SmallPtrSetImpl<Value*> *Sources;
-  SmallPtrSetImpl<Instruction*> *Sinks;
-  SmallPtrSetImpl<Instruction*> *SafeToPromote;
-  SmallPtrSetImpl<Instruction*> *SafeWrap;
 
   void ReplaceAllUsersOfWith(Value *From, Value *To);
   void PrepareWrappingAdds(void);
   void ExtendSources(void);
   void ConvertTruncs(void);
   void PromoteTree(void);
   void TruncateSinks(void);
   void Cleanup(void);
 
 public:
-  IRPromoter(LLVMContext &C, IntegerType *Ty, unsigned Width) :
-    Ctx(C), OrigTy(Ty), PromotedWidth(Width) {
+  IRPromoter(LLVMContext &C, IntegerType *Ty, unsigned Width,
+             SetVector<Value*> &visited, SetVector<Value*> &sources,
+             SetVector<Instruction*> &sinks,
+             SmallVectorImpl<Instruction*> &wrap) :
+    Ctx(C), OrigTy(Ty), PromotedWidth(Width), Visited(visited),
+    Sources(sources), Sinks(sinks), SafeWrap(wrap) {
     ExtTy = IntegerType::get(Ctx, PromotedWidth);
     assert(OrigTy->getPrimitiveSizeInBits() < ExtTy->getPrimitiveSizeInBits()
            && "Original type not smaller than extended type");
   }
 
-  void Mutate(SetVector<Value*> &Visited,
-              SmallPtrSetImpl<Value*> &Sources,
-              SmallPtrSetImpl<Instruction*> &Sinks,
-              SmallPtrSetImpl<Instruction*> &SafeToPromote,
-              SmallPtrSetImpl<Instruction*> &SafeWrap);
+  void Mutate();
 };
 
 class TypePromotion : public FunctionPass {
   unsigned TypeSize = 0;
   LLVMContext *Ctx = nullptr;
   unsigned RegisterBitWidth = 0;
   SmallPtrSet<Value*, 16> AllVisited;
   SmallPtrSet<Instruction*, 8> SafeToPromote;
-  SmallPtrSet<Instruction*, 4> SafeWrap;
+  SmallVector<Instruction*, 4> SafeWrap;
 
   // Does V have the same size result type as TypeSize.
   bool EqualTypeSize(Value *V);
   // Does V have the same size, or narrower, result type as TypeSize.
   bool LessOrEqualTypeSize(Value *V);
   // Does V have a result type that is wider than TypeSize.
   bool GreaterThanTypeSize(Value *V);
   // Does V have a result type that is narrower than TypeSize.
bool TypePromotion::isSafeWrap(Instruction *I) {
   } else if (Max.getBitWidth() > Total.getBitWidth()) {
     if (Total.zext(Max.getBitWidth()).ugt(Max))
       return false;
   } else if (Total.ugt(Max))
     return false;
 
   LLVM_DEBUG(dbgs() << "IR Promotion: Allowing safe overflow for "
              << *I << "\n");
-  SafeWrap.insert(I);
+  SafeWrap.push_back(I);
   return true;
 }
 
 bool TypePromotion::shouldPromote(Value *V) {
   if (!isa<IntegerType>(V->getType()) || isSink(V))
     return false;
 
   if (isSource(V))

 void IRPromoter::PrepareWrappingAdds() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Prepare wrapping adds.\n");
   IRBuilder<> Builder{Ctx};
 
   // For adds that safely wrap and use a negative immediate as operand 1, we
   // create an equivalent instruction using a positive immediate.
   // That positive immediate can then be zext along with all the other
   // immediates later.
-  for (auto *I : *SafeWrap) {
+  for (auto *I : SafeWrap) {
     if (I->getOpcode() != Instruction::Add)
       continue;
 
     LLVM_DEBUG(dbgs() << "IR Promotion: Adjusting " << *I << "\n");
     assert((isa<ConstantInt>(I->getOperand(1)) &&
             cast<ConstantInt>(I->getOperand(1))->isNegative()) &&
            "Wrapping should have a negative immediate as the second operand");
 
     auto Const = cast<ConstantInt>(I->getOperand(1));
     auto *NewConst = ConstantInt::get(Ctx, Const->getValue().abs());
     Builder.SetInsertPoint(I);
     Value *NewVal = Builder.CreateSub(I->getOperand(0), NewConst);
     if (auto *NewInst = dyn_cast<Instruction>(NewVal)) {
       NewInst->copyIRFlags(I);
       NewInsts.insert(NewInst);
     }
     InstsToRemove.insert(I);
     I->replaceAllUsesWith(NewVal);
     LLVM_DEBUG(dbgs() << "IR Promotion: New equivalent: " << *NewVal << "\n");
   }
   for (auto *I : NewInsts)
-    Visited->insert(I);
+    Visited.insert(I);
 }
 
 void IRPromoter::ExtendSources() {
   IRBuilder<> Builder{Ctx};
 
   auto InsertZExt = [&](Value *V, Instruction *InsertPt) {
     assert(V->getType() != ExtTy && "zext already extends to i32");
     LLVM_DEBUG(dbgs() << "IR Promotion: Inserting ZExt for " << *V << "\n");
if (auto *I = dyn_cast<Instruction>(ZExt)) {
       NewInsts.insert(I);
     }
 
     ReplaceAllUsersOfWith(V, ZExt);
   };
 
   // Now, insert extending instructions between the sources and their users.
   LLVM_DEBUG(dbgs() << "IR Promotion: Promoting sources:\n");
-  for (auto V : *Sources) {
+  for (auto V : Sources) {
     LLVM_DEBUG(dbgs() << " - " << *V << "\n");
     if (auto *I = dyn_cast<Instruction>(V))
       InsertZExt(I, I);
     else if (auto *Arg = dyn_cast<Argument>(V)) {
       BasicBlock &BB = Arg->getParent()->front();
       InsertZExt(Arg, &*BB.getFirstInsertionPt());
     } else {
       llvm_unreachable("unhandled source that needs extending");
     }
     Promoted.insert(V);
   }
 }
 
 void IRPromoter::PromoteTree() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Mutating the tree..\n");
 
   IRBuilder<> Builder{Ctx};
 
   // Mutate the types of the instructions within the tree. Here we handle
   // constant operands.
-  for (auto *V : *Visited) {
-    if (Sources->count(V))
+  for (auto *V : Visited) {
+    if (Sources.count(V))
       continue;
 
     auto *I = cast<Instruction>(V);
-    if (Sinks->count(I))
+    if (Sinks.count(I))
       continue;
 
     for (unsigned i = 0, e = I->getNumOperands(); i < e; ++i) {
       Value *Op = I->getOperand(i);
       if ((Op->getType() == ExtTy) || !isa<IntegerType>(Op->getType()))
         continue;
 
       if (auto *Const = dyn_cast<ConstantInt>(Op)) {
void IRPromoter::TruncateSinks() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Fixing up the sinks:\n");
 
   IRBuilder<> Builder{Ctx};
 
   auto InsertTrunc = [&](Value *V, Type *TruncTy) -> Instruction* {
     if (!isa<Instruction>(V) || !isa<IntegerType>(V->getType()))
       return nullptr;
 
-    if ((!Promoted.count(V) && !NewInsts.count(V)) || Sources->count(V))
+    if ((!Promoted.count(V) && !NewInsts.count(V)) || Sources.count(V))
       return nullptr;
 
     LLVM_DEBUG(dbgs() << "IR Promotion: Creating " << *TruncTy << " Trunc for "
                << *V << "\n");
     Builder.SetInsertPoint(cast<Instruction>(V));
     auto *Trunc = dyn_cast<Instruction>(Builder.CreateTrunc(V, TruncTy));
     if (Trunc)
       NewInsts.insert(Trunc);
     return Trunc;
   };
 
   // Fix up any stores or returns that use the results of the promoted
   // chain.
-  for (auto I : *Sinks) {
+  for (auto I : Sinks) {
     LLVM_DEBUG(dbgs() << "IR Promotion: For Sink: " << *I << "\n");
 
     // Handle calls separately as we need to iterate over arg operands.
     if (auto *Call = dyn_cast<CallInst>(I)) {
       for (unsigned i = 0; i < Call->getNumArgOperands(); ++i) {
         Value *Arg = Call->getArgOperand(i);
         Type *Ty = TruncTysMap[Call][i];
         if (Instruction *Trunc = InsertTrunc(Arg, Ty)) {
for (auto I : Sinks) {
     }
   }
 }
 
 void IRPromoter::Cleanup() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Cleanup..\n");
   // Some zexts will now have become redundant, along with their trunc
   // operands, so remove them
-  for (auto V : *Visited) {
+  for (auto V : Visited) {
     if (!isa<ZExtInst>(V))
       continue;
 
     auto ZExt = cast<ZExtInst>(V);
     if (ZExt->getDestTy() != ExtTy)
       continue;
 
     Value *Src = ZExt->getOperand(0);
for (auto *I : InstsToRemove) {
     I->eraseFromParent();
   }
 }
 
 void IRPromoter::ConvertTruncs() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Converting truncs..\n");
   IRBuilder<> Builder{Ctx};
 
-  for (auto *V : *Visited) {
-    if (!isa<TruncInst>(V) || Sources->count(V))
+  for (auto *V : Visited) {
+    if (!isa<TruncInst>(V) || Sources.count(V))
       continue;
 
     auto *Trunc = cast<TruncInst>(V);
     Builder.SetInsertPoint(Trunc);
     IntegerType *SrcTy = cast<IntegerType>(Trunc->getOperand(0)->getType());
     IntegerType *DestTy = cast<IntegerType>(TruncTysMap[Trunc][0]);
 
     unsigned NumBits = DestTy->getScalarSizeInBits();
     ConstantInt *Mask =
       ConstantInt::get(SrcTy, APInt::getMaxValue(NumBits).getZExtValue());
     Value *Masked = Builder.CreateAnd(Trunc->getOperand(0), Mask);
 
     if (auto *I = dyn_cast<Instruction>(Masked))
       NewInsts.insert(I);
 
     ReplaceAllUsersOfWith(Trunc, Masked);
   }
 }
 
-void IRPromoter::Mutate(SetVector<Value*> &Visited,
-                        SmallPtrSetImpl<Value*> &Sources,
-                        SmallPtrSetImpl<Instruction*> &Sinks,
-                        SmallPtrSetImpl<Instruction*> &SafeToPromote,
-                        SmallPtrSetImpl<Instruction*> &SafeWrap) {
+void IRPromoter::Mutate() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Promoting use-def chains from "
              << OrigTy->getBitWidth() << " to " << PromotedWidth << "-bits\n");
 
-  this->Visited = &Visited;
-  this->Sources = &Sources;
-  this->Sinks = &Sinks;
-  this->SafeToPromote = &SafeToPromote;
-  this->SafeWrap = &SafeWrap;
-
   // Cache original types of the values that will likely need truncating
   for (auto *I : Sinks) {
     if (auto *Call = dyn_cast<CallInst>(I)) {
       for (unsigned i = 0; i < Call->getNumArgOperands(); ++i) {
         Value *Arg = Call->getArgOperand(i);
         TruncTysMap[Call].push_back(Arg->getType());
       }
     } else if (auto *Switch = dyn_cast<SwitchInst>(I))
bool TypePromotion::TryToPromote(Value *V, unsigned PromotedWidth) {
 
   if (!isSupportedValue(V) || !shouldPromote(V) || !isLegalToPromote(V))
     return false;
 
   LLVM_DEBUG(dbgs() << "IR Promotion: TryToPromote: " << *V << ", from "
              << TypeSize << " bits to " << PromotedWidth << "\n");
 
   SetVector<Value*> WorkList;
-  SmallPtrSet<Value*, 8> Sources;
-  SmallPtrSet<Instruction*, 4> Sinks;
+  SetVector<Value*> Sources;
+  SetVector<Instruction*> Sinks;
   SetVector<Value*> CurrentVisited;
   WorkList.insert(V);
 
   // Return true if V was added to the worklist as a supported instruction,
   // if it was already visited, or if we don't need to explore it (e.g.
   // pointer values and GEPs), and false otherwise.
   auto AddLegalInst = [&](Value *V) {
     if (CurrentVisited.count(V))
bool TypePromotion::TryToPromote(Value *V, unsigned PromotedWidth) {
   // DAG optimisations should be able to handle these cases better, especially
   // for function arguments.
   if (ToPromote < 2 || (Blocks.size() == 1 && (NonFreeArgs > SafeWrap.size())))
     return false;
 
   if (ToPromote < 2)
     return false;
 
-  IRPromoter Promoter(*Ctx, cast<IntegerType>(OrigTy), PromotedWidth);
-  Promoter.Mutate(CurrentVisited, Sources, Sinks, SafeToPromote, SafeWrap);
+  IRPromoter Promoter(*Ctx, cast<IntegerType>(OrigTy), PromotedWidth,
+                      CurrentVisited, Sources, Sinks, SafeWrap);
+  Promoter.Mutate();
   return true;
 }
 
 bool TypePromotion::runOnFunction(Function &F) {
   if (skipFunction(F) || DisablePromotion)
     return false;
 
   LLVM_DEBUG(dbgs() << "IR Promotion: Running on " << F.getName() << "\n");