[AssumeBundles] Add API to fill a map from operand bundles of an llvm.assume.

[Tyker]

Summary

This patch adds a new way to query operand bundles of an llvm.assume that is much better suited to some users like the Attributor that need to do many queries on the operand bundles of llvm.assume. Some modifications of the IR like replaceAllUsesWith can cause information in the map to be outdated, so this API is more suited to analysis passes and passes that don't make modification that could invalidate the map.

[jdoerfert]

Sorry for the delay.

Two nits but otherwise LGTM.

if (Bundles.End - Bundles.Begin <compare> SOMETHIG)

Might deserve a helper function or two.

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

180

static or in the anonymous namespace.

llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 #define LLVM_TRANSFORMS_UTILS_ASSUMEBUILDER_H
 
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/ADT/DenseMap.h"
 
 namespace llvm {
 
 /// Build a call to llvm.assume to preserve informations that can be derived
 /// from the given instruction.
 /// If no information derived from \p I, this call returns null.
 /// The returned instruction is not inserted anywhere.
 CallInst *BuildAssumeFromInst(const Instruction *I, Module *M);

 inline bool hasAttributeInAssume(CallInst &AssumeCI, Value *IsOn,
                                  Attribute::AttrKind Kind,
                                  uint64_t *ArgVal = nullptr,
                                  AssumeQuery AQR = AssumeQuery::Highest) {
   return hasAttributeInAssume(
       AssumeCI, IsOn, Attribute::getNameFromAttrKind(Kind), ArgVal, AQR);
 }
 
-/// TODO: Add an function to create/fill a map from the bundle when users intend
-/// to make many different queries on the same bundles. to be used for example
-/// in the Attributor.
+template<> struct DenseMapInfo<Attribute::AttrKind> {
+  static constexpr auto MaxValue = std::numeric_limits<
+      std::underlying_type<Attribute::AttrKind>::type>::max();
+  static Attribute::AttrKind getEmptyKey() {
+    return static_cast<Attribute::AttrKind>(MaxValue);
+  }
+  static Attribute::AttrKind getTombstoneKey() {
+    return static_cast<Attribute::AttrKind>(MaxValue - 1);
+  }
+  static unsigned getHashValue(Attribute::AttrKind AK) {
+    return hash_combine(AK);
+  }
+  static bool isEqual(Attribute::AttrKind LHS, Attribute::AttrKind RHS) {
+    return LHS == RHS;
+  }
+};
+
+/// The map Key contains the Value on for which the attribute is valid and
+/// the Attribute that is valid for that value.
+/// If the Attribute is not on any value, the Value is nullptr.
+using RetainedKnowledgeKey = std::pair<Value *, Attribute::AttrKind>;
+
+struct MinMax {
+  unsigned Min;
+  unsigned Max;
+};
+
+using RetainedKnowledgeMap = DenseMap<RetainedKnowledgeKey, MinMax>;
+
+/// Insert into the map all the informations contained in the operand bundles of
+/// the llvm.assume. This should be used instead of hasAttributeInAssume when
+/// 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);
 
 //===----------------------------------------------------------------------===//
 // Utilities for testing
 //===----------------------------------------------------------------------===//
 
 /// This pass will try to build an llvm.assume for every instruction in the
 /// function. Its main purpose is testing.
 struct AssumeBuilderPass : public PassInfoMixin<AssumeBuilderPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 } // namespace llvm
 
 #endif

llvm/lib/Transforms/Utils/KnowledgeRetention.cpp

 } // namespace
 
 CallInst *llvm::BuildAssumeFromInst(const Instruction *I, Module *M) {
   AssumeBuilderState Builder(M);
   Builder.addInstruction(I);
   return Builder.build();
 }
 
