[DropUnnecessaryAssumes] Make the ephemeral value check more precise

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;
 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

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 @@ define void @type_test(ptr %x) {
   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
+}