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