Plumb AssumeInst through operand bundle apis [nfc]

Follow up to a6d2a8d.  This covers all the public interfaces of the bundle related code.  I tried to cleanup the internals where the changes were obvious, but there's definitely more room for improvement.

Author: preames

llvm/include/llvm/Analysis/AssumeBundleQueries.h

@@ -16,6 +16,7 @@
 
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace llvm {
@@ -39,12 +40,12 @@ enum AssumeBundleArg {
 ///
 /// Return true iff the queried attribute was found.
 /// If ArgVal is set. the argument will be stored to ArgVal.
-bool hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn, StringRef AttrName,
+bool hasAttributeInAssume(AssumeInst &Assume, Value *IsOn, StringRef AttrName,
                           uint64_t *ArgVal = nullptr);
-inline bool hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
+inline bool hasAttributeInAssume(AssumeInst &Assume, Value *IsOn,
                                  Attribute::AttrKind Kind,
                                  uint64_t *ArgVal = nullptr) {
-  return hasAttributeInAssume(AssumeCI, IsOn,
+  return hasAttributeInAssume(Assume, IsOn,
                               Attribute::getNameFromAttrKind(Kind), ArgVal);
 }
 
@@ -87,7 +88,7 @@ using RetainedKnowledgeMap =
 /// many queries are going to be made on the same llvm.assume.
 /// String attributes are not inserted in the map.
 /// If the IR changes the map will be outdated.
-void fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result);
+void fillMapFromAssume(AssumeInst &Assume, RetainedKnowledgeMap &Result);
 
 /// Represent one information held inside an operand bundle of an llvm.assume.
 /// AttrKind is the property that holds.
@@ -121,13 +122,13 @@ struct RetainedKnowledge {
 
 /// Retreive the information help by Assume on the operand at index Idx.
 /// Assume should be an llvm.assume and Idx should be in the operand bundle.
-RetainedKnowledge getKnowledgeFromOperandInAssume(CallInst &Assume,
+RetainedKnowledge getKnowledgeFromOperandInAssume(AssumeInst &Assume,
                                                   unsigned Idx);
 
 /// Retreive the information help by the Use U of an llvm.assume. the use should
 /// be in the operand bundle.
 inline RetainedKnowledge getKnowledgeFromUseInAssume(const Use *U) {
-  return getKnowledgeFromOperandInAssume(*cast<CallInst>(U->getUser()),
+  return getKnowledgeFromOperandInAssume(*cast<AssumeInst>(U->getUser()),
                                          U->getOperandNo());
 }
 
@@ -142,7 +143,7 @@ constexpr StringRef IgnoreBundleTag = "ignore";
 ///
 /// the argument to the call of llvm.assume may still be useful even if the
 /// function returned true.
-bool isAssumeWithEmptyBundle(CallInst &Assume);
+bool isAssumeWithEmptyBundle(AssumeInst &Assume);
 
 /// Return a valid Knowledge associated to the Use U if its Attribute kind is
 /// in AttrKinds.
@@ -168,7 +169,7 @@ RetainedKnowledge getKnowledgeValidInContext(
 
 /// This extracts the Knowledge from an element of an operand bundle.
 /// This is mostly for use in the assume builder.
-RetainedKnowledge getKnowledgeFromBundle(CallInst &Assume,
+RetainedKnowledge getKnowledgeFromBundle(AssumeInst &Assume,
                                          const CallBase::BundleOpInfo &BOI);
 
 } // namespace llvm

llvm/include/llvm/Transforms/Scalar/GVN.h

@@ -35,6 +35,7 @@
 namespace llvm {
 
 class AAResults;
+class AssumeInst;
 class AssumptionCache;
 class BasicBlock;
 class BranchInst;
@@ -309,7 +310,7 @@ class GVN : public PassInfoMixin<GVN> {
   // Helper functions of redundant load elimination
   bool processLoad(LoadInst *L);
   bool processNonLocalLoad(LoadInst *L);
-  bool processAssumeIntrinsic(IntrinsicInst *II);
+  bool processAssumeIntrinsic(AssumeInst *II);
 
   /// Given a local dependency (Def or Clobber) determine if a value is
   /// available for the load.  Returns true if an value is known to be

llvm/include/llvm/Transforms/Utils/AssumeBundleBuilder.h

@@ -30,7 +30,7 @@ class DominatorTree;
 /// from the given instruction.
 /// If no information derived from \p I, this call returns null.
 /// The returned instruction is not inserted anywhere.
-IntrinsicInst *buildAssumeFromInst(Instruction *I);
+AssumeInst *buildAssumeFromInst(Instruction *I);
 
 /// Calls BuildAssumeFromInst and if the resulting llvm.assume is valid insert
 /// if before I. This is usually what need to be done to salvage the knowledge
@@ -44,10 +44,10 @@ void salvageKnowledge(Instruction *I, AssumptionCache *AC = nullptr,
 
 /// Build and return a new assume created from the provided knowledge
 /// if the knowledge in the assume is fully redundant this will return nullptr
-IntrinsicInst *buildAssumeFromKnowledge(ArrayRef<RetainedKnowledge> Knowledge,
-                                        Instruction *CtxI,
-                                        AssumptionCache *AC = nullptr,
-                                        DominatorTree *DT = nullptr);
+AssumeInst *buildAssumeFromKnowledge(ArrayRef<RetainedKnowledge> Knowledge,
+                                     Instruction *CtxI,
+                                     AssumptionCache *AC = nullptr,
+                                     DominatorTree *DT = nullptr);
 
 /// This pass attempts to minimize the number of assume without loosing any
 /// information.
@@ -66,7 +66,7 @@ struct AssumeBuilderPass : public PassInfoMixin<AssumeBuilderPass> {
 /// canonicalize the RetainedKnowledge RK. it is assumed that RK is part of
 /// Assume. This will return an empty RetainedKnowledge if the knowledge is
 /// useless.
-RetainedKnowledge simplifyRetainedKnowledge(CallBase *Assume,
+RetainedKnowledge simplifyRetainedKnowledge(AssumeInst *Assume,
                                             RetainedKnowledge RK,
                                             AssumptionCache *AC,
                                             DominatorTree *DT);

llvm/lib/Analysis/AssumeBundleQueries.cpp

@@ -33,20 +33,15 @@ static bool bundleHasArgument(const CallBase::BundleOpInfo &BOI, unsigned Idx) {
   return BOI.End - BOI.Begin > Idx;
 }
 
-static Value *getValueFromBundleOpInfo(CallInst &Assume,
+static Value *getValueFromBundleOpInfo(AssumeInst &Assume,
                                        const CallBase::BundleOpInfo &BOI,
                                        unsigned Idx) {
   assert(bundleHasArgument(BOI, Idx) && "index out of range");
   return (Assume.op_begin() + BOI.Begin + Idx)->get();
 }
 
-bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
+bool llvm::hasAttributeInAssume(AssumeInst &Assume, Value *IsOn,
                                 StringRef AttrName, uint64_t *ArgVal) {
-  assert(isa<IntrinsicInst>(AssumeCI) &&
-         "this function is intended to be used on llvm.assume");
-  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
-  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
-         "this function is intended to be used on llvm.assume");
   assert(Attribute::isExistingAttribute(AttrName) &&
          "this attribute doesn't exist");
   assert((ArgVal == nullptr || Attribute::doesAttrKindHaveArgument(
@@ -72,10 +67,7 @@ bool llvm::hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
   return false;
 }
 
-void llvm::fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result) {
-  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
-  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
-         "this function is intended to be used on llvm.assume");
+void llvm::fillMapFromAssume(AssumeInst &Assume, RetainedKnowledgeMap &Result) {
   for (auto &Bundles : Assume.bundle_op_infos()) {
     std::pair<Value *, Attribute::AttrKind> Key{
         nullptr, Attribute::getAttrKindFromName(Bundles.Tag->getKey())};
@@ -104,7 +96,7 @@ void llvm::fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result) {
 }
 
 RetainedKnowledge
-llvm::getKnowledgeFromBundle(CallInst &Assume,
+llvm::getKnowledgeFromBundle(AssumeInst &Assume,
                              const CallBase::BundleOpInfo &BOI) {
   RetainedKnowledge Result;
   Result.AttrKind = Attribute::getAttrKindFromName(BOI.Tag->getKey());
@@ -124,19 +116,13 @@ llvm::getKnowledgeFromBundle(CallInst &Assume,
   return Result;
 }
 
-RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(CallInst &AssumeCI,
+RetainedKnowledge llvm::getKnowledgeFromOperandInAssume(AssumeInst &Assume,
                                                         unsigned Idx) {
-  IntrinsicInst &Assume = cast<IntrinsicInst>(AssumeCI);
-  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
-         "this function is intended to be used on llvm.assume");
   CallBase::BundleOpInfo BOI = Assume.getBundleOpInfoForOperand(Idx);
-  return getKnowledgeFromBundle(AssumeCI, BOI);
+  return getKnowledgeFromBundle(Assume, BOI);
 }
 
-bool llvm::isAssumeWithEmptyBundle(CallInst &CI) {
-  IntrinsicInst &Assume = cast<IntrinsicInst>(CI);
-  assert(Assume.getIntrinsicID() == Intrinsic::assume &&
-         "this function is intended to be used on llvm.assume");
+bool llvm::isAssumeWithEmptyBundle(AssumeInst &Assume) {
   return none_of(Assume.bundle_op_infos(),
                  [](const CallBase::BundleOpInfo &BOI) {
                    return BOI.Tag->getKey() != IgnoreBundleTag;
@@ -158,7 +144,7 @@ llvm::getKnowledgeFromUse(const Use *U,
   if (!Bundle)
     return RetainedKnowledge::none();
   RetainedKnowledge RK =
-      getKnowledgeFromBundle(*cast<CallInst>(U->getUser()), *Bundle);
+      getKnowledgeFromBundle(*cast<AssumeInst>(U->getUser()), *Bundle);
   if (llvm::is_contained(AttrKinds, RK.AttrKind))
     return RK;
   return RetainedKnowledge::none();
@@ -176,7 +162,7 @@ llvm::getKnowledgeForValue(const Value *V,
     return RetainedKnowledge::none();
   if (AC) {
     for (AssumptionCache::ResultElem &Elem : AC->assumptionsFor(V)) {
-      IntrinsicInst *II = cast_or_null<IntrinsicInst>(Elem.Assume);
+      auto *II = cast_or_null<AssumeInst>(Elem.Assume);
       if (!II || Elem.Index == AssumptionCache::ExprResultIdx)
         continue;
       if (RetainedKnowledge RK = getKnowledgeFromBundle(
@@ -197,7 +183,7 @@ llvm::getKnowledgeForValue(const Value *V,
     if (!Bundle)
       continue;
     if (RetainedKnowledge RK =
-            getKnowledgeFromBundle(*cast<CallInst>(U.getUser()), *Bundle))
+            getKnowledgeFromBundle(*cast<AssumeInst>(U.getUser()), *Bundle))
       if (is_contained(AttrKinds, RK.AttrKind) &&
           Filter(RK, cast<Instruction>(U.getUser()), Bundle)) {
         NumUsefullAssumeQueries++;

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -1910,9 +1910,8 @@ void InformationCache::initializeInformationCache(const Function &CF,
       // Calls are interesting on their own, additionally:
       // For `llvm.assume` calls we also fill the KnowledgeMap as we find them.
       // For `must-tail` calls we remember the caller and callee.
-      if (IntrinsicInst *Assume = dyn_cast<IntrinsicInst>(&I)) {
-        if (Assume->getIntrinsicID() == Intrinsic::assume)
-          fillMapFromAssume(*Assume, KnowledgeMap);
+      if (auto *Assume = dyn_cast<AssumeInst>(&I)) {
+        fillMapFromAssume(*Assume, KnowledgeMap);
       } else if (cast<CallInst>(I).isMustTailCall()) {
         FI.ContainsMustTailCall = true;
         if (const Function *Callee = cast<CallInst>(I).getCalledFunction())

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -1557,8 +1557,8 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     /// argument and remove the assume if it becomes useless.
     /// always returns nullptr for use as a return values.
     auto RemoveConditionFromAssume = [&](Instruction *Assume) -> Instruction * {
-      assert(isa<IntrinsicInst>(Assume));
-      if (isAssumeWithEmptyBundle(*cast<IntrinsicInst>(II)))
+      assert(isa<AssumeInst>(Assume));
+      if (isAssumeWithEmptyBundle(*cast<AssumeInst>(II)))
         return eraseInstFromFunction(CI);
       replaceUse(II->getOperandUse(0), ConstantInt::getTrue(II->getContext()));
       return nullptr;
@@ -1615,11 +1615,11 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     if (EnableKnowledgeRetention &&
         match(IIOperand, m_Cmp(Pred, m_Value(A), m_Zero())) &&
         Pred == CmpInst::ICMP_NE && A->getType()->isPointerTy()) {
-      if (IntrinsicInst *Replacement = buildAssumeFromKnowledge(
+      if (auto *Replacement = buildAssumeFromKnowledge(
               {RetainedKnowledge{Attribute::NonNull, 0, A}}, Next, &AC, &DT)) {
 
         Replacement->insertBefore(Next);
-        AC.registerAssumption(cast<AssumeInst>(Replacement));
+        AC.registerAssumption(Replacement);
         return RemoveConditionFromAssume(II);
       }
     }
@@ -1647,11 +1647,11 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
           /// the offset.
           RetainedKnowledge RK{Attribute::Alignment,
                                (unsigned)MinAlign(Offset, AlignMask + 1), A};
-          if (IntrinsicInst *Replacement =
+          if (auto *Replacement =
                   buildAssumeFromKnowledge(RK, Next, &AC, &DT)) {
 
             Replacement->insertAfter(II);
-            AC.registerAssumption(cast<AssumeInst>(Replacement));
+            AC.registerAssumption(Replacement);
           }
           return RemoveConditionFromAssume(II);
         }
@@ -1662,13 +1662,16 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     if (EnableKnowledgeRetention && II->hasOperandBundles()) {
       for (unsigned Idx = 0; Idx < II->getNumOperandBundles(); Idx++) {
         auto &BOI = II->bundle_op_info_begin()[Idx];
-        RetainedKnowledge RK = llvm::getKnowledgeFromBundle(*II, BOI);
+        RetainedKnowledge RK =
+          llvm::getKnowledgeFromBundle(cast<AssumeInst>(*II), BOI);
         if (BOI.End - BOI.Begin > 2)
           continue; // Prevent reducing knowledge in an align with offset since
                     // extracting a RetainedKnowledge form them looses offset
                     // information
-        RetainedKnowledge CanonRK = llvm::simplifyRetainedKnowledge(
-            II, RK, &getAssumptionCache(), &getDominatorTree());
+        RetainedKnowledge CanonRK =
+          llvm::simplifyRetainedKnowledge(cast<AssumeInst>(II), RK,
+                                          &getAssumptionCache(),
+                                          &getDominatorTree());
         if (CanonRK == RK)
           continue;
         if (!CanonRK) {
@@ -1694,7 +1697,7 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
     // then this one is redundant, and should be removed.
     KnownBits Known(1);
     computeKnownBits(IIOperand, Known, 0, II);
-    if (Known.isAllOnes() && isAssumeWithEmptyBundle(*II))
+    if (Known.isAllOnes() && isAssumeWithEmptyBundle(cast<AssumeInst>(*II)))
       return eraseInstFromFunction(*II);
 
     // Update the cache of affected values for this assumption (we might be

llvm/lib/Transforms/Scalar/GVN.cpp

@@ -1602,9 +1602,7 @@ static bool hasUsersIn(Value *V, BasicBlock *BB) {
   return false;
 }
 
-bool GVN::processAssumeIntrinsic(IntrinsicInst *IntrinsicI) {
-  assert(IntrinsicI->getIntrinsicID() == Intrinsic::assume &&
-         "This function can only be called with llvm.assume intrinsic");
+bool GVN::processAssumeIntrinsic(AssumeInst *IntrinsicI) {
   Value *V = IntrinsicI->getArgOperand(0);
 
   if (ConstantInt *Cond = dyn_cast<ConstantInt>(V)) {
@@ -2185,9 +2183,8 @@ bool GVN::processInstruction(Instruction *I) {
     }
   }
 
-  if (IntrinsicInst *IntrinsicI = dyn_cast<IntrinsicInst>(I))
-    if (IntrinsicI->getIntrinsicID() == Intrinsic::assume)
-      return processAssumeIntrinsic(IntrinsicI);
+  if (auto *Assume = dyn_cast<AssumeInst>(I))
+    return processAssumeIntrinsic(Assume);
 
   if (LoadInst *Load = dyn_cast<LoadInst>(I)) {
     if (processLoad(Load))