[DropUnnecessaryAssumes] Add support for operand bundles

[nikic]

This extends the DropUnnecessaryAssumes pass to also handle operand bundle assumes. For this purpose, export the affected value analysis for operand bundles from AssumptionCache.

If the bundle only affects ephemeral values, drop it. If all bundles on an assume are dropped, drop the whole assume.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Nikita Popov (nikic)

Changes

This extends the DropUnnecessaryAssumes pass to also handle operand bundle assumes. For this purpose, export the affected value analysis for operand bundles from AssumptionCache.

If the bundle only affects ephemeral values, drop it. If all bundles on an assume are dropped, drop the whole assume.

---

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

4 Files Affected:

• (modified) llvm/include/llvm/Analysis/AssumptionCache.h (+6)
• (modified) llvm/lib/Analysis/AssumptionCache.cpp (+20-11)
• (modified) llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp (+47-9)
• (modified) llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll (+42-5)

diff --git a/llvm/include/llvm/Analysis/AssumptionCache.h b/llvm/include/llvm/Analysis/AssumptionCache.h
index 1b026ef76a45e..d5c56b023584a 100644
--- a/llvm/include/llvm/Analysis/AssumptionCache.h
+++ b/llvm/include/llvm/Analysis/AssumptionCache.h
@@ -28,6 +28,7 @@
 namespace llvm {
 
 class AssumeInst;
+struct OperandBundleUse;
 class Function;
 class raw_ostream;
 class TargetTransformInfo;
@@ -165,6 +166,11 @@ class AssumptionCache {
 
     return AVI->second;
   }
+
+  /// Determine which values are affected by this assume operand bundle.
+  static void
+  findValuesAffectedByOperandBundle(OperandBundleUse Bundle,
+                                    function_ref<void(Value *)> InsertAffected);
 };
 
 /// A function analysis which provides an \c AssumptionCache.
diff --git a/llvm/lib/Analysis/AssumptionCache.cpp b/llvm/lib/Analysis/AssumptionCache.cpp
index 980a891266e50..3ad5d47b5a143 100644
--- a/llvm/lib/Analysis/AssumptionCache.cpp
+++ b/llvm/lib/Analysis/AssumptionCache.cpp
@@ -53,6 +53,22 @@ AssumptionCache::getOrInsertAffectedValues(Value *V) {
   return AffectedValues[AffectedValueCallbackVH(V, this)];
 }
 
+void AssumptionCache::findValuesAffectedByOperandBundle(
+    OperandBundleUse Bundle, function_ref<void(Value *)> InsertAffected) {
+  auto AddAffectedVal = [&](Value *V) {
+    if (isa<Argument>(V) || isa<GlobalValue>(V) || isa<Instruction>(V))
+      InsertAffected(V);
+  };
+
+  if (Bundle.getTagName() == "separate_storage") {
+    assert(Bundle.Inputs.size() == 2 && "separate_storage must have two args");
+    AddAffectedVal(getUnderlyingObject(Bundle.Inputs[0]));
+    AddAffectedVal(getUnderlyingObject(Bundle.Inputs[1]));
+  } else if (Bundle.Inputs.size() > ABA_WasOn &&
+             Bundle.getTagName() != IgnoreBundleTag)
+    AddAffectedVal(Bundle.Inputs[ABA_WasOn]);
+}
+
 static void
 findAffectedValues(CallBase *CI, TargetTransformInfo *TTI,
                    SmallVectorImpl<AssumptionCache::ResultElem> &Affected) {
@@ -69,17 +85,10 @@ findAffectedValues(CallBase *CI, TargetTransformInfo *TTI,
     }
   };
 
-  for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++) {
-    OperandBundleUse Bundle = CI->getOperandBundleAt(Idx);
-    if (Bundle.getTagName() == "separate_storage") {
-      assert(Bundle.Inputs.size() == 2 &&
-             "separate_storage must have two args");
-      AddAffectedVal(getUnderlyingObject(Bundle.Inputs[0]), Idx);
-      AddAffectedVal(getUnderlyingObject(Bundle.Inputs[1]), Idx);
-    } else if (Bundle.Inputs.size() > ABA_WasOn &&
-               Bundle.getTagName() != IgnoreBundleTag)
-      AddAffectedVal(Bundle.Inputs[ABA_WasOn], Idx);
-  }
+  for (unsigned Idx = 0; Idx != CI->getNumOperandBundles(); Idx++)
+    AssumptionCache::findValuesAffectedByOperandBundle(
+        CI->getOperandBundleAt(Idx),
+        [&](Value *V) { Affected.push_back({V, Idx}); });
 
   Value *Cond = CI->getArgOperand(0);
   findValuesAffectedByCondition(Cond, /*IsAssume=*/true, InsertAffected);
diff --git a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
index c2e58ba393553..72cc98ea6f72d 100644
--- a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
+++ b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
@@ -16,6 +16,17 @@
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
+static bool
+affectedValuesAreEphemeral(const SmallPtrSetImpl<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())));
+}
+
 PreservedAnalyses
 DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
   AssumptionCache &AC = FAM.getResult<AssumptionAnalysis>(F);
@@ -26,8 +37,41 @@ DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
     if (!Assume)
       continue;
 
