Coverage Report

Created: 2019-07-24 05:18

/Users/buildslave/jenkins/workspace/clang-stage2-coverage-R/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
Line
Count
Source (jump to first uncovered line)
1
//===- AMDGPURewriteOutArgumentsPass.cpp - Create struct returns ----------===//
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
/// \file This pass attempts to replace out argument usage with a return of a
10
/// struct.
11
///
12
/// We can support returning a lot of values directly in registers, but
13
/// idiomatic C code frequently uses a pointer argument to return a second value
14
/// rather than returning a struct by value. GPU stack access is also quite
15
/// painful, so we want to avoid that if possible. Passing a stack object
16
/// pointer to a function also requires an additional address expansion code
17
/// sequence to convert the pointer to be relative to the kernel's scratch wave
18
/// offset register since the callee doesn't know what stack frame the incoming
19
/// pointer is relative to.
20
///
21
/// The goal is to try rewriting code that looks like this:
22
///
23
///  int foo(int a, int b, int* out) {
24
///     *out = bar();
25
///     return a + b;
26
/// }
27
///
28
/// into something like this:
29
///
30
///  std::pair<int, int> foo(int a, int b) {
31
///     return std::make_pair(a + b, bar());
32
/// }
33
///
34
/// Typically the incoming pointer is a simple alloca for a temporary variable
35
/// to use the API, which if replaced with a struct return will be easily SROA'd
36
/// out when the stub function we create is inlined
37
///
38
/// This pass introduces the struct return, but leaves the unused pointer
39
/// arguments and introduces a new stub function calling the struct returning
40
/// body. DeadArgumentElimination should be run after this to clean these up.
41
//
42
//===----------------------------------------------------------------------===//
43
44
#include "AMDGPU.h"
45
#include "Utils/AMDGPUBaseInfo.h"
46
#include "llvm/Analysis/MemoryDependenceAnalysis.h"
47
#include "llvm/ADT/DenseMap.h"
48
#include "llvm/ADT/STLExtras.h"
49
#include "llvm/ADT/SmallSet.h"
50
#include "llvm/ADT/SmallVector.h"
51
#include "llvm/ADT/Statistic.h"
52
#include "llvm/Analysis/MemoryLocation.h"
53
#include "llvm/IR/Argument.h"
54
#include "llvm/IR/Attributes.h"
55
#include "llvm/IR/BasicBlock.h"
56
#include "llvm/IR/Constants.h"
57
#include "llvm/IR/DataLayout.h"
58
#include "llvm/IR/DerivedTypes.h"
59
#include "llvm/IR/Function.h"
60
#include "llvm/IR/IRBuilder.h"
61
#include "llvm/IR/Instructions.h"
62
#include "llvm/IR/Module.h"
63
#include "llvm/IR/Type.h"
64
#include "llvm/IR/Use.h"
65
#include "llvm/IR/User.h"
66
#include "llvm/IR/Value.h"
67
#include "llvm/Pass.h"
68
#include "llvm/Support/Casting.h"
69
#include "llvm/Support/CommandLine.h"
70
#include "llvm/Support/Debug.h"
71
#include "llvm/Support/raw_ostream.h"
72
#include <cassert>
73
#include <utility>
74
75
#define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
76
77
using namespace llvm;
78
79
static cl::opt<bool> AnyAddressSpace(
80
  "amdgpu-any-address-space-out-arguments",
81
  cl::desc("Replace pointer out arguments with "
82
           "struct returns for non-private address space"),
83
  cl::Hidden,
84
  cl::init(false));
85
86
static cl::opt<unsigned> MaxNumRetRegs(
87
  "amdgpu-max-return-arg-num-regs",
88
  cl::desc("Approximately limit number of return registers for replacing out arguments"),
89
  cl::Hidden,
90
  cl::init(16));
91
92
STATISTIC(NumOutArgumentsReplaced,
93
          "Number out arguments moved to struct return values");
94
STATISTIC(NumOutArgumentFunctionsReplaced,
95
          "Number of functions with out arguments moved to struct return values");
