[SimplifyCFG] Hoist common code with canonical orderingusers/guy-david/simplifycfg-stable-iterator

3 files changed, 203 insertions, 53 deletions

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index ed2a5c292fa5..9aeb353d0edb 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -1850,6 +1850,71 @@ static bool isSafeCheapLoadStore(const Instruction *I,
          getLoadStoreAlignment(I) < Value::MaximumAlignment;
 }
 
+/// Try to find an instruction in a range that matches the given instruction.
+/// This allows us to find matching instructions even when they appear in
+/// different orders across blocks, particularly for pure/readonly functions.
+static Instruction *findMatchingInstruction(Instruction *I,
+                                           BasicBlock::iterator Start,
+                                           BasicBlock::iterator End,
+                                           SmallPtrSetImpl<Instruction *> &AlreadyMatched) {
+  // First check if the current position matches
+  if (Start != End && !AlreadyMatched.count(&*Start) &&
+      areIdenticalUpToCommutativity(I, &*Start))
+    return &*Start;
+
+  // For calls to pure/readonly functions or other safe-to-reorder instructions,
+  // search the entire range for a match
+  auto *CB = dyn_cast<CallBase>(I);
+  bool CanSearch = false;
+
+  if (CB && (CB->onlyReadsMemory() || CB->doesNotAccessMemory()) &&
+      CB->doesNotThrow()) {
+    // Pure/readonly calls can be reordered
+    CanSearch = true;
+  } else if (!I->mayReadFromMemory() && !I->mayWriteToMemory() &&
+             !I->mayHaveSideEffects() && !isa<AllocaInst>(I) &&
+             isGuaranteedToTransferExecutionToSuccessor(I)) {
+    // Other instructions without side effects can also be reordered
+    CanSearch = true;
+  }
+
+  if (!CanSearch)
+    return nullptr;
+
+  // Search the entire range for a matching instruction
+  auto SearchIter = Start;
+  while (SearchIter != End) {
+    Instruction *CandidateI = &*SearchIter;
+
+    // Skip already matched instructions
+    if (AlreadyMatched.count(CandidateI)) {
+      ++SearchIter;
+      continue;
+    }
+
+    // Debug output for pure calls
+    if (CB && CandidateI) {
+      if (auto *CandidateCB = dyn_cast<CallBase>(CandidateI)) {
+        // dbgs() << "  Checking candidate: " << *CandidateI << "\n";
+        // dbgs() << "    areIdentical? " << areIdenticalUpToCommutativity(I, CandidateI) << "\n";
+      }
+    }
+
+    if (areIdenticalUpToCommutativity(I, CandidateI)) {
+      // Found a match - for pure functions, we can always match them
+      // The hoisting logic will handle dependencies correctly
+      if (CB) {
+        // dbgs() << "    Found matching instruction!\n";
+      }
+      return CandidateI;
+    }
+
+    ++SearchIter;
+  }
+
+  return nullptr;
+}
+
 /// Hoist any common code in the successor blocks up into the block. This
 /// function guarantees that BB dominates all successors. If AllInstsEqOnly is
 /// given, only perform hoisting in case all successors blocks contain matching
