[DropUnnecessaryAssumes] Make the ephemeral value check more precise (#160700)

The initial implementation used a very crude check where a value was
considered ephemeral if it has only one use. This is insufficient if
there are multiple assumes acting on the same value, or in more complex
cases like cyclic phis.

Generalize this to a more typical ephemeral value check, i.e. make sure
that all transitive users are in assumes, while stopping at
side-effecting instructions.
diff --git a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
index c215228..89980d5 100644
--- a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
+++ b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/DropUnnecessaryAssumes.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -17,13 +18,48 @@
 using namespace llvm::PatternMatch;
 
 static bool affectedValuesAreEphemeral(ArrayRef<Value *> Affected) {
-  // If all the affected uses have only one use (part of the assume), then
-  // the assume does not provide useful information. Note that additional
-  // users may appear as a result of inlining and CSE, so we should only
-  // make this assumption late in the optimization pipeline.
-  // TODO: Handle dead cyclic usages.
-  // TODO: Handle multiple dead assumes on the same value.
-  return all_of(Affected, match_fn(m_OneUse(m_Value())));
+  // Check whether all the uses are ephemeral, i.e. recursively only used
+  // by assumes. In that case, the assume does not provide useful information.
+  // Note that additional users may appear as a result of inlining and CSE,
+  // so we should only make this assumption late in the optimization pipeline.
+  SmallSetVector<Instruction *, 32> Worklist;
+  auto AddUsers = [&](Value *V) {
+    for (User *U : V->users()) {
+      // Bail out if we need to inspect too many users.
+      if (Worklist.size() >= 32)
+        return false;
+      Worklist.insert(cast<Instruction>(U));
+    }
+    return true;
+  };
+
+  for (Value *V : Affected) {
+    // Do not handle assumes on globals for now. The use list for them may
+    // contain uses in other functions.
+    if (!isa<Instruction, Argument>(V))
+      return false;
+
+    if (!AddUsers(V))
+      return false;
+  }
+
+  for (unsigned Idx = 0; Idx < Worklist.size(); ++Idx) {
+    Instruction *I = Worklist[Idx];
+
+    // Use in assume is ephemeral.
+    if (isa<AssumeInst>(I))
+      continue;
+
+    // Use in side-effecting instruction is non-ephemeral.
+    if (I->mayHaveSideEffects() || I->isTerminator())
+      return false;
+
+    // Otherwise, recursively look at the users.
+    if (!AddUsers(I))
+      return false;
+  }
+
+  return true;
 }
 
 PreservedAnalyses
diff --git a/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll b/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
index e2a9b4e..8a6f60b 100644
--- a/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
+++ b/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
@@ -1,6 +1,9 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
 ; RUN: opt -S -passes=drop-unnecessary-assumes < %s | FileCheck %s
 
+declare void @use(i32 %x)
+declare i32 @get()
+
 define void @basic_dead(i32 %x) {
 ; CHECK-LABEL: define void @basic_dead(
 ; CHECK-SAME: i32 [[X:%.*]]) {
@@ -180,3 +183,136 @@
   call void @llvm.assume(i1 %test)
   ret void
 }
+
+define void @multiple_dead_conds(i32 %x) {
+; CHECK-LABEL: define void @multiple_dead_conds(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  %cond1 = icmp sge i32 %x, 0
+  call void @llvm.assume(i1 %cond1)
+  %cond2 = icmp ne i32 %x, 64
+  call void @llvm.assume(i1 %cond2)
+  ret void
+}
+
+define void @multiple_dead_bundles(ptr %x) {
+; CHECK-LABEL: define void @multiple_dead_bundles(
+; CHECK-SAME: ptr [[X:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["align"(ptr %x, i64 8), "nonnull"(ptr %x)]
+  ret void
+}
+
+; The assume is eliminated, but currently leaves behind a dead cycle.
+define void @dead_cycle(i1 %loop.cond) {
+; CHECK-LABEL: define void @dead_cycle(
+; CHECK-SAME: i1 [[LOOP_COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    br i1 [[LOOP_COND]], label %[[LOOP]], label %[[EXIT:.*]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ]
+  %cond = icmp ne i32 %iv, 64
+  call void @llvm.assume(i1 %cond)
+  %iv.next = add i32 %iv, 1
+  br i1 %loop.cond, label %loop, label %exit
+
+exit:
+  ret void
+}
+
+define void @use_in_side_effect(i32 %x) {
+; CHECK-LABEL: define void @use_in_side_effect(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[COND:%.*]] = icmp sge i32 [[X]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    call void @use(i32 [[X]])
+; CHECK-NEXT:    ret void
+;
+  %cond = icmp sge i32 %x, 0
+  call void @llvm.assume(i1 %cond)
+  call void @use(i32 %x)
+  ret void
+}
+
+define void @indirect_use_in_side_effect(i32 %x) {
+; CHECK-LABEL: define void @indirect_use_in_side_effect(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[COND:%.*]] = icmp sge i32 [[X]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[X]], 1
+; CHECK-NEXT:    call void @use(i32 [[ADD]])
+; CHECK-NEXT:    ret void
+;
+  %cond = icmp sge i32 %x, 0
+  call void @llvm.assume(i1 %cond)
+  %add = add i32 %x, 1
+  call void @use(i32 %add)
+  ret void
+}
+
+; The affected value itself has a side effect, but we can still drop the
+; assume.
+define void @affected_value_has_side_effect() {
+; CHECK-LABEL: define void @affected_value_has_side_effect() {
+; CHECK-NEXT:    [[X:%.*]] = call i32 @get()
+; CHECK-NEXT:    ret void
+;
+  %x = call i32 @get()
+  %cond = icmp sge i32 %x, 0
+  call void @llvm.assume(i1 %cond)
+  ret void
+}
+
+define i32 @affected_value_has_side_effect_and_is_used() {
+; CHECK-LABEL: define i32 @affected_value_has_side_effect_and_is_used() {
+; CHECK-NEXT:    [[X:%.*]] = call i32 @get()
+; CHECK-NEXT:    [[COND:%.*]] = icmp sge i32 [[X]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  %x = call i32 @get()
+  %cond = icmp sge i32 %x, 0
+  call void @llvm.assume(i1 %cond)
+  ret i32 %x
+}
+
+@g = external global i8
+@g2 = external global i8
+
+; Assumes on globals are currently not supported.
+define void @assume_on_global() {
+; CHECK-LABEL: define void @assume_on_global() {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr @g, i64 8) ]
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["align"(ptr @g, i64 8)]
+  ret void
+}
+
+define void @assume_on_global_used_in_other_func() {
+; CHECK-LABEL: define void @assume_on_global_used_in_other_func() {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr @g2, i64 8) ]
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["align"(ptr @g2, i64 8)]
+  ret void
+}
+
+define ptr @other_func() {
+; CHECK-LABEL: define ptr @other_func() {
+; CHECK-NEXT:    ret ptr @g2
+;
+  ret ptr @g2
+}