1//==- AMDGPUArgumentrUsageInfo.h - Function Arg Usage Info -------*- 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//===----------------------------------------------------------------------===//
8
9#ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H
10#define LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H
11
12#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
13#include "llvm/ADT/DenseMap.h"
14#include "llvm/CodeGen/Register.h"
15#include "llvm/IR/PassManager.h"
16#include "llvm/Pass.h"
17#include "llvm/PassRegistry.h"
18#include <variant>
19
20namespace llvm {
21
22void initializeAMDGPUArgumentUsageInfoWrapperLegacyPass(PassRegistry &);
23
24class Function;
25class LLT;
26class raw_ostream;
27class TargetRegisterClass;
28class TargetRegisterInfo;
29
30struct ArgDescriptor {
31private:
32 friend struct AMDGPUFunctionArgInfo;
33 friend class AMDGPUArgumentUsageInfo;
34
35 std::variant<std::monostate, MCRegister, unsigned> Val;
36
37 // Bitmask to locate argument within the register.
38 unsigned Mask;
39
40public:
41 ArgDescriptor(unsigned Mask = ~0u) : Mask(Mask) {}
42
43 static ArgDescriptor createRegister(Register Reg, unsigned Mask = ~0u) {
44 ArgDescriptor Ret(Mask);
45 Ret.Val = Reg.asMCReg();
46 return Ret;
47 }
48
49 static ArgDescriptor createStack(unsigned Offset, unsigned Mask = ~0u) {
50 ArgDescriptor Ret(Mask);
51 Ret.Val = Offset;
52 return Ret;
53 }
54
55 static ArgDescriptor createArg(const ArgDescriptor &Arg, unsigned Mask) {
56 // Copy the descriptor, then change the mask.
57 ArgDescriptor Ret(Arg);
58 Ret.Mask = Mask;
59 return Ret;
60 }
61
62 bool isSet() const { return !std::holds_alternative<std::monostate>(v: Val); }
63
64 explicit operator bool() const {
65 return isSet();
66 }
67
68 bool isRegister() const { return std::holds_alternative<MCRegister>(v: Val); }
69
70 MCRegister getRegister() const { return std::get<MCRegister>(v: Val); }
71
72 unsigned getStackOffset() const { return std::get<unsigned>(v: Val); }
73
74 unsigned getMask() const {
75 // None of the target SGPRs or VGPRs are expected to have a 'zero' mask.
76 assert(Mask && "Invalid mask.");
77 return Mask;
78 }
79
80 bool isMasked() const {
81 return Mask != ~0u;
82 }
83
84 void print(raw_ostream &OS, const TargetRegisterInfo *TRI = nullptr) const;
85};
86
87inline raw_ostream &operator<<(raw_ostream &OS, const ArgDescriptor &Arg) {
88 Arg.print(OS);
89 return OS;
90}
91
92struct KernArgPreloadDescriptor : public ArgDescriptor {
93 KernArgPreloadDescriptor() = default;
94 SmallVector<MCRegister> Regs;
95};
96
97struct AMDGPUFunctionArgInfo {
98 // clang-format off
99 enum PreloadedValue {
100 // SGPRS:
101 PRIVATE_SEGMENT_BUFFER = 0,
102 DISPATCH_PTR = 1,
103 QUEUE_PTR = 2,
104 KERNARG_SEGMENT_PTR = 3,
105 DISPATCH_ID = 4,
106 FLAT_SCRATCH_INIT = 5,
107 LDS_KERNEL_ID = 6, // LLVM internal, not part of the ABI
108 WORKGROUP_ID_X = 10, // Also used for cluster ID X.
109 WORKGROUP_ID_Y = 11, // Also used for cluster ID Y.
110 WORKGROUP_ID_Z = 12, // Also used for cluster ID Z.
111 PRIVATE_SEGMENT_WAVE_BYTE_OFFSET = 14,
112 IMPLICIT_BUFFER_PTR = 15,
113 IMPLICIT_ARG_PTR = 16,
114 PRIVATE_SEGMENT_SIZE = 17,
115 CLUSTER_WORKGROUP_ID_X = 21,
116 CLUSTER_WORKGROUP_ID_Y = 22,
117 CLUSTER_WORKGROUP_ID_Z = 23,
118 CLUSTER_WORKGROUP_MAX_ID_X = 24,
119 CLUSTER_WORKGROUP_MAX_ID_Y = 25,
120 CLUSTER_WORKGROUP_MAX_ID_Z = 26,
121 CLUSTER_WORKGROUP_MAX_FLAT_ID = 27,
122
123 // VGPRS:
124 WORKITEM_ID_X = 28,
125 WORKITEM_ID_Y = 29,
126 WORKITEM_ID_Z = 30,
127 FIRST_VGPR_VALUE = WORKITEM_ID_X
128 };
129 // clang-format on
130
131 // Kernel input registers setup for the HSA ABI in allocation order.
132
133 // User SGPRs in kernels
134 // XXX - Can these require argument spills?
135 ArgDescriptor PrivateSegmentBuffer;
136 ArgDescriptor DispatchPtr;
137 ArgDescriptor QueuePtr;
138 ArgDescriptor KernargSegmentPtr;
139 ArgDescriptor DispatchID;
140 ArgDescriptor FlatScratchInit;
141 ArgDescriptor PrivateSegmentSize;
142 ArgDescriptor LDSKernelId;
143
144 // System SGPRs in kernels.
145 ArgDescriptor WorkGroupIDX;
146 ArgDescriptor WorkGroupIDY;
147 ArgDescriptor WorkGroupIDZ;
148 ArgDescriptor WorkGroupInfo;
149 ArgDescriptor PrivateSegmentWaveByteOffset;
150
151 // Pointer with offset from kernargsegmentptr to where special ABI arguments
152 // are passed to callable functions.
153 ArgDescriptor ImplicitArgPtr;
154
155 // Input registers for non-HSA ABI
156 ArgDescriptor ImplicitBufferPtr;
157
158 // VGPRs inputs. For entry functions these are either v0, v1 and v2 or packed
159 // into v0, 10 bits per dimension if packed-tid is set.
160 ArgDescriptor WorkItemIDX;
161 ArgDescriptor WorkItemIDY;
162 ArgDescriptor WorkItemIDZ;
163
164 // Map the index of preloaded kernel arguments to its descriptor.
165 SmallDenseMap<int, KernArgPreloadDescriptor> PreloadKernArgs{};
166 // The first user SGPR allocated for kernarg preloading.
167 Register FirstKernArgPreloadReg;
168
169 std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
170 getPreloadedValue(PreloadedValue Value) const;
171
172 static AMDGPUFunctionArgInfo fixedABILayout();
173};
174
175class AMDGPUArgumentUsageInfo {
176private:
177 DenseMap<const Function *, AMDGPUFunctionArgInfo> ArgInfoMap;
178
179public:
180 static const AMDGPUFunctionArgInfo ExternFunctionInfo;
181 static const AMDGPUFunctionArgInfo FixedABIFunctionInfo;
182
183 void print(raw_ostream &OS, const Module *M = nullptr) const;
184
185 void clear() { ArgInfoMap.clear(); }
186
187 void setFuncArgInfo(const Function &F, const AMDGPUFunctionArgInfo &ArgInfo) {
188 ArgInfoMap[&F] = ArgInfo;
189 }
190
191 const AMDGPUFunctionArgInfo &lookupFuncArgInfo(const Function &F) const;
192
193 bool invalidate(Module &M, const PreservedAnalyses &PA,
194 ModuleAnalysisManager::Invalidator &Inv);
195};
196
197class AMDGPUArgumentUsageInfoWrapperLegacy : public ImmutablePass {
198 std::unique_ptr<AMDGPUArgumentUsageInfo> AUIP;
199
200public:
201 static char ID;
202
203 AMDGPUArgumentUsageInfoWrapperLegacy() : ImmutablePass(ID) {
204 initializeAMDGPUArgumentUsageInfoWrapperLegacyPass(
205 *PassRegistry::getPassRegistry());
206 }
207
208 AMDGPUArgumentUsageInfo &getArgUsageInfo() { return *AUIP; }
209 const AMDGPUArgumentUsageInfo &getArgUsageInfo() const { return *AUIP; }
210
211 void getAnalysisUsage(AnalysisUsage &AU) const override {
212 AU.setPreservesAll();
213 }
214
215 bool doInitialization(Module &M) override {
216 AUIP = std::make_unique<AMDGPUArgumentUsageInfo>();
217 return false;
218 }
219
220 bool doFinalization(Module &M) override {
221 AUIP->clear();
222 return false;
223 }
224
225 void print(raw_ostream &OS, const Module *M = nullptr) const override {
226 AUIP->print(OS, M);
227 }
228};
229
230class AMDGPUArgumentUsageAnalysis
231 : public AnalysisInfoMixin<AMDGPUArgumentUsageAnalysis> {
232 friend AnalysisInfoMixin<AMDGPUArgumentUsageAnalysis>;
233 static AnalysisKey Key;
234
235public:
236 using Result = AMDGPUArgumentUsageInfo;
237
238 AMDGPUArgumentUsageInfo run(Module &M, ModuleAnalysisManager &);
239};
240
241} // end namespace llvm
242
243#endif
244