@@ -1940,6 +2005,9 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(Instruction *TI,
 
   bool Changed = false;
 
+  // Track matched instructions to avoid matching the same instruction twice
+  SmallPtrSet<Instruction *, 8> AlreadyMatched;
+
   for (;;) {
     auto *SuccIterPairBegin = SuccIterPairs.begin();
     auto &BB1ItrPair = *SuccIterPairBegin++;
@@ -1949,19 +2017,60 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(Instruction *TI,
 
     Instruction *I1 = &*BB1ItrPair.first;
 
+    // Try to find matching instructions in other blocks, even if they're not
+    // at the exact same position (for pure/readonly calls)
+    SmallVector<Instruction *, 8> OtherInsts;
     bool AllInstsAreIdentical = true;
     bool HasTerminator = I1->isTerminator();
-    for (auto &SuccIter : OtherSuccIterRange) {
-      Instruction *I2 = &*SuccIter;
-      HasTerminator |= I2->isTerminator();
-      if (AllInstsAreIdentical && (!areIdenticalUpToCommutativity(I1, I2) ||
-                                   MMRAMetadata(*I1) != MMRAMetadata(*I2)))
+
+    // Check each successor block for a matching instruction
+    for (auto &SuccIterPair : iterator_range(SuccIterPairBegin, SuccIterPairs.end())) {
+      auto &SuccIter = SuccIterPair.first;
+      Instruction *CurrentI = &*SuccIter;
+
+      // First check for exact match at current position
+      bool IsMatch = !AlreadyMatched.count(CurrentI) &&
+                     areIdenticalUpToCommutativity(I1, CurrentI) &&
+                     MMRAMetadata(*I1) == MMRAMetadata(*CurrentI);
+
+      Instruction *MatchedI = nullptr;
+
+      if (IsMatch) {
+        // Exact match at current position
+        MatchedI = CurrentI;
+      } else if (!I1->isTerminator()) {
+        // Try to find a match elsewhere in the block
+        BasicBlock *SuccBB = CurrentI->getParent();
+        // For pure functions, search the entire block since they can be reordered
+        MatchedI = findMatchingInstruction(I1, SuccBB->begin(), SuccBB->end(), AlreadyMatched);
+
+        // Debug: If we're looking for a pure call but didn't find a match
+        if (!MatchedI) {
+          if (auto *CB1 = dyn_cast<CallBase>(I1)) {
+            if ((CB1->onlyReadsMemory() || CB1->doesNotAccessMemory()) &&
+                CB1->doesNotThrow()) {
+              // We have a pure call in I1 but couldn't find a match
+              // This might be because our search isn't comprehensive enough
+              // dbgs() << "SimplifyCFG: Failed to find match for pure call: " << *I1 << "\n";
+              // dbgs() << "  in block: " << SuccBB->getName() << "\n";
+            }
+          }
+        }
+      }
+
+      if (MatchedI) {
+        OtherInsts.push_back(MatchedI);
+        // Mark this instruction as matched so we don't match it again
+        AlreadyMatched.insert(MatchedI);
+        // Note: We don't advance the iterator here if we found a match at a
+        // different position - that will be handled later
+      } else {
         AllInstsAreIdentical = false;
-    }
+        OtherInsts.push_back(CurrentI);
+      }
 
-    SmallVector<Instruction *, 8> OtherInsts;
-    for (auto &SuccIter : OtherSuccIterRange)
-      OtherInsts.push_back(&*SuccIter);
+      HasTerminator |= CurrentI->isTerminator();
+    }
 
     // If we are hoisting the terminator instruction, don't move one (making a
     // broken BB), instead clone it, and remove BI.
@@ -1982,13 +2091,10 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(Instruction *TI,
       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.
+          all_of(OtherInsts, [=](Instruction *I2) {
+            // For instructions found at different positions, we need to check
+            // if they can be safely hoisted considering what's been skipped
+            unsigned SkipFlagsBB2 = 0; // TODO: need to track skip flags per instruction
             return isSafeToHoistInstr(I2, SkipFlagsBB2) &&
                    shouldHoistCommonInstructions(I1, I2, TTI);
           });
@@ -2004,8 +2110,21 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(Instruction *TI,
       // and leave any that were not hoisted behind (by calling moveBefore
       // rather than moveBeforePreserving).
       I1->moveBefore(TI->getIterator());
-      for (auto &SuccIter : OtherSuccIterRange) {
-        Instruction *I2 = &*SuccIter++;
+
+      // Replace uses and erase the matched instructions
+      size_t InstIdx = 0;
+      for (auto &SuccIterPair : iterator_range(SuccIterPairBegin, SuccIterPairs.end())) {
+        Instruction *I2 = OtherInsts[InstIdx];
+
+        // Check if this was a match found at a different position
+        if (I2 != &*SuccIterPair.first) {
+          // This instruction was found out of order - we need to handle it specially
+          // The iterator should only advance if we used the instruction at current position
+        } else {
+          // Normal case - advance the iterator
+          ++SuccIterPair.first;
+        }
+
         assert(I2 != I1);
         if (!I2->use_empty())
           I2->replaceAllUsesWith(I1);
@@ -2023,6 +2142,8 @@ bool SimplifyCFGOpt::hoistCommonCodeFromSuccessors(Instruction *TI,
         // location is the merged locations of the original instructions.
         I1->applyMergedLocation(I1->getDebugLoc(), I2->getDebugLoc());
         I2->eraseFromParent();
+
+        ++InstIdx;
       }
       if (!Changed)
         NumHoistCommonCode += SuccIterPairs.size();
diff --git a/llvm/test/Transforms/SimplifyCFG/dont-hoist-deoptimize.ll b/llvm/test/Transforms/SimplifyCFG/dont-hoist-deoptimize.ll
index 6bfa967ff23e..158de317afc0 100644
--- a/llvm/test/Transforms/SimplifyCFG/dont-hoist-deoptimize.ll
+++ b/llvm/test/Transforms/SimplifyCFG/dont-hoist-deoptimize.ll
@@ -9,15 +9,14 @@ define void @widget(i1 %arg) {
 ; CHECK-LABEL: @widget(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[TMP:%.*]] = trunc i64 5 to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 0 to i32
 ; CHECK-NEXT:    br i1 [[ARG:%.*]], label [[BB1:%.*]], label [[BB4:%.*]]
 ; CHECK:       bb1:
-; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 0 to i32
 ; CHECK-NEXT:    [[TMP3:%.*]] = trunc i64 0 to i32
 ; CHECK-NEXT:    call void (...) @llvm.experimental.deoptimize.isVoid(i32 13) #[[ATTR0:[0-9]+]] [ "deopt"() ]
 ; CHECK-NEXT:    ret void
 ; CHECK:       bb4:
 ; CHECK-NEXT:    [[TMP6:%.*]] = trunc i64 1 to i32
-; CHECK-NEXT:    [[TMP7:%.*]] = trunc i64 0 to i32
 ; CHECK-NEXT:    call void (...) @llvm.experimental.deoptimize.isVoid(i32 13) #[[ATTR0]] [ "deopt"() ]
 ; CHECK-NEXT:    ret void
 ;
diff --git a/llvm/test/Transforms/SimplifyCFG/mmra.ll b/llvm/test/Transforms/SimplifyCFG/mmra.ll
index 667065747137..14cc55b1ce0b 100644
--- a/llvm/test/Transforms/SimplifyCFG/mmra.ll
+++ b/llvm/test/Transforms/SimplifyCFG/mmra.ll
@@ -1,7 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
-; RUN: opt -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S %s | FileCheck %s
-
-; RUN: opt -passes='simplifycfg<sink-common-insts;hoist-common-insts>,verify' -S %s | FileCheck %s
+; RUN: opt -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S %s | FileCheck %s --check-prefixes=CHECK,CHECK-PRESERVE
+; RUN: opt -passes='simplifycfg<sink-common-insts;hoist-common-insts>,verify' -S %s | FileCheck %s --check-prefixes=CHECK,CHECK-NO-PRESERVE
 
 declare void @clobber1()
 declare void @clobber2()
@@ -40,20 +39,34 @@ exit:
 }
 
 define void @hoist_store(ptr %arg, i1 %c) {
-; CHECK-LABEL: define void @hoist_store(
-; CHECK-SAME: ptr [[ARG:%.*]], i1 [[C:%.*]]) {
-; CHECK-NEXT:  bb:
-; CHECK-NEXT:    br i1 [[C]], label [[THEN:%.*]], label [[ELSE:%.*]]
-; CHECK:       then:
-; CHECK-NEXT:    store ptr null, ptr [[ARG]], align 8
-; CHECK-NEXT:    call void @clobber1()
-; CHECK-NEXT:    br label [[EXIT:%.*]]
-; CHECK:       else:
-; CHECK-NEXT:    store ptr null, ptr [[ARG]], align 8, !mmra [[META0]]
-; CHECK-NEXT:    call void @clobber2()
-; CHECK-NEXT:    br label [[EXIT]]
-; CHECK:       exit:
-; CHECK-NEXT:    ret void
+; CHECK-PRESERVE-LABEL: define void @hoist_store(
+; CHECK-PRESERVE-SAME: ptr [[ARG:%.*]], i1 [[C:%.*]]) {
+; CHECK-PRESERVE-NEXT:  bb:
+; CHECK-PRESERVE-NEXT:    br i1 [[C]], label [[THEN:%.*]], label [[ELSE:%.*]]
+; CHECK-PRESERVE:       then:
+; CHECK-PRESERVE-NEXT:    store ptr null, ptr [[ARG]], align 8
+; CHECK-PRESERVE-NEXT:    call void @clobber1()
+; CHECK-PRESERVE-NEXT:    br label [[EXIT:%.*]]
+; CHECK-PRESERVE:       else:
+; CHECK-PRESERVE-NEXT:    store ptr null, ptr [[ARG]], align 8, !mmra [[META0]]
+; CHECK-PRESERVE-NEXT:    call void @clobber2()
+; CHECK-PRESERVE-NEXT:    br label [[EXIT]]
+; CHECK-PRESERVE:       exit:
+; CHECK-PRESERVE-NEXT:    ret void
+;
+; CHECK-NO-PRESERVE-LABEL: define void @hoist_store(
+; CHECK-NO-PRESERVE-SAME: ptr [[ARG:%.*]], i1 [[C:%.*]]) {
+; CHECK-NO-PRESERVE-NEXT:  bb:
+; CHECK-NO-PRESERVE-NEXT:    store ptr null, ptr [[ARG]], align 8, !mmra [[META1:![0-9]+]]
+; CHECK-NO-PRESERVE-NEXT:    br i1 [[C]], label [[THEN:%.*]], label [[ELSE:%.*]]
+; CHECK-NO-PRESERVE:       then:
+; CHECK-NO-PRESERVE-NEXT:    call void @clobber1()
+; CHECK-NO-PRESERVE-NEXT:    br label [[EXIT:%.*]]
+; CHECK-NO-PRESERVE:       else:
+; CHECK-NO-PRESERVE-NEXT:    call void @clobber2()
+; CHECK-NO-PRESERVE-NEXT:    br label [[EXIT]]
+; CHECK-NO-PRESERVE:       exit:
+; CHECK-NO-PRESERVE-NEXT:    ret void
 ;
 bb:
   br i1 %c, label %then, label %else
@@ -108,21 +121,35 @@ exit:
 }
 
 define ptr @hoist_load(ptr %arg, i1 %c) {
-; CHECK-LABEL: define ptr @hoist_load(
-; CHECK-SAME: ptr [[ARG:%.*]], i1 [[C:%.*]]) {
-; CHECK-NEXT:  bb:
-; CHECK-NEXT:    br i1 [[C]], label [[THEN:%.*]], label [[ELSE:%.*]]
-; CHECK:       then:
-; CHECK-NEXT:    [[L1:%.*]] = load ptr, ptr [[ARG]], align 8
-; CHECK-NEXT:    call void @clobber1()
-; CHECK-NEXT:    br label [[EXIT:%.*]]
-; CHECK:       else:
-; CHECK-NEXT:    [[L2:%.*]] = load ptr, ptr [[ARG]], align 8, !mmra [[META0]]
-; CHECK-NEXT:    call void @clobber2()
-; CHECK-NEXT:    br label [[EXIT]]
-; CHECK:       exit:
-; CHECK-NEXT:    [[P:%.*]] = phi ptr [ [[L1]], [[THEN]] ], [ [[L2]], [[ELSE]] ]
-; CHECK-NEXT:    ret ptr [[P]]
+; CHECK-PRESERVE-LABEL: define ptr @hoist_load(
+; CHECK-PRESERVE-SAME: ptr [[ARG:%.*]], i1 [[C:%.*]]) {
+; CHECK-PRESERVE-NEXT:  bb:
+; CHECK-PRESERVE-NEXT:    br i1 [[C]], label [[THEN:%.*]], label [[ELSE:%.*]]
+; CHECK-PRESERVE:       then:
+; CHECK-PRESERVE-NEXT:    [[L1:%.*]] = load ptr, ptr [[ARG]], align 8
+; CHECK-PRESERVE-NEXT:    call void @clobber1()
+; CHECK-PRESERVE-NEXT:    br label [[EXIT:%.*]]
+; CHECK-PRESERVE:       else:
+; CHECK-PRESERVE-NEXT:    [[L2:%.*]] = load ptr, ptr [[ARG]], align 8, !mmra [[META0]]
+; CHECK-PRESERVE-NEXT:    call void @clobber2()
+; CHECK-PRESERVE-NEXT:    br label [[EXIT]]
+; CHECK-PRESERVE:       exit:
+; CHECK-PRESERVE-NEXT:    [[P:%.*]] = phi ptr [ [[L1]], [[THEN]] ], [ [[L2]], [[ELSE]] ]
+; CHECK-PRESERVE-NEXT:    ret ptr [[P]]
+;
+; CHECK-NO-PRESERVE-LABEL: define ptr @hoist_load(
+; CHECK-NO-PRESERVE-SAME: ptr [[ARG:%.*]], i1 [[C:%.*]]) {
+; CHECK-NO-PRESERVE-NEXT:  bb:
+; CHECK-NO-PRESERVE-NEXT:    [[L1:%.*]] = load ptr, ptr [[ARG]], align 8, !mmra [[META1]]
+; CHECK-NO-PRESERVE-NEXT:    br i1 [[C]], label [[THEN:%.*]], label [[ELSE:%.*]]
+; CHECK-NO-PRESERVE:       then:
+; CHECK-NO-PRESERVE-NEXT:    call void @clobber1()
+; CHECK-NO-PRESERVE-NEXT:    br label [[EXIT:%.*]]
+; CHECK-NO-PRESERVE:       else:
+; CHECK-NO-PRESERVE-NEXT:    call void @clobber2()
+; CHECK-NO-PRESERVE-NEXT:    br label [[EXIT]]
+; CHECK-NO-PRESERVE:       exit:
+; CHECK-NO-PRESERVE-NEXT:    ret ptr [[L1]]
 ;
 bb:
   br i1 %c, label %then, label %else
@@ -146,5 +173,8 @@ exit:
 !0 = !{!"foo", !"bar"}
 
 ;.
-; CHECK: [[META0]] = !{!"foo", !"bar"}
+; CHECK-PRESERVE: [[META0]] = !{!"foo", !"bar"}
+;.
+; CHECK-NO-PRESERVE: [[META0]] = !{!"foo", !"bar"}
+; CHECK-NO-PRESERVE: [[META1]] = !{}
 ;.