[RISCV] Fix mgather -> riscv.masked.strided.load combine not extending indices (...
[llvm-project.git] / llvm / lib / Transforms / Utils / AssumeBundleBuilder.cpp
blobefa8e874b955e0a05e94e46867ae70de313d93bb
1 //===- AssumeBundleBuilder.cpp - tools to preserve informations -*- 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 //===----------------------------------------------------------------------===//
9 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
10 #include "llvm/ADT/DepthFirstIterator.h"
11 #include "llvm/ADT/MapVector.h"
12 #include "llvm/ADT/Statistic.h"
13 #include "llvm/Analysis/AssumeBundleQueries.h"
14 #include "llvm/Analysis/AssumptionCache.h"
15 #include "llvm/Analysis/ValueTracking.h"
16 #include "llvm/IR/Dominators.h"
17 #include "llvm/IR/Function.h"
18 #include "llvm/IR/InstIterator.h"
19 #include "llvm/IR/IntrinsicInst.h"
20 #include "llvm/IR/Module.h"
21 #include "llvm/IR/Operator.h"
22 #include "llvm/Support/CommandLine.h"
23 #include "llvm/Support/DebugCounter.h"
24 #include "llvm/Transforms/Utils/Local.h"
26 using namespace llvm;
28 namespace llvm {
29 cl::opt<bool> ShouldPreserveAllAttributes(
30 "assume-preserve-all", cl::init(false), cl::Hidden,
31 cl::desc("enable preservation of all attrbitues. even those that are "
32 "unlikely to be usefull"));
34 cl::opt<bool> EnableKnowledgeRetention(
35 "enable-knowledge-retention", cl::init(false), cl::Hidden,
36 cl::desc(
37 "enable preservation of attributes throughout code transformation"));
38 } // namespace llvm
40 #define DEBUG_TYPE "assume-builder"
42 STATISTIC(NumAssumeBuilt, "Number of assume built by the assume builder");
43 STATISTIC(NumBundlesInAssumes, "Total number of Bundles in the assume built");
44 STATISTIC(NumAssumesMerged,
45 "Number of assume merged by the assume simplify pass");
46 STATISTIC(NumAssumesRemoved,
47 "Number of assume removed by the assume simplify pass");
49 DEBUG_COUNTER(BuildAssumeCounter, "assume-builder-counter",
50 "Controls which assumes gets created");
52 namespace {
54 bool isUsefullToPreserve(Attribute::AttrKind Kind) {
55 switch (Kind) {
56 case Attribute::NonNull:
57 case Attribute::NoUndef:
58 case Attribute::Alignment:
59 case Attribute::Dereferenceable:
60 case Attribute::DereferenceableOrNull:
61 case Attribute::Cold:
62 return true;
63 default:
64 return false;
68 /// This function will try to transform the given knowledge into a more
69 /// canonical one. the canonical knowledge maybe the given one.
70 RetainedKnowledge canonicalizedKnowledge(RetainedKnowledge RK,
71 const DataLayout &DL) {
72 switch (RK.AttrKind) {
73 default:
74 return RK;
75 case Attribute::NonNull:
76 RK.WasOn = getUnderlyingObject(RK.WasOn);
77 return RK;
78 case Attribute::Alignment: {
79 Value *V = RK.WasOn->stripInBoundsOffsets([&](const Value *Strip) {
80 if (auto *GEP = dyn_cast<GEPOperator>(Strip))
81 RK.ArgValue =
82 MinAlign(RK.ArgValue, GEP->getMaxPreservedAlignment(DL).value());
83 });
84 RK.WasOn = V;
85 return RK;
87 case Attribute::Dereferenceable:
88 case Attribute::DereferenceableOrNull: {
89 int64_t Offset = 0;
90 Value *V = GetPointerBaseWithConstantOffset(RK.WasOn, Offset, DL,
91 /*AllowNonInBounds*/ false);
92 if (Offset < 0)
93 return RK;
94 RK.ArgValue = RK.ArgValue + Offset;
95 RK.WasOn = V;
98 return RK;
101 /// This class contain all knowledge that have been gather while building an
102 /// llvm.assume and the function to manipulate it.
103 struct AssumeBuilderState {
104 Module *M;
106 using MapKey = std::pair<Value *, Attribute::AttrKind>;
107 SmallMapVector<MapKey, uint64_t, 8> AssumedKnowledgeMap;
108 Instruction *InstBeingModified = nullptr;
109 AssumptionCache* AC = nullptr;
110 DominatorTree* DT = nullptr;
112 AssumeBuilderState(Module *M, Instruction *I = nullptr,
113 AssumptionCache *AC = nullptr, DominatorTree *DT = nullptr)
114 : M(M), InstBeingModified(I), AC(AC), DT(DT) {}
116 bool tryToPreserveWithoutAddingAssume(RetainedKnowledge RK) {
117 if (!InstBeingModified || !RK.WasOn)
118 return false;
119 bool HasBeenPreserved = false;
120 Use* ToUpdate = nullptr;
121 getKnowledgeForValue(
122 RK.WasOn, {RK.AttrKind}, AC,
123 [&](RetainedKnowledge RKOther, Instruction *Assume,
124 const CallInst::BundleOpInfo *Bundle) {
125 if (!isValidAssumeForContext(Assume, InstBeingModified, DT))
126 return false;
127 if (RKOther.ArgValue >= RK.ArgValue) {
128 HasBeenPreserved = true;
129 return true;
130 } else if (isValidAssumeForContext(InstBeingModified, Assume, DT)) {
131 HasBeenPreserved = true;
132 IntrinsicInst *Intr = cast<IntrinsicInst>(Assume);
133 ToUpdate = &Intr->op_begin()[Bundle->Begin + ABA_Argument];
134 return true;
136 return false;
138 if (ToUpdate)
139 ToUpdate->set(
140 ConstantInt::get(Type::getInt64Ty(M->getContext()), RK.ArgValue));
141 return HasBeenPreserved;
144 bool isKnowledgeWorthPreserving(RetainedKnowledge RK) {
145 if (!RK)
146 return false;
147 if (!RK.WasOn)
148 return true;
149 if (RK.WasOn->getType()->isPointerTy()) {
150 Value *UnderlyingPtr = getUnderlyingObject(RK.WasOn);
151 if (isa<AllocaInst>(UnderlyingPtr) || isa<GlobalValue>(UnderlyingPtr))
152 return false;
154 if (auto *Arg = dyn_cast<Argument>(RK.WasOn)) {
155 if (Arg->hasAttribute(RK.AttrKind) &&
156 (!Attribute::isIntAttrKind(RK.AttrKind) ||
157 Arg->getAttribute(RK.AttrKind).getValueAsInt() >= RK.ArgValue))
158 return false;
159 return true;
161 if (auto *Inst = dyn_cast<Instruction>(RK.WasOn))
162 if (wouldInstructionBeTriviallyDead(Inst)) {
163 if (RK.WasOn->use_empty())
164 return false;
165 Use *SingleUse = RK.WasOn->getSingleUndroppableUse();
166 if (SingleUse && SingleUse->getUser() == InstBeingModified)
167 return false;
169 return true;
172 void addKnowledge(RetainedKnowledge RK) {
173 RK = canonicalizedKnowledge(RK, M->getDataLayout());
175 if (!isKnowledgeWorthPreserving(RK))
176 return;
178 if (tryToPreserveWithoutAddingAssume(RK))
179 return;
180 MapKey Key{RK.WasOn, RK.AttrKind};
181 auto Lookup = AssumedKnowledgeMap.find(Key);
182 if (Lookup == AssumedKnowledgeMap.end()) {
183 AssumedKnowledgeMap[Key] = RK.ArgValue;
184 return;
186 assert(((Lookup->second == 0 && RK.ArgValue == 0) ||
187 (Lookup->second != 0 && RK.ArgValue != 0)) &&
188 "inconsistent argument value");
190 /// This is only desirable because for all attributes taking an argument
191 /// higher is better.
192 Lookup->second = std::max(Lookup->second, RK.ArgValue);
195 void addAttribute(Attribute Attr, Value *WasOn) {
196 if (Attr.isTypeAttribute() || Attr.isStringAttribute() ||
197 (!ShouldPreserveAllAttributes &&
198 !isUsefullToPreserve(Attr.getKindAsEnum())))
199 return;
200 uint64_t AttrArg = 0;
201 if (Attr.isIntAttribute())
202 AttrArg = Attr.getValueAsInt();
203 addKnowledge({Attr.getKindAsEnum(), AttrArg, WasOn});
206 void addCall(const CallBase *Call) {
207 auto addAttrList = [&](AttributeList AttrList, unsigned NumArgs) {
208 for (unsigned Idx = 0; Idx < NumArgs; Idx++)
209 for (Attribute Attr : AttrList.getParamAttrs(Idx)) {
210 bool IsPoisonAttr = Attr.hasAttribute(Attribute::NonNull) ||
211 Attr.hasAttribute(Attribute::Alignment);
212 if (!IsPoisonAttr || Call->isPassingUndefUB(Idx))
213 addAttribute(Attr, Call->getArgOperand(Idx));
215 for (Attribute Attr : AttrList.getFnAttrs())
216 addAttribute(Attr, nullptr);
218 addAttrList(Call->getAttributes(), Call->arg_size());
219 if (Function *Fn = Call->getCalledFunction())
220 addAttrList(Fn->getAttributes(), Fn->arg_size());
223 AssumeInst *build() {
224 if (AssumedKnowledgeMap.empty())
225 return nullptr;
226 if (!DebugCounter::shouldExecute(BuildAssumeCounter))
227 return nullptr;
228 Function *FnAssume = Intrinsic::getDeclaration(M, Intrinsic::assume);
229 LLVMContext &C = M->getContext();
230 SmallVector<OperandBundleDef, 8> OpBundle;
231 for (auto &MapElem : AssumedKnowledgeMap) {
232 SmallVector<Value *, 2> Args;
233 if (MapElem.first.first)
234 Args.push_back(MapElem.first.first);
236 /// This is only valid because for all attribute that currently exist a
237 /// value of 0 is useless. and should not be preserved.
238 if (MapElem.second)
239 Args.push_back(ConstantInt::get(Type::getInt64Ty(M->getContext()),
240 MapElem.second));
241 OpBundle.push_back(OperandBundleDefT<Value *>(
242 std::string(Attribute::getNameFromAttrKind(MapElem.first.second)),
243 Args));
244 NumBundlesInAssumes++;
246 NumAssumeBuilt++;
247 return cast<AssumeInst>(CallInst::Create(
248 FnAssume, ArrayRef<Value *>({ConstantInt::getTrue(C)}), OpBundle));
251 void addAccessedPtr(Instruction *MemInst, Value *Pointer, Type *AccType,
252 MaybeAlign MA) {
253 unsigned DerefSize = MemInst->getModule()
254 ->getDataLayout()
255 .getTypeStoreSize(AccType)
256 .getKnownMinValue();
257 if (DerefSize != 0) {
258 addKnowledge({Attribute::Dereferenceable, DerefSize, Pointer});
259 if (!NullPointerIsDefined(MemInst->getFunction(),
260 Pointer->getType()->getPointerAddressSpace()))
261 addKnowledge({Attribute::NonNull, 0u, Pointer});
263 if (MA.valueOrOne() > 1)
264 addKnowledge({Attribute::Alignment, MA.valueOrOne().value(), Pointer});
267 void addInstruction(Instruction *I) {
268 if (auto *Call = dyn_cast<CallBase>(I))
269 return addCall(Call);
270 if (auto *Load = dyn_cast<LoadInst>(I))
271 return addAccessedPtr(I, Load->getPointerOperand(), Load->getType(),
272 Load->getAlign());
273 if (auto *Store = dyn_cast<StoreInst>(I))
274 return addAccessedPtr(I, Store->getPointerOperand(),
275 Store->getValueOperand()->getType(),
276 Store->getAlign());
277 // TODO: Add support for the other Instructions.
278 // TODO: Maybe we should look around and merge with other llvm.assume.
282 } // namespace
284 AssumeInst *llvm::buildAssumeFromInst(Instruction *I) {
285 if (!EnableKnowledgeRetention)
286 return nullptr;
287 AssumeBuilderState Builder(I->getModule());
288 Builder.addInstruction(I);
289 return Builder.build();
292 bool llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC,
293 DominatorTree *DT) {
294 if (!EnableKnowledgeRetention || I->isTerminator())
295 return false;
296 bool Changed = false;
297 AssumeBuilderState Builder(I->getModule(), I, AC, DT);
298 Builder.addInstruction(I);
299 if (auto *Intr = Builder.build()) {
300 Intr->insertBefore(I);
301 Changed = true;
302 if (AC)
303 AC->registerAssumption(Intr);
305 return Changed;
308 AssumeInst *
309 llvm::buildAssumeFromKnowledge(ArrayRef<RetainedKnowledge> Knowledge,
310 Instruction *CtxI, AssumptionCache *AC,
311 DominatorTree *DT) {
312 AssumeBuilderState Builder(CtxI->getModule(), CtxI, AC, DT);
313 for (const RetainedKnowledge &RK : Knowledge)
314 Builder.addKnowledge(RK);
315 return Builder.build();
318 RetainedKnowledge llvm::simplifyRetainedKnowledge(AssumeInst *Assume,
319 RetainedKnowledge RK,
320 AssumptionCache *AC,
321 DominatorTree *DT) {
322 AssumeBuilderState Builder(Assume->getModule(), Assume, AC, DT);
323 RK = canonicalizedKnowledge(RK, Assume->getModule()->getDataLayout());
325 if (!Builder.isKnowledgeWorthPreserving(RK))
326 return RetainedKnowledge::none();
328 if (Builder.tryToPreserveWithoutAddingAssume(RK))
329 return RetainedKnowledge::none();
330 return RK;
333 namespace {
335 struct AssumeSimplify {
336 Function &F;
337 AssumptionCache &AC;
338 DominatorTree *DT;
339 LLVMContext &C;
340 SmallDenseSet<IntrinsicInst *> CleanupToDo;
341 StringMapEntry<uint32_t> *IgnoreTag;
342 SmallDenseMap<BasicBlock *, SmallVector<IntrinsicInst *, 4>, 8> BBToAssume;
343 bool MadeChange = false;
345 AssumeSimplify(Function &F, AssumptionCache &AC, DominatorTree *DT,
346 LLVMContext &C)
347 : F(F), AC(AC), DT(DT), C(C),
348 IgnoreTag(C.getOrInsertBundleTag(IgnoreBundleTag)) {}
350 void buildMapping(bool FilterBooleanArgument) {
351 BBToAssume.clear();
352 for (Value *V : AC.assumptions()) {
353 if (!V)
354 continue;
355 IntrinsicInst *Assume = cast<IntrinsicInst>(V);
356 if (FilterBooleanArgument) {
357 auto *Arg = dyn_cast<ConstantInt>(Assume->getOperand(0));
358 if (!Arg || Arg->isZero())
359 continue;
361 BBToAssume[Assume->getParent()].push_back(Assume);
364 for (auto &Elem : BBToAssume) {
365 llvm::sort(Elem.second,
366 [](const IntrinsicInst *LHS, const IntrinsicInst *RHS) {
367 return LHS->comesBefore(RHS);
372 /// Remove all asumes in CleanupToDo if there boolean argument is true and
373 /// ForceCleanup is set or the assume doesn't hold valuable knowledge.
374 void RunCleanup(bool ForceCleanup) {
375 for (IntrinsicInst *Assume : CleanupToDo) {
376 auto *Arg = dyn_cast<ConstantInt>(Assume->getOperand(0));
377 if (!Arg || Arg->isZero() ||
378 (!ForceCleanup &&
379 !isAssumeWithEmptyBundle(cast<AssumeInst>(*Assume))))
380 continue;
381 MadeChange = true;
382 if (ForceCleanup)
383 NumAssumesMerged++;
384 else
385 NumAssumesRemoved++;
386 Assume->eraseFromParent();
388 CleanupToDo.clear();
391 /// Remove knowledge stored in assume when it is already know by an attribute
392 /// or an other assume. This can when valid update an existing knowledge in an
393 /// attribute or an other assume.
394 void dropRedundantKnowledge() {
395 struct MapValue {
396 IntrinsicInst *Assume;
397 uint64_t ArgValue;
398 CallInst::BundleOpInfo *BOI;
400 buildMapping(false);
401 SmallDenseMap<std::pair<Value *, Attribute::AttrKind>,
402 SmallVector<MapValue, 2>, 16>
403 Knowledge;
404 for (BasicBlock *BB : depth_first(&F))
405 for (Value *V : BBToAssume[BB]) {
406 if (!V)
407 continue;
408 IntrinsicInst *Assume = cast<IntrinsicInst>(V);
409 for (CallInst::BundleOpInfo &BOI : Assume->bundle_op_infos()) {
410 auto RemoveFromAssume = [&]() {
411 CleanupToDo.insert(Assume);
412 if (BOI.Begin != BOI.End) {
413 Use *U = &Assume->op_begin()[BOI.Begin + ABA_WasOn];
414 U->set(UndefValue::get(U->get()->getType()));
416 BOI.Tag = IgnoreTag;
418 if (BOI.Tag == IgnoreTag) {
419 CleanupToDo.insert(Assume);
420 continue;
422 RetainedKnowledge RK =
423 getKnowledgeFromBundle(cast<AssumeInst>(*Assume), BOI);
424 if (auto *Arg = dyn_cast_or_null<Argument>(RK.WasOn)) {
425 bool HasSameKindAttr = Arg->hasAttribute(RK.AttrKind);
426 if (HasSameKindAttr)
427 if (!Attribute::isIntAttrKind(RK.AttrKind) ||
428 Arg->getAttribute(RK.AttrKind).getValueAsInt() >=
429 RK.ArgValue) {
430 RemoveFromAssume();
431 continue;
433 if (isValidAssumeForContext(
434 Assume, &*F.getEntryBlock().getFirstInsertionPt()) ||
435 Assume == &*F.getEntryBlock().getFirstInsertionPt()) {
436 if (HasSameKindAttr)
437 Arg->removeAttr(RK.AttrKind);
438 Arg->addAttr(Attribute::get(C, RK.AttrKind, RK.ArgValue));
439 MadeChange = true;
440 RemoveFromAssume();
441 continue;
444 auto &Lookup = Knowledge[{RK.WasOn, RK.AttrKind}];
445 for (MapValue &Elem : Lookup) {
446 if (!isValidAssumeForContext(Elem.Assume, Assume, DT))
447 continue;
448 if (Elem.ArgValue >= RK.ArgValue) {
449 RemoveFromAssume();
450 continue;
451 } else if (isValidAssumeForContext(Assume, Elem.Assume, DT)) {
452 Elem.Assume->op_begin()[Elem.BOI->Begin + ABA_Argument].set(
453 ConstantInt::get(Type::getInt64Ty(C), RK.ArgValue));
454 MadeChange = true;
455 RemoveFromAssume();
456 continue;
459 Lookup.push_back({Assume, RK.ArgValue, &BOI});
464 using MergeIterator = SmallVectorImpl<IntrinsicInst *>::iterator;
466 /// Merge all Assumes from Begin to End in and insert the resulting assume as
467 /// high as possible in the basicblock.
468 void mergeRange(BasicBlock *BB, MergeIterator Begin, MergeIterator End) {
469 if (Begin == End || std::next(Begin) == End)
470 return;
471 /// Provide no additional information so that AssumeBuilderState doesn't
472 /// try to do any punning since it already has been done better.
473 AssumeBuilderState Builder(F.getParent());
475 /// For now it is initialized to the best value it could have
476 Instruction *InsertPt = BB->getFirstNonPHI();
477 if (isa<LandingPadInst>(InsertPt))
478 InsertPt = InsertPt->getNextNode();
479 for (IntrinsicInst *I : make_range(Begin, End)) {
480 CleanupToDo.insert(I);
481 for (CallInst::BundleOpInfo &BOI : I->bundle_op_infos()) {
482 RetainedKnowledge RK =
483 getKnowledgeFromBundle(cast<AssumeInst>(*I), BOI);
484 if (!RK)
485 continue;
486 Builder.addKnowledge(RK);
487 if (auto *I = dyn_cast_or_null<Instruction>(RK.WasOn))
488 if (I->getParent() == InsertPt->getParent() &&
489 (InsertPt->comesBefore(I) || InsertPt == I))
490 InsertPt = I->getNextNode();
494 /// Adjust InsertPt if it is before Begin, since mergeAssumes only
495 /// guarantees we can place the resulting assume between Begin and End.
496 if (InsertPt->comesBefore(*Begin))
497 for (auto It = (*Begin)->getIterator(), E = InsertPt->getIterator();
498 It != E; --It)
499 if (!isGuaranteedToTransferExecutionToSuccessor(&*It)) {
500 InsertPt = It->getNextNode();
501 break;
503 auto *MergedAssume = Builder.build();
504 if (!MergedAssume)
505 return;
506 MadeChange = true;
507 MergedAssume->insertBefore(InsertPt);
508 AC.registerAssumption(MergedAssume);
511 /// Merge assume when they are in the same BasicBlock and for all instruction
512 /// between them isGuaranteedToTransferExecutionToSuccessor returns true.
513 void mergeAssumes() {
514 buildMapping(true);
516 SmallVector<MergeIterator, 4> SplitPoints;
517 for (auto &Elem : BBToAssume) {
518 SmallVectorImpl<IntrinsicInst *> &AssumesInBB = Elem.second;
519 if (AssumesInBB.size() < 2)
520 continue;
521 /// AssumesInBB is already sorted by order in the block.
523 BasicBlock::iterator It = AssumesInBB.front()->getIterator();
524 BasicBlock::iterator E = AssumesInBB.back()->getIterator();
525 SplitPoints.push_back(AssumesInBB.begin());
526 MergeIterator LastSplit = AssumesInBB.begin();
527 for (; It != E; ++It)
528 if (!isGuaranteedToTransferExecutionToSuccessor(&*It)) {
529 for (; (*LastSplit)->comesBefore(&*It); ++LastSplit)
531 if (SplitPoints.back() != LastSplit)
532 SplitPoints.push_back(LastSplit);
534 SplitPoints.push_back(AssumesInBB.end());
535 for (auto SplitIt = SplitPoints.begin();
536 SplitIt != std::prev(SplitPoints.end()); SplitIt++) {
537 mergeRange(Elem.first, *SplitIt, *(SplitIt + 1));
539 SplitPoints.clear();
544 bool simplifyAssumes(Function &F, AssumptionCache *AC, DominatorTree *DT) {
545 AssumeSimplify AS(F, *AC, DT, F.getContext());
547 /// Remove knowledge that is already known by a dominating other assume or an
548 /// attribute.
549 AS.dropRedundantKnowledge();
551 /// Remove assume that are empty.
552 AS.RunCleanup(false);
554 /// Merge assume in the same basicblock when possible.
555 AS.mergeAssumes();
557 /// Remove assume that were merged.
558 AS.RunCleanup(true);
559 return AS.MadeChange;
562 } // namespace
564 PreservedAnalyses AssumeSimplifyPass::run(Function &F,
565 FunctionAnalysisManager &AM) {
566 if (!EnableKnowledgeRetention)
567 return PreservedAnalyses::all();
568 if (!simplifyAssumes(F, &AM.getResult<AssumptionAnalysis>(F),
569 AM.getCachedResult<DominatorTreeAnalysis>(F)))
570 return PreservedAnalyses::all();
571 PreservedAnalyses PA;
572 PA.preserveSet<CFGAnalyses>();
573 return PA;
576 PreservedAnalyses AssumeBuilderPass::run(Function &F,
577 FunctionAnalysisManager &AM) {
578 AssumptionCache *AC = &AM.getResult<AssumptionAnalysis>(F);
579 DominatorTree* DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
580 bool Changed = false;
581 for (Instruction &I : instructions(F))
582 Changed |= salvageKnowledge(&I, AC, DT);
583 if (!Changed)
584 PreservedAnalyses::all();
585 PreservedAnalyses PA;
586 PA.preserveSet<CFGAnalyses>();
587 return PA;