202 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H10#define LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H11 12#include "MCTargetDesc/AMDGPUMCTargetDesc.h"13#include "llvm/ADT/DenseMap.h"14#include "llvm/CodeGen/Register.h"15#include "llvm/Pass.h"16#include <variant>17 18namespace llvm {19 20class Function;21class LLT;22class raw_ostream;23class TargetRegisterClass;24class TargetRegisterInfo;25 26struct ArgDescriptor {27private:28 friend struct AMDGPUFunctionArgInfo;29 friend class AMDGPUArgumentUsageInfo;30 31 std::variant<std::monostate, MCRegister, unsigned> Val;32 33 // Bitmask to locate argument within the register.34 unsigned Mask;35 36public:37 ArgDescriptor(unsigned Mask = ~0u) : Mask(Mask) {}38 39 static ArgDescriptor createRegister(Register Reg, unsigned Mask = ~0u) {40 ArgDescriptor Ret(Mask);41 Ret.Val = Reg.asMCReg();42 return Ret;43 }44 45 static ArgDescriptor createStack(unsigned Offset, unsigned Mask = ~0u) {46 ArgDescriptor Ret(Mask);47 Ret.Val = Offset;48 return Ret;49 }50 51 static ArgDescriptor createArg(const ArgDescriptor &Arg, unsigned Mask) {52 // Copy the descriptor, then change the mask.53 ArgDescriptor Ret(Arg);54 Ret.Mask = Mask;55 return Ret;56 }57 58 bool isSet() const { return !std::holds_alternative<std::monostate>(Val); }59 60 explicit operator bool() const {61 return isSet();62 }63 64 bool isRegister() const { return std::holds_alternative<MCRegister>(Val); }65 66 MCRegister getRegister() const { return std::get<MCRegister>(Val); }67 68 unsigned getStackOffset() const { return std::get<unsigned>(Val); }69 70 unsigned getMask() const {71 // None of the target SGPRs or VGPRs are expected to have a 'zero' mask.72 assert(Mask && "Invalid mask.");73 return Mask;74 }75 76 bool isMasked() const {77 return Mask != ~0u;78 }79 80 void print(raw_ostream &OS, const TargetRegisterInfo *TRI = nullptr) const;81};82 83inline raw_ostream &operator<<(raw_ostream &OS, const ArgDescriptor &Arg) {84 Arg.print(OS);85 return OS;86}87 88struct KernArgPreloadDescriptor : public ArgDescriptor {89 KernArgPreloadDescriptor() = default;90 SmallVector<MCRegister> Regs;91};92 93struct AMDGPUFunctionArgInfo {94 // clang-format off95 enum PreloadedValue {96 // SGPRS:97 PRIVATE_SEGMENT_BUFFER = 0,98 DISPATCH_PTR = 1,99 QUEUE_PTR = 2,100 KERNARG_SEGMENT_PTR = 3,101 DISPATCH_ID = 4,102 FLAT_SCRATCH_INIT = 5,103 LDS_KERNEL_ID = 6, // LLVM internal, not part of the ABI104 WORKGROUP_ID_X = 10, // Also used for cluster ID X.105 WORKGROUP_ID_Y = 11, // Also used for cluster ID Y.106 WORKGROUP_ID_Z = 12, // Also used for cluster ID Z.107 PRIVATE_SEGMENT_WAVE_BYTE_OFFSET = 14,108 IMPLICIT_BUFFER_PTR = 15,109 IMPLICIT_ARG_PTR = 16,110 PRIVATE_SEGMENT_SIZE = 17,111 CLUSTER_WORKGROUP_ID_X = 21,112 CLUSTER_WORKGROUP_ID_Y = 22,113 CLUSTER_WORKGROUP_ID_Z = 23,114 CLUSTER_WORKGROUP_MAX_ID_X = 24,115 CLUSTER_WORKGROUP_MAX_ID_Y = 25,116 CLUSTER_WORKGROUP_MAX_ID_Z = 26,117 CLUSTER_WORKGROUP_MAX_FLAT_ID = 27,118 119 // VGPRS:120 WORKITEM_ID_X = 28,121 WORKITEM_ID_Y = 29,122 WORKITEM_ID_Z = 30,123 FIRST_VGPR_VALUE = WORKITEM_ID_X124 };125 // clang-format on126 127 // Kernel input registers setup for the HSA ABI in allocation order.128 129 // User SGPRs in kernels130 // 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 arguments148 // are passed to callable functions.149 ArgDescriptor ImplicitArgPtr;150 151 // Input registers for non-HSA ABI152 ArgDescriptor ImplicitBufferPtr;153 154 // VGPRs inputs. For entry functions these are either v0, v1 and v2 or packed155 // 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 // The first user SGPR allocated for kernarg preloading.163 Register FirstKernArgPreloadReg;164 165 std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>166 getPreloadedValue(PreloadedValue Value) const;167 168 static AMDGPUFunctionArgInfo fixedABILayout();169};170 171class AMDGPUArgumentUsageInfo : public ImmutablePass {172private:173 DenseMap<const Function *, AMDGPUFunctionArgInfo> ArgInfoMap;174 175public:176 static char ID;177 178 static const AMDGPUFunctionArgInfo ExternFunctionInfo;179 static const AMDGPUFunctionArgInfo FixedABIFunctionInfo;180 181 AMDGPUArgumentUsageInfo() : ImmutablePass(ID) { }182 183 void getAnalysisUsage(AnalysisUsage &AU) const override {184 AU.setPreservesAll();185 }186 187 bool doInitialization(Module &M) override;188 bool doFinalization(Module &M) override;189 190 void print(raw_ostream &OS, const Module *M = nullptr) const override;191 192 void setFuncArgInfo(const Function &F, const AMDGPUFunctionArgInfo &ArgInfo) {193 ArgInfoMap[&F] = ArgInfo;194 }195 196 const AMDGPUFunctionArgInfo &lookupFuncArgInfo(const Function &F) const;197};198 199} // end namespace llvm200 201#endif202