+Value *getValueFromBundleOpInfo(IntrinsicInst &Assume,
+                                const CallBase::BundleOpInfo &BOI,
+                                unsigned Idx) {
+  assert(BOI.End - BOI.Begin > Idx && "index out of range");
+  return (Assume.op_begin() + BOI.Begin + Idx)->get();
+}
+
 #ifndef NDEBUG
 
 static bool isExistingAttribute(StringRef Name) {
   return StringSwitch<bool>(Name)
 #define GET_ATTR_NAMES
 #define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) .Case(#DISPLAY_NAME, true)
 #include "llvm/IR/Attributes.inc"
       .Default(false);
else
     Lookup = std::prev(
         llvm::upper_bound(Assume.bundle_op_infos(), AttrName,
                           [](StringRef LHS, const CallBase::BundleOpInfo &BOI) {
                             assert(isExistingAttribute(BOI.Tag->getKey()) &&
                                    "this attribute doesn't exist");
                             return LHS < BOI.Tag->getKey();
                           }));
 
-  auto getValueFromBundleOpInfo = [&Assume](const CallBase::BundleOpInfo &BOI,
-                                            unsigned Idx) {
-    assert(BOI.End - BOI.Begin > Idx && "index out of range");
-    return (Assume.op_begin() + BOI.Begin + Idx)->get();
-  };
-
   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(*Lookup, BOIE_WasOn) == IsOn)
+      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);
     }
   }
 
   if (Lookup->End - Lookup->Begin < BOIE_Argument)
     return true;
   if (ArgVal)
-    *ArgVal =
-        cast<ConstantInt>(getValueFromBundleOpInfo(*Lookup, BOIE_Argument))
+    *ArgVal = cast<ConstantInt>(
+                  getValueFromBundleOpInfo(Assume, *Lookup, BOIE_Argument))
                   ->getZExtValue();
   return true;
 }
 
+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");
+  for (auto &Bundles : Assume.bundle_op_infos()) {
+    std::pair<Value *, Attribute::AttrKind> Key{
+        nullptr, Attribute::getAttrKindFromName(Bundles.Tag->getKey())};
+    if (Bundles.End - Bundles.Begin > BOIE_WasOn)
+      Key.first = getValueFromBundleOpInfo(Assume, Bundles, BOIE_WasOn);
+
+    if (Key.first == nullptr && Key.second == Attribute::None)
+      continue;
+    if (Bundles.End - Bundles.Begin <= BOIE_Argument) {
+      Result[Key] = {0, 0};
+      continue;
+    }
+    unsigned Val = cast<ConstantInt>(
+                       getValueFromBundleOpInfo(Assume, Bundles, BOIE_Argument))
+                       ->getZExtValue();
+    auto Lookup = Result.find(Key);
+    if (Lookup == Result.end()) {
+      Result[Key] = {Val, Val};
+      continue;
+    }
+    Lookup->second.Min = std::min(Val, Lookup->second.Min);
+    Lookup->second.Max = std::max(Val, Lookup->second.Max);
+  }
+}
+
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   for (Instruction &I : instructions(F))
     if (Instruction *Assume = BuildAssumeFromInst(&I))
       Assume->insertBefore(&I);
   return PreservedAnalyses::all();
 }

llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp

static void RunTest(
   unsigned Idx = 0;
   for (Instruction &I : (*Mod->getFunction("test")->begin())) {
     if (Idx < Tests.size())
       Tests[Idx].second(&I);
     Idx++;
   }
 }
 
