|
23 | 23 | #include "llvm/IR/IRBuilder.h"
|
24 | 24 | #include "llvm/Pass.h"
|
25 | 25 |
|
26 |
| -#include <algorithm> |
27 |
| - |
28 | 26 | using namespace llvm;
|
29 |
| -using namespace std; |
30 | 27 |
|
31 | 28 | namespace {
|
32 | 29 | class PromotePointerKernArgsToGlobal : public FunctionPass {
|
33 | 30 | // TODO: query the address space robustly.
|
34 | 31 | static constexpr unsigned int GenericAddrSpace{0u};
|
35 | 32 | static constexpr unsigned int GlobalAddrSpace{1u};
|
| 33 | + |
| 34 | + void createPromotableCast(IRBuilder<>& Builder, Value *From, Value *To) { |
| 35 | + From->replaceAllUsesWith(To); |
| 36 | + |
| 37 | + Value *FToG = Builder.CreateAddrSpaceCast( |
| 38 | + From, |
| 39 | + cast<PointerType>(From->getType()) |
| 40 | + ->getElementType()->getPointerTo(GlobalAddrSpace)); |
| 41 | + Value *GToF = Builder.CreateAddrSpaceCast(FToG, From->getType()); |
| 42 | + |
| 43 | + To->replaceAllUsesWith(GToF); |
| 44 | + } |
| 45 | + |
| 46 | + void maybePromoteUse(IRBuilder<>& Builder, Instruction *UI) { |
| 47 | + if (!UI) |
| 48 | + return; |
| 49 | + |
| 50 | + Builder.SetInsertPoint(UI->getNextNonDebugInstruction()); |
| 51 | + |
| 52 | + Value *Tmp = Builder.CreateBitCast(UndefValue::get(UI->getType()), |
| 53 | + UI->getType()); |
| 54 | + createPromotableCast(Builder, UI, Tmp); |
| 55 | + } |
| 56 | + |
| 57 | + void promoteArgument(IRBuilder<>& Builder, Argument *Arg) { |
| 58 | + Value *Tmp = Builder.CreateBitCast(UndefValue::get(Arg->getType()), |
| 59 | + Arg->getType()); |
| 60 | + createPromotableCast(Builder, Arg, Tmp); |
| 61 | + } |
36 | 62 | public:
|
37 | 63 | static char ID;
|
38 | 64 | PromotePointerKernArgsToGlobal() : FunctionPass{ID} {}
|
39 | 65 |
|
40 | 66 | bool runOnFunction(Function &F) override
|
41 | 67 | {
|
42 |
| - if (F.getCallingConv() != CallingConv::AMDGPU_KERNEL) return false; |
| 68 | + if (F.getCallingConv() != CallingConv::AMDGPU_KERNEL) |
| 69 | + return false; |
| 70 | + |
| 71 | + SmallVector<Argument *, 8> PromotableArgs; |
| 72 | + SmallVector<User *, 8> PromotableUses; |
| 73 | + for (auto &&Arg : F.args()) { |
| 74 | + for (auto &&U : Arg.users()) { |
| 75 | + if (!U->getType()->isPointerTy()) |
| 76 | + continue; |
| 77 | + if (U->getType()->getPointerAddressSpace() != GenericAddrSpace) |
| 78 | + continue; |
43 | 79 |
|
44 |
| - SmallVector<Argument *, 8> PtrArgs; |
45 |
| - for_each(F.arg_begin(), F.arg_end(), [&](Argument &Arg) { |
46 |
| - if (!Arg.getType()->isPointerTy()) return; |
47 |
| - if (Arg.getType()->getPointerAddressSpace() != GenericAddrSpace) { |
48 |
| - return; |
| 80 | + PromotableUses.push_back(U); |
49 | 81 | }
|
50 | 82 |
|
51 |
| - PtrArgs.push_back(&Arg); |
52 |
| - }); |
| 83 | + if (!Arg.getType()->isPointerTy()) |
| 84 | + continue; |
| 85 | + if (Arg.getType()->getPointerAddressSpace() != GenericAddrSpace) |
| 86 | + continue; |
53 | 87 |
|
54 |
| - if (PtrArgs.empty()) return false; |
| 88 | + PromotableArgs.push_back(&Arg); |
| 89 | + } |
55 | 90 |
|
56 |
| - static IRBuilder<> Builder{F.getContext()}; |
57 |
| - Builder.SetInsertPoint(&F.getEntryBlock().front()); |
| 91 | + if (PromotableArgs.empty() && PromotableUses.empty()) |
| 92 | + return false; |
58 | 93 |
|
59 |
| - for_each(PtrArgs.begin(), PtrArgs.end(), [](Argument *PArg) { |
60 |
| - Argument Tmp{PArg->getType(), PArg->getName()}; |
61 |
| - PArg->replaceAllUsesWith(&Tmp); |
62 |
| - |
63 |
| - Value *FToG = Builder.CreateAddrSpaceCast( |
64 |
| - PArg, |
65 |
| - cast<PointerType>(PArg->getType()) |
66 |
| - ->getElementType()->getPointerTo(GlobalAddrSpace)); |
67 |
| - Value *GToF = Builder.CreateAddrSpaceCast(FToG, PArg->getType()); |
68 |
| - |
69 |
| - Tmp.replaceAllUsesWith(GToF); |
70 |
| - }); |
| 94 | + static IRBuilder<> Builder{F.getContext()}; |
| 95 | + for (auto &&PU : PromotableUses) |
| 96 | + maybePromoteUse(Builder, dyn_cast<Instruction>(PU)); |
71 | 97 |
|
| 98 | + Builder.SetInsertPoint(&F.getEntryBlock().front()); |
| 99 | + for (auto &&Arg : PromotableArgs) |
| 100 | + promoteArgument(Builder, Arg); |
72 | 101 | return true;
|
73 | 102 | }
|
74 | 103 | };
|
|
0 commit comments