96
97
namespace {
98
99
class AMDGPURewriteOutArguments : public FunctionPass {
100
private:
101
  const DataLayout *DL = nullptr;
102
  MemoryDependenceResults *MDA = nullptr;
103
104
  bool checkArgumentUses(Value &Arg) const;
105
  bool isOutArgumentCandidate(Argument &Arg) const;
106
107
#ifndef NDEBUG
108
  bool isVec3ToVec4Shuffle(Type *Ty0, Type* Ty1) const;
109
#endif
110
111
public:
112
  static char ID;
113
114
2
  AMDGPURewriteOutArguments() : FunctionPass(ID) {}
115
116
2
  void getAnalysisUsage(AnalysisUsage &AU) const override {
117
2
    AU.addRequired<MemoryDependenceWrapperPass>();
118
2
    FunctionPass::getAnalysisUsage(AU);
119
2
  }
120
121
  bool doInitialization(Module &M) override;
122
  bool runOnFunction(Function &F) override;
123
};
124
125
} // end anonymous namespace
126
127
101k
INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
128
101k
                      "AMDGPU Rewrite Out Arguments", false, false)
129
101k
INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
130
101k
INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
131
                    "AMDGPU Rewrite Out Arguments", false, false)
132
133
char AMDGPURewriteOutArguments::ID = 0;
134
135
68
bool AMDGPURewriteOutArguments::checkArgumentUses(Value &Arg) const {
136
68
  const int MaxUses = 10;
137
68
  int UseCount = 0;
138
68
139
69
  for (Use &U : Arg.uses()) {
140
69
    StoreInst *SI = dyn_cast<StoreInst>(U.getUser());
141
69
    if (UseCount > MaxUses)
142
0
      return false;
143
69
144
69
    if (!SI) {
145
21
      auto *BCI = dyn_cast<BitCastInst>(U.getUser());
146
21
      if (!BCI || 
!BCI->hasOneUse()19
)
147
2
        return false;
148
19
149
19
      // We don't handle multiple stores currently, so stores to aggregate
150
19
      // pointers aren't worth the trouble since they are canonically split up.
151
19
      Type *DestEltTy = BCI->getType()->getPointerElementType();
152
19
      if (DestEltTy->isAggregateType())
153
2
        return false;
154
17
155
17
      // We could handle these if we had a convenient way to bitcast between
156
17
      // them.
157
17
      Type *SrcEltTy = Arg.getType()->getPointerElementType();
158
17
      if (SrcEltTy->isArrayTy())
159
0
        return false;
160
17
161
17
      // Special case handle structs with single members. It is useful to handle
162
17
      // some casts between structs and non-structs, but we can't bitcast
163
17
      // directly between them.  directly bitcast between them.  Blender uses
164
17
      // some casts that look like { <3 x float> }* to <4 x float>*
165
17
      if ((SrcEltTy->isStructTy() && 
(SrcEltTy->getStructNumElements() != 1)10
))
166
2
        return false;
167
15
168
15
      // Clang emits OpenCL 3-vector type accesses with a bitcast to the
169
15
      // equivalent 4-element vector and accesses that, and we're looking for
170
15
      // this pointer cast.
171
15
      if (DL->getTypeAllocSize(SrcEltTy) != DL->getTypeAllocSize(DestEltTy))
172
3
        return false;
173
12
174
12
      return checkArgumentUses(*BCI);
175
12
    }
176
48
177
48
    if (!SI->isSimple() ||
178
48
        U.getOperandNo() != StoreInst::getPointerOperandIndex())
179
1
      return false;
180
47
181
47
    ++UseCount;
182
47
  }
183
68
184
68
  // Skip unused arguments.
185
68
  
return UseCount > 046
;
186
68
}
187
188
81
bool AMDGPURewriteOutArguments::isOutArgumentCandidate(Argument &Arg) const {
189
81
  const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
190
81
  PointerType *ArgTy = dyn_cast<PointerType>(Arg.getType());
191
81
192
81
  // TODO: It might be useful for any out arguments, not just privates.
193
81
  if (!ArgTy || 
(61
ArgTy->getAddressSpace() != DL->getAllocaAddrSpace()61
&&
194
61
                 
!AnyAddressSpace3
) ||
195
81
      
Arg.hasByValAttr()60
||
Arg.hasStructRetAttr()57
||
196
81
      
DL->getTypeStoreSize(ArgTy->getPointerElementType()) > MaxOutArgSizeBytes57
) {
197
25
    return false;
198
25
  }
199
56
200
56
  return checkArgumentUses(Arg);
201
56
}
202
203
2
bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
204
2
  DL = &M.getDataLayout();
