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/Analysis/AssumptionCache.h"
18#include "llvm/Analysis/ValueTracking.h"
19#include "llvm/IR/Instruction.h"
20#include "llvm/IR/Instructions.h"
21#include "llvm/IR/IntrinsicInst.h"
22#include "llvm/IR/PatternMatch.h"
23#include "llvm/Support/KnownBits.h"
24#include "llvm/Transforms/Scalar.h"
25#include "llvm/Transforms/Utils/Local.h"
26
27using namespace llvm;
28using namespace llvm::PatternMatch;
29
30static bool tryToImproveAlign(
31 const DataLayout &DL, Instruction *I,
32 function_ref<Align(Value *PtrOp, Align OldAlign, Align PrefAlign)> Fn) {
33
34 if (auto *PtrOp = getLoadStorePointerOperand(V: I)) {
35 Align OldAlign = getLoadStoreAlignment(I);
36 Align PrefAlign = DL.getPrefTypeAlign(Ty: getLoadStoreType(I));
37
38 Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign);
39 if (NewAlign > OldAlign) {
40 setLoadStoreAlignment(I, NewAlign);
41 return true;
42 }
43 }
44
45 Value *PtrOp;
46 const APInt *Const;
47 if (match(V: I, P: m_And(L: m_PtrToIntOrAddr(Op: m_Value(V&: PtrOp)), R: m_APInt(Res&: Const)))) {
48 Align ActualAlign = Fn(PtrOp, Align(1), Align(1));
49 if (Const->ult(RHS: ActualAlign.value())) {
50 I->replaceAllUsesWith(V: Constant::getNullValue(Ty: I->getType()));
51 return true;
52 }
53 if (Const->uge(
54 RHS: APInt::getBitsSetFrom(numBits: Const->getBitWidth(), loBit: Log2(A: ActualAlign)))) {
55 I->replaceAllUsesWith(V: I->getOperand(i: 0));
56 return true;
57 }
58 }
59 if (match(V: I, P: m_Trunc(Op: m_PtrToIntOrAddr(Op: m_Value(V&: PtrOp))))) {
60 Align ActualAlign = Fn(PtrOp, Align(1), Align(1));
61 if (Log2(A: ActualAlign) >= I->getType()->getScalarSizeInBits()) {
62 I->replaceAllUsesWith(V: Constant::getNullValue(Ty: I->getType()));
63 return true;
64 }
65 }
66
67 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: I);
68 if (!II)
69 return false;
70
71 // TODO: Handle more memory intrinsics.
72 switch (II->getIntrinsicID()) {
73 case Intrinsic::masked_load:
74 case Intrinsic::masked_store: {
75 unsigned PtrOpIdx = II->getIntrinsicID() == Intrinsic::masked_load ? 0 : 1;
76 Value *PtrOp = II->getArgOperand(i: PtrOpIdx);
77 Type *Type = II->getIntrinsicID() == Intrinsic::masked_load
78 ? II->getType()
79 : II->getArgOperand(i: 0)->getType();
80
81 Align OldAlign = II->getParamAlign(ArgNo: PtrOpIdx).valueOrOne();
82 Align PrefAlign = DL.getPrefTypeAlign(Ty: Type);
83 Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign);
84 if (NewAlign <= OldAlign)
85 return false;
86
87 II->addParamAttr(ArgNo: PtrOpIdx,
88 Attr: Attribute::getWithAlignment(Context&: II->getContext(), Alignment: NewAlign));
89 return true;
90 }
91 default:
92 return false;
93 }
94}
95
96bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) {
97 const DataLayout &DL = F.getDataLayout();
98 bool Changed = false;
99
100 // Enforce preferred type alignment if possible. We do this as a separate
101 // pass first, because it may improve the alignments we infer below.
102 for (BasicBlock &BB : F) {
103 for (Instruction &I : BB) {
104 Changed |= tryToImproveAlign(
105 DL, I: &I, Fn: [&](Value *PtrOp, Align OldAlign, Align PrefAlign) {
106 if (PrefAlign > OldAlign)
107 return std::max(a: OldAlign,
108 b: tryEnforceAlignment(V: PtrOp, PrefAlign, DL));
109 return OldAlign;
110 });
111 }
112 }
113
114 // Compute alignment from known bits.
115 auto InferFromKnownBits = [&](Instruction &I, Value *PtrOp) {
116 KnownBits Known = computeKnownBits(V: PtrOp, DL, AC: &AC, CxtI: &I, DT: &DT);
117 unsigned TrailZ =
118 std::min(a: Known.countMinTrailingZeros(), b: +Value::MaxAlignmentExponent);
119 return Align(1ull << std::min(a: Known.getBitWidth() - 1, b: TrailZ));
120 };
121
122 // Propagate alignment between loads and stores that originate from the
123 // same base pointer.
124 DenseMap<Value *, Align> BestBasePointerAligns;
125 auto InferFromBasePointer = [&](Value *PtrOp, Align LoadStoreAlign) {
126 APInt OffsetFromBase(DL.getIndexTypeSizeInBits(Ty: PtrOp->getType()), 0);
127 PtrOp = PtrOp->stripAndAccumulateConstantOffsets(DL, Offset&: OffsetFromBase, AllowNonInbounds: true);
128 // Derive the base pointer alignment from the load/store alignment
129 // and the offset from the base pointer.
130 Align BasePointerAlign =
131 commonAlignment(A: LoadStoreAlign, Offset: OffsetFromBase.getLimitedValue());
132
133 auto [It, Inserted] =
134 BestBasePointerAligns.try_emplace(Key: PtrOp, Args&: BasePointerAlign);
135 if (!Inserted) {
136 // If the stored base pointer alignment is better than the
137 // base pointer alignment we derived, we may be able to use it
138 // to improve the load/store alignment. If not, store the
139 // improved base pointer alignment for future iterations.
140 if (It->second > BasePointerAlign) {
141 Align BetterLoadStoreAlign =
142 commonAlignment(A: It->second, Offset: OffsetFromBase.getLimitedValue());
143 return BetterLoadStoreAlign;
144 }
145 It->second = BasePointerAlign;
146 }
147 return LoadStoreAlign;
148 };
149
150 for (BasicBlock &BB : F) {
151 // We need to reset the map for each block because alignment information
152 // can only be propagated from instruction A to B if A dominates B.
153 // This is because control flow (and exception throwing) could be dependent
154 // on the address (and its alignment) at runtime. Some sort of dominator
155 // tree approach could be better, but doing a simple forward pass through a
156 // single basic block is correct too.
157 BestBasePointerAligns.clear();
158
159 for (Instruction &I : BB) {
160 Changed |= tryToImproveAlign(
161 DL, I: &I, Fn: [&](Value *PtrOp, Align OldAlign, Align PrefAlign) {
162 return std::max(a: InferFromKnownBits(I, PtrOp),
163 b: InferFromBasePointer(PtrOp, OldAlign));
164 });
165 }
166 }
167
168 return Changed;
169}
170
171PreservedAnalyses InferAlignmentPass::run(Function &F,
172 FunctionAnalysisManager &AM) {
173 AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(IR&: F);
174 DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(IR&: F);
175 inferAlignment(F, AC, DT);
176 // Changes to alignment shouldn't invalidated analyses.
177 return PreservedAnalyses::all();
178}
179