[llvm] [JumpThreading] Preserve assume instructions to maintain constraint information (PR #176600)

[Bu Le via llvm-commits]

https://github.com/lebu26 updated https://github.com/llvm/llvm-project/pull/176600

>From 965a42d79d2984a47a045ae9cebf35ac4c99b4fc Mon Sep 17 00:00:00 2001
From: Bu Le <cityubule at gmail.com>
Date: Sun, 18 Jan 2026 06:13:29 +0800
Subject: [PATCH] [JumpThreading] Preserve assume instructions to maintain
 constraint information

When jump threading optimizes a conditional branch in a basic block, it may
replace the condition variable with a constant. However, if the basic block
contains an assume() intrinsic that constrains this variable, we should not
replace the condition in the assume instruction.

This is important because downstream basic blocks that are dominated by the
current block may need to analyze the original constraint on the variable for
their own optimizations. Replacing the condition with assume(true) would break
the data flow for such analyses.

Add test case: assume-edge-dom1.ll
This test verifies that JumpThreading does not modify assume instructions when
the condition variable is optimized away in other branches.
---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp  |  7 ++
 .../JumpThreading/assume-edge-dom1.ll         | 71 +++++++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 llvm/test/Transforms/JumpThreading/assume-edge-dom1.ll

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index e8b4ba5f5106f..9533edde035c7 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -411,6 +411,13 @@ static bool replaceFoldableUses(Instruction *Cond, Value *ToVal,
     // of BB, where we know Cond is ToVal.
     if (!isGuaranteedToTransferExecutionToSuccessor(&I))
       break;
+    // Do not replace assume instruction's condition because other basic blocks
+    // dominated by this one may need to analyze the original constraint on the
+    // condition variable.
+    if (auto *II = llvm::dyn_cast<llvm::IntrinsicInst>(&I)) {
+      if (II->getIntrinsicID() == llvm::Intrinsic::assume)
+        break;
+    }
     Changed |= I.replaceUsesOfWith(Cond, ToVal);
   }
   if (Cond->use_empty() && !Cond->mayHaveSideEffects()) {
diff --git a/llvm/test/Transforms/JumpThreading/assume-edge-dom1.ll b/llvm/test/Transforms/JumpThreading/assume-edge-dom1.ll
new file mode 100644
index 0000000000000..1064567698907
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/assume-edge-dom1.ll
@@ -0,0 +1,71 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt %s -passes=jump-threading -S | FileCheck %s
+;
+; Test: Preserve assume information in dominated basic blocks
+;
+; This test ensures that when a basic block contains an assume() that constrains
+; a variable (e.g., assume(x == 0)), the JumpThreading pass does not replace it
+; with assume(true). If it did, dominated basic blocks that depend on analyzing
+; the variable's constraints would lose the ability to find the original definition
+; and could miss optimization opportunities or produce incorrect code.
+;
+; The key scenario:
+; 1. assume_bb contains: call void @llvm.assume(i1 %check_val)
+; 2. check_bb (dominated by assume_bb) branches on the same %check_val
+; 3. JumpThreading should NOT rewrite the assume to assume(true) in assume_bb
+;    because it would break the data flow for analyses in other dominated blocks
+
+declare void @side_effect(i32)
+
+; Function Attrs: inaccessiblememonly nocallback nofree nosync nounwind willreturn
+declare void @llvm.assume(i1 noundef) #0
+
+define i64 @test_preserve_assume_info(i32 %unused, i1 %outer_cond) {
+; CHECK-LABEL: define i64 @test_preserve_assume_info(
+; CHECK-SAME: i32 [[UNUSED:%.*]], i1 [[OUTER_COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CHECK_VAL:%.*]] = icmp eq i32 [[UNUSED]], 0
+; CHECK-NEXT:    br i1 [[OUTER_COND]], label %[[EXIT:.*]], label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CHECK_VAL]])
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  %check_val = icmp eq i32 %unused, 0
+  br i1 %outer_cond, label %merge, label %cond_true
+
+cond_true:
+  br label %merge
+
+merge:
+  br label %assume_bb
+
+assume_bb:
+  ; This assume constrains %check_val. JumpThreading should NOT replace it with
+  ; assume(true) because dominated blocks need to analyze %check_val's definition.
+  call void @llvm.assume(i1 %check_val)
+  br i1 %check_val, label %merge1, label %use_bb1
+
+use_bb1:
+  call void @side_effect(i32 99998)
+  br label %merge1
+
+merge1:
+  br label %check_bb
+
+check_bb:
+  ; This block is dominated by assume_bb. It needs the original constraint on
+  ; %check_val to be preserved in assume_bb for correct analysis.
+  br i1 %check_val, label %exit, label %use_bb2
+
+use_bb2:
+  ; This call should be removed by JumpThreading because %check_val is false here.
+  ; This only works correctly if the assume in assume_bb is preserved properly.
+  call void @side_effect(i32 99999)
+  br label %exit
+
+exit:
+  ret i64 0
+}
+
+attributes #0 = { inaccessiblememonly nocallback nofree nosync nounwind willreturn }