205
2
  return false;
206
2
}
207
208
#ifndef NDEBUG
209
bool AMDGPURewriteOutArguments::isVec3ToVec4Shuffle(Type *Ty0, Type* Ty1) const {
210
  VectorType *VT0 = dyn_cast<VectorType>(Ty0);
211
  VectorType *VT1 = dyn_cast<VectorType>(Ty1);
212
  if (!VT0 || !VT1)
213
    return false;
214
215
  if (VT0->getNumElements() != 3 ||
216
      VT1->getNumElements() != 4)
217
    return false;
218
219
  return DL->getTypeSizeInBits(VT0->getElementType()) ==
220
         DL->getTypeSizeInBits(VT1->getElementType());
221
}
222
#endif
223
224
60
bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
225
60
  if (skipFunction(F))
226
1
    return false;
227
59
228
59
  // TODO: Could probably handle variadic functions.
229
59
  if (F.isVarArg() || F.hasStructRetAttr() ||
230
59
      
AMDGPU::isEntryFunctionCC(F.getCallingConv())58
)
231
1
    return false;
232
58
233
58
  MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
234
58
235
58
  unsigned ReturnNumRegs = 0;
236
58
  SmallSet<int, 4> OutArgIndexes;
237
58
  SmallVector<Type *, 4> ReturnTypes;
238
58
  Type *RetTy = F.getReturnType();
239
58
  if (!RetTy->isVoidTy()) {
240
9
    ReturnNumRegs = DL->getTypeStoreSize(RetTy) / 4;
241
9
242
9
    if (ReturnNumRegs >= MaxNumRetRegs)
243
1
      return false;
244
8
245
8
    ReturnTypes.push_back(RetTy);
246
8
  }
247
58
248
58
  SmallVector<Argument *, 4> OutArgs;
249
81
  for (Argument &Arg : F.args()) {
250
81
    if (isOutArgumentCandidate(Arg)) {
251
44
      LLVM_DEBUG(dbgs() << "Found possible out argument " << Arg
252
44
                        << " in function " << F.getName() << '\n');
253
44
      OutArgs.push_back(&Arg);
254
44
    }
255
81
  }
256
57
257
57
  if (OutArgs.empty())
258
17
    return false;
259
40
260
40
  using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
261
40
262
40
  DenseMap<ReturnInst *, ReplacementVec> Replacements;
263
40
264
40
  SmallVector<ReturnInst *, 4> Returns;
265
49
  for (BasicBlock &BB : F) {
266
49
    if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back()))
267
43
      Returns.push_back(RI);
268
49
  }
269
40
270
40
  if (Returns.empty())
271
0
    return false;
272
40
273
40
  bool Changing;
