1//===- InferAlignment.cpp -------------------------------------------------===//
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// Infer alignment for load, stores and other memory operations based on
10// trailing zero known bits information.
11//
12//===----------------------------------------------------------------------===//
13
14#include "llvm/Transforms/Scalar/InferAlignment.h"
15#include "llvm/ADT/APInt.h"
16#include "llvm/ADT/STLFunctionalExtras.h"
17#include "llvm/ADT/ScopedHashTable.h"
18#include "llvm/Analysis/AssumptionCache.h"
19#include "llvm/Analysis/DomConditionCache.h"
20#include "llvm/Analysis/SimplifyQuery.h"
21#include "llvm/Analysis/ValueTracking.h"
22#include "llvm/IR/Instruction.h"
23#include "llvm/IR/Instructions.h"
24#include "llvm/IR/IntrinsicInst.h"
25#include "llvm/IR/PatternMatch.h"
26#include "llvm/Support/KnownBits.h"
27#include "llvm/Transforms/Scalar.h"
28#include "llvm/Transforms/Utils/Local.h"
29#include <list>
30
31using namespace llvm;
32using namespace llvm::PatternMatch;
33
34static bool tryToImproveAlign(
35 const DataLayout &DL, Instruction *I,
36 function_ref<Align(Value *PtrOp, Align OldAlign, Align PrefAlign)> Fn) {
37
38 if (auto *PtrOp = getLoadStorePointerOperand(V: I)) {
39 Align OldAlign = getLoadStoreAlignment(I);
40 Align PrefAlign = DL.getPrefTypeAlign(Ty: getLoadStoreType(I));
41
42 Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign);
43 if (NewAlign > OldAlign) {
44 setLoadStoreAlignment(I, NewAlign);
45 return true;
46 }
47 }
48
49 Value *PtrOp;
50 const APInt *Const;
51 if (match(V: I, P: m_And(L: m_PtrToIntOrAddr(Op: m_Value(V&: PtrOp)), R: m_APInt(Res&: Const)))) {
52 Align ActualAlign = Fn(PtrOp, Align(1), Align(1));
53 if (Const->ult(RHS: ActualAlign.value())) {
54 I->replaceAllUsesWith(V: Constant::getNullValue(Ty: I->getType()));
55 return true;
56 }
57 if (Const->uge(
58 RHS: APInt::getBitsSetFrom(numBits: Const->getBitWidth(), loBit: Log2(A: ActualAlign)))) {
59 I->replaceAllUsesWith(V: I->getOperand(i: 0));
60 return true;
61 }
62 }
63 if (match(V: I, P: m_Trunc(Op: m_PtrToIntOrAddr(Op: m_Value(V&: PtrOp))))) {
64 Align ActualAlign = Fn(PtrOp, Align(1), Align(1));
65 if (Log2(A: ActualAlign) >= I->getType()->getScalarSizeInBits()) {
66 I->replaceAllUsesWith(V: Constant::getNullValue(Ty: I->getType()));
67 return true;
68 }
69 }
70
71 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: I);
72 if (!II)
73 return false;
74
75 if (!isa<MemIntrinsic>(Val: II) &&
76 II->getIntrinsicID() != Intrinsic::masked_load &&
77 II->getIntrinsicID() != Intrinsic::masked_store)
78 return false;
79
80 bool Changed = false;
81 for (unsigned ArgNo = 0; ArgNo != II->arg_size(); ++ArgNo) {
82 Value *Arg = II->getArgOperand(i: ArgNo);
83 if (!Arg->getType()->isPointerTy())
84 continue;
85
86 // The align attribute only proves the alignment if passing poison is UB
87 // (e.g. noundef): otherwise a zero-length memset may pass an unaligned,
88 // poison pointer, so it must not seed the base pointer alignment.
89 Align OldAlign = II->getParamAlign(ArgNo).valueOrOne();
90 Align KnownAlign = II->isPassingUndefUB(ArgNo) ? OldAlign : Align(1);
91 Align NewAlign = Fn(Arg, KnownAlign, Align(1));
92 if (NewAlign <= OldAlign)
93 continue;
94
95 II->addParamAttr(ArgNo,
96 Attr: Attribute::getWithAlignment(Context&: II->getContext(), Alignment: NewAlign));
97 Changed = true;
98 }
99 return Changed;
100}
101
102using ScopedHT =
103 ScopedHashTable<Value *, Align, DenseMapInfo<Value *>, BumpPtrAllocator>;
104struct AlignmentScope {
105 // If BB is nullptr, the BB is processed.
106 BasicBlock *BB;
107 DomTreeNode::const_iterator Iter;
108 DomTreeNode::const_iterator End;
109 ScopedHT::ScopeTy Scope;
110
111 AlignmentScope(DomTreeNode *N, ScopedHT &Table)
112 : BB(N->getBlock()), Iter(N->begin()), End(N->end()), Scope(Table) {}
113};
114
115bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) {
116 const DataLayout &DL = F.getDataLayout();
117 bool Changed = false;
118
119 DomConditionCache DC;
120 // Register all conditions before querying known bits, because function block
121 // order does not guarantee dominators precede their dominated blocks.
122 for (BasicBlock &BB : F)
123 if (auto *BI = dyn_cast<CondBrInst>(Val: BB.getTerminator()))
124 DC.registerBranch(BI);
125
126 // Enforce preferred type alignment if possible. We do this as a separate
127 // pass first, because it may improve the alignments we infer below.
128 for (BasicBlock &BB : F) {
129 for (Instruction &I : BB) {
130 Changed |= tryToImproveAlign(
131 DL, I: &I, Fn: [&](Value *PtrOp, Align OldAlign, Align PrefAlign) {
132 if (PrefAlign > OldAlign)
133 return std::max(a: OldAlign,
134 b: tryEnforceAlignment(V: PtrOp, PrefAlign, DL));
135 return OldAlign;
136 });
137 }
138 }
139
140 SimplifyQuery SQ(DL, /*TLI=*/nullptr, &DT, &AC, /*CxtI=*/nullptr,
141 /*UseInstrInfo=*/true, /*CanUseUndef=*/true, &DC);
142
143 // Compute alignment from known bits.
144 auto InferFromKnownBits = [&](Instruction &I, Value *PtrOp) {
145 KnownBits Known = computeKnownBits(V: PtrOp, Q: SQ.getWithInstruction(I: &I));
146 unsigned TrailZ =
147 std::min(a: Known.countMinTrailingZeros(), b: +Value::MaxAlignmentExponent);
148 return Align::fromLog2(Log2Value: std::min(a: Known.getBitWidth() - 1, b: TrailZ));
149 };
150
151 // Propagate alignment between loads and stores that originate from the
152 // same base pointer.
153 ScopedHT BestBasePointerAligns;
154 auto InferFromBasePointer = [&](Value *PtrOp, Align LoadStoreAlign) {
155 APInt OffsetFromBase(DL.getIndexTypeSizeInBits(Ty: PtrOp->getType()), 0);
156 PtrOp = PtrOp->stripAndAccumulateConstantOffsets(DL, Offset&: OffsetFromBase, AllowNonInbounds: true);
157 // Derive the base pointer alignment from the load/store alignment
158 // and the offset from the base pointer.
159 Align BasePointerAlign =
160 commonAlignment(A: LoadStoreAlign, Offset: OffsetFromBase.getLimitedValue());
161
162 if (auto BestAlign = BestBasePointerAligns.lookup(Key: PtrOp);
163 BestAlign != Align()) {
164 // If the stored base pointer alignment is better than the
165 // base pointer alignment we derived, we may be able to use it
166 // to improve the load/store alignment. If not, store the
167 // improved base pointer alignment for future iterations.
168 if (BestAlign > BasePointerAlign) {
169 Align BetterLoadStoreAlign =
170 commonAlignment(A: BestAlign, Offset: OffsetFromBase.getLimitedValue());
171 return BetterLoadStoreAlign;
172 }
173 }
174
175 BestBasePointerAligns.insert(Key: PtrOp, Val: BasePointerAlign);
176 return LoadStoreAlign;
177 };
178
179 // AlignmentScope is unmovable.
180 std::list<AlignmentScope> Stack;
181 Stack.emplace_back(args: DT.getRootNode(), args&: BestBasePointerAligns);
182 while (!Stack.empty()) {
183 AlignmentScope &Top = Stack.back();
184 if (Top.BB) {
185 for (Instruction &I : *Top.BB) {
186 Changed |= tryToImproveAlign(
187 DL, I: &I, Fn: [&](Value *PtrOp, Align OldAlign, Align PrefAlign) {
188 return std::max(a: InferFromKnownBits(I, PtrOp),
189 b: InferFromBasePointer(PtrOp, OldAlign));
190 });
191 }
192 Top.BB = nullptr;
193 }
194
195 if (Top.Iter != Top.End)
196 Stack.emplace_back(args: *Top.Iter++, args&: BestBasePointerAligns);
197 else
198 Stack.pop_back();
199 }
200
201 return Changed;
202}
203
204PreservedAnalyses InferAlignmentPass::run(Function &F,
205 FunctionAnalysisManager &AM) {
206 AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(IR&: F);
207 DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(IR&: F);
208 inferAlignment(F, AC, DT);
209 // Changes to alignment shouldn't invalidated analyses.
210 return PreservedAnalyses::all();
211}
212