1//===-- NVPTXPromoteParamAlign.cpp - Promote .param alignment ------------===//
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// Increase the .param-space alignment of NVPTX arguments and return values so
10// their loads and stores can be vectorized. On every defined function:
11//
12// 1. Give each byval param an explicit ABI `align` for its pointee type
13// (capped at the PTX max). byval already implies this, but the alignment
14// is otherwise invisible to IR alignment analyses.
15// 2. For a local function whose every use is a type-compatible direct call,
16// we control all call sites and raise aggregate/byval param and return
17// alignment to at least 16 (for 128-bit vectorization). This is recorded
18// as `stackalign` and mirrored onto the calls.
19// 3. Propagate the result onto byval loads at a known offset, since `align`
20// and `stackalign` aren't both picked up by IR alignment analyses.
21//
22// (2) runs before (1) so byval `align` still matches between caller and callee
23// while stackalign is mirrored onto the calls.
24//
25//===----------------------------------------------------------------------===//
26
27#include "NVPTX.h"
28#include "NVPTXUtilities.h"
29#include "llvm/ADT/Sequence.h"
30#include "llvm/IR/Attributes.h"
31#include "llvm/IR/Function.h"
32#include "llvm/IR/Instructions.h"
33#include "llvm/IR/Module.h"
34#include "llvm/InitializePasses.h"
35#include "llvm/Pass.h"
36#include "llvm/Support/Debug.h"
37#include <optional>
38#include <queue>
39
40#define DEBUG_TYPE "nvptx-promote-param-align"
41
42using namespace llvm;
43
44namespace {
45class NVPTXPromoteParamAlignLegacyPass : public ModulePass {
46 bool runOnModule(Module &M) override;
47
48public:
49 static char ID;
50 NVPTXPromoteParamAlignLegacyPass() : ModulePass(ID) {}
51 StringRef getPassName() const override {
52 return "Promote alignment of parameters and return values (NVPTX)";
53 }
54};
55} // namespace
56
57char NVPTXPromoteParamAlignLegacyPass::ID = 0;
58
59INITIALIZE_PASS(NVPTXPromoteParamAlignLegacyPass, "nvptx-promote-param-align",
60 "Promote alignment of parameters and return values (NVPTX)",
61 false, false)
62
63// Return true if the attributes that determine an NVPTX .param slot's layout
64// match.
65static bool layoutAttrsMatch(AttributeSet CalleeAttrs, AttributeSet CallAttrs) {
66 if (CalleeAttrs.getByValType() != CallAttrs.getByValType() ||
67 CalleeAttrs.getStackAlignment() != CallAttrs.getStackAlignment())
68 return false;
69
70 // `align` only affects the layout for byval parameters.
71 return !CalleeAttrs.getByValType() ||
72 CalleeAttrs.getAlignment() == CallAttrs.getAlignment();
73}
74
75static bool callSiteMatchesCalleeABI(const CallBase &CB, const Function &F) {
76 const AttributeList CalleeAttrs = F.getAttributes();
77 const AttributeList CallAttrs = CB.getAttributes();
78
79 if (!layoutAttrsMatch(CalleeAttrs: CalleeAttrs.getRetAttrs(), CallAttrs: CallAttrs.getRetAttrs()))
80 return false;
81
82 return all_of(Range: seq(Size: F.arg_size()), P: [&](size_t I) {
83 return layoutAttrsMatch(CalleeAttrs: CalleeAttrs.getParamAttrs(ArgNo: I),
84 CallAttrs: CallAttrs.getParamAttrs(ArgNo: I));
85 });
86}
87
88// Promotable if the function is local and every use is an ABI-compatible direct
89// call, so we control every call site and can raise alignment on both sides.
90static bool canPromoteParamAlign(Function &F) {
91 if (F.isDeclaration() || !F.hasLocalLinkage())
92 return false;
93
94 if (F.hasAddressTaken(/*Users=*/nullptr, /*IgnoreCallbackUses=*/false,
95 /*IgnoreAssumeLikeCalls=*/true,
96 /*IgnoreLLVMUsed=*/IngoreLLVMUsed: true))
97 return false;
98
99 return all_of(Range: F.users(), P: [&](const User *U) {
100 const auto *CB = dyn_cast<CallBase>(Val: U);
101 if (!CB || CB->getCalledOperand() != &F)
102 return true;
103 return CB->getFunctionType() == F.getFunctionType() &&
104 callSiteMatchesCalleeABI(CB: *CB, F);
105 });
106}
107
108// Raise the alignment of every load reachable from a byval pointer at a known
109// constant offset. Must stay in sync with the param load/store in LowerCall.
110static bool propagateAlignmentToLoads(Value *Val, Align NewAlign,
111 const DataLayout &DL) {
112 struct Load {
113 LoadInst *Inst;
114 uint64_t Offset;
115 };
116
117 struct LoadContext {
118 Value *InitialVal;
119 uint64_t Offset;
120 };
121
122 SmallVector<Load> Loads;
123 std::queue<LoadContext> Worklist;
124 Worklist.push(x: {.InitialVal: Val, .Offset: 0});
125
126 while (!Worklist.empty()) {
127 LoadContext Ctx = Worklist.front();
128 Worklist.pop();
129
130 for (User *CurUser : Ctx.InitialVal->users()) {
131 if (auto *I = dyn_cast<LoadInst>(Val: CurUser))
132 Loads.push_back(Elt: {.Inst: I, .Offset: Ctx.Offset});
133 else if (isa<BitCastInst>(Val: CurUser) || isa<AddrSpaceCastInst>(Val: CurUser))
134 Worklist.push(x: {.InitialVal: cast<Instruction>(Val: CurUser), .Offset: Ctx.Offset});
135 else if (auto *I = dyn_cast<GetElementPtrInst>(Val: CurUser)) {
136 APInt OffsetAccumulated =
137 APInt::getZero(numBits: DL.getIndexTypeSizeInBits(Ty: I->getType()));
138
139 if (!I->accumulateConstantOffset(DL, Offset&: OffsetAccumulated))
140 continue;
141
142 uint64_t OffsetLimit = -1;
143 uint64_t Offset = OffsetAccumulated.getLimitedValue(Limit: OffsetLimit);
144 assert(Offset != OffsetLimit && "Expect Offset less than UINT64_MAX");
145
146 Worklist.push(x: {.InitialVal: I, .Offset: Ctx.Offset + Offset});
147 }
148 }
149 }
150
151 bool Changed = false;
152 for (Load &CurLoad : Loads) {
153 Align NewLoadAlign = commonAlignment(A: NewAlign, Offset: CurLoad.Offset);
154 if (NewLoadAlign > CurLoad.Inst->getAlign()) {
155 CurLoad.Inst->setAlignment(NewLoadAlign);
156 Changed = true;
157 }
158 }
159 return Changed;
160}
161
162// Bump an alignment up to at least 16 (for 128-bit vectorization), or nullopt
163// if it's already large enough.
164static MaybeAlign getPromotedParamAlign(Align CurrentAlign) {
165 const Align PromotedAlign = std::max(a: CurrentAlign, b: Align(16));
166 if (PromotedAlign > CurrentAlign)
167 return PromotedAlign;
168 return std::nullopt;
169}
170
171static bool promoteParamAlign(Function &F) {
172 if (!canPromoteParamAlign(F))
173 return false;
174
175 LLVMContext &Ctx = F.getContext();
176 const DataLayout &DL = F.getDataLayout();
177
178 // Promoted (arg index, new alignment) pairs, to mirror onto call sites.
179 SmallVector<std::pair<unsigned, Align>, 8> PromotedParams;
180 MaybeAlign PromotedRet;
181
182 // Promote aggregate and byval parameters.
183 for (Argument &Arg : F.args()) {
184 const bool IsByVal = Arg.hasByValAttr();
185 Type *ArgTy = IsByVal ? Arg.getParamByValType() : Arg.getType();
186 if (ArgTy->isEmptyTy() || (!IsByVal && !shouldPassAsArray(Ty: ArgTy)))
187 continue;
188
189 // An explicit stackalign already wins at emission time, nothing to promote.
190 if (Arg.getParamStackAlign())
191 continue;
192 const unsigned ArgNo = Arg.getArgNo();
193
194 // `align` only applies to byval (pointer) args, not by-value aggregates.
195 Align CurrentAlign = getPTXParamTypeAlign(ArgTy, DL);
196 if (IsByVal)
197 CurrentAlign = std::max(a: CurrentAlign, b: Arg.getParamAlign().valueOrOne());
198 const MaybeAlign PromotedAlign = getPromotedParamAlign(CurrentAlign);
199 if (!PromotedAlign)
200 continue;
201
202 LLVM_DEBUG(dbgs() << "Promoting alignment of " << Arg << " to "
203 << PromotedAlign->value() << '\n');
204 Arg.addAttr(Attr: Attribute::getWithStackAlignment(Context&: Ctx, Alignment: *PromotedAlign));
205 PromotedParams.emplace_back(Args: ArgNo, Args: *PromotedAlign);
206 }
207
208 // Promote an aggregate return value.
209 Type *RetTy = F.getReturnType();
210 if (shouldPassAsArray(Ty: RetTy) && !RetTy->isEmptyTy() &&
211 !F.getAttributes().getRetStackAlignment()) {
212 const MaybeAlign PromotedAlign =
213 getPromotedParamAlign(CurrentAlign: getPTXParamTypeAlign(ArgTy: RetTy, DL));
214 if (PromotedAlign) {
215 F.addRetAttr(Attr: Attribute::getWithStackAlignment(Context&: Ctx, Alignment: *PromotedAlign));
216 PromotedRet = *PromotedAlign;
217 }
218 }
219
220 if (PromotedParams.empty() && !PromotedRet)
221 return false;
222
223 // Mirror the promotion onto every direct call site so both sides agree on the
224 // .param layout. canPromoteParamAlign already verified they're
225 // ABI-compatible.
226 for (User *U : F.users()) {
227 auto *CB = dyn_cast<CallBase>(Val: U);
228 if (!CB || CB->getCalledOperand() != &F)
229 continue;
230
231 for (const auto &[ArgNo, PromotedAlign] : PromotedParams)
232 CB->addParamAttr(ArgNo,
233 Attr: Attribute::getWithStackAlignment(Context&: Ctx, Alignment: PromotedAlign));
234 if (PromotedRet)
235 CB->addRetAttr(Attr: Attribute::getWithStackAlignment(Context&: Ctx, Alignment: *PromotedRet));
236
237 assert(callSiteMatchesCalleeABI(*CB, F) &&
238 "mirroring must preserve call-site/callee ABI compatibility");
239 }
240
241 return true;
242}
243
244// Spell out each byval parameter's ABI alignment as an explicit `align` (step 1
245// above). Runs after promoteParamAlign, which needs byval `align` to still
246// match between callers and callees.
247static bool setByValParamABIAlign(Function &F) {
248 if (F.isDeclaration())
249 return false;
250
251 LLVMContext &Ctx = F.getContext();
252 const DataLayout &DL = F.getDataLayout();
253 bool Changed = false;
254 for (Argument &Arg : F.args()) {
255 if (!Arg.hasByValAttr())
256 continue;
257 Type *ETy = Arg.getParamByValType();
258 if (ETy->isEmptyTy())
259 continue;
260 const Align ABIAlign = getPTXParamTypeAlign(ArgTy: ETy, DL);
261 if (Arg.getParamAlign().valueOrOne() >= ABIAlign)
262 continue;
263 Arg.removeAttr(Kind: Attribute::Alignment);
264 Arg.addAttr(Attr: Attribute::getWithAlignment(Context&: Ctx, Alignment: ABIAlign));
265 Changed = true;
266 }
267 return Changed;
268}
269
270// Propagate each byval parameter's .param alignment onto its constant-offset
271// loads (step 3 above). Runs after promoteParamAlign so the promoted
272// `stackalign` is included, and on every function so kernels and external
273// functions benefit too.
274static bool propagateByValParamLoadAlign(Function &F) {
275 if (F.isDeclaration())
276 return false;
277
278 const DataLayout &DL = F.getDataLayout();
279 bool Changed = false;
280 for (Argument &Arg : F.args()) {
281 if (!Arg.hasByValAttr())
282 continue;
283 Type *ETy = Arg.getParamByValType();
284 if (ETy->isEmptyTy())
285 continue;
286 const unsigned ParamIdx = Arg.getArgNo() + AttributeList::FirstArgIndex;
287 const Align ParamAlign = getDeviceByValParamAlign(F: &F, ArgTy: ETy, AttrIdx: ParamIdx, DL);
288 Changed |= propagateAlignmentToLoads(Val: &Arg, NewAlign: ParamAlign, DL);
289 }
290 return Changed;
291}
292
293static bool promoteParamAlignModule(Module &M) {
294 bool Changed = false;
295 for (Function &F : M) {
296 // Order matters (see the file header): promote, normalize `align`, then
297 // propagate to loads.
298 Changed |= promoteParamAlign(F);
299 Changed |= setByValParamABIAlign(F);
300 Changed |= propagateByValParamLoadAlign(F);
301 }
302 return Changed;
303}
304
305bool NVPTXPromoteParamAlignLegacyPass::runOnModule(Module &M) {
306 return promoteParamAlignModule(M);
307}
308
309ModulePass *llvm::createNVPTXPromoteParamAlignPass() {
310 return new NVPTXPromoteParamAlignLegacyPass();
311}
312
313PreservedAnalyses NVPTXPromoteParamAlignPass::run(Module &M,
314 ModuleAnalysisManager &AM) {
315 return promoteParamAlignModule(M) ? PreservedAnalyses::none()
316 : PreservedAnalyses::all();
317}
318