1 //===-- NVPTXLowerAlloca.cpp - Make alloca to use local memory =====--===// 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 // For all alloca instructions, and add a pair of cast to local address for 10 // each of them. For example, 11 // 12 // %A = alloca i32 13 // store i32 0, i32* %A ; emits st.u32 14 // 15 // will be transformed to 16 // 17 // %A = alloca i32 18 // %Local = addrspacecast i32* %A to i32 addrspace(5)* 19 // %Generic = addrspacecast i32 addrspace(5)* %A to i32* 20 // store i32 0, i32 addrspace(5)* %Generic ; emits st.local.u32 21 // 22 // And we will rely on NVPTXInferAddressSpaces to combine the last two 23 // instructions. 24 // 25 //===----------------------------------------------------------------------===// 26 27 #include "MCTargetDesc/NVPTXBaseInfo.h" 28 #include "NVPTX.h" 29 #include "llvm/IR/Function.h" 30 #include "llvm/IR/Instructions.h" 31 #include "llvm/IR/Type.h" 32 #include "llvm/Pass.h" 33 34 using namespace llvm; 35 36 namespace llvm { 37 void initializeNVPTXLowerAllocaPass(PassRegistry &); 38 } 39 40 namespace { 41 class NVPTXLowerAlloca : public FunctionPass { 42 bool runOnFunction(Function &F) override; 43 44 public: 45 static char ID; // Pass identification, replacement for typeid 46 NVPTXLowerAlloca() : FunctionPass(ID) {} 47 StringRef getPassName() const override { 48 return "convert address space of alloca'ed memory to local"; 49 } 50 }; 51 } // namespace 52 53 char NVPTXLowerAlloca::ID = 1; 54 55 INITIALIZE_PASS(NVPTXLowerAlloca, "nvptx-lower-alloca", "Lower Alloca", false, 56 false) 57 58 // ============================================================================= 59 // Main function for this pass. 60 // ============================================================================= 61 bool NVPTXLowerAlloca::runOnFunction(Function &F) { 62 if (skipFunction(F)) 63 return false; 64 65 bool Changed = false; 66 for (auto &BB : F) 67 for (auto &I : BB) { 68 if (auto allocaInst = dyn_cast<AllocaInst>(&I)) { 69 Changed = true; 70 71 PointerType *AllocInstPtrTy = 72 cast<PointerType>(allocaInst->getType()->getScalarType()); 73 unsigned AllocAddrSpace = AllocInstPtrTy->getAddressSpace(); 74 assert((AllocAddrSpace == ADDRESS_SPACE_GENERIC || 75 AllocAddrSpace == ADDRESS_SPACE_LOCAL) && 76 "AllocaInst can only be in Generic or Local address space for " 77 "NVPTX."); 78 79 Instruction *AllocaInLocalAS = allocaInst; 80 auto ETy = allocaInst->getAllocatedType(); 81 82 // We need to make sure that LLVM has info that alloca needs to go to 83 // ADDRESS_SPACE_LOCAL for InferAddressSpace pass. 84 // 85 // For allocas in ADDRESS_SPACE_LOCAL, we add addrspacecast to 86 // ADDRESS_SPACE_LOCAL and back to ADDRESS_SPACE_GENERIC, so that 87 // the alloca's users still use a generic pointer to operate on. 88 // 89 // For allocas already in ADDRESS_SPACE_LOCAL, we just need 90 // addrspacecast to ADDRESS_SPACE_GENERIC. 91 if (AllocAddrSpace == ADDRESS_SPACE_GENERIC) { 92 auto ASCastToLocalAS = new AddrSpaceCastInst( 93 allocaInst, 94 PointerType::get(ETy->getContext(), ADDRESS_SPACE_LOCAL), ""); 95 ASCastToLocalAS->insertAfter(allocaInst->getIterator()); 96 AllocaInLocalAS = ASCastToLocalAS; 97 } 98 99 auto AllocaInGenericAS = new AddrSpaceCastInst( 100 AllocaInLocalAS, 101 PointerType::get(ETy->getContext(), ADDRESS_SPACE_GENERIC), ""); 102 AllocaInGenericAS->insertAfter(AllocaInLocalAS->getIterator()); 103 104 for (Use &AllocaUse : llvm::make_early_inc_range(allocaInst->uses())) { 105 // Check Load, Store, GEP, and BitCast Uses on alloca and make them 106 // use the converted generic address, in order to expose non-generic 107 // addrspacecast to NVPTXInferAddressSpaces. For other types 108 // of instructions this is unnecessary and may introduce redundant 109 // address cast. 110 auto LI = dyn_cast<LoadInst>(AllocaUse.getUser()); 111 if (LI && LI->getPointerOperand() == allocaInst && 112 !LI->isVolatile()) { 113 LI->setOperand(LI->getPointerOperandIndex(), AllocaInGenericAS); 114 continue; 115 } 116 auto SI = dyn_cast<StoreInst>(AllocaUse.getUser()); 117 if (SI && SI->getPointerOperand() == allocaInst && 118 !SI->isVolatile()) { 119 SI->setOperand(SI->getPointerOperandIndex(), AllocaInGenericAS); 120 continue; 121 } 122 auto GI = dyn_cast<GetElementPtrInst>(AllocaUse.getUser()); 123 if (GI && GI->getPointerOperand() == allocaInst) { 124 GI->setOperand(GI->getPointerOperandIndex(), AllocaInGenericAS); 125 continue; 126 } 127 auto BI = dyn_cast<BitCastInst>(AllocaUse.getUser()); 128 if (BI && BI->getOperand(0) == allocaInst) { 129 BI->setOperand(0, AllocaInGenericAS); 130 continue; 131 } 132 } 133 } 134 } 135 return Changed; 136 } 137 138 FunctionPass *llvm::createNVPTXLowerAllocaPass() { 139 return new NVPTXLowerAlloca(); 140 } 141