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