xref: /llvm-project/llvm/unittests/Analysis/AssumeBundleQueriesTest.cpp (revision 813f438baaa9638529023b2218875e01ea037735)
1 //===- AssumeBundleQueriesTest.cpp ------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "llvm/Analysis/AssumptionCache.h"
10 #include "llvm/Analysis/AssumeBundleQueries.h"
11 #include "llvm/AsmParser/Parser.h"
12 #include "llvm/IR/CallSite.h"
13 #include "llvm/IR/LLVMContext.h"
14 #include "llvm/IR/IntrinsicInst.h"
15 #include "llvm/Support/Regex.h"
16 #include "llvm/Support/SourceMgr.h"
17 #include "llvm/Support/CommandLine.h"
18 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
19 #include "gtest/gtest.h"
20 #include <random>
21 
22 using namespace llvm;
23 
24 extern cl::opt<bool> ShouldPreserveAllAttributes;
25 extern cl::opt<bool> EnableKnowledgeRetention;
26 
27 static void RunTest(
28     StringRef Head, StringRef Tail,
29     std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
30         &Tests) {
31   for (auto &Elem : Tests) {
32     std::string IR;
33     IR.append(Head.begin(), Head.end());
34     IR.append(Elem.first.begin(), Elem.first.end());
35     IR.append(Tail.begin(), Tail.end());
36     LLVMContext C;
37     SMDiagnostic Err;
38     std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
39     if (!Mod)
40       Err.print("AssumeQueryAPI", errs());
41     Elem.second(&*(Mod->getFunction("test")->begin()->begin()));
42   }
43 }
44 
45 bool hasMatchesExactlyAttributes(IntrinsicInst *Assume, Value *WasOn,
46                                     StringRef AttrToMatch) {
47   Regex Reg(AttrToMatch);
48   SmallVector<StringRef, 1> Matches;
49   for (StringRef Attr : {
50 #define GET_ATTR_NAMES
51 #define ATTRIBUTE_ALL(ENUM_NAME, DISPLAY_NAME) StringRef(#DISPLAY_NAME),
52 #include "llvm/IR/Attributes.inc"
53        }) {
54     bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
55     if (ShouldHaveAttr != hasAttributeInAssume(*Assume, WasOn, Attr))
56       return false;
57   }
58   return true;
59 }
60 
61 bool hasTheRightValue(IntrinsicInst *Assume, Value *WasOn,
62                             Attribute::AttrKind Kind, unsigned Value, bool Both,
63                             AssumeQuery AQ = AssumeQuery::Highest) {
64   if (!Both) {
65     uint64_t ArgVal = 0;
66     if (!hasAttributeInAssume(*Assume, WasOn, Kind, &ArgVal, AQ))
67       return false;
68     if (ArgVal != Value)
69       return false;
70     return true;
71   }
72   uint64_t ArgValLow = 0;
73   uint64_t ArgValHigh = 0;
74   bool ResultLow = hasAttributeInAssume(*Assume, WasOn, Kind, &ArgValLow,
75                                         AssumeQuery::Lowest);
76   bool ResultHigh = hasAttributeInAssume(*Assume, WasOn, Kind, &ArgValHigh,
77                                          AssumeQuery::Highest);
78   if (ResultLow != ResultHigh || ResultHigh == false)
79     return false;
80   if (ArgValLow != Value || ArgValLow != ArgValHigh)
81     return false;
82   return true;
83 }
84 
85 TEST(AssumeQueryAPI, hasAttributeInAssume) {
86   EnableKnowledgeRetention.setValue(true);
87   StringRef Head =
88       "declare void @llvm.assume(i1)\n"
89       "declare void @func(i32*, i32*)\n"
90       "declare void @func1(i32*, i32*, i32*, i32*)\n"
91       "declare void @func_many(i32*) \"no-jump-tables\" nounwind "
92       "\"less-precise-fpmad\" willreturn norecurse\n"
93       "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3) {\n";
94   StringRef Tail = "ret void\n"
95                    "}";
96   std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
97       Tests;
98   Tests.push_back(std::make_pair(
99       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
100       "8 noalias %P1)\n",
101       [](Instruction *I) {
102         IntrinsicInst *Assume = buildAssumeFromInst(I);
103         Assume->insertBefore(I);
104         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(0),
105                                        "(nonnull|align|dereferenceable)"));
106         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
107                                        "(align)"));
108         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
109                                Attribute::AttrKind::Dereferenceable, 16, true));
110         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
111                                Attribute::AttrKind::Alignment, 4, true));
112         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
113                                Attribute::AttrKind::Alignment, 4, true));
114       }));
115   Tests.push_back(std::make_pair(
116       "call void @func1(i32* nonnull align 32 dereferenceable(48) %P, i32* "
117       "nonnull "
118       "align 8 dereferenceable(28) %P, i32* nonnull align 64 "
119       "dereferenceable(4) "
120       "%P, i32* nonnull align 16 dereferenceable(12) %P)\n",
121       [](Instruction *I) {
122         IntrinsicInst *Assume = buildAssumeFromInst(I);
123         Assume->insertBefore(I);
124         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(0),
125                                        "(nonnull|align|dereferenceable)"));
126         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
127                                        "(nonnull|align|dereferenceable)"));
128         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(2),
129                                        "(nonnull|align|dereferenceable)"));
130         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(3),
131                                        "(nonnull|align|dereferenceable)"));
132         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
133                                Attribute::AttrKind::Dereferenceable, 48, false,
134                                AssumeQuery::Highest));
135         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
136                                Attribute::AttrKind::Alignment, 64, false,
137                                AssumeQuery::Highest));
138         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(1),
139                                Attribute::AttrKind::Alignment, 64, false,
140                                AssumeQuery::Highest));
141         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
142                                Attribute::AttrKind::Dereferenceable, 4, false,
143                                AssumeQuery::Lowest));
144         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
145                                Attribute::AttrKind::Alignment, 8, false,
146                                AssumeQuery::Lowest));
147         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(1),
148                                Attribute::AttrKind::Alignment, 8, false,
149                                AssumeQuery::Lowest));
150       }));
151   Tests.push_back(std::make_pair(
152       "call void @func_many(i32* align 8 %P1) cold\n", [](Instruction *I) {
153         ShouldPreserveAllAttributes.setValue(true);
154         IntrinsicInst *Assume = buildAssumeFromInst(I);
155         Assume->insertBefore(I);
156         ASSERT_TRUE(hasMatchesExactlyAttributes(
157             Assume, nullptr,
158             "(align|no-jump-tables|less-precise-fpmad|"
159             "nounwind|norecurse|willreturn|cold)"));
160         ShouldPreserveAllAttributes.setValue(false);
161       }));
162   Tests.push_back(
163       std::make_pair("call void @llvm.assume(i1 true)\n", [](Instruction *I) {
164         IntrinsicInst *Assume = cast<IntrinsicInst>(I);
165         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, nullptr, ""));
166       }));
167   Tests.push_back(std::make_pair(
168       "call void @func1(i32* readnone align 32 "
169       "dereferenceable(48) noalias %P, i32* "
170       "align 8 dereferenceable(28) %P1, i32* align 64 "
171       "dereferenceable(4) "
172       "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
173       [](Instruction *I) {
174         IntrinsicInst *Assume = buildAssumeFromInst(I);
175         Assume->insertBefore(I);
176         ASSERT_TRUE(hasMatchesExactlyAttributes(
177             Assume, I->getOperand(0),
178             "(align|dereferenceable)"));
179         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
180                                        "(align|dereferenceable)"));
181         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(2),
182                                        "(align|dereferenceable)"));
183         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(3),
184                                        "(nonnull|align|dereferenceable)"));
185         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
186                                Attribute::AttrKind::Alignment, 32, true));
187         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
188                                Attribute::AttrKind::Dereferenceable, 48, true));
189         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(1),
190                                Attribute::AttrKind::Dereferenceable, 28, true));
191         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(1),
192                                Attribute::AttrKind::Alignment, 8, true));
193         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(2),
194                                Attribute::AttrKind::Alignment, 64, true));
195         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(2),
196                                Attribute::AttrKind::Dereferenceable, 4, true));
197         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(3),
198                                Attribute::AttrKind::Alignment, 16, true));
199         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(3),
200                                Attribute::AttrKind::Dereferenceable, 12, true));
201       }));
202 
203   Tests.push_back(std::make_pair(
204       "call void @func1(i32* readnone align 32 "
205       "dereferenceable(48) noalias %P, i32* "
206       "align 8 dereferenceable(28) %P1, i32* align 64 "
207       "dereferenceable(4) "
208       "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
209       [](Instruction *I) {
210         IntrinsicInst *Assume = buildAssumeFromInst(I);
211         Assume->insertBefore(I);
212         I->getOperand(1)->dropDroppableUses();
213         I->getOperand(2)->dropDroppableUses();
214         I->getOperand(3)->dropDroppableUses();
215         ASSERT_TRUE(hasMatchesExactlyAttributes(
216             Assume, I->getOperand(0),
217             "(align|dereferenceable)"));
218         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(1),
219                                        ""));
220         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(2),
221                                        ""));
222         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, I->getOperand(3),
223                                        ""));
224         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
225                                Attribute::AttrKind::Alignment, 32, true));
226         ASSERT_TRUE(hasTheRightValue(Assume, I->getOperand(0),
227                                Attribute::AttrKind::Dereferenceable, 48, true));
228       }));
229   Tests.push_back(std::make_pair(
230       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
231       "8 noalias %P1)\n",
232       [](Instruction *I) {
233         IntrinsicInst *Assume = buildAssumeFromInst(I);
234         Assume->insertBefore(I);
235         Value *New = I->getFunction()->getArg(3);
236         Value *Old = I->getOperand(0);
237         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, New, ""));
238         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, Old,
239                                        "(nonnull|align|dereferenceable)"));
240         Old->replaceAllUsesWith(New);
241         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, New,
242                                        "(nonnull|align|dereferenceable)"));
243         ASSERT_TRUE(hasMatchesExactlyAttributes(Assume, Old, ""));
244       }));
245   RunTest(Head, Tail, Tests);
246 }
247 
248 static bool FindExactlyAttributes(RetainedKnowledgeMap &Map, Value *WasOn,
249                                  StringRef AttrToMatch) {
250   Regex Reg(AttrToMatch);
251   SmallVector<StringRef, 1> Matches;
252   for (StringRef Attr : {
253 #define GET_ATTR_NAMES
254 #define ATTRIBUTE_ENUM(ENUM_NAME, DISPLAY_NAME) StringRef(#DISPLAY_NAME),
255 #include "llvm/IR/Attributes.inc"
256        }) {
257     bool ShouldHaveAttr = Reg.match(Attr, &Matches) && Matches[0] == Attr;
258 
259     if (ShouldHaveAttr != (Map.find(RetainedKnowledgeKey{WasOn, Attribute::getAttrKindFromName(Attr)}) != Map.end()))
260       return false;
261   }
262   return true;
263 }
264 
265 static bool MapHasRightValue(RetainedKnowledgeMap &Map, IntrinsicInst *II,
266                              RetainedKnowledgeKey Key, MinMax MM) {
267   auto LookupIt = Map.find(Key);
268   return (LookupIt != Map.end()) && (LookupIt->second[II].Min == MM.Min) &&
269          (LookupIt->second[II].Max == MM.Max);
270 }
271 
272 TEST(AssumeQueryAPI, fillMapFromAssume) {
273   EnableKnowledgeRetention.setValue(true);
274   StringRef Head =
275       "declare void @llvm.assume(i1)\n"
276       "declare void @func(i32*, i32*)\n"
277       "declare void @func1(i32*, i32*, i32*, i32*)\n"
278       "declare void @func_many(i32*) \"no-jump-tables\" nounwind "
279       "\"less-precise-fpmad\" willreturn norecurse\n"
280       "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3) {\n";
281   StringRef Tail = "ret void\n"
282                    "}";
283   std::vector<std::pair<StringRef, llvm::function_ref<void(Instruction *)>>>
284       Tests;
285   Tests.push_back(std::make_pair(
286       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
287       "8 noalias %P1)\n",
288       [](Instruction *I) {
289         IntrinsicInst *Assume = buildAssumeFromInst(I);
290         Assume->insertBefore(I);
291 
292         RetainedKnowledgeMap Map;
293         fillMapFromAssume(*Assume, Map);
294         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(0),
295                                        "(nonnull|align|dereferenceable)"));
296         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(1),
297                                        "(align)"));
298         ASSERT_TRUE(MapHasRightValue(
299             Map, Assume, {I->getOperand(0), Attribute::Dereferenceable}, {16, 16}));
300         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
301                                {4, 4}));
302         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
303                                {4, 4}));
304       }));
305   Tests.push_back(std::make_pair(
306       "call void @func1(i32* nonnull align 32 dereferenceable(48) %P, i32* "
307       "nonnull "
308       "align 8 dereferenceable(28) %P, i32* nonnull align 64 "
309       "dereferenceable(4) "
310       "%P, i32* nonnull align 16 dereferenceable(12) %P)\n",
311       [](Instruction *I) {
312         IntrinsicInst *Assume = buildAssumeFromInst(I);
313         Assume->insertBefore(I);
314 
315         RetainedKnowledgeMap Map;
316         fillMapFromAssume(*Assume, Map);
317 
318         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(0),
319                                        "(nonnull|align|dereferenceable)"));
320         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(1),
321                                        "(nonnull|align|dereferenceable)"));
322         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(2),
323                                        "(nonnull|align|dereferenceable)"));
324         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(3),
325                                        "(nonnull|align|dereferenceable)"));
326         ASSERT_TRUE(MapHasRightValue(
327             Map, Assume, {I->getOperand(0), Attribute::Dereferenceable}, {4, 48}));
328         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
329                                {8, 64}));
330       }));
331   Tests.push_back(std::make_pair(
332       "call void @func_many(i32* align 8 %P1) cold\n", [](Instruction *I) {
333         ShouldPreserveAllAttributes.setValue(true);
334         IntrinsicInst *Assume = buildAssumeFromInst(I);
335         Assume->insertBefore(I);
336 
337         RetainedKnowledgeMap Map;
338         fillMapFromAssume(*Assume, Map);
339 
340         ASSERT_TRUE(FindExactlyAttributes(
341             Map, nullptr, "(nounwind|norecurse|willreturn|cold)"));
342         ShouldPreserveAllAttributes.setValue(false);
343       }));
344   Tests.push_back(
345       std::make_pair("call void @llvm.assume(i1 true)\n", [](Instruction *I) {
346         RetainedKnowledgeMap Map;
347         fillMapFromAssume(*cast<IntrinsicInst>(I), Map);
348 
349         ASSERT_TRUE(FindExactlyAttributes(Map, nullptr, ""));
350         ASSERT_TRUE(Map.empty());
351       }));
352   Tests.push_back(std::make_pair(
353       "call void @func1(i32* readnone align 32 "
354       "dereferenceable(48) noalias %P, i32* "
355       "align 8 dereferenceable(28) %P1, i32* align 64 "
356       "dereferenceable(4) "
357       "%P2, i32* nonnull align 16 dereferenceable(12) %P3)\n",
358       [](Instruction *I) {
359         IntrinsicInst *Assume = buildAssumeFromInst(I);
360         Assume->insertBefore(I);
361 
362         RetainedKnowledgeMap Map;
363         fillMapFromAssume(*Assume, Map);
364 
365         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(0),
366                                     "(align|dereferenceable)"));
367         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(1),
368                                     "(align|dereferenceable)"));
369         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(2),
370                                        "(align|dereferenceable)"));
371         ASSERT_TRUE(FindExactlyAttributes(Map, I->getOperand(3),
372                                        "(nonnull|align|dereferenceable)"));
373         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(0), Attribute::Alignment},
374                                {32, 32}));
375         ASSERT_TRUE(MapHasRightValue(
376             Map, Assume, {I->getOperand(0), Attribute::Dereferenceable}, {48, 48}));
377         ASSERT_TRUE(MapHasRightValue(
378             Map, Assume, {I->getOperand(1), Attribute::Dereferenceable}, {28, 28}));
379         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(1), Attribute::Alignment},
380                                {8, 8}));
381         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(2), Attribute::Alignment},
382                                {64, 64}));
383         ASSERT_TRUE(MapHasRightValue(
384             Map, Assume, {I->getOperand(2), Attribute::Dereferenceable}, {4, 4}));
385         ASSERT_TRUE(MapHasRightValue(Map, Assume, {I->getOperand(3), Attribute::Alignment},
386                                {16, 16}));
387         ASSERT_TRUE(MapHasRightValue(
388             Map, Assume, {I->getOperand(3), Attribute::Dereferenceable}, {12, 12}));
389       }));
390 
391   /// Keep this test last as it modifies the function.
392   Tests.push_back(std::make_pair(
393       "call void @func(i32* nonnull align 4 dereferenceable(16) %P, i32* align "
394       "8 noalias %P1)\n",
395       [](Instruction *I) {
396         IntrinsicInst *Assume = buildAssumeFromInst(I);
397         Assume->insertBefore(I);
398 
399         RetainedKnowledgeMap Map;
400         fillMapFromAssume(*Assume, Map);
401 
402         Value *New = I->getFunction()->getArg(3);
403         Value *Old = I->getOperand(0);
404         ASSERT_TRUE(FindExactlyAttributes(Map, New, ""));
405         ASSERT_TRUE(FindExactlyAttributes(Map, Old,
406                                        "(nonnull|align|dereferenceable)"));
407         Old->replaceAllUsesWith(New);
408         Map.clear();
409         fillMapFromAssume(*Assume, Map);
410         ASSERT_TRUE(FindExactlyAttributes(Map, New,
411                                        "(nonnull|align|dereferenceable)"));
412         ASSERT_TRUE(FindExactlyAttributes(Map, Old, ""));
413       }));
414   RunTest(Head, Tail, Tests);
415 }
416 
417 static void RunRandTest(uint64_t Seed, int Size, int MinCount, int MaxCount,
418                         unsigned MaxValue) {
419   LLVMContext C;
420   SMDiagnostic Err;
421 
422   std::random_device dev;
423   std::mt19937 Rng(Seed);
424   std::uniform_int_distribution<int> DistCount(MinCount, MaxCount);
425   std::uniform_int_distribution<unsigned> DistValue(0, MaxValue);
426   std::uniform_int_distribution<unsigned> DistAttr(0,
427                                                    Attribute::EndAttrKinds - 1);
428 
429   std::unique_ptr<Module> Mod = std::make_unique<Module>("AssumeQueryAPI", C);
430   if (!Mod)
431     Err.print("AssumeQueryAPI", errs());
432 
433   std::vector<Type *> TypeArgs;
434   for (int i = 0; i < (Size * 2); i++)
435     TypeArgs.push_back(Type::getInt32PtrTy(C));
436   FunctionType *FuncType =
437       FunctionType::get(Type::getVoidTy(C), TypeArgs, false);
438 
439   Function *F =
440       Function::Create(FuncType, GlobalValue::ExternalLinkage, "test", &*Mod);
441   BasicBlock *BB = BasicBlock::Create(C);
442   BB->insertInto(F);
443   Instruction *Ret = ReturnInst::Create(C);
444   BB->getInstList().insert(BB->begin(), Ret);
445   Function *FnAssume = Intrinsic::getDeclaration(Mod.get(), Intrinsic::assume);
446 
447   std::vector<Argument *> ShuffledArgs;
448   std::vector<bool> HasArg;
449   for (auto &Arg : F->args()) {
450     ShuffledArgs.push_back(&Arg);
451     HasArg.push_back(false);
452   }
453 
454   std::shuffle(ShuffledArgs.begin(), ShuffledArgs.end(), Rng);
455 
456   std::vector<OperandBundleDef> OpBundle;
457   OpBundle.reserve(Size);
458   std::vector<Value *> Args;
459   Args.reserve(2);
460   for (int i = 0; i < Size; i++) {
461     int count = DistCount(Rng);
462     int value = DistValue(Rng);
463     int attr = DistAttr(Rng);
464     std::string str;
465     raw_string_ostream ss(str);
466     ss << Attribute::getNameFromAttrKind(
467         static_cast<Attribute::AttrKind>(attr));
468     Args.clear();
469 
470     if (count > 0) {
471       Args.push_back(ShuffledArgs[i]);
472       HasArg[i] = true;
473     }
474     if (count > 1)
475       Args.push_back(ConstantInt::get(Type::getInt32Ty(C), value));
476 
477     OpBundle.push_back(OperandBundleDef{ss.str().c_str(), std::move(Args)});
478   }
479 
480   auto *Assume = cast<IntrinsicInst>(IntrinsicInst::Create(
481       FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}), OpBundle));
482   Assume->insertBefore(&F->begin()->front());
483   RetainedKnowledgeMap Map;
484   fillMapFromAssume(*Assume, Map);
485   for (int i = 0; i < (Size * 2); i++) {
486     if (!HasArg[i])
487       continue;
488     RetainedKnowledge K =
489         getKnowledgeFromUseInAssume(&*ShuffledArgs[i]->use_begin());
490     auto LookupIt = Map.find(RetainedKnowledgeKey{K.WasOn, K.AttrKind});
491     ASSERT_TRUE(LookupIt != Map.end());
492     MinMax MM = LookupIt->second[Assume];
493     ASSERT_TRUE(MM.Min == MM.Max);
494     ASSERT_TRUE(MM.Min == K.ArgValue);
495   }
496 }
497 
498 TEST(AssumeQueryAPI, getKnowledgeFromUseInAssume) {
499   // // For Fuzzing
500   // std::random_device dev;
501   // std::mt19937 Rng(dev());
502   // while (true) {
503   //   unsigned Seed = Rng();
504   //   dbgs() << Seed << "\n";
505   //   RunRandTest(Seed, 100000, 0, 2, 100);
506   // }
507   RunRandTest(23456, 4, 0, 2, 100);
508   RunRandTest(560987, 25, -3, 2, 100);
509 
510   // Large bundles can lead to special cases. this is why this test is soo
511   // large.
512   RunRandTest(9876789, 100000, -0, 7, 100);
513 }
514 
515 TEST(AssumeQueryAPI, AssumptionCache) {
516   LLVMContext C;
517   SMDiagnostic Err;
518   std::unique_ptr<Module> Mod = parseAssemblyString(
519       "declare void @llvm.assume(i1)\n"
520       "define void @test(i32* %P, i32* %P1, i32* %P2, i32* %P3, i1 %B) {\n"
521       "call void @llvm.assume(i1 true) [\"nonnull\"(i32* %P), \"align\"(i32* "
522       "%P2, i32 4), \"align\"(i32* %P, i32 8)]\n"
523       "call void @llvm.assume(i1 %B) [\"test\"(i32* %P1), "
524       "\"dereferenceable\"(i32* %P, i32 4)]\n"
525       "ret void\n}\n",
526       Err, C);
527   if (!Mod)
528     Err.print("AssumeQueryAPI", errs());
529   Function *F = Mod->getFunction("test");
530   BasicBlock::iterator First = F->begin()->begin();
531   BasicBlock::iterator Second = F->begin()->begin();
532   Second++;
533   AssumptionCacheTracker ACT;
534   AssumptionCache &AC = ACT.getAssumptionCache(*F);
535   auto AR = AC.assumptionsFor(F->getArg(3));
536   ASSERT_EQ(AR.size(), 0u);
537   AR = AC.assumptionsFor(F->getArg(1));
538   ASSERT_EQ(AR.size(), 1u);
539   ASSERT_EQ(AR[0].Index, 0u);
540   ASSERT_EQ(AR[0].Assume, &*Second);
541   AR = AC.assumptionsFor(F->getArg(2));
542   ASSERT_EQ(AR.size(), 1u);
543   ASSERT_EQ(AR[0].Index, 1u);
544   ASSERT_EQ(AR[0].Assume, &*First);
545   AR = AC.assumptionsFor(F->getArg(0));
546   ASSERT_EQ(AR.size(), 3u);
547   llvm::sort(AR,
548              [](const auto &L, const auto &R) { return L.Index < R.Index; });
549   ASSERT_EQ(AR[0].Assume, &*First);
550   ASSERT_EQ(AR[0].Index, 0u);
551   ASSERT_EQ(AR[1].Assume, &*Second);
552   ASSERT_EQ(AR[1].Index, 1u);
553   ASSERT_EQ(AR[2].Assume, &*First);
554   ASSERT_EQ(AR[2].Index, 2u);
555   AR = AC.assumptionsFor(F->getArg(4));
556   ASSERT_EQ(AR.size(), 1u);
557   ASSERT_EQ(AR[0].Assume, &*Second);
558   ASSERT_EQ(AR[0].Index, AssumptionCache::ExprResultIdx);
559   AC.unregisterAssumption(cast<CallInst>(&*Second));
560   AR = AC.assumptionsFor(F->getArg(1));
561   ASSERT_EQ(AR.size(), 0u);
562   AR = AC.assumptionsFor(F->getArg(0));
563   ASSERT_EQ(AR.size(), 3u);
564   llvm::sort(AR,
565              [](const auto &L, const auto &R) { return L.Index < R.Index; });
566   ASSERT_EQ(AR[0].Assume, &*First);
567   ASSERT_EQ(AR[0].Index, 0u);
568   ASSERT_EQ(AR[1].Assume, nullptr);
569   ASSERT_EQ(AR[1].Index, 1u);
570   ASSERT_EQ(AR[2].Assume, &*First);
571   ASSERT_EQ(AR[2].Index, 2u);
572   AR = AC.assumptionsFor(F->getArg(2));
573   ASSERT_EQ(AR.size(), 1u);
574   ASSERT_EQ(AR[0].Index, 1u);
575   ASSERT_EQ(AR[0].Assume, &*First);
576 }
577