xref: /llvm-project/llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h (revision 2e5c2982819625d84e0b61aea0ec00de859f0e95)
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 
17 namespace llvm {
18 
19 class Function;
20 class LLT;
21 class raw_ostream;
22 class TargetRegisterClass;
23 class TargetRegisterInfo;
24 
25 struct ArgDescriptor {
26 private:
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 
41 public:
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 
93 inline raw_ostream &operator<<(raw_ostream &OS, const ArgDescriptor &Arg) {
94   Arg.print(OS);
95   return OS;
96 }
97 
98 struct KernArgPreloadDescriptor : public ArgDescriptor {
99   KernArgPreloadDescriptor() {}
100   SmallVector<MCRegister> Regs;
101 };
102 
103 struct 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 
174 class AMDGPUArgumentUsageInfo : public ImmutablePass {
175 private:
176   DenseMap<const Function *, AMDGPUFunctionArgInfo> ArgInfoMap;
177 
178 public:
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