274
40
275
78
  do {
276
78
    Changing = false;
277
78
278
78
    // Keep retrying if we are able to successfully eliminate an argument. This
279
78
    // helps with cases with multiple arguments which may alias, such as in a
280
78
    // sincos implemntation. If we have 2 stores to arguments, on the first
281
78
    // attempt the MDA query will succeed for the second store but not the
282
78
    // first. On the second iteration we've removed that out clobbering argument
283
78
    // (by effectively moving it into another function) and will find the second
284
78
    // argument is OK to move.
285
88
    for (Argument *OutArg : OutArgs) {
286
88
      bool ThisReplaceable = true;
287
88
      SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
288
88
289
88
      Type *ArgTy = OutArg->getType()->getPointerElementType();
290
88
291
88
      // Skip this argument if converting it will push us over the register
292
88
      // count to return limit.
293
88
294
88
      // TODO: This is an approximation. When legalized this could be more. We
295
88
      // can ask TLI for exactly how many.
296
88
      unsigned ArgNumRegs = DL->getTypeStoreSize(ArgTy) / 4;
297
88
      if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
298
0
        continue;
299
88
300
88
      // An argument is convertible only if all exit blocks are able to replace
301
88
      // it.
302
90
      
for (ReturnInst *RI : Returns)88
{
303
90
        BasicBlock *BB = RI->getParent();
304
90
305
90
        MemDepResult Q = MDA->getPointerDependencyFrom(MemoryLocation(OutArg),
306
90
                                                       true, BB->end(), BB, RI);
307
90
        StoreInst *SI = nullptr;
308
90
        if (Q.isDef())
309
43
          SI = dyn_cast<StoreInst>(Q.getInst());
310
90
311
90
        if (SI) {
312
43
          LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
313
43
          ReplaceableStores.emplace_back(RI, SI);
314
47
        } else {
315
47
          ThisReplaceable = false;
316
47
          break;
317
47
        }
318
90
      }
319
88
320
88
      if (!ThisReplaceable)
321
47
        continue; // Try the next argument candidate.
322
41
323
43
      
for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores)41
{
324
43
        Value *ReplVal = Store.second->getValueOperand();
325
43
326
43
        auto &ValVec = Replacements[Store.first];
327
43
        if (llvm::find_if(ValVec,
328
43
              [OutArg](const std::pair<Argument *, Value *> &Entry) {
329
5
                 return Entry.first == OutArg;}) != ValVec.end()) {
330
1
          LLVM_DEBUG(dbgs()
331
1
                     << "Saw multiple out arg stores" << *OutArg << '\n');
332
1
          // It is possible to see stores to the same argument multiple times,
333
1
          // but we expect these would have been optimized out already.
334
1
          ThisReplaceable = false;
335
1
          break;
336
1
        }
337
42
338
42
        ValVec.emplace_back(OutArg, ReplVal);
339
42
        Store.second->eraseFromParent();
340
42
      }
341
41
342
41
      if (ThisReplaceable) {
343
40
        ReturnTypes.push_back(ArgTy);
344
40
        OutArgIndexes.insert(OutArg->getArgNo());
345
40
        ++NumOutArgumentsReplaced;
346
40
        Changing = true;
347
40
      }
348
41
    }
349
78
  } while (Changing);
350
40
351
40
  if (Replacements.empty())
352
4
    return false;
353
36
354
36
  LLVMContext &Ctx = F.getParent()->getContext();
355
36
  StructType *NewRetTy = StructType::create(Ctx, ReturnTypes, F.getName());
356
36
357
36
  FunctionType *NewFuncTy = FunctionType::get(NewRetTy,
358
36
                                              F.getFunctionType()->params(),
359
36
                                              F.isVarArg());
360
36
361
36
  LLVM_DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
362
36
363
36
  Function *NewFunc = Function::Create(NewFuncTy, Function::PrivateLinkage,
364
36
                                       F.getName() + ".body");
365
36
  F.getParent()->getFunctionList().insert(F.getIterator(), NewFunc);
366
36
  NewFunc->copyAttributesFrom(&F);
367
36
  NewFunc->setComdat(F.getComdat());
368
36
369
36
  // We want to preserve the function and param attributes, but need to strip
370
36
  // off any return attributes, e.g. zeroext doesn't make sense with a struct.
371
36
  NewFunc->stealArgumentListFrom(F);
372
36
373
36
  AttrBuilder RetAttrs;
374
36
  RetAttrs.addAttribute(Attribute::SExt);
375
36
  RetAttrs.addAttribute(Attribute::ZExt);
376
36
  RetAttrs.addAttribute(Attribute::NoAlias);
377
36
  NewFunc->removeAttributes(AttributeList::ReturnIndex, RetAttrs);
378
36
  // TODO: How to preserve metadata?
379
36
380
36
  // Move the body of the function into the new rewritten function, and replace
381
36
  // this function with a stub.
382
36
  NewFunc->getBasicBlockList().splice(NewFunc->begin(), F.getBasicBlockList());
383
36
384
38
  for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
385
38
    ReturnInst *RI = Replacement.first;
386
38
    IRBuilder<> B(RI);
387
38
    B.SetCurrentDebugLocation(RI->getDebugLoc());
388
38
389
38
    int RetIdx = 0;
