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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Bu Le (lebu26)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/176600.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/JumpThreading.cpp (+7) 
- (added) llvm/test/Transforms/JumpThreading/assume-edge-dom1.ll (+62) 


``````````diff
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..f1ed831c38c3e
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/assume-edge-dom1.ll
@@ -0,0 +1,62 @@
+; 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, downstream 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) {
+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
+}
+
+; CHECK-NOT: @side_effect(i32 99999)
+attributes #0 = { inaccessiblememonly nocallback nofree nosync nounwind willreturn }

``````````

</details>


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