-void AssertMatchesExactlyAttributes(CallInst *Assume, Value *WasOn,
+static void AssertMatchesExactlyAttributes(CallInst *Assume, Value *WasOn,
                                     StringRef AttrToMatch) {
   Regex Reg(AttrToMatch);
   SmallVector<StringRef, 1> Matches;
   for (StringRef Attr : {
 #define GET_ATTR_NAMES
 #define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) StringRef(#DISPLAY_NAME),
 #include "llvm/IR/Attributes.inc"
        }) {
     bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
     if (ShouldHaveAttr != hasAttributeInAssume(*Assume, WasOn, Attr)) {
       ASSERT_TRUE(false);
     }
   }
 }
 
-void AssertHasTheRightValue(CallInst *Assume, Value *WasOn,
+static void AssertHasTheRightValue(CallInst *Assume, Value *WasOn,
                             Attribute::AttrKind Kind, unsigned Value, bool Both,
                             AssumeQuery AQ = AssumeQuery::Highest) {
   if (!Both) {
     uint64_t ArgVal = 0;
     ASSERT_TRUE(hasAttributeInAssume(*Assume, WasOn, Kind, &ArgVal, AQ));
     ASSERT_EQ(ArgVal, Value);
     return;
   }
   uint64_t ArgValLow = 0;
   uint64_t ArgValHigh = 0;
   bool ResultLow = hasAttributeInAssume(*Assume, WasOn, Kind, &ArgValLow,
                                         AssumeQuery::Lowest);
   bool ResultHigh = hasAttributeInAssume(*Assume, WasOn, Kind, &ArgValHigh,
                                          AssumeQuery::Highest);
   if (ResultLow != ResultHigh || ResultHigh == false) {
     ASSERT_TRUE(false);
   }
   if (ArgValLow != Value || ArgValLow != ArgValHigh) {
     ASSERT_TRUE(false);
   }
 }
 
-TEST(AssumeQueryAPI, Basic) {
+TEST(AssumeQueryAPI, hasAttributeInAssume) {
   StringRef Head =
       "declare void @llvm.assume(i1)\n"
       "declare void @func(i32*, i32*)\n"
       "declare void @func1(i32*, i32*, i32*, i32*)\n"
       "declare void @func_many(i32*) \"no-jump-tables\" nounwind "
       "\"less-precise-fpmad\" willreturn norecurse\n"
       "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3) {\n";
   StringRef Tail = "ret void\n"
Tests.push_back(std::make_pair(
                                        "(nonnull|align|dereferenceable)");
         Old->replaceAllUsesWith(New);
         AssertMatchesExactlyAttributes(Assume, New,
                                        "(nonnull|align|dereferenceable)");
         AssertMatchesExactlyAttributes(Assume, Old, "");
       }));
   RunTest(Head, Tail, Tests);
 }
+
+static void AssertFindExactlyAttributes(RetainedKnowledgeMap &Map, Value *WasOn,
+                                 StringRef AttrToMatch) {
+  Regex Reg(AttrToMatch);
+  SmallVector<StringRef, 1> Matches;
+  for (StringRef Attr : {
+#define GET_ATTR_NAMES
+#define ATTRIBUTE_ENUM(ENUM_NAME, DISPLAY_NAME) StringRef(#DISPLAY_NAME),
+#include "llvm/IR/Attributes.inc"
+       }) {
+    bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
+
+    if (ShouldHaveAttr != (Map.find(RetainedKnowledgeKey{WasOn, Attribute::getAttrKindFromName(Attr)}) != Map.end())) {
+      ASSERT_TRUE(false);
+    }
+  }
+}
+
+static void AssertMapHasRightValue(RetainedKnowledgeMap &Map,
+                                   RetainedKnowledgeKey Key, MinMax MM) {
+  auto LookupIt = Map.find(Key);
+  ASSERT_TRUE(LookupIt != Map.end());
+  ASSERT_TRUE(LookupIt->second.Min == MM.Min);
+  ASSERT_TRUE(LookupIt->second.Max == MM.Max);
+}
+
+TEST(AssumeQueryAPI, fillMapFromAssume) {
+    StringRef Head =
+      "declare void @llvm.assume(i1)\n"
+      "declare void @func(i32*, i32*)\n"
+      "declare void @func1(i32*, i32*, i32*, i32*)\n"
+      "declare void @func_many(i32*) \"no-jump-tables\" nounwind "
+      "\"less-precise-fpmad\" willreturn norecurse\n"
+      "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3) {\n";
+  StringRef Tail = "ret void\n"
+                   "}";
+  std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
+      Tests;
+  Tests.push_back(std::make_pair(
+      "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
+      "8 noalias %P1)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+
+        RetainedKnowledgeMap Map;
+        fillMapFromAssume(*Assume, Map);
+        AssertFindExactlyAttributes(Map, I->getOperand(0),
+                                       "(nonnull|align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, I->getOperand(1),
+                                       "(noalias|align)");
+        AssertMapHasRightValue(
+            Map, {I->getOperand(0), Attribute::Dereferenceable}, {16, 16});
+        AssertMapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+                               {4, 4});
+        AssertMapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+                               {4, 4});
+      }));
+  Tests.push_back(std::make_pair(
+      "call void @func1(i32* nonnull align 32 dereferenceable(48) %P, i32* "
+      "nonnull "
+      "align 8 dereferenceable(28) %P, i32* nonnull align 64 "
+      "dereferenceable(4) "
+      "%P, i32* nonnull align 16 dereferenceable(12) %P)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+
+        RetainedKnowledgeMap Map;
+        fillMapFromAssume(*Assume, Map);
+
+        AssertFindExactlyAttributes(Map, I->getOperand(0),
+                                       "(nonnull|align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, I->getOperand(1),
+                                       "(nonnull|align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, I->getOperand(2),
+                                       "(nonnull|align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, I->getOperand(3),
+                                       "(nonnull|align|dereferenceable)");
+        AssertMapHasRightValue(
+            Map, {I->getOperand(0), Attribute::Dereferenceable}, {4, 48});
+        AssertMapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+                               {8, 64});
+      }));
+  Tests.push_back(std::make_pair(
+      "call void @func_many(i32* align 8 %P1) cold\n", [](Instruction *I) {
+        ShouldPreserveAllAttributes.setValue(true);
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+
+        RetainedKnowledgeMap Map;
+        fillMapFromAssume(*Assume, Map);
+
+        AssertFindExactlyAttributes(
+            Map, nullptr, "(nounwind|norecurse|willreturn|cold)");
+        ShouldPreserveAllAttributes.setValue(false);
+      }));
+  Tests.push_back(
+      std::make_pair("call void @llvm.assume(i1 true)\n", [](Instruction *I) {
+        RetainedKnowledgeMap Map;
+        fillMapFromAssume(*cast<CallInst>(I), Map);
+
+        AssertFindExactlyAttributes(Map, nullptr, "");
+        ASSERT_TRUE(Map.empty());
+      }));
+  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);
+
+        RetainedKnowledgeMap Map;
+        fillMapFromAssume(*Assume, Map);
+
+        AssertFindExactlyAttributes(Map, I->getOperand(0),
+                                    "(readnone|align|dereferenceable|noalias)");
+        AssertFindExactlyAttributes(Map, I->getOperand(1),
+                                    "(align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, I->getOperand(2),
+                                       "(align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, I->getOperand(3),
+                                       "(nonnull|align|dereferenceable)");
+        AssertMapHasRightValue(Map, {I->getOperand(0), Attribute::Alignment},
+                               {32, 32});
+        AssertMapHasRightValue(
+            Map, {I->getOperand(0), Attribute::Dereferenceable}, {48, 48});
+        AssertMapHasRightValue(
+            Map, {I->getOperand(0), Attribute::NoAlias}, {0, 0});
+        AssertMapHasRightValue(
+            Map, {I->getOperand(1), Attribute::Dereferenceable}, {28, 28});
+        AssertMapHasRightValue(Map, {I->getOperand(1), Attribute::Alignment},
+                               {8, 8});
+        AssertMapHasRightValue(Map, {I->getOperand(2), Attribute::Alignment},
+                               {64, 64});
+        AssertMapHasRightValue(
+            Map, {I->getOperand(2), Attribute::Dereferenceable}, {4, 4});
+        AssertMapHasRightValue(Map, {I->getOperand(3), Attribute::Alignment},
+                               {16, 16});
+        AssertMapHasRightValue(
+            Map, {I->getOperand(3), Attribute::Dereferenceable}, {12, 12});
+      }));
+
+  /// Keep this test last as it modifies the function.
+  Tests.push_back(std::make_pair(
+      "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
+      "8 noalias %P1)\n",
+      [](Instruction *I) {
+        CallInst *Assume = BuildAssumeFromInst(I);
+        Assume->insertBefore(I);
+
+        RetainedKnowledgeMap Map;
+        fillMapFromAssume(*Assume, Map);
+
+        Value *New = I->getFunction()->getArg(3);
+        Value *Old = I->getOperand(0);
+        AssertFindExactlyAttributes(Map, New, "");
+        AssertFindExactlyAttributes(Map, Old,
+                                       "(nonnull|align|dereferenceable)");
+        Old->replaceAllUsesWith(New);
+        Map.clear();
+        fillMapFromAssume(*Assume, Map);
+        AssertFindExactlyAttributes(Map, New,
+                                       "(nonnull|align|dereferenceable)");
+        AssertFindExactlyAttributes(Map, Old, "");
+      }));
+  RunTest(Head, Tail, Tests);
+}