llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

   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     PRIVATE_SEGMENT_SIZE = 17,
 118
 119     // VGPRS:
 120     WORKITEM_ID_X       = 18,
 121     WORKITEM_ID_Y       = 19,
 122     WORKITEM_ID_Z       = 20,
 123     FIRST_VGPR_VALUE    = WORKITEM_ID_X
 124   };
 125   // clang-format on
 126
 127   // Kernel input registers setup for the HSA ABI in allocation order.
 128
 129   // User SGPRs in kernels
 130   // XXX - Can these require argument spills?
 131   ArgDescriptor PrivateSegmentBuffer;
 132   ArgDescriptor DispatchPtr;
 133   ArgDescriptor QueuePtr;
 134   ArgDescriptor KernargSegmentPtr;
 135   ArgDescriptor DispatchID;
 136   ArgDescriptor FlatScratchInit;
 137   ArgDescriptor PrivateSegmentSize;
 138   ArgDescriptor LDSKernelId;
 139
 140   // System SGPRs in kernels.
 141   ArgDescriptor WorkGroupIDX;
 142   ArgDescriptor WorkGroupIDY;
 143   ArgDescriptor WorkGroupIDZ;
 144   ArgDescriptor WorkGroupInfo;
 145   ArgDescriptor PrivateSegmentWaveByteOffset;
 146
 147   // Pointer with offset from kernargsegmentptr to where special ABI arguments
 148   // are passed to callable functions.
 149   ArgDescriptor ImplicitArgPtr;
 150
 151   // Input registers for non-HSA ABI
 152   ArgDescriptor ImplicitBufferPtr;
 153
 154   // VGPRs inputs. For entry functions these are either v0, v1 and v2 or packed
 155   // into v0, 10 bits per dimension if packed-tid is set.
 156   ArgDescriptor WorkItemIDX;
 157   ArgDescriptor WorkItemIDY;
 158   ArgDescriptor WorkItemIDZ;
 159
 160   // Map the index of preloaded kernel arguments to its descriptor.
 161   SmallDenseMap<int, KernArgPreloadDescriptor> PreloadKernArgs{};
 162
 163   std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
 164   getPreloadedValue(PreloadedValue Value) const;
 165
 166   static AMDGPUFunctionArgInfo fixedABILayout();
 167 };
 168
 169 class AMDGPUArgumentUsageInfo : public ImmutablePass {
 170 private:
 171   DenseMap<const Function *, AMDGPUFunctionArgInfo> ArgInfoMap;
 172
 173 public:
 174   static char ID;
 175
 176   static const AMDGPUFunctionArgInfo ExternFunctionInfo;
 177   static const AMDGPUFunctionArgInfo FixedABIFunctionInfo;
 178
 179   AMDGPUArgumentUsageInfo() : ImmutablePass(ID) { }
 180
 181   void getAnalysisUsage(AnalysisUsage &AU) const override {
 182     AU.setPreservesAll();
 183   }
 184
 185   bool doInitialization(Module &M) override;
 186   bool doFinalization(Module &M) override;
 187
 188   void print(raw_ostream &OS, const Module *M = nullptr) const override;
 189
 190   void setFuncArgInfo(const Function &F, const AMDGPUFunctionArgInfo &ArgInfo) {
 191     ArgInfoMap[&F] = ArgInfo;
 192   }
 193
 194   const AMDGPUFunctionArgInfo &lookupFuncArgInfo(const Function &F) const;
 195 };
 196
 197 } // end namespace llvm
 198
 199 #endif