Basis of dropping uses in llvm.assume.

Summary: This patch adds the basic utilities to deal with dropable uses. dropable uses are uses that we rather drop than prevent transformations, for now they are limited to uses in llvm.assume.

Reviewers: jdoerfert, sstefan1

Reviewed By: jdoerfert

Subscribers: uenoku, lebedev.ri, mgorny, hiraditya, dexonsmith, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D73404

llvm/docs/LangRef.rst

@@ -2119,8 +2119,9 @@ An assume operand bundle has the form:
 
       "<tag>"([ <holds for value> [, <attribute argument>] ])
 
-* The tag of the operand bundle is the name of attribute that can be assumed
-  to hold.
+* The tag of the operand bundle is usually the name of attribute that can be
+  assumed to hold. It can also be `ignore`, this tag doesn't contain any
+  information and should be ignored.
 * The first argument if present is the value for which the attribute hold.
 * The second argument if present is an argument of the attribute.
 

llvm/include/llvm/IR/User.h

@@ -218,6 +218,11 @@ public:
     NumUserOperands = NumOps;
   }
 
+  /// A droppable user is a user for which uses can be dropped without affecting
+  /// correctness and should be dropped rather than preventing a transformation
+  /// from happening.
+  bool isDroppable() const;
+
   // ---------------------------------------------------------------------------
   // Operand Iterator interface...
   //

llvm/include/llvm/IR/Value.h

@@ -444,6 +444,34 @@ public:
   /// This is logically equivalent to getNumUses() >= N.
   bool hasNUsesOrMore(unsigned N) const;
 
+  /// Return true if there is exactly one user of this value that cannot be
+  /// dropped.
+  ///
+  /// This is specialized because it is a common request and does not require
+  /// traversing the whole use list.
+  Use *getSingleUndroppableUse();
+
+  /// Return true if there this value.
+  ///
+  /// This is specialized because it is a common request and does not require
+  /// traversing the whole use list.
+  bool hasNUndroppableUses(unsigned N) const;
+
+  /// Return true if this value has N users or more.
+  ///
+  /// This is logically equivalent to getNumUses() >= N.
+  bool hasNUndroppableUsesOrMore(unsigned N) const;
+
+  /// Remove every uses that can safely be removed.
+  ///
+  /// This will remove for example uses in llvm.assume.
+  /// This should be used when performing want to perform a tranformation but
+  /// some Droppable uses pervent it.
+  /// This function optionally takes a filter to only remove some droppable
+  /// uses.
+  void dropDroppableUses(llvm::function_ref<bool(const Use *)> ShouldDrop =
+                             [](const Use *) { return true; });
+
   /// Check if this value is used in the specified basic block.
   bool isUsedInBasicBlock(const BasicBlock *BB) const;
 

llvm/lib/IR/KnowledgeRetention.cpp

@@ -197,51 +197,27 @@ bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
   if (Assume.bundle_op_infos().empty())
     return false;
 
