[llvm] [SimpleLoopUnswitch] Fix bug in partially invariant unswitching (PR #222118)

[Ehsan Amiri via llvm-commits]

https://github.com/amehsan updated https://github.com/llvm/llvm-project/pull/222118

>From e25c5be1ec1bd99070da6d42e0da05474f155cb0 Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Thu, 3 Sep 2026 15:25:41 -0400
Subject: [PATCH 1/4] add the failing testcase

---
 .../Transforms/SimpleLoopUnswitch/pr219510.ll | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)
 create mode 100644 llvm/test/Transforms/SimpleLoopUnswitch/pr219510.ll

diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/pr219510.ll b/llvm/test/Transforms/SimpleLoopUnswitch/pr219510.ll
new file mode 100644
index 0000000000000..555bfc1d992cc
--- /dev/null
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/pr219510.ll
@@ -0,0 +1,34 @@
+; RUN: opt -passes='loop-mssa(simple-loop-unswitch<nontrivial>)' -S < %s | FileCheck %s 
+
+declare void @barrier() memory(read)
+declare void @clobber()
+
+; CHECK: call void @barrier()
+; CHECK-NEXT: %lv = load i32, ptr %ptr
+
+
+define i32 @bad_unswitch(ptr %ptr, i32 %N) {
+entry:
+  br label %loop.header
+loop.header:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  call void @barrier()
+  %lv = load i32, ptr %ptr
+  %sc = icmp eq i32 %lv, 100
+  br i1 %sc, label %noclobber, label %clobber
+
+noclobber:
+  br label %loop.latch
+
+clobber:
+  call void @clobber()
+  br label %loop.latch
+
+loop.latch:
+  %c = icmp ult i32 %iv, %N
+  %iv.next = add i32 %iv, 1
+  br i1 %c, label %loop.header, label %exit
+
+exit:
+  ret i32 10
+}

>From 1cd8f20ac73b7534dad1f7c291c9d3d8e97ba1eb Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Wed, 9 Sep 2026 13:56:24 -0400
Subject: [PATCH 2/4] actual bugfix

---
 llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 7f6a08454b749..2bf4fa309f830 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -2787,7 +2787,10 @@ static bool collectUnswitchCandidates(
     AddUnswitchCandidatesForInst(BI, BI->getCondition());
   }
 
-  if (MSSAU && !findOptionMDForLoop(&L, "llvm.loop.unswitch.partial.disable") &&
+  BasicBlock *Header = L.getHeader();
+  // Need to make sure the load instruction to be hoisted is always executed.
+  bool HeaderCondGuaranteedToExecute = isGuaranteedToTransferExecutionToSuccessor(Header->begin(), Header->getTerminator()->getIterator());
+  if (MSSAU && HeaderCondGuaranteedToExecute && !findOptionMDForLoop(&L, "llvm.loop.unswitch.partial.disable") &&
       !any_of(UnswitchCandidates, [&L](auto &TerminatorAndInvariants) {
          return TerminatorAndInvariants.TI == L.getHeader()->getTerminator();
        })) {
@@ -2797,11 +2800,11 @@ static bool collectUnswitchCandidates(
           dbgs() << "simple-loop-unswitch: Found partially invariant condition "
                  << *Info->InstToDuplicate[0] << "\n");
       PartialIVInfo = *Info;
-      PartialIVCondBranch = L.getHeader()->getTerminator();
+      PartialIVCondBranch = Header->getTerminator();
       TinyPtrVector<Value *> ValsToDuplicate;
       llvm::append_range(ValsToDuplicate, Info->InstToDuplicate);
       UnswitchCandidates.push_back(
-          {L.getHeader()->getTerminator(), std::move(ValsToDuplicate)});
+          {Header->getTerminator(), std::move(ValsToDuplicate)});
     }
   }
   return !UnswitchCandidates.empty();

>From e679baa2a3b25f177e9a0fac811e84c7b437a4ae Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Wed, 9 Sep 2026 14:51:18 -0400
Subject: [PATCH 3/4] fix code style

---
 llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 2bf4fa309f830..7331ee85b9983 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -2789,8 +2789,11 @@ static bool collectUnswitchCandidates(
 
   BasicBlock *Header = L.getHeader();
   // Need to make sure the load instruction to be hoisted is always executed.
-  bool HeaderCondGuaranteedToExecute = isGuaranteedToTransferExecutionToSuccessor(Header->begin(), Header->getTerminator()->getIterator());
-  if (MSSAU && HeaderCondGuaranteedToExecute && !findOptionMDForLoop(&L, "llvm.loop.unswitch.partial.disable") &&
+  bool HeaderCondGuaranteedToExecute =
+      isGuaranteedToTransferExecutionToSuccessor(
+          Header->begin(), Header->getTerminator()->getIterator());
+  if (MSSAU && HeaderCondGuaranteedToExecute &&
+      !findOptionMDForLoop(&L, "llvm.loop.unswitch.partial.disable") &&
       !any_of(UnswitchCandidates, [&L](auto &TerminatorAndInvariants) {
          return TerminatorAndInvariants.TI == L.getHeader()->getTerminator();
        })) {

>From aa1345e2fe34898e70fb0ccfdce3c4d18cb1a2fb Mon Sep 17 00:00:00 2001
From: Ehsan Amiri <ehsan.amiri at huawei.com>
Date: Wed, 9 Sep 2026 14:56:24 -0400
Subject: [PATCH 4/4] fix seemingly irrelevant code style issue that CI
 complains about

---
 llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 7331ee85b9983..061960de3eb66 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -2795,8 +2795,8 @@ static bool collectUnswitchCandidates(
   if (MSSAU && HeaderCondGuaranteedToExecute &&
       !findOptionMDForLoop(&L, "llvm.loop.unswitch.partial.disable") &&
       !any_of(UnswitchCandidates, [&L](auto &TerminatorAndInvariants) {
-         return TerminatorAndInvariants.TI == L.getHeader()->getTerminator();
-       })) {
+        return TerminatorAndInvariants.TI == L.getHeader()->getTerminator();
+      })) {
     MemorySSA *MSSA = MSSAU->getMemorySSA();
     if (auto Info = hasPartialIVCondition(L, MSSAThreshold, *MSSA, AA)) {
       LLVM_DEBUG(