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