[DropUnnecessaryAssumes] Add support for operand bundles

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.

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, GlobalValue, 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);

llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp

@@ -16,6 +16,16 @@
 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())));
+}
+
 PreservedAnalyses
 DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
   AssumptionCache &AC = FAM.getResult<AssumptionAnalysis>(F);
@@ -26,26 +36,64 @@ DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
     if (!Assume)
       continue;
 
-    // TODO: Handle assumes with operand bundles.
-    if (Assume->hasOperandBundles())
+    if (Assume->hasOperandBundles()) {
+      // Handle operand bundle assumptions.
+      SmallVector<WeakTrackingVH> DeadBundleArgs;
+      SmallVector<OperandBundleDef> KeptBundles;
+      unsigned NumBundles = Assume->getNumOperandBundles();
+      for (unsigned I = 0; I != NumBundles; ++I) {
+        auto IsDead = [](OperandBundleUse Bundle) {
+          // "ignore" operand bundles are always dead.
+          if (Bundle.getTagName() == "ignore")
+            return true;
+
+          // Bundles without arguments do not affect any specific values.
+          // Always keep them for now.
+          if (Bundle.Inputs.empty())
+            return false;
+
+          SmallVector<Value *> Affected;
+          AssumptionCache::findValuesAffectedByOperandBundle(
+              Bundle, [&](Value *A) { Affected.push_back(A); });
+
+          return affectedValuesAreEphemeral(Affected);
+        };
+
+        OperandBundleUse Bundle = Assume->getOperandBundleAt(I);
+        if (IsDead(Bundle))
+          append_range(DeadBundleArgs, Bundle.Inputs);
+        else
+          KeptBundles.emplace_back(Bundle);
+      }
+
+      if (KeptBundles.size() != NumBundles) {
+        if (KeptBundles.empty()) {
+          // 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;
+    }
 
     Value *Cond = Assume->getArgOperand(0);
     // Don't drop type tests, which have special semantics.
     if (match(Cond, m_Intrinsic<Intrinsic::type_test>()))
       continue;
 
-    SmallPtrSet<Value *, 8> Affected;
+    SmallVector<Value *> Affected;
     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()))))
+                                  [&](Value *A) { Affected.push_back(A); });
+
+    if (!affectedValuesAreEphemeral(Affected))
       continue;
 
     Assume->eraseFromParent();

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,93 @@ 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 @operand_bundle_no_args() {
+; CHECK-LABEL: define void @operand_bundle_no_args() {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "cold"() ]
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["cold"()]
+  ret void
+}
+
+; Can always drop ignore bundles, regardless of uses.
+define ptr @operand_bundle_ignore(ptr %x) {
+; CHECK-LABEL: define ptr @operand_bundle_ignore(
+; CHECK-SAME: ptr [[X:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(ptr [[X]]) ]
+; CHECK-NEXT:    ret ptr [[X]]
+;
+  call void @llvm.assume(i1 true) ["ignore"(), "ignore"(ptr %x), "nonnull"(ptr %x)]
+  ret ptr %x
+}
+
+define void @operand_bundle_separate_storage_both_dead(ptr %x, ptr %y) {
+; CHECK-LABEL: define void @operand_bundle_separate_storage_both_dead(
+; CHECK-SAME: ptr [[X:%.*]], ptr [[Y:%.*]]) {
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.assume(i1 true) ["separate_storage"(ptr %x, ptr %y)]
+  ret void
+}
+
+define ptr @operand_bundle_separate_storage_one_live1(ptr %x, ptr %y) {
+; CHECK-LABEL: define ptr @operand_bundle_separate_storage_one_live1(
+; CHECK-SAME: ptr [[X:%.*]], ptr [[Y:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "separate_storage"(ptr [[X]], ptr [[Y]]) ]
+; CHECK-NEXT:    ret ptr [[Y]]
+;
+  call void @llvm.assume(i1 true) ["separate_storage"(ptr %x, ptr %y)]
+  ret ptr %y
+}
+
+define ptr @operand_bundle_separate_storage_one_live2(ptr %x, ptr %y) {
+; CHECK-LABEL: define ptr @operand_bundle_separate_storage_one_live2(
+; CHECK-SAME: ptr [[X:%.*]], ptr [[Y:%.*]]) {
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "separate_storage"(ptr [[X]], ptr [[Y]]) ]
+; CHECK-NEXT:    ret ptr [[X]]
+;
+  call void @llvm.assume(i1 true) ["separate_storage"(ptr %x, ptr %y)]
+  ret ptr %x
+}
+
 define void @type_test(ptr %x) {
 ; CHECK-LABEL: define void @type_test(
 ; CHECK-SAME: ptr [[X:%.*]]) {