[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 d5eb50700b0b8472643585cc1a2dd52e5dbe5789 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 basic blocks 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 415136b612ac2..292b82c8f87e3 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -413,6 +413,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 }