Bug type: miscompilation
Component: InstCombine
Bug location: llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp: InstCombinerImpl::SimplifyDemandedBits
InstCombinerImpl::SimplifyDemandedUseBits
Issue title: [InstCombine] Miscompilation with llvm.assume in unreachable blocks
Labels: miscompilation
llvm:instcombine
Reported by: dtcxzyw

Reproducer: https://alive2.llvm.org/ce/z/zKvpXV
```
; bin/opt -passes=instcombine reduced.ll -S
define i16 @src(i16 %g, ptr %e, ptr %d) {
entry:
  %0 = load i64, ptr %d, align 8
  %conv = trunc i64 %0 to i16
  %tobool.not.i = icmp eq i16 %g, 0
  br i1 %tobool.not.i, label %i.exit, label %for.cond.preheader.i

for.cond.preheader.i:
  %cmp5.i = icmp ne i16 %g, %conv
  %conv6.i = zext i1 %cmp5.i to i32
  store i32 %conv6.i, ptr %e, align 4
  %cmp7.i = icmp eq i64 %0, 1
  call void @llvm.assume(i1 %cmp7.i)
  unreachable

i.exit:
  ret i16 %conv
}
```
```
define i16 @tgt(i16 %g, ptr %e, ptr %d) {
entry:
  %tobool.not.i = icmp eq i16 %g, 0
  br i1 %tobool.not.i, label %i.exit, label %for.cond.preheader.i

for.cond.preheader.i:                             ; preds = %entry
  unreachable

i.exit:                                           ; preds = %entry
  ret i16 1
}
```

llvm version: a8e1c3e1239604ac787b6a2d39b5278ddec8aa8a

cc @nikic

Comments: author: dtcxzyw
body: Reverting https://github.com/llvm/llvm-project/commit/11484cb817bcc2a6e2ef9572be982a1a5a4964ec fixes this issue.

author: nikic
body: I've applied two fixes for this. https://github.com/llvm/llvm-project/commit/86b37944a70229b07626e63bdb9a46b4bc3d1460 tries to stick closer to the current logic, while https://github.com/llvm/llvm-project/commit/05670b42f5b45710bfdba48dcb7e8c30c8c7478f is a more thorough fix on top of that.
author: dtcxzyw
body: Confirmed that the original case has been fixed. Thank you!

author: chapuni
body: @nikic Fixed, thanks!

Fix patch:
commit b558ac0eef57a3737b1e27844115fa91e0b32582
Author: Nikita Popov <npopov@redhat.com>
Date:   Tue Jul 2 10:08:45 2024 +0200

    [InstCombine] Fix context for multi-use demanded bits simplification
    
    When simplifying a multi-use root value, the demanded bits were
    reset to full, but we also need to reset the context extract. To
    make this convenient (without requiring by-value passing of
    SimplifyQuery), move the logic that that handles constants and
    dispatches to SimplifyDemandedUseBits/SimplifyMultipleUseDemandedBits
    into SimplifyDemandedBits. The SimplifyDemandedInstructionBits
    caller starts with full demanded bits and an appropriate context
    anyway.
    
    The different context instruction does mean that the ephemeral
    value protection no longer triggers in some cases, as the changes
    to assume tests show.
    
    An alternative, which I will explore in a followup, is to always
    use SimplifyMultipleUseDemandedBits() -- the previous root special
    case is only really intended for SimplifyDemandedInstructionBits(),
    which now no longer shares this code path.
    
    Fixes https://github.com/llvm/llvm-project/issues/97330.

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 318c455fd7ef..64fbcc80e0ed 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -545,10 +545,11 @@ public:
                                ConstantInt *&Less, ConstantInt *&Equal,
                                ConstantInt *&Greater);
 
-  /// Attempts to replace V with a simpler value based on the demanded
+  /// Attempts to replace I with a simpler value based on the demanded
   /// bits.
