1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
|
//===-- AMDGPULowerKernelAttributes.cpp ------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
/// \file This pass does attempts to make use of reqd_work_group_size metadata
/// to eliminate loads from the dispatch packet and to constant fold OpenCL
/// get_local_size-like functions.
//
//===----------------------------------------------------------------------===//
#include "AMDGPU.h"
#include "Utils/AMDGPUBaseInfo.h"
#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsAMDGPU.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/Pass.h"
#define DEBUG_TYPE "amdgpu-lower-kernel-attributes"
using namespace llvm;
namespace {
// Field offsets in hsa_kernel_dispatch_packet_t.
enum DispatchPackedOffsets {
WORKGROUP_SIZE_X = 4,
WORKGROUP_SIZE_Y = 6,
WORKGROUP_SIZE_Z = 8,
GRID_SIZE_X = 12,
GRID_SIZE_Y = 16,
GRID_SIZE_Z = 20
};
// Field offsets to implicit kernel argument pointer.
enum ImplicitArgOffsets {
HIDDEN_BLOCK_COUNT_X = 0,
HIDDEN_BLOCK_COUNT_Y = 4,
HIDDEN_BLOCK_COUNT_Z = 8,
HIDDEN_GROUP_SIZE_X = 12,
HIDDEN_GROUP_SIZE_Y = 14,
HIDDEN_GROUP_SIZE_Z = 16,
HIDDEN_REMAINDER_X = 18,
HIDDEN_REMAINDER_Y = 20,
HIDDEN_REMAINDER_Z = 22,
};
class AMDGPULowerKernelAttributes : public ModulePass {
public:
static char ID;
AMDGPULowerKernelAttributes() : ModulePass(ID) {}
bool runOnModule(Module &M) override;
StringRef getPassName() const override {
return "AMDGPU Kernel Attributes";
}
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.setPreservesAll();
}
};
Function *getBasePtrIntrinsic(Module &M, bool IsV5OrAbove) {
auto IntrinsicId = IsV5OrAbove ? Intrinsic::amdgcn_implicitarg_ptr
: Intrinsic::amdgcn_dispatch_ptr;
return Intrinsic::getDeclarationIfExists(&M, IntrinsicId);
}
} // end anonymous namespace
static void annotateGridSizeLoadWithRangeMD(LoadInst *Load,
uint32_t MaxNumGroups) {
if (MaxNumGroups == 0 || MaxNumGroups == std::numeric_limits<uint32_t>::max())
return;
if (!Load->getType()->isIntegerTy(32))
return;
// TODO: If there is existing range metadata, preserve it if it is stricter.
MDBuilder MDB(Load->getContext());
MDNode *Range = MDB.createRange(APInt(32, 1), APInt(32, MaxNumGroups + 1));
Load->setMetadata(LLVMContext::MD_range, Range);
}
static bool processUse(CallInst *CI, bool IsV5OrAbove) {
Function *F = CI->getParent()->getParent();
auto *MD = F->getMetadata("reqd_work_group_size");
const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;
const bool HasUniformWorkGroupSize =
F->getFnAttribute("uniform-work-group-size").getValueAsBool();
SmallVector<unsigned> MaxNumWorkgroups =
AMDGPU::getIntegerVecAttribute(*F, "amdgpu-max-num-workgroups",
/*Size=*/3, /*DefaultVal=*/0);
if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize &&
none_of(MaxNumWorkgroups, [](unsigned X) { return X != 0; }))
return false;
Value *BlockCounts[3] = {nullptr, nullptr, nullptr};
Value *GroupSizes[3] = {nullptr, nullptr, nullptr};
Value *Remainders[3] = {nullptr, nullptr, nullptr};
Value *GridSizes[3] = {nullptr, nullptr, nullptr};
const DataLayout &DL = F->getDataLayout();
// We expect to see several GEP users, casted to the appropriate type and
// loaded.
for (User *U : CI->users()) {
if (!U->hasOneUse())
continue;
int64_t Offset = 0;
auto *Load = dyn_cast<LoadInst>(U); // Load from ImplicitArgPtr/DispatchPtr?
auto *BCI = dyn_cast<BitCastInst>(U);
if (!Load && !BCI) {
if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
continue;
Load = dyn_cast<LoadInst>(*U->user_begin()); // Load from GEP?
BCI = dyn_cast<BitCastInst>(*U->user_begin());
}
if (BCI) {
if (!BCI->hasOneUse())
continue;
Load = dyn_cast<LoadInst>(*BCI->user_begin()); // Load from BCI?
}
if (!Load || !Load->isSimple())
continue;
unsigned LoadSize = DL.getTypeStoreSize(Load->getType());
// TODO: Handle merged loads.
if (IsV5OrAbove) { // Base is ImplicitArgPtr.
switch (Offset) {
case HIDDEN_BLOCK_COUNT_X:
if (LoadSize == 4) {
BlockCounts[0] = Load;
annotateGridSizeLoadWithRangeMD(Load, MaxNumWorkgroups[0]);
}
break;
case HIDDEN_BLOCK_COUNT_Y:
if (LoadSize == 4) {
BlockCounts[1] = Load;
annotateGridSizeLoadWithRangeMD(Load, MaxNumWorkgroups[1]);
}
break;
case HIDDEN_BLOCK_COUNT_Z:
if (LoadSize == 4) {
BlockCounts[2] = Load;
annotateGridSizeLoadWithRangeMD(Load, MaxNumWorkgroups[2]);
}
break;
case HIDDEN_GROUP_SIZE_X:
if (LoadSize == 2)
GroupSizes[0] = Load;
break;
case HIDDEN_GROUP_SIZE_Y:
if (LoadSize == 2)
GroupSizes[1] = Load;
break;
case HIDDEN_GROUP_SIZE_Z:
if (LoadSize == 2)
GroupSizes[2] = Load;
break;
case HIDDEN_REMAINDER_X:
if (LoadSize == 2)
Remainders[0] = Load;
break;
case HIDDEN_REMAINDER_Y:
if (LoadSize == 2)
Remainders[1] = Load;
break;
case HIDDEN_REMAINDER_Z:
if (LoadSize == 2)
Remainders[2] = Load;
break;
default:
break;
}
} else { // Base is DispatchPtr.
switch (Offset) {
case WORKGROUP_SIZE_X:
if (LoadSize == 2)
GroupSizes[0] = Load;
break;
case WORKGROUP_SIZE_Y:
if (LoadSize == 2)
GroupSizes[1] = Load;
break;
case WORKGROUP_SIZE_Z:
if (LoadSize == 2)
GroupSizes[2] = Load;
break;
case GRID_SIZE_X:
if (LoadSize == 4)
GridSizes[0] = Load;
break;
case GRID_SIZE_Y:
if (LoadSize == 4)
GridSizes[1] = Load;
break;
case GRID_SIZE_Z:
if (LoadSize == 4)
GridSizes[2] = Load;
break;
default:
break;
}
}
}
bool MadeChange = false;
if (IsV5OrAbove && HasUniformWorkGroupSize) {
// Under v5 __ockl_get_local_size returns the value computed by the expression:
//
// workgroup_id < hidden_block_count ? hidden_group_size : hidden_remainder
//
// For functions with the attribute uniform-work-group-size=true. we can evaluate
// workgroup_id < hidden_block_count as true, and thus hidden_group_size is returned
// for __ockl_get_local_size.
for (int I = 0; I < 3; ++I) {
Value *BlockCount = BlockCounts[I];
if (!BlockCount)
continue;
using namespace llvm::PatternMatch;
auto GroupIDIntrin =
I == 0 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>()
: (I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>()
: m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());
for (User *ICmp : BlockCount->users()) {
if (match(ICmp, m_SpecificICmp(ICmpInst::ICMP_ULT, GroupIDIntrin,
m_Specific(BlockCount)))) {
ICmp->replaceAllUsesWith(llvm::ConstantInt::getTrue(ICmp->getType()));
MadeChange = true;
}
}
}
// All remainders should be 0 with uniform work group size.
for (Value *Remainder : Remainders) {
if (!Remainder)
continue;
Remainder->replaceAllUsesWith(Constant::getNullValue(Remainder->getType()));
MadeChange = true;
}
} else if (HasUniformWorkGroupSize) { // Pre-V5.
// Pattern match the code used to handle partial workgroup dispatches in the
// library implementation of get_local_size, so the entire function can be
// constant folded with a known group size.
//
// uint r = grid_size - group_id * group_size;
// get_local_size = (r < group_size) ? r : group_size;
//
// If we have uniform-work-group-size (which is the default in OpenCL 1.2),
// the grid_size is required to be a multiple of group_size). In this case:
//
// grid_size - (group_id * group_size) < group_size
// ->
// grid_size < group_size + (group_id * group_size)
//
// (grid_size / group_size) < 1 + group_id
//
// grid_size / group_size is at least 1, so we can conclude the select
// condition is false (except for group_id == 0, where the select result is
// the same).
for (int I = 0; I < 3; ++I) {
Value *GroupSize = GroupSizes[I];
Value *GridSize = GridSizes[I];
if (!GroupSize || !GridSize)
continue;
using namespace llvm::PatternMatch;
auto GroupIDIntrin =
I == 0 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>()
: (I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>()
: m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());
for (User *U : GroupSize->users()) {
auto *ZextGroupSize = dyn_cast<ZExtInst>(U);
if (!ZextGroupSize)
continue;
for (User *UMin : ZextGroupSize->users()) {
if (match(UMin,
m_UMin(m_Sub(m_Specific(GridSize),
m_Mul(GroupIDIntrin, m_Specific(ZextGroupSize))),
m_Specific(ZextGroupSize)))) {
if (HasReqdWorkGroupSize) {
ConstantInt *KnownSize
= mdconst::extract<ConstantInt>(MD->getOperand(I));
UMin->replaceAllUsesWith(ConstantFoldIntegerCast(
KnownSize, UMin->getType(), false, DL));
} else {
UMin->replaceAllUsesWith(ZextGroupSize);
}
MadeChange = true;
}
}
}
}
}
// If reqd_work_group_size is set, we can replace work group size with it.
if (!HasReqdWorkGroupSize)
return MadeChange;
for (int I = 0; I < 3; I++) {
Value *GroupSize = GroupSizes[I];
if (!GroupSize)
continue;
ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(I));
GroupSize->replaceAllUsesWith(
ConstantFoldIntegerCast(KnownSize, GroupSize->getType(), false, DL));
MadeChange = true;
}
return MadeChange;
}
// TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
// TargetPassConfig for subtarget.
bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
bool MadeChange = false;
bool IsV5OrAbove =
AMDGPU::getAMDHSACodeObjectVersion(M) >= AMDGPU::AMDHSA_COV5;
Function *BasePtr = getBasePtrIntrinsic(M, IsV5OrAbove);
if (!BasePtr) // ImplicitArgPtr/DispatchPtr not used.
return false;
SmallPtrSet<Instruction *, 4> HandledUses;
for (auto *U : BasePtr->users()) {
CallInst *CI = cast<CallInst>(U);
if (HandledUses.insert(CI).second) {
if (processUse(CI, IsV5OrAbove))
MadeChange = true;
}
}
return MadeChange;
}
INITIALIZE_PASS_BEGIN(AMDGPULowerKernelAttributes, DEBUG_TYPE,
"AMDGPU Kernel Attributes", false, false)
INITIALIZE_PASS_END(AMDGPULowerKernelAttributes, DEBUG_TYPE,
"AMDGPU Kernel Attributes", false, false)
char AMDGPULowerKernelAttributes::ID = 0;
ModulePass *llvm::createAMDGPULowerKernelAttributesPass() {
return new AMDGPULowerKernelAttributes();
}
PreservedAnalyses
AMDGPULowerKernelAttributesPass::run(Function &F, FunctionAnalysisManager &AM) {
bool IsV5OrAbove =
AMDGPU::getAMDHSACodeObjectVersion(*F.getParent()) >= AMDGPU::AMDHSA_COV5;
Function *BasePtr = getBasePtrIntrinsic(*F.getParent(), IsV5OrAbove);
if (!BasePtr) // ImplicitArgPtr/DispatchPtr not used.
return PreservedAnalyses::all();
for (Instruction &I : instructions(F)) {
if (CallInst *CI = dyn_cast<CallInst>(&I)) {
if (CI->getCalledFunction() == BasePtr)
processUse(CI, IsV5OrAbove);
}
}
return PreservedAnalyses::all();
}
|