-    // TODO: Handle assumes with operand bundles.
-    if (Assume->hasOperandBundles())
+    SmallVector<WeakTrackingVH> DeadBundleArgs;
+    SmallVector<OperandBundleDef> KeptBundles;
+    unsigned NumBundles = Assume->getNumOperandBundles();
+    for (unsigned I = 0; I != NumBundles; ++I) {
+      // Handle operand bundle assumptions.
+      OperandBundleUse Bundle = Assume->getOperandBundleAt(I);
+      SmallPtrSet<Value *, 8> Affected;
+      AssumptionCache::findValuesAffectedByOperandBundle(
+          Bundle, [&](Value *A) { Affected.insert(A); });
+
+      if (affectedValuesAreEphemeral(Affected))
+        append_range(DeadBundleArgs, Bundle.Inputs);
+      else
+        KeptBundles.emplace_back(Bundle);
+    }
+
+    if (KeptBundles.size() != NumBundles) {
+      if (KeptBundles.size() == 0) {
+        // All operand bundles are dead, remove the whole assume.
+        Assume->eraseFromParent();
+      } else {
+        // Otherwise only drop the dead operand bundles.
+        CallBase *NewAssume =
+            CallBase::Create(Assume, KeptBundles, Assume->getIterator());
+        AC.registerAssumption(cast<AssumeInst>(NewAssume));
+        Assume->eraseFromParent();
+      }
+
+      RecursivelyDeleteTriviallyDeadInstructionsPermissive(DeadBundleArgs);
+      Changed = true;
+      continue;
+    }
+
+    // Ignore condition on assumes with operand bundles.
+    if (NumBundles != 0)
       continue;
 
     Value *Cond = Assume->getArgOperand(0);
@@ -39,13 +83,7 @@ DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
     findValuesAffectedByCondition(Cond, /*IsAssume=*/true,
                                   [&](Value *A) { Affected.insert(A); });
 
-    // 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.
-    if (!all_of(Affected, match_fn(m_OneUse(m_Value()))))
+    if (!affectedValuesAreEphemeral(Affected))
       continue;
 
     Assume->eraseFromParent();
diff --git a/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll b/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
index ea0d5d3fca8ff..f984b3db88337 100644
--- a/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
+++ b/llvm/test/Transforms/DropUnnecessaryAssumes/basic.ll
@@ -63,18 +63,17 @@ define i32 @multiple_live2(i32 %x, i32 %y) {
   ret i32 %y
 }
 
-define void @operand_bundle_dead(ptr %x) {
-; CHECK-LABEL: define void @operand_bundle_dead(
+define void @operand_bundle_one_dead(ptr %x) {
+; CHECK-LABEL: define void @operand_bundle_one_dead(
 ; CHECK-SAME: ptr [[X:%.*]]) {
-; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[X]], i64 8) ]
 ; CHECK-NEXT:    ret void
 ;
   call void @llvm.assume(i1 true) ["align"(ptr %x, i64 8)]
   ret void
 }
 
-define ptr @operand_bundle_live(ptr %x) {
-; CHECK-LABEL: define ptr @operand_bundle_live(
+define ptr @operand_bundle_one_live(ptr %x) {
+; CHECK-LABEL: define ptr @operand_bundle_one_live(
 ; CHECK-SAME: ptr [[X:%.*]]) {
 ; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[X]], i64 8) ]
 ; CHECK-NEXT:    ret ptr [[X]]
@@ -83,6 +82,44 @@ define ptr @operand_bundle_live(ptr %x) {
   ret ptr %x
 }
 
+define void @operand_bundle_multiple_dead(ptr %x, ptr %y) {
+; CHECK-LABEL: define void @operand_bundle_multiple_dead(
+; CHECK-SAME: ptr [[X:%.*]], ptr [[Y:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["align"(ptr %x, i64 8), "align"(ptr %y, i64 8)]
+  ret void
+}
+
+define ptr @operand_bundle_one_live_one_dead(ptr %x, ptr %y) {
+; CHECK-LABEL: define ptr @operand_bundle_one_live_one_dead(
+; CHECK-SAME: ptr [[X:%.*]], ptr [[Y:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[Y]], i64 8) ]
+; CHECK-NEXT:    ret ptr [[Y]]
+;
+  call void @llvm.assume(i1 true) ["align"(ptr %x, i64 8), "align"(ptr %y, i64 8)]
+  ret ptr %y
+}
+
+define i64 @operand_bundle_ignore_unaffected_operands(ptr %x, i64 %align) {
+; CHECK-LABEL: define i64 @operand_bundle_ignore_unaffected_operands(
+; CHECK-SAME: ptr [[X:%.*]], i64 [[ALIGN:%.*]]) {
+; CHECK-NEXT:    ret i64 [[ALIGN]]
+;
+  call void @llvm.assume(i1 true) ["align"(ptr %x, i64 %align)]
+  ret i64 %align
+}
+
+define void @operand_bundle_remove_dead_insts(ptr %x) {
+; CHECK-LABEL: define void @operand_bundle_remove_dead_insts(
+; CHECK-SAME: ptr [[X:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  %gep = getelementptr i8, ptr %x, i64 8
+  call void @llvm.assume(i1 true) ["align"(ptr %gep, i64 8)]
+  ret void
+}
+
 define void @type_test(ptr %x) {
 ; CHECK-LABEL: define void @type_test(
 ; CHECK-SAME: ptr [[X:%.*]]) {

[artagnon]

Thanks, this looks good from my end. I'll let @andjo403 and others have a look.

[andjo403]

Looks good to me

[fhahn]

LGTM, thanks