-  Value *SimplifyDemandedUseBits(Value *V, APInt DemandedMask, KnownBits &Known,
-                                 unsigned Depth, const SimplifyQuery &Q);
+  Value *SimplifyDemandedUseBits(Instruction *I, const APInt &DemandedMask,
+                                 KnownBits &Known, unsigned Depth,
+                                 const SimplifyQuery &Q);
   using InstCombiner::SimplifyDemandedBits;
   bool SimplifyDemandedBits(Instruction *I, unsigned Op,
                             const APInt &DemandedMask, KnownBits &Known,
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
index 6cf2e71363ab..b1d03786f321 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
@@ -91,8 +91,47 @@ bool InstCombinerImpl::SimplifyDemandedBits(Instruction *I, unsigned OpNo,
                                             KnownBits &Known, unsigned Depth,
                                             const SimplifyQuery &Q) {
   Use &U = I->getOperandUse(OpNo);
-  Value *NewVal = SimplifyDemandedUseBits(U.get(), DemandedMask, Known,
-                                          Depth, Q);
+  Value *V = U.get();
+  if (isa<Constant>(V)) {
+    llvm::computeKnownBits(V, Known, Depth, Q);
+    return false;
+  }
+
+  Known.resetAll();
+  if (DemandedMask.isZero()) {
+    // Not demanding any bits from V.
+    replaceUse(U, UndefValue::get(V->getType()));
+    return true;
+  }
+
+  if (Depth == MaxAnalysisRecursionDepth)
+    return false;
+
+  Instruction *VInst = dyn_cast<Instruction>(V);
+  if (!VInst) {
+    llvm::computeKnownBits(V, Known, Depth, Q);
+    return false;
+  }
+
+  Value *NewVal;
+  if (VInst->hasOneUse()) {
+    // If the instruction has one use, we can directly simplify it.
+    NewVal = SimplifyDemandedUseBits(VInst, DemandedMask, Known, Depth, Q);
+  } else if (Depth != 0) {
+    // If there are multiple uses of this instruction and we aren't at the root,
+    // then we can simplify VInst to some other value, but not modify the
+    // instruction.
+    NewVal =
+        SimplifyMultipleUseDemandedBits(VInst, DemandedMask, Known, Depth, Q);
+  } else {
+    // If this is the root being simplified, allow it to have multiple uses,
+    // just set the DemandedMask to all bits and reset the context instruction.
+    // This allows visitTruncInst (for example) to simplify the operand of a
+    // trunc without duplicating all the SimplifyDemandedUseBits() logic.
+    NewVal =
+        SimplifyDemandedUseBits(VInst, APInt::getAllOnes(Known.getBitWidth()),
+                                Known, Depth, Q.getWithInstruction(VInst));
+  }
   if (!NewVal) return false;
   if (Instruction* OpInst = dyn_cast<Instruction>(U))
     salvageDebugInfo(*OpInst);
@@ -124,50 +163,21 @@ bool InstCombinerImpl::SimplifyDemandedBits(Instruction *I, unsigned OpNo,
 /// operands based on the information about what bits are demanded. This returns
 /// some other non-null value if it found out that V is equal to another value
 /// in the context where the specified bits are demanded, but not for all users.
-Value *InstCombinerImpl::SimplifyDemandedUseBits(Value *V, APInt DemandedMask,
+Value *InstCombinerImpl::SimplifyDemandedUseBits(Instruction *I,
+                                                 const APInt &DemandedMask,
                                                  KnownBits &Known,
                                                  unsigned Depth,
                                                  const SimplifyQuery &Q) {
-  assert(V != nullptr && "Null pointer of Value???");
+  assert(I != nullptr && "Null pointer of Value???");
   assert(Depth <= MaxAnalysisRecursionDepth && "Limit Search Depth");
   uint32_t BitWidth = DemandedMask.getBitWidth();
-  Type *VTy = V->getType();
+  Type *VTy = I->getType();
   assert(
       (!VTy->isIntOrIntVectorTy() || VTy->getScalarSizeInBits() == BitWidth) &&
       Known.getBitWidth() == BitWidth &&
       "Value *V, DemandedMask and Known must have same BitWidth");
 
-  if (isa<Constant>(V)) {
-    llvm::computeKnownBits(V, Known, Depth, Q);
-    return nullptr;
-  }
-
-  Known.resetAll();
-  if (DemandedMask.isZero()) // Not demanding any bits from V.
-    return UndefValue::get(VTy);
-
-  if (Depth == MaxAnalysisRecursionDepth)
-    return nullptr;
-
-  Instruction *I = dyn_cast<Instruction>(V);
-  if (!I) {
-    llvm::computeKnownBits(V, Known, Depth, Q);
-    return nullptr;        // Only analyze instructions.
-  }
-
-  // If there are multiple uses of this value and we aren't at the root, then
-  // we can't do any simplifications of the operands, because DemandedMask
-  // only reflects the bits demanded by *one* of the users.
-  if (Depth != 0 && !I->hasOneUse())
-    return SimplifyMultipleUseDemandedBits(I, DemandedMask, Known, Depth, Q);
-
   KnownBits LHSKnown(BitWidth), RHSKnown(BitWidth);
-  // If this is the root being simplified, allow it to have multiple uses,
-  // just set the DemandedMask to all bits so that we can try to simplify the
-  // operands.  This allows visitTruncInst (for example) to simplify the
-  // operand of a trunc without duplicating all the logic below.
-  if (Depth == 0 && !V->hasOneUse())
-    DemandedMask.setAllBits();
 
   // Update flags after simplifying an operand based on the fact that some high
   // order bits are not demanded.
@@ -1105,13 +1115,13 @@ Value *InstCombinerImpl::SimplifyDemandedUseBits(Value *V, APInt DemandedMask,
     }
 
     if (!KnownBitsComputed)
-      llvm::computeKnownBits(V, Known, Depth, Q);
+      llvm::computeKnownBits(I, Known, Depth, Q);
     break;
   }
   }
 
-  if (V->getType()->isPointerTy()) {
-    Align Alignment = V->getPointerAlignment(DL);
+  if (I->getType()->isPointerTy()) {
+    Align Alignment = I->getPointerAlignment(DL);
     Known.Zero.setLowBits(Log2(Alignment));
   }
 
@@ -1119,13 +1129,14 @@ Value *InstCombinerImpl::SimplifyDemandedUseBits(Value *V, APInt DemandedMask,
   // constant. We can't directly simplify pointers as a constant because of
   // pointer provenance.
   // TODO: We could return `(inttoptr const)` for pointers.
-  if (!V->getType()->isPointerTy() && DemandedMask.isSubsetOf(Known.Zero | Known.One))
+  if (!I->getType()->isPointerTy() &&
+      DemandedMask.isSubsetOf(Known.Zero | Known.One))
     return Constant::getIntegerValue(VTy, Known.One);
 
   if (VerifyKnownBits) {
-    KnownBits ReferenceKnown = llvm::computeKnownBits(V, Depth, Q);
+    KnownBits ReferenceKnown = llvm::computeKnownBits(I, Depth, Q);
     if (Known != ReferenceKnown) {
-      errs() << "Mismatched known bits for " << *V << " in "
+      errs() << "Mismatched known bits for " << *I << " in "
              << I->getFunction()->getName() << "\n";
       errs() << "computeKnownBits(): " << ReferenceKnown << "\n";
       errs() << "SimplifyDemandedBits(): " << Known << "\n";


Regression tests:
file: llvm/test/Transforms/InstCombine/assume-inseltpoison.ll
commands: opt < %s -passes=instcombine -S
tests: test_name: PR40940
test_body: target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"

; Function Attrs: nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write)
declare void @llvm.assume(i1 noundef) #0

define i32 @PR40940(<4 x i8> %x) {
  %shuf = shufflevector <4 x i8> %x, <4 x i8> poison, <4 x i32> <i32 1, i32 1, i32 2, i32 3>
  %t2 = bitcast <4 x i8> %shuf to i32
  %t3 = icmp ult i32 %t2, 65536
  call void @llvm.assume(i1 %t3)
  ret i32 %t2
}

attributes #0 = { nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write) }

file: llvm/test/Transforms/InstCombine/assume.ll
commands: opt < %s -passes=instcombine -S
opt < %s -passes=instcombine --enable-knowledge-retention -S
opt < %s -passes=instcombine -S --try-experimental-debuginfo-iterators
opt < %s -passes=instcombine --enable-knowledge-retention -S --try-experimental-debuginfo-iterators
tests: test_name: PR40940
test_body: target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"

; Function Attrs: nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write)
declare void @llvm.assume(i1 noundef) #0

define i32 @PR40940(<4 x i8> %x) {
  %shuf = shufflevector <4 x i8> %x, <4 x i8> undef, <4 x i32> <i32 1, i32 1, i32 2, i32 3>
  %t2 = bitcast <4 x i8> %shuf to i32
  %t3 = icmp ult i32 %t2, 65536
  call void @llvm.assume(i1 %t3)
  ret i32 %t2
}

attributes #0 = { nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write) }

!llvm.dbg.cu = !{}
!llvm.module.flags = !{!0, !1, !2, !3}

!0 = !{i32 2, !"Dwarf Version", i32 4}
!1 = !{i32 2, !"Debug Info Version", i32 3}
!2 = !{i32 1, !"wchar_size", i32 4}
!3 = !{i32 8, !"PIC Level", i32 2}

file: llvm/test/Transforms/InstCombine/known-bits.ll
commands: opt -S -passes=instcombine < %s
tests: test_name: simplifydemanded_context
test_body: define i8 @simplifydemanded_context(i8 %x, i8 %y) {
  %and1 = and i8 %x, 1
  call void @dummy()
  %x.lobits = and i8 %x, 3
  %precond = icmp eq i8 %x.lobits, 0
  call void @llvm.assume(i1 %precond)
  %and2 = and i8 %and1, %y
  ret i8 %and2
}

declare void @dummy()

; Function Attrs: nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write)
declare void @llvm.assume(i1 noundef) #0

attributes #0 = { nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write) }

test_name: pr97330
test_body: define i16 @pr97330(i1 %c, ptr %p1, ptr %p2) {
entry:
  %v = load i64, ptr %p1, align 8
  %conv = trunc i64 %v to i16
  br i1 %c, label %exit, label %if

if:                                               ; preds = %entry
  %cmp = icmp ne i16 %conv, 1
  %conv2 = zext i1 %cmp to i32
  store i32 %conv2, ptr %p2, align 4
  %cmp2 = icmp eq i64 %v, 1
  call void @llvm.assume(i1 %cmp2)
  unreachable

exit:                                             ; preds = %entry
  ret i16 %conv
}

; Function Attrs: nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write)
declare void @llvm.assume(i1 noundef) #0

attributes #0 = { nocallback nofree nosync nounwind willreturn memory(inaccessiblemem: write) }