-  CallInst::bundle_op_iterator Lookup;
-
-  /// The right attribute can be found by binary search. After this finding the
-  /// right WasOn needs to be done via linear search.
-  /// Element have been ordered by argument value so the first we find is the
-  /// one we need.
-  if (AQR == AssumeQuery::Lowest)
-    Lookup =
-        llvm::lower_bound(Assume.bundle_op_infos(), AttrName,
-                          [](const CallBase::BundleOpInfo &BOI, StringRef RHS) {
-                            return BOI.Tag->getKey() < RHS;
-                          });
-  else
-    Lookup = std::prev(
-        llvm::upper_bound(Assume.bundle_op_infos(), AttrName,
-                          [](StringRef LHS, const CallBase::BundleOpInfo &BOI) {
-                            return LHS < BOI.Tag->getKey();
-                          }));
-
-  if (Lookup == Assume.bundle_op_info_end() ||
-      Lookup->Tag->getKey() != AttrName)
-    return false;
-  if (IsOn) {
-    assert((Lookup->End - Lookup->Begin > BOIE_WasOn) &&
-           "missing argument of attribute");
-    while (true) {
-      if (Lookup == Assume.bundle_op_info_end() ||
-          Lookup->Tag->getKey() != AttrName)
-        return false;
-      if (getValueFromBundleOpInfo(Assume, *Lookup, BOIE_WasOn) == IsOn)
-        break;
-      if (AQR == AssumeQuery::Highest &&
-          Lookup == Assume.bundle_op_info_begin())
-        return false;
-      Lookup = Lookup + (AQR == AssumeQuery::Lowest ? 1 : -1);
+  auto Loop = [&](auto &&Range) {
+    for (auto &BOI : Range) {
+      if (BOI.Tag->getKey() != AttrName)
+        continue;
+      if (IsOn && (BOI.End - BOI.Begin <= BOIE_WasOn ||
+                   IsOn != getValueFromBundleOpInfo(Assume, BOI, BOIE_WasOn)))
+        continue;
+      if (ArgVal) {
+        assert(BOI.End - BOI.Begin > BOIE_Argument);
+        *ArgVal = cast<ConstantInt>(
+                      getValueFromBundleOpInfo(Assume, BOI, BOIE_Argument))
+                      ->getZExtValue();
+      }
+      return true;
     }
-  }
+    return false;
+  };
 
-  if (Lookup->End - Lookup->Begin < BOIE_Argument)
-    return true;
-  if (ArgVal)
-    *ArgVal = cast<ConstantInt>(
-                  getValueFromBundleOpInfo(Assume, *Lookup, BOIE_Argument))
-                  ->getZExtValue();
-  return true;
+  if (AQR == AssumeQuery::Lowest)
+    return Loop(Assume.bundle_op_infos());
+  return Loop(reverse(Assume.bundle_op_infos()));
 }
 
 void llvm::fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result) {

llvm/lib/IR/User.cpp

@@ -9,6 +9,7 @@
 #include "llvm/IR/User.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/IntrinsicInst.h"
 
 namespace llvm {
 class BasicBlock;
@@ -105,6 +106,12 @@ MutableArrayRef<uint8_t> User::getDescriptor() {
       reinterpret_cast<uint8_t *>(DI) - DI->SizeInBytes, DI->SizeInBytes);
 }
 
+bool User::isDroppable() const {
+  if (const auto *Intr = dyn_cast<IntrinsicInst>(this))
+    return Intr->getIntrinsicID() == Intrinsic::assume;
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 //                         User operator new Implementations
 //===----------------------------------------------------------------------===//

llvm/lib/IR/Value.cpp

@@ -141,6 +141,51 @@ bool Value::hasNUsesOrMore(unsigned N) const {
   return hasNItemsOrMore(use_begin(), use_end(), N);
 }
 
+static bool isUnDroppableUser(const User *U) { return !U->isDroppable(); }
+
+Use *Value::getSingleUndroppableUse() {
+  Use *Result = nullptr;
+  for (Use &U : uses()) {
+    if (!U.getUser()->isDroppable()) {
+      if (Result)
+        return nullptr;
+      Result = &U;
+    }
+  }
+  return Result;
+}
+
+bool Value::hasNUndroppableUses(unsigned int N) const {
+  return hasNItems(user_begin(), user_end(), N, isUnDroppableUser);
+}
+
+bool Value::hasNUndroppableUsesOrMore(unsigned int N) const {
+  return hasNItemsOrMore(user_begin(), user_end(), N, isUnDroppableUser);
+}
+
+void Value::dropDroppableUses(
+    llvm::function_ref<bool(const Use *)> ShouldDrop) {
+  SmallVector<Use *, 8> ToBeEdited;
+  for (Use &U : uses())
+    if (U.getUser()->isDroppable() && ShouldDrop(&U))
+      ToBeEdited.push_back(&U);
+  for (Use *U : ToBeEdited) {
+    U->removeFromList();
+    if (auto *Assume = dyn_cast<IntrinsicInst>(U->getUser())) {
+      assert(Assume->getIntrinsicID() == Intrinsic::assume);
+      unsigned OpNo = U->getOperandNo();
+      if (OpNo == 0)
+        Assume->setOperand(0, ConstantInt::getTrue(Assume->getContext()));
+      else {
+        Assume->setOperand(OpNo, UndefValue::get(U->get()->getType()));
+        CallInst::BundleOpInfo &BOI = Assume->getBundleOpInfoForOperand(OpNo);
+        BOI.Tag = getContext().pImpl->getOrInsertBundleTag("ignore");
+      }
+    } else
+      llvm_unreachable("unkown droppable use");
+  }
+}
+
 bool Value::isUsedInBasicBlock(const BasicBlock *BB) const {
   // This can be computed either by scanning the instructions in BB, or by
   // scanning the use list of this Value. Both lists can be very long, but

llvm/lib/IR/Verifier.cpp

@@ -4322,7 +4322,8 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
     break;
   case Intrinsic::assume: {
     for (auto &Elem : Call.bundle_op_infos()) {
-      Assert(Attribute::isExistingAttribute(Elem.Tag->getKey()),
+      Assert(Elem.Tag->getKey() == "ignore" ||
+                 Attribute::isExistingAttribute(Elem.Tag->getKey()),
              "tags must be valid attribute names");
       Assert(Elem.End - Elem.Begin <= 2, "to many arguments");
       Attribute::AttrKind Kind =

llvm/unittests/IR/KnowledgeRetentionTest.cpp

@@ -25,21 +25,17 @@ static void RunTest(
     StringRef Head, StringRef Tail,
     std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
         &Tests) {
-  std::string IR;
-  IR.append(Head.begin(), Head.end());
-  for (auto &Elem : Tests)
+  for (auto &Elem : Tests) {
+    std::string IR;
+    IR.append(Head.begin(), Head.end());
     IR.append(Elem.first.begin(), Elem.first.end());
-  IR.append(Tail.begin(), Tail.end());
-  LLVMContext C;
-  SMDiagnostic Err;
-  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
-  if (!Mod)
-    Err.print("AssumeQueryAPI", errs());
-  unsigned Idx = 0;
-  for (Instruction &I : (*Mod->getFunction("test")->begin())) {
-    if (Idx < Tests.size())
-      Tests[Idx].second(&I);
-    Idx++;
+    IR.append(Tail.begin(), Tail.end());
+    LLVMContext C;
+    SMDiagnostic Err;
+    std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+    if (!Mod)
+      Err.print("AssumeQueryAPI", errs());
+    Elem.second(&*(Mod->getFunction("test")->begin()->begin()));
   }
 }
 
@@ -199,7 +195,29 @@ TEST(AssumeQueryAPI, hasAttributeInAssume) {
                                Attribute::AttrKind::Dereferenceable, 12, true);
       }));
 
-  /// Keep this test last as it modifies the function.
+  Tests.push_back(std::make_pair(
+      "call void @func1(i32* readnone align 32 "
+      "dereferenceable(48) noalias %P, i32* "
+      "align 8 dereferenceable(28) %P1, i32* align 64 "
+      "dereferenceable(4) "
+      "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+        I->getOperand(1)->dropDroppableUses();
+        I->getOperand(2)->dropDroppableUses();
+        I->getOperand(3)->dropDroppableUses();
+        AssertMatchesExactlyAttributes(
+            Assume, I->getOperand(0),
+            "(readnone|align|dereferenceable|noalias)");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(1), "");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(2), "");
+        AssertMatchesExactlyAttributes(Assume, I->getOperand(3), "");
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Alignment, 32, true);
+        AssertHasTheRightValue(Assume, I->getOperand(0),
+                               Attribute::AttrKind::Dereferenceable, 48, true);
+      }));
   Tests.push_back(std::make_pair(
       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
       "8 noalias %P1)\n",