[SimplifyCFG] Fix hoisting problem in SimplifyCFG

Author: RouzbehPaktinat

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -1582,19 +1582,29 @@ hoistLockstepIdenticalDPValues(Instruction *TI, Instruction *I1,
   }
 }
 
+// Hash instructions based on following factors:
+// 1- Instruction Opcode
+// 2- Instruction type
+// 3- Instruction operands
+llvm::hash_code getHash(Instruction *Instr) {
+  std::vector<Value *> operands(Instr->op_begin(), Instr->op_end());
+  return llvm::hash_combine(
+      Instr->getOpcode(), Instr->getType(),
+      hash_combine_range(operands.begin(), operands.end()));
+}
+
 /// Hoist any common code in the successor blocks up into the block. This
 /// function guarantees that BB dominates all successors. If EqTermsOnly is
 /// given, only perform hoisting in case both blocks only contain a terminator.
 /// In that case, only the original BI will be replaced and selects for PHIs are
 /// added.
 bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(BasicBlock *BB,
                                                    bool EqTermsOnly) {
-  // This does very trivial matching, with limited scanning, to find identical
-  // instructions in the two blocks. In particular, we don't want to get into
-  // O(N1*N2*...) situations here where Ni are the sizes of these successors. As
-  // such, we currently just scan for obviously identical instructions in an
-  // identical order, possibly separated by the same number of non-identical
-  // instructions.
+  // We first sort successors based on the number of instructions each block
+  // holds. Then for each successor we make a hashmap from its instructions,
+  // except for the first successor. After that, we iterate over the
+  // instructions of the first successor. If we find identical instructions from
+  // every other successor, we hoist all of them into the predeccessor.
   unsigned int SuccSize = succ_size(BB);
   if (SuccSize < 2)
     return false;
@@ -1608,10 +1618,21 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(BasicBlock *BB,
 
   auto *TI = BB->getTerminator();
 
+  SmallVector<BasicBlock *> SuccessorBlocks;
+  for (auto *Succ : successors(BB))
+    SuccessorBlocks.push_back(Succ);
+
+  // Sort successor blocks based on the number of instructions.
+  // This is because we always want to iterate over instructions
+  // of the smallest block.
+  llvm::stable_sort(SuccessorBlocks, [](BasicBlock *BB1, BasicBlock *BB2) {
+    return BB1->sizeWithoutDebug() < BB2->sizeWithoutDebug();
+  });
+
   // The second of pair is a SkipFlags bitmask.
   using SuccIterPair = std::pair<BasicBlock::iterator, unsigned>;
   SmallVector<SuccIterPair, 8> SuccIterPairs;
-  for (auto *Succ : successors(BB)) {
+  for (auto *Succ : SuccessorBlocks) {
     BasicBlock::iterator SuccItr = Succ->begin();
     if (isa<PHINode>(*SuccItr))
       return false;
@@ -1645,77 +1666,124 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(BasicBlock *BB,
   }
 
   bool Changed = false;
+  auto *SuccIterPairBegin = SuccIterPairs.begin();
+  SuccIterPairBegin++;
+  auto OtherSuccIterPairRange =
+      iterator_range(SuccIterPairBegin, SuccIterPairs.end());
+  auto OtherSuccIterRange = make_first_range(OtherSuccIterPairRange);
+  using InstrFlagPair = std::pair<Instruction *, unsigned>;
+  SmallVector<DenseMap<llvm::hash_code, InstrFlagPair>, 2> OtherSuccessorsHash;
+
+  for (auto BBItrPair : OtherSuccIterRange) {
+    // Fill the hashmap for every other successor
+    DenseMap<llvm::hash_code, InstrFlagPair> hashMap;
+    unsigned skipFlag = 0;
+    Instruction *I = nullptr;
+    do {
+      I = &*BBItrPair;
+      skipFlag |= skippedInstrFlags(I);
+      hashMap[getHash(I)] = InstrFlagPair(I, skipFlag);
+      BBItrPair++;
+    } while (!I->isTerminator());
+    OtherSuccessorsHash.push_back(hashMap);
+  }
 
+  // Keep track of instructions skipped in the first successor
+  unsigned SkipFlagsBB1 = 0;
+  bool SameLevelHoist = true;
   for (;;) {
     auto *SuccIterPairBegin = SuccIterPairs.begin();
     auto &BB1ItrPair = *SuccIterPairBegin++;
     auto OtherSuccIterPairRange =
         iterator_range(SuccIterPairBegin, SuccIterPairs.end());
-    auto OtherSuccIterRange = make_first_range(OtherSuccIterPairRange);
-
     Instruction *I1 = &*BB1ItrPair.first;
-
-    // Skip debug info if it is not identical.
-    bool AllDbgInstsAreIdentical = all_of(OtherSuccIterRange, [I1](auto &Iter) {
-      Instruction *I2 = &*Iter;
-      return I1->isIdenticalToWhenDefined(I2);
-    });
-    if (!AllDbgInstsAreIdentical) {
-      while (isa<DbgInfoIntrinsic>(I1))
-        I1 = &*++BB1ItrPair.first;
-      for (auto &SuccIter : OtherSuccIterRange) {
-        Instruction *I2 = &*SuccIter;
-        while (isa<DbgInfoIntrinsic>(I2))
-          I2 = &*++SuccIter;
+    bool HasIdenticalInst = true;
+
+    // Check if there are identical instructions in all other successors
+    for (auto &map : OtherSuccessorsHash) {
+      Instruction *I2 = map[getHash(I1)].first;
+      // We might face with same hash values for different instructions.
+      // If that happens, ignore the instruction.
+      if (!I2 || !I1->isIdenticalTo(I2)) {
+        HasIdenticalInst = false;
+        break;
       }
     }
 
-    bool AllInstsAreIdentical = true;
-    bool HasTerminator = I1->isTerminator();
-    for (auto &SuccIter : OtherSuccIterRange) {
-      Instruction *I2 = &*SuccIter;
-      HasTerminator |= I2->isTerminator();
-      if (AllInstsAreIdentical && !I1->isIdenticalToWhenDefined(I2))
-        AllInstsAreIdentical = false;
+    if (!HasIdenticalInst) {
+      if (NumSkipped >= HoistCommonSkipLimit)
+        return Changed;
+      SkipFlagsBB1 |= skippedInstrFlags(I1);
+      if (SameLevelHoist) {
+        for (auto &SuccIterPair : OtherSuccIterPairRange) {
+          Instruction *I = &*SuccIterPair.first++;
+          SuccIterPair.second |= skippedInstrFlags(I);
+        }
+      }
+      NumSkipped++;
+      if (I1->isTerminator())
+        return Changed;
+      ++BB1ItrPair.first;
+      continue;
     }
 
     SmallVector<Instruction *, 8> OtherInsts;
-    for (auto &SuccIter : OtherSuccIterRange)
-      OtherInsts.push_back(&*SuccIter);
+    if (SameLevelHoist) { 
+      for (auto &SuccIterPair : OtherSuccIterPairRange)
+        OtherInsts.push_back(&*(SuccIterPair.first));
+    } else {
+      for (auto &map : OtherSuccessorsHash)
+        OtherInstrs.push_back(map[getHash(I1)].first);
+    }
 
     // If we are hoisting the terminator instruction, don't move one (making a
     // broken BB), instead clone it, and remove BI.
-    if (HasTerminator) {
+    if (I1->isTerminator()) {
       // Even if BB, which contains only one unreachable instruction, is ignored
       // at the beginning of the loop, we can hoist the terminator instruction.
       // If any instructions remain in the block, we cannot hoist terminators.
-      if (NumSkipped || !AllInstsAreIdentical) {
+      if (NumSkipped) {
         hoistLockstepIdenticalDPValues(TI, I1, OtherInsts);
         return Changed;
       }
-
-      return hoistSuccIdenticalTerminatorToSwitchOrIf(TI, I1, OtherInsts) ||
-             Changed;
+      SmallVector<Instruction *, 8> Insts;
+      for (auto &map : OtherSuccessorsHash) {
+        Instruction *I2 = map[getHash(I1)].first;
+        // BB holding I2 should only contain the branch instruction
+        auto itr = I2->getParent()->instructionsWithoutDebug();
+        if (&*itr.begin() != I2)
+          return Changed;
+        Insts.push_back(I2);
+      }
+      return hoistSuccIdenticalTerminatorToSwitchOrIf(TI, I1, Insts) || Changed;
     }
 
-    if (AllInstsAreIdentical) {
-      unsigned SkipFlagsBB1 = BB1ItrPair.second;
-      AllInstsAreIdentical =
-          isSafeToHoistInstr(I1, SkipFlagsBB1) &&
-          all_of(OtherSuccIterPairRange, [=](const auto &Pair) {
-            Instruction *I2 = &*Pair.first;
-            unsigned SkipFlagsBB2 = Pair.second;
-            // Even if the instructions are identical, it may not
-            // be safe to hoist them if we have skipped over
-            // instructions with side effects or their operands
-            // weren't hoisted.
-            return isSafeToHoistInstr(I2, SkipFlagsBB2) &&
-                   shouldHoistCommonInstructions(I1, I2, TTI);
-          });
+    bool SafeToHoist = isSafeToHoistInstr(I1, SkipFlagsBB1);
+    unsigned index = 0;
+    for (auto &SuccIterPair : OtherSuccIterPairRange) {
+      Instruction *I2 = OtherSuccessorsHash[index][getHash(I1)].first;
+      // If instructions of all successors are at the same level, use the
+      // skipFlag of its BB, i.e., SameLevelHoist. Otherwise, use the skipFlag
+      // that was calculated initially for this instruction in the hashmap
+      if (SameLevelHoist && I2 == (&*(SuccIterPair.first))) {
+        SafeToHoist = SafeToHoist &&
+                      isSafeToHoistInstr(I2, SuccIterPair.second) &&
+                      shouldHoistCommonInstructions(I1, I2, TTI);
+      } else {
+        unsigned skipFlag = OtherSuccessorsHash[index][getHash(I1)].second;
+        SafeToHoist = SafeToHoist && isSafeToHoistInstr(I2, skipFlag) &&
+                      shouldHoistCommonInstructions(I1, I2, TTI);
+        SameLevelHoist = false;
+      }
+      index++;
     }
 
-    if (AllInstsAreIdentical) {
+    if (SafeToHoist) {
       BB1ItrPair.first++;
+      if (SameLevelHoist) {
+        for (auto &SuccIterPair : OtherSuccIterPairRange)
+          SuccIterPair.first++;
+      }
       if (isa<DbgInfoIntrinsic>(I1)) {
         // The debug location is an integral part of a debug info intrinsic
         // and can't be separated from it or replaced.  Instead of attempting
@@ -1725,8 +1793,8 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(BasicBlock *BB,
         // leave any that were not hoisted behind (by calling moveBefore
         // rather than moveBeforePreserving).
         I1->moveBefore(TI);
-        for (auto &SuccIter : OtherSuccIterRange) {
-          auto *I2 = &*SuccIter++;
+        for (auto &map : OtherSuccessorsHash) {
+          Instruction *I2 = map[getHash(I1)].first;
           assert(isa<DbgInfoIntrinsic>(I2));
           I2->moveBefore(TI);
         }
@@ -1739,8 +1807,8 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(BasicBlock *BB,
         // leave any that were not hoisted behind (by calling moveBefore
         // rather than moveBeforePreserving).
         I1->moveBefore(TI);
-        for (auto &SuccIter : OtherSuccIterRange) {
-          Instruction *I2 = &*SuccIter++;
+        for (auto &map : OtherSuccessorsHash) {
+          Instruction *I2 = map[getHash(I1)].first;
           assert(I2 != I1);
           if (!I2->use_empty())
             I2->replaceAllUsesWith(I1);
@@ -1764,9 +1832,12 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(BasicBlock *BB,
       // We are about to skip over a pair of non-identical instructions. Record
       // if any have characteristics that would prevent reordering instructions
       // across them.
-      for (auto &SuccIterPair : SuccIterPairs) {
-        Instruction *I = &*SuccIterPair.first++;
-        SuccIterPair.second |= skippedInstrFlags(I);
+      SkipFlagsBB1 |= skippedInstrFlags(I1);
+      if (SameLevelHoist) {
+        for (auto &SuccIterPair : OtherSuccIterPairRange) { // update flags
+          Instruction *I = &*SuccIterPair.first;
+          SuccIterPair.second |= skippedInstrFlags(I);
+        }
       }
       ++NumSkipped;
     }
@@ -1810,7 +1881,6 @@ bool SimplifyCFGOpt::hoistSuccIdenticalTerminatorToSwitchOrIf(
         Value *BB2V = PN.getIncomingValueForBlock(OtherSuccTI->getParent());
         if (BB1V == BB2V)
           continue;
-
         // In the case of an if statement, check for
         // passingValueIsAlwaysUndefined here because we would rather eliminate
         // undefined control flow then converting it to a select.
@@ -1882,20 +1952,16 @@ bool SimplifyCFGOpt::hoistSuccIdenticalTerminatorToSwitchOrIf(
       }
     }
   }
-
   SmallVector<DominatorTree::UpdateType, 4> Updates;
-
   // Update any PHI nodes in our new successors.
   for (BasicBlock *Succ : successors(BB1)) {
     AddPredecessorToBlock(Succ, TIParent, BB1);
     if (DTU)
       Updates.push_back({DominatorTree::Insert, TIParent, Succ});
   }
-
   if (DTU)
     for (BasicBlock *Succ : successors(TI))
       Updates.push_back({DominatorTree::Delete, TIParent, Succ});
-
   EraseTerminatorAndDCECond(TI);
   if (DTU)
     DTU->applyUpdates(Updates);