diff --git a/llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h b/llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h --- a/llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h +++ b/llvm/include/llvm/Transforms/Utils/KnowledgeRetention.h @@ -19,6 +19,7 @@ #include "llvm/IR/Attributes.h" #include "llvm/IR/Instruction.h" #include "llvm/IR/PassManager.h" +#include "llvm/ADT/DenseMap.h" namespace llvm { @@ -58,9 +59,41 @@ 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 { + static constexpr auto MaxValue = std::numeric_limits< + std::underlying_type::type>::max(); + static Attribute::AttrKind getEmptyKey() { + return static_cast(MaxValue); + } + static Attribute::AttrKind getTombstoneKey() { + return static_cast(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; + +struct MinMax { + unsigned Min; + unsigned Max; +}; + +using RetainedKnowledgeMap = DenseMap; + +/// 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 diff --git a/llvm/lib/Transforms/Utils/KnowledgeRetention.cpp b/llvm/lib/Transforms/Utils/KnowledgeRetention.cpp --- a/llvm/lib/Transforms/Utils/KnowledgeRetention.cpp +++ b/llvm/lib/Transforms/Utils/KnowledgeRetention.cpp @@ -171,6 +171,18 @@ return Builder.build(); } +static bool BundleHasArguement(const CallBase::BundleOpInfo &BOI, + unsigned Idx) { + return BOI.End - BOI.Begin > Idx; +} + +static Value *getValueFromBundleOpInfo(IntrinsicInst &Assume, + const CallBase::BundleOpInfo &BOI, + unsigned Idx) { + assert(BundleHasArguement(BOI, Idx) && "index out of range"); + return (Assume.op_begin() + BOI.Begin + Idx)->get(); +} + #ifndef NDEBUG static bool isExistingAttribute(StringRef Name) { @@ -219,12 +231,6 @@ 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; @@ -235,7 +241,7 @@ 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()) @@ -247,12 +253,41 @@ if (Lookup->End - Lookup->Begin < BOIE_Argument) return true; if (ArgVal) - *ArgVal = - cast(getValueFromBundleOpInfo(*Lookup, BOIE_Argument)) - ->getZExtValue(); + *ArgVal = cast( + getValueFromBundleOpInfo(Assume, *Lookup, BOIE_Argument)) + ->getZExtValue(); return true; } +void llvm::fillMapFromAssume(CallInst &AssumeCI, RetainedKnowledgeMap &Result) { + IntrinsicInst &Assume = cast(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 Key{ + nullptr, Attribute::getAttrKindFromName(Bundles.Tag->getKey())}; + if (BundleHasArguement(Bundles, BOIE_WasOn)) + Key.first = getValueFromBundleOpInfo(Assume, Bundles, BOIE_WasOn); + + if (Key.first == nullptr && Key.second == Attribute::None) + continue; + if (!BundleHasArguement(Bundles, BOIE_Argument)) { + Result[Key] = {0, 0}; + continue; + } + unsigned Val = cast( + 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)) diff --git a/llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp b/llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp --- a/llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp +++ b/llvm/unittests/Transforms/Utils/KnowledgeRetentionTest.cpp @@ -41,7 +41,7 @@ } } -void AssertMatchesExactlyAttributes(CallInst *Assume, Value *WasOn, +static void AssertMatchesExactlyAttributes(CallInst *Assume, Value *WasOn, StringRef AttrToMatch) { Regex Reg(AttrToMatch); SmallVector Matches; @@ -57,7 +57,7 @@ } } -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) { @@ -80,7 +80,7 @@ } } -TEST(AssumeQueryAPI, Basic) { +TEST(AssumeQueryAPI, hasAttributeInAssume) { StringRef Head = "declare void @llvm.assume(i1)\n" "declare void @func(i32*, i32*)\n" @@ -216,3 +216,174 @@ })); RunTest(Head, Tail, Tests); } + +static void AssertFindExactlyAttributes(RetainedKnowledgeMap &Map, Value *WasOn, + StringRef AttrToMatch) { + Regex Reg(AttrToMatch); + SmallVector 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>> + 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(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); +}