1//===----------------------------------------------------------------------===//
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#include "AMDGPU.h"
10#include "AMDGPUArgumentUsageInfo.h"
11#include "AMDGPUTargetMachine.h"
12#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
13#include "SIRegisterInfo.h"
14#include "llvm/Support/NativeFormatting.h"
15#include "llvm/Support/raw_ostream.h"
16
17using namespace llvm;
18
19#define DEBUG_TYPE "amdgpu-argument-reg-usage-info"
20
21INITIALIZE_PASS(AMDGPUArgumentUsageInfo, DEBUG_TYPE,
22                "Argument Register Usage Information Storage", false, true)
23
24void ArgDescriptor::print(raw_ostream &OS,
25                          const TargetRegisterInfo *TRI) const {
26  if (!isSet()) {
27    OS << "<not set>\n";
28    return;
29  }
30
31  if (isRegister())
32    OS << "Reg " << printReg(getRegister(), TRI);
33  else
34    OS << "Stack offset " << getStackOffset();
35
36  if (isMasked()) {
37    OS << " & ";
38    llvm::write_hex(OS, Mask, llvm::HexPrintStyle::PrefixLower);
39  }
40
41  OS << '\n';
42}
43
44char AMDGPUArgumentUsageInfo::ID = 0;
45
46const AMDGPUFunctionArgInfo AMDGPUArgumentUsageInfo::ExternFunctionInfo{};
47
48// Hardcoded registers from fixed function ABI
49const AMDGPUFunctionArgInfo AMDGPUArgumentUsageInfo::FixedABIFunctionInfo
50  = AMDGPUFunctionArgInfo::fixedABILayout();
51
52bool AMDGPUArgumentUsageInfo::doInitialization(Module &M) {
53  return false;
54}
55
56bool AMDGPUArgumentUsageInfo::doFinalization(Module &M) {
57  ArgInfoMap.clear();
58  return false;
59}
60
61void AMDGPUArgumentUsageInfo::print(raw_ostream &OS, const Module *M) const {
62  for (const auto &FI : ArgInfoMap) {
63    OS << "Arguments for " << FI.first->getName() << '\n'
64       << "  PrivateSegmentBuffer: " << FI.second.PrivateSegmentBuffer
65       << "  DispatchPtr: " << FI.second.DispatchPtr
66       << "  QueuePtr: " << FI.second.QueuePtr
67       << "  KernargSegmentPtr: " << FI.second.KernargSegmentPtr
68       << "  DispatchID: " << FI.second.DispatchID
69       << "  FlatScratchInit: " << FI.second.FlatScratchInit
70       << "  PrivateSegmentSize: " << FI.second.PrivateSegmentSize
71       << "  WorkGroupIDX: " << FI.second.WorkGroupIDX
72       << "  WorkGroupIDY: " << FI.second.WorkGroupIDY
73       << "  WorkGroupIDZ: " << FI.second.WorkGroupIDZ
74       << "  WorkGroupInfo: " << FI.second.WorkGroupInfo
75       << "  PrivateSegmentWaveByteOffset: "
76          << FI.second.PrivateSegmentWaveByteOffset
77       << "  ImplicitBufferPtr: " << FI.second.ImplicitBufferPtr
78       << "  ImplicitArgPtr: " << FI.second.ImplicitArgPtr
79       << "  WorkItemIDX " << FI.second.WorkItemIDX
80       << "  WorkItemIDY " << FI.second.WorkItemIDY
81       << "  WorkItemIDZ " << FI.second.WorkItemIDZ
82       << '\n';
83  }
84}
85
86std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
87AMDGPUFunctionArgInfo::getPreloadedValue(
88    AMDGPUFunctionArgInfo::PreloadedValue Value) const {
89  switch (Value) {
90  case AMDGPUFunctionArgInfo::PRIVATE_SEGMENT_BUFFER: {
91    return std::make_tuple(PrivateSegmentBuffer ? &PrivateSegmentBuffer
92                                                : nullptr,
93                           &AMDGPU::SGPR_128RegClass, LLT::vector(4, 32));
94  }
95  case AMDGPUFunctionArgInfo::IMPLICIT_BUFFER_PTR:
96    return std::make_tuple(ImplicitBufferPtr ? &ImplicitBufferPtr : nullptr,
97                           &AMDGPU::SGPR_64RegClass,
98                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
99  case AMDGPUFunctionArgInfo::WORKGROUP_ID_X:
100    return std::make_tuple(WorkGroupIDX ? &WorkGroupIDX : nullptr,
101                           &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
102  case AMDGPUFunctionArgInfo::WORKGROUP_ID_Y:
103    return std::make_tuple(WorkGroupIDY ? &WorkGroupIDY : nullptr,
104                           &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
105  case AMDGPUFunctionArgInfo::WORKGROUP_ID_Z:
106    return std::make_tuple(WorkGroupIDZ ? &WorkGroupIDZ : nullptr,
107                           &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
108  case AMDGPUFunctionArgInfo::PRIVATE_SEGMENT_WAVE_BYTE_OFFSET:
109    return std::make_tuple(
110        PrivateSegmentWaveByteOffset ? &PrivateSegmentWaveByteOffset : nullptr,
111        &AMDGPU::SGPR_32RegClass, LLT::scalar(32));
112  case AMDGPUFunctionArgInfo::KERNARG_SEGMENT_PTR:
113    return std::make_tuple(KernargSegmentPtr ? &KernargSegmentPtr : nullptr,
114                           &AMDGPU::SGPR_64RegClass,
115                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
116  case AMDGPUFunctionArgInfo::IMPLICIT_ARG_PTR:
117    return std::make_tuple(ImplicitArgPtr ? &ImplicitArgPtr : nullptr,
118                           &AMDGPU::SGPR_64RegClass,
119                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
120  case AMDGPUFunctionArgInfo::DISPATCH_ID:
121    return std::make_tuple(DispatchID ? &DispatchID : nullptr,
122                           &AMDGPU::SGPR_64RegClass, LLT::scalar(64));
123  case AMDGPUFunctionArgInfo::FLAT_SCRATCH_INIT:
124    return std::make_tuple(FlatScratchInit ? &FlatScratchInit : nullptr,
125                           &AMDGPU::SGPR_64RegClass, LLT::scalar(64));
126  case AMDGPUFunctionArgInfo::DISPATCH_PTR:
127    return std::make_tuple(DispatchPtr ? &DispatchPtr : nullptr,
128                           &AMDGPU::SGPR_64RegClass,
129                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
130  case AMDGPUFunctionArgInfo::QUEUE_PTR:
131    return std::make_tuple(QueuePtr ? &QueuePtr : nullptr,
132                           &AMDGPU::SGPR_64RegClass,
133                           LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64));
134  case AMDGPUFunctionArgInfo::WORKITEM_ID_X:
135    return std::make_tuple(WorkItemIDX ? &WorkItemIDX : nullptr,
136                           &AMDGPU::VGPR_32RegClass, LLT::scalar(32));
137  case AMDGPUFunctionArgInfo::WORKITEM_ID_Y:
138    return std::make_tuple(WorkItemIDY ? &WorkItemIDY : nullptr,
139                           &AMDGPU::VGPR_32RegClass, LLT::scalar(32));
140  case AMDGPUFunctionArgInfo::WORKITEM_ID_Z:
141    return std::make_tuple(WorkItemIDZ ? &WorkItemIDZ : nullptr,
142                           &AMDGPU::VGPR_32RegClass, LLT::scalar(32));
143  }
144  llvm_unreachable("unexpected preloaded value type");
145}
146
147constexpr AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
148  AMDGPUFunctionArgInfo AI;
149  AI.PrivateSegmentBuffer
150    = ArgDescriptor::createRegister(AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3);
151  AI.DispatchPtr = ArgDescriptor::createRegister(AMDGPU::SGPR4_SGPR5);
152  AI.QueuePtr = ArgDescriptor::createRegister(AMDGPU::SGPR6_SGPR7);
153
154  // Do not pass kernarg segment pointer, only pass increment version in its
155  // place.
156  AI.ImplicitArgPtr = ArgDescriptor::createRegister(AMDGPU::SGPR8_SGPR9);
157  AI.DispatchID = ArgDescriptor::createRegister(AMDGPU::SGPR10_SGPR11);
158
159  // Skip FlatScratchInit/PrivateSegmentSize
160  AI.WorkGroupIDX = ArgDescriptor::createRegister(AMDGPU::SGPR12);
161  AI.WorkGroupIDY = ArgDescriptor::createRegister(AMDGPU::SGPR13);
162  AI.WorkGroupIDZ = ArgDescriptor::createRegister(AMDGPU::SGPR14);
163
164  const unsigned Mask = 0x3ff;
165  AI.WorkItemIDX = ArgDescriptor::createRegister(AMDGPU::VGPR31, Mask);
166  AI.WorkItemIDY = ArgDescriptor::createRegister(AMDGPU::VGPR31, Mask << 10);
167  AI.WorkItemIDZ = ArgDescriptor::createRegister(AMDGPU::VGPR31, Mask << 20);
168  return AI;
169}
170
171const AMDGPUFunctionArgInfo &
172AMDGPUArgumentUsageInfo::lookupFuncArgInfo(const Function &F) const {
173  auto I = ArgInfoMap.find(&F);
174  if (I == ArgInfoMap.end()) {
175    if (AMDGPUTargetMachine::EnableFixedFunctionABI)
176      return FixedABIFunctionInfo;
177
178    // Without the fixed ABI, we assume no function has special inputs.
179    assert(F.isDeclaration());
180    return ExternFunctionInfo;
181  }
182
183  return I->second;
184}
185