[DropUnnecessaryAssumes] Make the ephemeral value check more precise

[nikic]

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.

Interestingly, this doesn't seem to actually matter in practice (no difference on llvm-opt-benchmark).

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

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.

Interestingly, this doesn't seem to actually matter in practice (no difference on llvm-opt-benchmark).

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp (+39-7)
• (modified) llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll (+80)

diff --git a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
index c215228b480d2..d4fb420211cb0 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,44 @@ 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<User *, 16> Worklist;
+  auto AddUser = [&](User *U) {
+    // Bail out if we need to inspect too many users.
+    if (Worklist.size() >= 32)
+      return false;
+    Worklist.insert(U);
+    return true;
+  };
+
+  for (Value *V : Affected)
+    for (User *U : V->users())
+      if (!AddUser(U))
+        return false;
+
+  for (unsigned Idx = 0; Idx < Worklist.size(); ++Idx) {
+    auto *I = dyn_cast<Instruction>(Worklist[Idx]);
+    if (!I)
+      return false;
+
+    // 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.
+    for (User *NestedU : I->users())
+      if (!AddUser(NestedU))
+        return false;
+  }
+
+  return true;
 }
 
 PreservedAnalyses
diff --git a/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll b/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
index e2a9b4eea2c7d..a61dc2dcc4b9f 100644
--- a/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
+++ b/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
@@ -1,6 +1,8 @@
 ; 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 @dummy(i32 %x)
+
 define void @basic_dead(i32 %x) {
 ; CHECK-LABEL: define void @basic_dead(
 ; CHECK-SAME: i32 [[X:%.*]]) {
@@ -180,3 +182,81 @@ 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 @dummy(i32 [[X]])
+; CHECK-NEXT:    ret void
+;
+  %cond = icmp sge i32 %x, 0
+  call void @llvm.assume(i1 %cond)
+  call void @dummy(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 @dummy(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 @dummy(i32 %add)
+  ret void
+}

[artagnon]

Interestingly, this doesn't seem to actually matter in practice (no difference on llvm-opt-benchmark).

Not sure how I feel about this. I've personally abandoned patches in the past due to zero-diff (or even small diff) on llvm-opt-benchmark. The patch is certainly nice for completeness though?

[nikic]

Interestingly, this doesn't seem to actually matter in practice (no difference on llvm-opt-benchmark).

Not sure how I feel about this. I've personally abandoned patches in the past due to zero-diff (or even small diff) on llvm-opt-benchmark. The patch is certainly nice for completeness though?

Generally agreed. I do think it's valuable to have this to make sure the pass behavior is consistent (i.e. the pass shouldn't break just because we added a second dead assume).

[artagnon]

Okay, let's use the opportunity to think about some ideas then: could you kindly think about collectUsersRecursively?

[artagnon]

After some more thought, I think a collectUsersRecursively is probably not the best idea. LGTM, thanks!