390
38
    Value *NewRetVal = UndefValue::get(NewRetTy);
391
38
392
38
    Value *RetVal = RI->getReturnValue();
393
38
    if (RetVal)
394
7
      NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
395
38
396
42
    for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second) {
397
42
      Argument *Arg = ReturnPoint.first;
398
42
      Value *Val = ReturnPoint.second;
399
42
      Type *EltTy = Arg->getType()->getPointerElementType();
400
42
      if (Val->getType() != EltTy) {
401
13
        Type *EffectiveEltTy = EltTy;
402
13
        if (StructType *CT = dyn_cast<StructType>(EltTy)) {
403
8
          assert(CT->getNumElements() == 1);
404
8
          EffectiveEltTy = CT->getElementType(0);
405
8
        }
406
13
407
13
        if (DL->getTypeSizeInBits(EffectiveEltTy) !=
408
13
            DL->getTypeSizeInBits(Val->getType())) {
409
8
          assert(isVec3ToVec4Shuffle(EffectiveEltTy, Val->getType()));
410
8
          Val = B.CreateShuffleVector(Val, UndefValue::get(Val->getType()),
411
8
                                      { 0, 1, 2 });
412
8
        }
413
13
414
13
        Val = B.CreateBitCast(Val, EffectiveEltTy);
415
13
416
13
        // Re-create single element composite.
417
13
        if (EltTy != EffectiveEltTy)
418
8
          Val = B.CreateInsertValue(UndefValue::get(EltTy), Val, 0);
419
13
      }
420
42
421
42
      NewRetVal = B.CreateInsertValue(NewRetVal, Val, RetIdx++);
422
42
    }
423
38
424
38
    if (RetVal)
425
7
      RI->setOperand(0, NewRetVal);
426
31
    else {
427
31
      B.CreateRet(NewRetVal);
428
31
      RI->eraseFromParent();
429
31
    }
430
38
  }
431
36
432
36
  SmallVector<Value *, 16> StubCallArgs;
433
54
  for (Argument &Arg : F.args()) {
434
54
    if (OutArgIndexes.count(Arg.getArgNo())) {
435
40
      // It's easier to preserve the type of the argument list. We rely on
436
40
      // DeadArgumentElimination to take care of these.
437
40
      StubCallArgs.push_back(UndefValue::get(Arg.getType()));
438
40
    } else {
439
14
      StubCallArgs.push_back(&Arg);
440
14
    }
441
54
  }
442
36
443
36
  BasicBlock *StubBB = BasicBlock::Create(Ctx, "", &F);
444
36
  IRBuilder<> B(StubBB);
445
36
  CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
446
36
447
36
  int RetIdx = RetTy->isVoidTy() ? 
029
:
17
;
448
54
  for (Argument &Arg : F.args()) {
449
54
    if (!OutArgIndexes.count(Arg.getArgNo()))
450
14
      continue;
451
40
452
40
    PointerType *ArgType = cast<PointerType>(Arg.getType());
453
40
454
40
    auto *EltTy = ArgType->getElementType();
455
40
    unsigned Align = Arg.getParamAlignment();
456
40
    if (Align == 0)
457
39
      Align = DL->getABITypeAlignment(EltTy);
458
40
459
40
    Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
460
40
    Type *PtrTy = Val->getType()->getPointerTo(ArgType->getAddressSpace());
461
40
462
40
    // We can peek through bitcasts, so the type may not match.
463
40
    Value *PtrVal = B.CreateBitCast(&Arg, PtrTy);
464
40
465
40
    B.CreateAlignedStore(Val, PtrVal, Align);
466
40
  }
467
36
468
36
  if (!RetTy->isVoidTy()) {
469
7
    B.CreateRet(B.CreateExtractValue(StubCall, 0));
470
29
  } else {
471
29
    B.CreateRetVoid();
472
29
  }
473
36
474
36
  // The function is now a stub we want to inline.
475
36
  F.addFnAttr(Attribute::AlwaysInline);
476
36
477
36
  ++NumOutArgumentFunctionsReplaced;
478
36
  return true;
479
36
}
480
481
0
FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {
482
0
  return new AMDGPURewriteOutArguments();
483
0
}