[BypassSlowDivision] Refactor fast division insertion logic (NFC)

[n.bozhenov]

Summary

The most important goal of the patch is to break large insertFastDiv function into separate pieces, so that later a different fast insertion logic can be implemented using some of these pieces.

[jlebar]

lib/Transforms/Utils/BypassSlowDivision.cpp

139 ↗

(On Diff #88213)

Can we add a "// anonymous namespace" comment?

346 ↗

(On Diff #88213)

I'm totally behind the notion of creating classes that encapsulate our temporary state. But I'm not wild about the way we use FastDivInsertionTask as basically a bucket of *mutable* state. We have had really nasty bugs in LLVM with similar designs where we forget to reset one piece of the mutable state. (I think Chandler was saying that some pass effectively had a global cost cap instead of a per-function cost cap, because they just forgot to set the cost to 0 at the end of the function.)

I went through a few iterations, and the one that seems best to me right now is to get rid of FastDivInserter, which is already very simple, and make FastDivInsertionTask a one-shot class.

DivCacheTy Cache;
while (Next != nullptr) {
  Instruction *I = Next;
  Next = Next->getNextNode();

  if (Value *Replacement = FastDivInsertionTask(I, Cache, BypassWidths).getReplacement()) {
    I->RAUW(Replacement);
    MadeChange |= true;
  }
}

WDYT?

[jlebar]

lib/Transforms/Utils/BypassSlowDivision.cpp

80 ↗

(On Diff #88213)

This is actually SlowInstr or InstrToReplace or SlowDivOrRem -- we should be careful here and elsewhere not to call things divs that may be rems.

92 ↗

(On Diff #88213)

I would feel a lot more comfortable if our constructor initialized all of these variables. Otherwise it seems like we're asking for trouble.

98 ↗

(On Diff #88213)

"Runtime" as in "runtime check" is usually one word.

[n.bozhenov]

lib/Transforms/Utils/BypassSlowDivision.cpp

346 ↗

(On Diff #88213)

Hi Justin,

Thanks for reviewing the code so quickly. And I really like the code sample you suggested. If I understand correctly, your main concern is that in my code the Task object gets reused for different instructions during the pass, isn't it? If this is the case, then moving the code from FastDivInserter into FastDivInsertionTask constructor is indeed one of the possible solutions.

Another approach is to create a new Task object for each instruction inside an FastDivInserter object and never expose it outside. Something like this:

// In bypassSlowDivision:
FastDivInserter FDI;
while (Next != nullptr) {
  Instruction *I = Next;
  Next = Next->getNextNode();
  MadeChange |= FDI.tryReplaceSlowDiv(I);
}

// In FastDivInserter::tryReplaceSlowDiv:
FastDivInsertionTask Task(I, BypassWidths);
if (!Task.isSlowDivision())
  return false;

Value *R = Task.getReplacement(Cache);
if (!R)
  return false;

I->RAUW(R);
return true;

Currently I'm playing with these two approaches and will come with an updated patch by tomorrow. Please, comment if you're strongly in favour of one of these approaches.

[jlebar]

lib/Transforms/Utils/BypassSlowDivision.cpp

346 ↗

(On Diff #88213)

If I understand correctly, your main concern is that in my code the Task object gets reused for different instructions during the pass, isn't it?

Yes.

I'm also a little concerned about the mutable state inside of the Task object used for passing between functions, but I wanted to get this part figured out first.

Another approach is to create a new Task object for each instruction inside an FastDivInserter object and never expose it outside

Right. If FastDivInserter has nontrivial complexity that it doesn't make sense to push into FastDivInsertionTask or the outer pass function body, this might make sense. When I looked at it, I didn't think it did.

[n.bozhenov]

lib/Transforms/Utils/BypassSlowDivision.cpp

346 ↗

(On Diff #88213)

I have updated the patch. I have got rid of FastDivInserter class. And the Task is not reused any more, but a new object is constructed for each instruction.

As for mutable fields used for passing data between functions, it's possible to get rid of them as well. We could define

struct IncomingDivRemPair {
  BasicBlock *BB; //< PHINode predecessor for the following values.
  Value *Quotient;
  Value *Remainder;
};

and return such a structure by value from createSlowBB and createFastBB methods. Later, we could pass a pair of such structures into createDivRemPhiNodes to make the latter both more flexible and easier to understand. The only reason I haven't done this yet is that I don't like very much returning structures by value. Do you think it is worth doing?

[n.bozhenov]

lib/Transforms/Utils/BypassSlowDivision.cpp

80 ↗

(On Diff #88213)

During this optimization we generally use 'slow division' to refer to both Div and Rem operations. I believe that strictly following your suggestion would be overkill. For example, the file itself and its main entry point are named BypassSlowDivision.cpp and llvm::bypassSlowDivision. I'm not sure if we should rename them into something like bypassSlowDivisionOrRemainder.

To make the naming somewhat less confusing I suggest using the full word (Division) to refer to both Div and Rem operations, and use the abbreviations (Div/Rem) to refer to specific operation kinds. According to that, SlowDiv should be renamed into SlowDivision, FastDivInsertionTask into FastDivisionInsertionTask, and isDivisionOp into isDivOp. What do you think?

[jlebar]

lib/Transforms/Utils/BypassSlowDivision.cpp

80 ↗

(On Diff #88552)

Do we really need this field, since it's just SlowDiv->getOpcode()?

82 ↗

(On Diff #88552)

Similarly, do we need this field, since it can be derived from SlowDiv?

You could have a member function for it if it's annoying to get N separate times. I think a member function would be better because it makes it explicit that this is SlowDiv->getType().

84 ↗

(On Diff #88552)

Similarly for the MainBB field.

86 ↗

(On Diff #88552)

Similarly for the Dividend and Divisor fields.

98 ↗

(On Diff #88552)

FastBB and SlowBB, to match the functions right above?

118 ↗

(On Diff #88552)

Can some of these be made private and/or deleted now?

126 ↗

(On Diff #88552)

Consider using inline initializers for these, otherwise it's easy to add a member variable and forget to initialize it.

class Foo {
  void* ptr = nullptr;
};

128 ↗

(On Diff #88552)

Can we do this in the initialization list?

146 ↗

(On Diff #88552)

Perhaps

// Skip division on vector types; only optimize integer instructions.
auto* SlowType = dyn_cast<IntegerType>(I->getType());
if (!SlowType) return;
auto BI = BypassWidths.find(SlowType->getBitWidth());

(Make SlowType a local variable, use dyn_cast instead of checking twice, and use auto for the iterator type. Also get rid of the comma splice in the comment and add a period.)

148 ↗

(On Diff #88552)

Nit, please add periods at the end of sentences.

166 ↗

(On Diff #88552)

previously-computed

177 ↗

(On Diff #88552)

auto

181 ↗

(On Diff #88552)

It seems like much less of a layering violation to me if insertFastDiv (possibly renamed to createFastDivOrRem) would return Optional<std::pair<Value*, Value*>>. Then we would update the cache here.

185 ↗

(On Diff #88552)

s/Cache/cache/

196 ↗

(On Diff #88552)

slow div and rem operations

214 ↗

(On Diff #88552)

fast div and rem operations

224 ↗

(On Diff #88552)

While we're here, probably should add missing word: "because *this* optimization only handles positive numbers"

253 ↗

(On Diff #88552)

"Creates a runtime check to test whether both the divisor and dividend fit into BypassType"?

255 ↗

(On Diff #88552)

Nit, a linebreak without any intervening whitespace like we have here isn't a meaningful punctuation. Please flow as a single paragraph, or split into two paragraphs by inserting a blank line.

282 ↗

(On Diff #88552)

This sentence runs on. We could remove "and the longer-slower div/rem instruction otherwise." Otherwise, can you split into two sentences?

295 ↗

(On Diff #88552)

Maybe "Split the basic block before the div/rem."

297 ↗

(On Diff #88552)

This line could use a comment, I think.

306 ↗

(On Diff #88552)

Perhaps we should update this comment.

324 ↗

(On Diff #88552)

I know I wrote it this way initially, but now that I see it...MadeChange = true? :)

80 ↗

(On Diff #88213)

During this optimization we generally use 'slow division' to refer to both Div and Rem operations. I believe that strictly following your suggestion would be overkill. For example, the file itself and its main entry point are named BypassSlowDivision.cpp and llvm::bypassSlowDivision. I'm not sure if we should rename them into something like bypassSlowDivisionOrRemainder.

I agree that renaming the file and pass would be overkill. But I think we can still improve readability by being just a little less strict about the suggestion.

Even if we had a comment explaining that "division" == "div or mod" at the top of the file, that would be us inventing a new naming idiom, which is going to add to the cognitive cost to users of understanding our code. And realistically speaking, a large fraction of people aren't going to notice this comment at all, in which case they're just going to be lost.

In contrast, SlowDivOrRem is the same number of characters as SlowDivision, but is unambiguous. I'd also be OK with SlowInstr if you prefer that.

Again I don't think we need to be totally strict about the rule. I don't think that FastDivInsertionTask is ambiguous, for example, although perhaps a better name is FastDivRemInsertionTask would more closely match what it does (it always inserts both instrs). But the SlowDiv member here is, I think, ambiguous, because some of the members here *are* necessarily divs. Similarly, I think createDivRuntimeCheck would be less ambiguous as insertOperandRuntimeCheck or maybe something like insertFastSlowBranch, and insertFastDiv would be better named insertFastDivAndRem.

346 ↗

(On Diff #88213)

Oh yes, I like this much better.

The only reason I haven't done this yet is that I don't like very much returning structures by value.

If it makes you feel any better, the calling convention returns large structs "by outparam". :)

[n.bozhenov]

lib/Transforms/Utils/BypassSlowDivision.cpp

84 ↗

(On Diff #88552)

Things get more difficult with MainBB. After splitting the basic block it is not obvious any more what basic block SlowDivOrRem->getParent() will return. So, I believe it's better to keep MainBB member in the class.

181 ↗

(On Diff #88552)

Right. I have moved all the caching logic into this routine. As a result, I could get rid of Cache field in the class. Now it's only passed here as a method argument.

306 ↗

(On Diff #88552)

Not sure what's wrong with the comment.

[jlebar]

Sorry to keep going back and forth on this. I'm almost happy, but still have some meaty comments.

lib/Transforms/Utils/BypassSlowDivision.cpp

94 ↗

(On Diff #88939)

I think we could still be more explicit about which variables have to do with division specifically and which have to do with div/rem. Right now it's confusing that IsSlowDivision actually means "should we try to generate a fast div/rem for SlowDivOrRem?". Similarly it's confusing that isSignedDiv means "is SlowDivOrRem a signed div or rem?".

Maybe rename "IsSlowDivision" to "IsValid" or "IsValidTask"? And isSignedDiv to isSignedOp to match isDivisionOp?

106 ↗

(On Diff #88939)

isSignedOp()?

174 ↗

(On Diff #88939)

It's kind of confusing that DivRemResult is "nullable". While the result of insertFastDivAndRem is nullable, the value inside of the cache must not be null.

Seems like it would be clearer if we made DivRemResult not nullable and then returned an optional<DivRemResult> from insertFastDivAndRem.

181 ↗

(On Diff #88939)

Perhaps just

CacheI = Cache.Insert({Key, DivRem}).first

? Not sure we need the comment if you write it like this, either.

188 ↗

(On Diff #88939)

Nit, perhaps clearer as a ternary:

return isDivisionOp() ? Value.Quotient : Value.Remainder;

193 ↗

(On Diff #88939)

"IncomingDivRemPair" is kind a confusing name, because the values are "outgoing" from one block and "incoming" to another.

It's also subtly different from DivRemResult.

It's going to be true that at least one of the Value*s inside our IncomingDivRemPair is an Instruction, right? If so, maybe we should

• Rename DivRemResult to DivRemPair
• Get rid of IncomingDivRemPair
• Add a DivRemPair::basicBlock() method that gets the BB from its Values (assuming that at least one is an Instruction).

84 ↗

(On Diff #88552)

sgtm. It's a lot more clear now.

306 ↗

(On Diff #88552)

I could have sworn it used to say "identifies division instructions", but now it says "DIV/REM" instructions, which is all I wanted. We're good here.

[n.bozhenov]

lib/Transforms/Utils/BypassSlowDivision.cpp

193 ↗

(On Diff #88939)

I do agree that DivRemResult and IncomingDivRemPair have something in common and it's tempting to get rid of one of them. On the other hand, they serve different purposes and have no complex logic to share. As a result, I believe these two structures should be kept separate one from another. Not sure about if their names are perfect, though.

DivRemResult is a type that is used to cache results of bypassing. Usually, it is a pair of PHINodes. However, in D29897 there is a special case when it is a pair of ZExt instructions.

IncomingDivRemPair structures are used to create two PHINodes in createDivRemPhiNodes. Each of the IncomingDivRemPair arguments describes one of the incoming values for the PHINodes. Again, in D29897 there is a special case when Quotient = Zero and Remainder = Dividend. In this case it would be impossible to derive information about PHINode predecessor from DivRemPair. So, basic block information should be explicitely included somehow into createDivRemPhiNodes arguments.

Moreover, DivOpInfo and DivRemResult are auxilliary types for caching only. They contain only information significant for caching and it is extremely unlikely that any of these types will be extended in future. On the other hand, one may want to pass some additional information to createDivRemPhiNodes in future (e.g. some sort of metadata) and add additional fields to IncomingDivRemPair. It is not a problem only if DivRemResult and IncomingDivRemPair are two different types.

[jlebar]

lib/Transforms/Utils/BypassSlowDivision.cpp

45 ↗

(On Diff #89065)

Now we can remove this constructor, right?

172 ↗

(On Diff #89065)

It's not really a "div rem" -- it's the *result* of the div and rem, which is a quotient plus a remainder. Maybe just QuotRem or Result?

284 ↗

(On Diff #89065)

return None here and elsewhere.

193 ↗

(On Diff #88939)

Okay, how about we just name them based on what they contain, QuotRemPair and QuotRemWithBB?

[n.bozhenov]

lib/Transforms/Utils/BypassSlowDivision.cpp

45 ↗

(On Diff #89065)

Yep.

172 ↗

(On Diff #89065)

Makes sense.

284 ↗

(On Diff #89065)

Right. Thanks.

193 ↗

(On Diff #88939)

Definitely, such names have the advantage of being absolutely unambiguous.

[jlebar]

lgtm with comments updated. \o/

lib/Transforms/Utils/BypassSlowDivision.cpp

40 ↗

(On Diff #89133)

Please update comment. This is now just a pair, it's not the result of bypassing.

51 ↗

(On Diff #89133)

Please update comment. This is now just two values plus a BB. If you want, you can say what the BB is supposed to represent. We should also say how this should be used, rather than describing how it is currently used in the algorithm. Perhaps:

A quotient and remainder, plus a BB from which they logically "originate". If you use Quotient or Remainder in a Phi node, you should use BB as its corresponding predecessor.

[n.bozhenov]

In D29896#681722, @jlebar wrote:

lgtm with comments updated. \o/

Thank you very much for reviewing the patch and very helpful suggestions.

llvm/trunk/lib/Transforms/Utils/BypassSlowDivision.cpp

struct DivOpInfo {
     bool SignedOp;
     Value *Dividend;
     Value *Divisor;
 
     DivOpInfo(bool InSignedOp, Value *InDividend, Value *InDivisor)
       : SignedOp(InSignedOp), Dividend(InDividend), Divisor(InDivisor) {}
   };
 
-  struct DivPhiNodes {
-    PHINode *Quotient;
-    PHINode *Remainder;
+  struct QuotRemPair {
+    Value *Quotient;
+    Value *Remainder;
 
-    DivPhiNodes(PHINode *InQuotient, PHINode *InRemainder)
+    QuotRemPair(Value *InQuotient, Value *InRemainder)
         : Quotient(InQuotient), Remainder(InRemainder) {}
   };
+
+  /// A quotient and remainder, plus a BB from which they logically "originate".
+  /// If you use Quotient or Remainder in a Phi node, you should use BB as its
+  /// corresponding predecessor.
+  struct QuotRemWithBB {
+    BasicBlock *BB = nullptr;
+    Value *Quotient = nullptr;
+    Value *Remainder = nullptr;
+  };
 }
 
 namespace llvm {
   template<>
   struct DenseMapInfo<DivOpInfo> {
     static bool isEqual(const DivOpInfo &Val1, const DivOpInfo &Val2) {
       return Val1.SignedOp == Val2.SignedOp &&
              Val1.Dividend == Val2.Dividend &&
struct DenseMapInfo<DivOpInfo> {
 
     static unsigned getHashValue(const DivOpInfo &Val) {
       return (unsigned)(reinterpret_cast<uintptr_t>(Val.Dividend) ^
                         reinterpret_cast<uintptr_t>(Val.Divisor)) ^
                         (unsigned)Val.SignedOp;
     }
   };
 
-  typedef DenseMap<DivOpInfo, DivPhiNodes> DivCacheTy;
+  typedef DenseMap<DivOpInfo, QuotRemPair> DivCacheTy;
+  typedef DenseMap<unsigned, unsigned> BypassWidthsTy;
 }
 
-// insertFastDiv - Substitutes the div/rem instruction with code that checks the
-// value of the operands and uses a shorter-faster div/rem instruction when
-// possible and the longer-slower div/rem instruction otherwise.
-static bool insertFastDiv(Instruction *I, IntegerType *BypassType,
-                          bool UseDivOp, bool UseSignedOp,
-                          DivCacheTy &PerBBDivCache) {
-  Function *F = I->getParent()->getParent();
-  // Get instruction operands
-  Value *Dividend = I->getOperand(0);
-  Value *Divisor = I->getOperand(1);
+namespace {
+class FastDivInsertionTask {
+  bool IsValidTask = false;
+  Instruction *SlowDivOrRem = nullptr;
+  IntegerType *BypassType = nullptr;
+  BasicBlock *MainBB = nullptr;
+
+  QuotRemWithBB createSlowBB(BasicBlock *Successor);
+  QuotRemWithBB createFastBB(BasicBlock *Successor);
+  QuotRemPair createDivRemPhiNodes(QuotRemWithBB &LHS, QuotRemWithBB &RHS,
+                                   BasicBlock *PhiBB);
+  Value *insertOperandRuntimeCheck();
+  Optional<QuotRemPair> insertFastDivAndRem();
+
+  bool isSignedOp() {
+    return SlowDivOrRem->getOpcode() == Instruction::SDiv ||
+           SlowDivOrRem->getOpcode() == Instruction::SRem;
+  }
+  bool isDivisionOp() {
+    return SlowDivOrRem->getOpcode() == Instruction::SDiv ||
+           SlowDivOrRem->getOpcode() == Instruction::UDiv;
+  }
+  Type *getSlowType() { return SlowDivOrRem->getType(); }
+
+public:
+  FastDivInsertionTask(Instruction *I, const BypassWidthsTy &BypassWidths);
+  Value *getReplacement(DivCacheTy &Cache);
+};
+} // anonymous namespace
 
-  if (isa<ConstantInt>(Divisor)) {
-    // Division by a constant should have been been solved and replaced earlier
-    // in the pipeline.
-    return false;
-  }
+FastDivInsertionTask::FastDivInsertionTask(Instruction *I,
+                                           const BypassWidthsTy &BypassWidths) {
+  switch (I->getOpcode()) {
+  case Instruction::UDiv:
+  case Instruction::SDiv:
+  case Instruction::URem:
+  case Instruction::SRem:
+    SlowDivOrRem = I;
+    break;
+  default:
+    // I is not a div/rem operation.
+    return;
+  }
+
+  // Skip division on vector types. Only optimize integer instructions.
+  IntegerType *SlowType = dyn_cast<IntegerType>(SlowDivOrRem->getType());
+  if (!SlowType)
+    return;
 
-  // If the numerator is a constant, bail if it doesn't fit into BypassType.
-  if (ConstantInt *ConstDividend = dyn_cast<ConstantInt>(Dividend))
-    if (ConstDividend->getValue().getActiveBits() > BypassType->getBitWidth())
-      return false;
+  // Skip if this bitwidth is not bypassed.
+  auto BI = BypassWidths.find(SlowType->getBitWidth());
+  if (BI == BypassWidths.end())
+    return;
 
-  // Basic Block is split before divide
-  BasicBlock *MainBB = &*I->getParent();
-  BasicBlock *SuccessorBB = MainBB->splitBasicBlock(I);
-
-  // Add new basic block for slow divide operation
-  BasicBlock *SlowBB =
-      BasicBlock::Create(F->getContext(), "", MainBB->getParent(), SuccessorBB);
-  SlowBB->moveBefore(SuccessorBB);
-  IRBuilder<> SlowBuilder(SlowBB, SlowBB->begin());
-  Value *SlowQuotientV;
-  Value *SlowRemainderV;
-  if (UseSignedOp) {
-    SlowQuotientV = SlowBuilder.CreateSDiv(Dividend, Divisor);
-    SlowRemainderV = SlowBuilder.CreateSRem(Dividend, Divisor);
+  // Get type for div/rem instruction with bypass bitwidth.
+  IntegerType *BT = IntegerType::get(I->getContext(), BI->second);
+  BypassType = BT;
+
+  // The original basic block.
+  MainBB = I->getParent();
+
+  // The instruction is indeed a slow div or rem operation.
+  IsValidTask = true;
+}
+
+/// Reuses previously-computed dividend or remainder from the current BB if
+/// operands and operation are identical. Otherwise calls insertFastDivAndRem to
+/// perform the optimization and caches the resulting dividend and remainder.
+/// If no replacement can be generated, nullptr is returned.
+Value *FastDivInsertionTask::getReplacement(DivCacheTy &Cache) {
+  // First, make sure that the task is valid.
+  if (!IsValidTask)
+    return nullptr;
+
+  // Then, look for a value in Cache.
+  Value *Dividend = SlowDivOrRem->getOperand(0);
+  Value *Divisor = SlowDivOrRem->getOperand(1);
+  DivOpInfo Key(isSignedOp(), Dividend, Divisor);
+  auto CacheI = Cache.find(Key);
+
+  if (CacheI == Cache.end()) {
+    // If previous instance does not exist, try to insert fast div.
+    Optional<QuotRemPair> OptResult = insertFastDivAndRem();
+    // Bail out if insertFastDivAndRem has failed.
+    if (!OptResult)
+      return nullptr;
+    CacheI = Cache.insert({Key, *OptResult}).first;
+  }
+
+  QuotRemPair &Value = CacheI->second;
+  return isDivisionOp() ? Value.Quotient : Value.Remainder;
+}
+
+/// Add new basic block for slow div and rem operations and put it before
+/// SuccessorBB.
+QuotRemWithBB FastDivInsertionTask::createSlowBB(BasicBlock *SuccessorBB) {
+  QuotRemWithBB DivRemPair;
+  DivRemPair.BB = BasicBlock::Create(MainBB->getParent()->getContext(), "",
+                                     MainBB->getParent(), SuccessorBB);
+  IRBuilder<> Builder(DivRemPair.BB, DivRemPair.BB->begin());
+
+  Value *Dividend = SlowDivOrRem->getOperand(0);
+  Value *Divisor = SlowDivOrRem->getOperand(1);
+
+  if (isSignedOp()) {
+    DivRemPair.Quotient = Builder.CreateSDiv(Dividend, Divisor);
+    DivRemPair.Remainder = Builder.CreateSRem(Dividend, Divisor);
   } else {
-    SlowQuotientV = SlowBuilder.CreateUDiv(Dividend, Divisor);
-    SlowRemainderV = SlowBuilder.CreateURem(Dividend, Divisor);
+    DivRemPair.Quotient = Builder.CreateUDiv(Dividend, Divisor);
+    DivRemPair.Remainder = Builder.CreateURem(Dividend, Divisor);
   }
-  SlowBuilder.CreateBr(SuccessorBB);
 
-  // Add new basic block for fast divide operation
-  BasicBlock *FastBB =
-      BasicBlock::Create(F->getContext(), "", MainBB->getParent(), SuccessorBB);
-  FastBB->moveBefore(SlowBB);
-  IRBuilder<> FastBuilder(FastBB, FastBB->begin());
-  Value *ShortDivisorV = FastBuilder.CreateCast(Instruction::Trunc, Divisor,
-                                                BypassType);
-  Value *ShortDividendV = FastBuilder.CreateCast(Instruction::Trunc, Dividend,
-                                                 BypassType);
-
-  // udiv/urem because optimization only handles positive numbers
-  Value *ShortQuotientV = FastBuilder.CreateUDiv(ShortDividendV, ShortDivisorV);
-  Value *ShortRemainderV = FastBuilder.CreateURem(ShortDividendV,
-                                                  ShortDivisorV);
-  Value *FastQuotientV = FastBuilder.CreateCast(Instruction::ZExt,
-                                                ShortQuotientV,
-                                                Dividend->getType());
-  Value *FastRemainderV = FastBuilder.CreateCast(Instruction::ZExt,
-                                                 ShortRemainderV,
-                                                 Dividend->getType());
-  FastBuilder.CreateBr(SuccessorBB);
-
-  // Phi nodes for result of div and rem
-  IRBuilder<> SuccessorBuilder(SuccessorBB, SuccessorBB->begin());
-  PHINode *QuoPhi = SuccessorBuilder.CreatePHI(I->getType(), 2);
-  QuoPhi->addIncoming(SlowQuotientV, SlowBB);
-  QuoPhi->addIncoming(FastQuotientV, FastBB);
-  PHINode *RemPhi = SuccessorBuilder.CreatePHI(I->getType(), 2);
-  RemPhi->addIncoming(SlowRemainderV, SlowBB);
-  RemPhi->addIncoming(FastRemainderV, FastBB);
-
-  // Replace I with appropriate phi node
-  if (UseDivOp)
-    I->replaceAllUsesWith(QuoPhi);
-  else
-    I->replaceAllUsesWith(RemPhi);
-  I->eraseFromParent();
+  Builder.CreateBr(SuccessorBB);
+  return DivRemPair;
+}
 
-  // Combine operands into a single value with OR for value testing below
-  MainBB->getInstList().back().eraseFromParent();
-  IRBuilder<> MainBuilder(MainBB, MainBB->end());
+/// Add new basic block for fast div and rem operations and put it before
+/// SuccessorBB.
+QuotRemWithBB FastDivInsertionTask::createFastBB(BasicBlock *SuccessorBB) {
+  QuotRemWithBB DivRemPair;
+  DivRemPair.BB = BasicBlock::Create(MainBB->getParent()->getContext(), "",
+                                     MainBB->getParent(), SuccessorBB);
+  IRBuilder<> Builder(DivRemPair.BB, DivRemPair.BB->begin());
+
+  Value *Dividend = SlowDivOrRem->getOperand(0);
+  Value *Divisor = SlowDivOrRem->getOperand(1);
+  Value *ShortDivisorV =
+      Builder.CreateCast(Instruction::Trunc, Divisor, BypassType);
+  Value *ShortDividendV =
+      Builder.CreateCast(Instruction::Trunc, Dividend, BypassType);
+
+  // udiv/urem because this optimization only handles positive numbers.
+  Value *ShortQV = Builder.CreateUDiv(ShortDividendV, ShortDivisorV);
+  Value *ShortRV = Builder.CreateURem(ShortDividendV, ShortDivisorV);
+  DivRemPair.Quotient =
+      Builder.CreateCast(Instruction::ZExt, ShortQV, getSlowType());
+  DivRemPair.Remainder =
+      Builder.CreateCast(Instruction::ZExt, ShortRV, getSlowType());
+  Builder.CreateBr(SuccessorBB);
+
+  return DivRemPair;
+}
+
+/// Creates Phi nodes for result of Div and Rem.
+QuotRemPair FastDivInsertionTask::createDivRemPhiNodes(QuotRemWithBB &LHS,
+                                                       QuotRemWithBB &RHS,
+                                                       BasicBlock *PhiBB) {
+  IRBuilder<> Builder(PhiBB, PhiBB->begin());
+  PHINode *QuoPhi = Builder.CreatePHI(getSlowType(), 2);
+  QuoPhi->addIncoming(LHS.Quotient, LHS.BB);
+  QuoPhi->addIncoming(RHS.Quotient, RHS.BB);
+  PHINode *RemPhi = Builder.CreatePHI(getSlowType(), 2);
+  RemPhi->addIncoming(LHS.Remainder, LHS.BB);
+  RemPhi->addIncoming(RHS.Remainder, RHS.BB);
+  return QuotRemPair(QuoPhi, RemPhi);
+}
+
+/// Creates a runtime check to test whether both the divisor and dividend fit
+/// into BypassType. The check is inserted at the end of MainBB. True return
+/// value means that the operands fit.
+Value *FastDivInsertionTask::insertOperandRuntimeCheck() {
+  IRBuilder<> Builder(MainBB, MainBB->end());
+  Value *Dividend = SlowDivOrRem->getOperand(0);
+  Value *Divisor = SlowDivOrRem->getOperand(1);
 
   // We should have bailed out above if the divisor is a constant, but the
   // dividend may still be a constant.  Set OrV to our non-constant operands
   // OR'ed together.
   assert(!isa<ConstantInt>(Divisor));
 
   Value *OrV;
   if (!isa<ConstantInt>(Dividend))
-    OrV = MainBuilder.CreateOr(Dividend, Divisor);
+    OrV = Builder.CreateOr(Dividend, Divisor);
   else
     OrV = Divisor;
 
   // BitMask is inverted to check if the operands are
   // larger than the bypass type
   uint64_t BitMask = ~BypassType->getBitMask();
-  Value *AndV = MainBuilder.CreateAnd(OrV, BitMask);
+  Value *AndV = Builder.CreateAnd(OrV, BitMask);
 
-  // Compare operand values and branch
-  Value *ZeroV = ConstantInt::getSigned(Dividend->getType(), 0);
-  Value *CmpV = MainBuilder.CreateICmpEQ(AndV, ZeroV);
-  MainBuilder.CreateCondBr(CmpV, FastBB, SlowBB);
-
-  // Cache phi nodes to be used later in place of other instances
-  // of div or rem with the same sign, dividend, and divisor
-  DivOpInfo Key(UseSignedOp, Dividend, Divisor);
-  DivPhiNodes Value(QuoPhi, RemPhi);
-  PerBBDivCache.insert(std::pair<DivOpInfo, DivPhiNodes>(Key, Value));
-  return true;
-}
-
-// reuseOrInsertFastDiv - Reuses previously computed dividend or remainder from
-// the current BB if operands and operation are identical. Otherwise calls
-// insertFastDiv to perform the optimization and caches the resulting dividend
-// and remainder.
-static bool reuseOrInsertFastDiv(Instruction *I, IntegerType *BypassType,
-                                 bool UseDivOp, bool UseSignedOp,
-                                 DivCacheTy &PerBBDivCache) {
-  // Get instruction operands
-  DivOpInfo Key(UseSignedOp, I->getOperand(0), I->getOperand(1));
-  DivCacheTy::iterator CacheI = PerBBDivCache.find(Key);
-
-  if (CacheI == PerBBDivCache.end()) {
-    // If previous instance does not exist, insert fast div
-    return insertFastDiv(I, BypassType, UseDivOp, UseSignedOp, PerBBDivCache);
-  }
-
-  // Replace operation value with previously generated phi node
-  DivPhiNodes &Value = CacheI->second;
-  if (UseDivOp) {
-    // Replace all uses of div instruction with quotient phi node
-    I->replaceAllUsesWith(Value.Quotient);
-  } else {
-    // Replace all uses of rem instruction with remainder phi node
-    I->replaceAllUsesWith(Value.Remainder);
-  }
+  // Compare operand values
+  Value *ZeroV = ConstantInt::getSigned(getSlowType(), 0);
+  return Builder.CreateICmpEQ(AndV, ZeroV);
 
-  // Remove redundant operation
-  I->eraseFromParent();
-  return true;
+/// Substitutes the div/rem instruction with code that checks the value of the
+/// operands and uses a shorter-faster div/rem instruction when possible.
+Optional<QuotRemPair> FastDivInsertionTask::insertFastDivAndRem() {
+  Value *Dividend = SlowDivOrRem->getOperand(0);
+  Value *Divisor = SlowDivOrRem->getOperand(1);
+
+  if (isa<ConstantInt>(Divisor)) {
+    // Keep division by a constant for DAGCombiner.
+    return None;
+  }
+
+  // If the numerator is a constant, bail if it doesn't fit into BypassType.
+  if (ConstantInt *ConstDividend = dyn_cast<ConstantInt>(Dividend))
+    if (ConstDividend->getValue().getActiveBits() > BypassType->getBitWidth())
+      return None;
+
+  // Split the basic block before the div/rem.
+  BasicBlock *SuccessorBB = MainBB->splitBasicBlock(SlowDivOrRem);
+  // Remove the unconditional branch from MainBB to SuccessorBB.
+  MainBB->getInstList().back().eraseFromParent();
+  QuotRemWithBB Fast = createFastBB(SuccessorBB);
+  QuotRemWithBB Slow = createSlowBB(SuccessorBB);
+  QuotRemPair Result = createDivRemPhiNodes(Fast, Slow, SuccessorBB);
+  Value *CmpV = insertOperandRuntimeCheck();
+  IRBuilder<> Builder(MainBB, MainBB->end());
+  Builder.CreateCondBr(CmpV, Fast.BB, Slow.BB);
+  return Result;
 }
 
-// bypassSlowDivision - This optimization identifies DIV instructions in a BB
-// that can be profitably bypassed and carried out with a shorter, faster
-// divide.
-bool llvm::bypassSlowDivision(
-    BasicBlock *BB, const DenseMap<unsigned int, unsigned int> &BypassWidths) {
-  DivCacheTy DivCache;
+/// This optimization identifies DIV/REM instructions in a BB that can be
+/// profitably bypassed and carried out with a shorter, faster divide.
+bool llvm::bypassSlowDivision(BasicBlock *BB,
+                              const BypassWidthsTy &BypassWidths) {
+  DivCacheTy PerBBDivCache;
 
   bool MadeChange = false;
   Instruction* Next = &*BB->begin();
   while (Next != nullptr) {
     // We may add instructions immediately after I, but we want to skip over
     // them.
     Instruction* I = Next;
     Next = Next->getNextNode();
 
-    // Get instruction details
-    unsigned Opcode = I->getOpcode();
-    bool UseDivOp = Opcode == Instruction::SDiv || Opcode == Instruction::UDiv;
-    bool UseRemOp = Opcode == Instruction::SRem || Opcode == Instruction::URem;
-    bool UseSignedOp = Opcode == Instruction::SDiv ||
-                       Opcode == Instruction::SRem;
-
-    // Only optimize div or rem ops
-    if (!UseDivOp && !UseRemOp)
-      continue;
-
-    // Skip division on vector types, only optimize integer instructions
-    if (!I->getType()->isIntegerTy())
-      continue;
-
-    // Get bitwidth of div/rem instruction
-    IntegerType *T = cast<IntegerType>(I->getType());
-    unsigned int bitwidth = T->getBitWidth();
-
-    // Continue if bitwidth is not bypassed
-    DenseMap<unsigned int, unsigned int>::const_iterator BI = BypassWidths.find(bitwidth);
-    if (BI == BypassWidths.end())
-      continue;
-
-    // Get type for div/rem instruction with bypass bitwidth
-    IntegerType *BT = IntegerType::get(I->getContext(), BI->second);
-
-    MadeChange |= reuseOrInsertFastDiv(I, BT, UseDivOp, UseSignedOp, DivCache);
+    FastDivInsertionTask Task(I, BypassWidths);
+    if (Value *Replacement = Task.getReplacement(PerBBDivCache)) {
+      I->replaceAllUsesWith(Replacement);
+      I->eraseFromParent();
+      MadeChange = true;
+    }
   }
 
   // Above we eagerly create divs and rems, as pairs, so that we can efficiently
   // create divrem machine instructions.  Now erase any unused divs / rems so we
   // don't leave extra instructions sitting around.
-  for (auto &KV : DivCache)
-    for (Instruction *Phi : {KV.second.Quotient, KV.second.Remainder})
-      RecursivelyDeleteTriviallyDeadInstructions(Phi);
+  for (auto &KV : PerBBDivCache)
+    for (Value *V : {KV.second.Quotient, KV.second.Remainder})
+      RecursivelyDeleteTriviallyDeadInstructions(V);
 
   return MadeChange;
 }