[llvm] [CVP] Stop CVP constant propagation from destroying `llvm.assume` (PR #183688)

[via llvm-commits]

https://github.com/lijinpei-amd updated https://github.com/llvm/llvm-project/pull/183688

>From c0cac23d1c529735dab937557b3e4e8be7633a94 Mon Sep 17 00:00:00 2001
From: Li Jinpei <jinpli at amd.com>
Date: Thu, 26 Feb 2026 17:40:23 +0800
Subject: [PATCH] [CVP] Stop CVP constant propagation from destroying
 `llvm.assume`

For following code sequence:
```
  %cmp = icmp ugt i64 %idx, 1
  tail call void @llvm.assume(i1 %cmp)
```
CVP may replace all use of `%cmp` with `true`, and the information about
`%idx` is lost. This commit stop CVP from doing that, so that later pass
can exploit the `llvm.assume`.
---
 .../Scalar/CorrelatedValuePropagation.cpp      | 15 +++++++++++----
 .../CorrelatedValuePropagation/conflict.ll     |  3 ++-
 .../CorrelatedValuePropagation/icmp.ll         | 18 ++++++++++++++++++
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp b/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
index f4897034569c6..8eaf385065737 100644
--- a/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
+++ b/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
@@ -333,10 +333,17 @@ static bool constantFoldCmp(CmpInst *Cmp, LazyValueInfo *LVI) {
   if (!Res)
     return false;
 
-  ++NumCmps;
-  Cmp->replaceAllUsesWith(Res);
-  Cmp->eraseFromParent();
-  return true;
+  bool Changed = Cmp->replaceUsesWithIf(
+      Res, [](Use &U) { return !isa<AssumeInst>(U.getUser()); });
+  if (Cmp->use_empty()) {
+    Cmp->eraseFromParent();
+    Changed = true;
+  }
+
+  if (Changed)
+    ++NumCmps;
+
+  return Changed;
 }
 
 static bool processCmp(CmpInst *Cmp, LazyValueInfo *LVI) {
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/conflict.ll b/llvm/test/Transforms/CorrelatedValuePropagation/conflict.ll
index 7bba0565d9988..88f4aa12e90f2 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/conflict.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/conflict.ll
@@ -58,7 +58,8 @@ define i8 @test3(i8 %a) {
 ; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i8 [[A:%.*]], 5
 ; CHECK-NEXT:    br i1 [[CMP1]], label [[DEAD:%.*]], label [[EXIT:%.*]]
 ; CHECK:       dead:
-; CHECK-NEXT:    call void @llvm.assume(i1 false)
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i8 [[A]], 3
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP2]])
 ; CHECK-NEXT:    ret i8 5
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i8 0
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll b/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll
index 7fb1e0718795d..c87b02700694d 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll
@@ -6,6 +6,7 @@ target triple = "x86_64-apple-macosx10.10.0"
 
 declare void @check1(i1) #1
 declare void @check2(i1) #1
+declare void @check3(i64)
 declare void @llvm.assume(i1)
 
 ; Make sure we propagate the value of %tmp35 to the true/false cases
@@ -1657,3 +1658,20 @@ if.true:
 if.false:
   ret void
 }
+
+define void @test_assume_not_removed(i64 %idx) {
+; CHECK-LABEL: @test_assume_not_removed(
+; CHECK-NEXT:    [[IDX1:%.*]] = add nuw i64 [[IDX:%.*]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i64 [[IDX1]], 1
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    tail call void @check3(i64 [[IDX1]])
+; CHECK-NEXT:    tail call void @check1(i1 true)
+; CHECK-NEXT:    ret void
+;
+  %idx1 = add i64 %idx, 1
+  %cmp = icmp ugt i64 %idx1, 1
+  tail call void @llvm.assume(i1 %cmp)
+  tail call void @check3(i64 %idx1)
+  tail call void @check1(i1 %cmp)
+  ret void
+}