From 32d6ef36d6c0da2d0b2165531cf9fcdc4dedc95a Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo@fhahn.com>
Date: Thu, 19 May 2022 17:01:11 +0100
Subject: [PATCH] [SimpleLoopUnswitch] Skip trivial selects during trivial
 unswitching.

Update the remaining places in unswitchTrivialBranch to properly skip
trivial selects.

Fixes #55526.
---
 llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp  | 13 ++++++-----
 .../trivial-unswitch-logical-and-or.ll             | 26 ++++++++++++++++++++++
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 8be8138..b48b628 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -460,12 +460,12 @@ static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
   // some input conditions to the branch.
   bool FullUnswitch = false;
 
-  if (L.isLoopInvariant(BI.getCondition())) {
-    Invariants.push_back(BI.getCondition());
+  Value *Cond = skipTrivialSelect(BI.getCondition());
+  if (L.isLoopInvariant(Cond)) {
+    Invariants.push_back(Cond);
     FullUnswitch = true;
   } else {
-    if (auto *CondInst =
-            dyn_cast<Instruction>(skipTrivialSelect(BI.getCondition())))
+    if (auto *CondInst = dyn_cast<Instruction>(Cond))
       Invariants = collectHomogenousInstGraphLoopInvariants(L, *CondInst, LI);
     if (Invariants.empty()) {
       LLVM_DEBUG(dbgs() << "   Couldn't find invariant inputs!\n");
@@ -499,7 +499,6 @@ static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
   // is a graph of `or` operations, or the exit block is along the false edge
   // and the condition is a graph of `and` operations.
   if (!FullUnswitch) {
-    Value *Cond = skipTrivialSelect(BI.getCondition());
     if (ExitDirection ? !match(Cond, m_LogicalOr())
                       : !match(Cond, m_LogicalAnd())) {
       LLVM_DEBUG(dbgs() << "   Branch condition is in improper form for "
@@ -566,6 +565,7 @@ static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
     // its successors.
     OldPH->getInstList().splice(OldPH->end(), BI.getParent()->getInstList(),
                                 BI);
+    BI.setCondition(Cond);
     if (MSSAU) {
       // Temporarily clone the terminator, to make MSSA update cheaper by
       // separating "insert edge" updates from "remove edge" ones.
@@ -1040,7 +1040,8 @@ static bool unswitchAllTrivialConditions(Loop &L, DominatorTree &DT,
     // Don't bother trying to unswitch past an unconditional branch or a branch
     // with a constant value. These should be removed by simplifycfg prior to
     // running this pass.
-    if (!BI->isConditional() || isa<Constant>(BI->getCondition()))
+    if (!BI->isConditional() ||
+        isa<Constant>(skipTrivialSelect(BI->getCondition())))
       return Changed;
 
     // Found a trivial condition candidate: non-foldable conditional branch. If
diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch-logical-and-or.ll b/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch-logical-and-or.ll
index 2af8dec..64d8742 100644
--- a/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch-logical-and-or.ll
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch-logical-and-or.ll
@@ -248,3 +248,29 @@ do_something:
 loop_exit:
   ret i32 0
 }
+
+; Test case for PR55526.
+define void @test_pr55526(i16 %a) {
+; CHECK-LABEL: @test_pr55526(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp ne i16 [[A:%.*]], 0
+; CHECK-NEXT:    br i1 [[TOBOOL]], label [[ENTRY_SPLIT:%.*]], label [[EXIT:%.*]]
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 true, i1 true, i1 false
+; CHECK-NEXT:    br label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tobool = icmp ne i16 %a, 0
+  br label %loop
+
+loop:
+  %sel = select i1 %tobool, i1 true, i1 false
+  br i1 %sel, label %loop, label %exit
+
+exit:
+  ret void
+}
-- 
2.7.4