66//
77// ===----------------------------------------------------------------------===//
88//
9- // For all alloca instructions, and add a pair of cast to local address for
10- // each of them. For example,
9+ // Change the Module's DataLayout to have the local address space for alloca's.
10+ // Change the address space of each alloca to local and add an addrspacecast to
11+ // generic address space. For example,
1112//
1213// %A = alloca i32
1314// store i32 0, i32* %A ; emits st.u32
1415//
1516// will be transformed to
1617//
17- // %A = alloca i32
18- // %Local = addrspacecast i32* %A to i32 addrspace(5)*
18+ // %A = alloca i32, addrspace(5)
1919// %Generic = addrspacecast i32 addrspace(5)* %A to i32*
2020// store i32 0, i32 addrspace(5)* %Generic ; emits st.local.u32
2121//
2424//
2525// ===----------------------------------------------------------------------===//
2626
27- #include " MCTargetDesc/NVPTXBaseInfo .h"
27+ #include " llvm/Support/NVPTXAddrSpace .h"
2828#include " NVPTX.h"
29+ #include " llvm/ADT/SmallVector.h"
2930#include " llvm/IR/Function.h"
31+ #include " llvm/IR/IRBuilder.h"
3032#include " llvm/IR/Instructions.h"
33+ #include " llvm/IR/IntrinsicInst.h"
34+ #include " llvm/IR/Module.h"
3135#include " llvm/IR/Type.h"
3236#include " llvm/Pass.h"
3337
3438using namespace llvm ;
39+ using namespace NVPTXAS ;
3540
3641namespace {
3742class NVPTXLowerAlloca : public FunctionPass {
3843 bool runOnFunction (Function &F) override ;
44+ bool doInitialization (Module &M) override ;
3945
4046public:
4147 static char ID; // Pass identification, replacement for typeid
@@ -58,77 +64,55 @@ bool NVPTXLowerAlloca::runOnFunction(Function &F) {
5864 if (skipFunction (F))
5965 return false ;
6066
61- bool Changed = false ;
67+ SmallVector<AllocaInst *, 16 > Allocas ;
6268 for (auto &BB : F)
63- for (auto &I : BB) {
64- if (
auto allocaInst = dyn_cast<AllocaInst>(&I)) {
65- Changed = true ;
69+ for (auto &I : BB)
70+ if (
auto *Alloca = dyn_cast<AllocaInst>(&I))
71+ if (Alloca->getAddressSpace () != ADDRESS_SPACE_LOCAL)
72+ Allocas.push_back (Alloca);
73+
74+ if (Allocas.empty ())
75+ return false ;
6676
67- PointerType *AllocInstPtrTy =
68- cast<PointerType>(allocaInst->getType ()->getScalarType ());
69- unsigned AllocAddrSpace = AllocInstPtrTy->getAddressSpace ();
70- assert ((AllocAddrSpace == ADDRESS_SPACE_GENERIC ||
71- AllocAddrSpace == ADDRESS_SPACE_LOCAL) &&
72- " AllocaInst can only be in Generic or Local address space for "
73- " NVPTX."
77+ for (AllocaInst *Alloca : Allocas) {
78+ auto *NewAlloca = new AllocaInst (
79+ Alloca->getAllocatedType (), ADDRESS_SPACE_LOCAL, Alloca->getArraySize (),
80+ Alloca->getAlign (), Alloca->getName ());
81+ auto *Cast = new AddrSpaceCastInst (
82+ NewAlloca,
83+ PointerType::get (Alloca->getAllocatedType ()->getContext (),
84+ ADDRESS_SPACE_GENERIC),
85+ " "
86+ Cast->insertBefore (Alloca->getIterator ());
87+ NewAlloca->insertBefore (Cast->getIterator ());
88+ for (auto &U : llvm::make_early_inc_range (Alloca->uses ())) {
89+ auto *II = dyn_cast<IntrinsicInst>(U.getUser ());
90+ if (!II || (II->getIntrinsicID () != Intrinsic::lifetime_start &&
91+ II->getIntrinsicID () != Intrinsic::lifetime_end))
92+ continue ;
7493
75- Instruction *AllocaInLocalAS = allocaInst;
76- auto ETy = allocaInst->getAllocatedType ();
94+ IRBuilder<> Builder (II);
95+ Builder.CreateIntrinsic (II->getIntrinsicID (), {NewAlloca->getType ()},
96+ {NewAlloca});
97+ II->eraseFromParent ();
98+ }
7799
78- // We need to make sure that LLVM has info that alloca needs to go to
79- // ADDRESS_SPACE_LOCAL for InferAddressSpace pass.
80- //
81- // For allocas in ADDRESS_SPACE_LOCAL, we add addrspacecast to
82- // ADDRESS_SPACE_LOCAL and back to ADDRESS_SPACE_GENERIC, so that
83- // the alloca's users still use a generic pointer to operate on.
84- //
85- // For allocas already in ADDRESS_SPACE_LOCAL, we just need
86- // addrspacecast to ADDRESS_SPACE_GENERIC.
87- if (AllocAddrSpace == ADDRESS_SPACE_GENERIC) {
88- auto ASCastToLocalAS = new AddrSpaceCastInst (
89- allocaInst,
90- PointerType::get (ETy->getContext (), ADDRESS_SPACE_LOCAL), " "
91- ASCastToLocalAS->insertAfter (allocaInst->getIterator ());
92- AllocaInLocalAS = ASCastToLocalAS;
93- }
100+ Alloca->replaceAllUsesWith (Cast);
101+ Alloca->eraseFromParent ();
102+ }
103+ return true ;
104+ }
94105
95- auto AllocaInGenericAS = new AddrSpaceCastInst (
96- AllocaInLocalAS,
97- PointerType::get (ETy->getContext (), ADDRESS_SPACE_GENERIC), " "
98- AllocaInGenericAS->insertAfter (AllocaInLocalAS->getIterator ());
106+ bool NVPTXLowerAlloca::doInitialization (Module &M) {
107+ const auto &DL = M.getDataLayout ();
108+ if (DL.getAllocaAddrSpace () == ADDRESS_SPACE_LOCAL)
109+ return false ;
110+ auto DLStr = DL.getStringRepresentation ();
99111
100- for (Use &AllocaUse : llvm::make_early_inc_range (allocaInst->uses ())) {
101- // Check Load, Store, GEP, and BitCast Uses on alloca and make them
102- // use the converted generic address, in order to expose non-generic
103- // addrspacecast to NVPTXInferAddressSpaces. For other types
104- // of instructions this is unnecessary and may introduce redundant
105- // address cast.
106- auto LI = dyn_cast<LoadInst>(AllocaUse.getUser ());
107- if (LI && LI->getPointerOperand () == allocaInst &&
108- !LI->isVolatile ()) {
109- LI->setOperand (LI->getPointerOperandIndex (), AllocaInGenericAS);
110- continue ;
111- }
112- auto SI = dyn_cast<StoreInst>(AllocaUse.getUser ());
113- if (SI && SI->getPointerOperand () == allocaInst &&
114- !SI->isVolatile ()) {
115- SI->setOperand (SI->getPointerOperandIndex (), AllocaInGenericAS);
116- continue ;
117- }
118- auto GI = dyn_cast<GetElementPtrInst>(AllocaUse.getUser ());
119- if (GI && GI->getPointerOperand () == allocaInst) {
120- GI->setOperand (GI->getPointerOperandIndex (), AllocaInGenericAS);
121- continue ;
122- }
123- auto BI = dyn_cast<BitCastInst>(AllocaUse.getUser ());
124- if (BI && BI->getOperand (0 ) == allocaInst) {
125- BI->setOperand (0 , AllocaInGenericAS);
126- continue ;
127- }
128- }
129- }
130- }
131- return Changed;
112+ auto AddrSpaceStr = " A" std::to_string (ADDRESS_SPACE_LOCAL);
113+ assert (!StringRef (DLStr).contains (" A" " DataLayout should not contain A"
114+ M.setDataLayout (DLStr.empty () ? AddrSpaceStr : DLStr + " -"
115+ return true ;
132116}
133117
134118FunctionPass *llvm::createNVPTXLowerAllocaPass () {
0 commit comments