137 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// For all alloca instructions, and add a pair of cast to local address for10// each of them. For example,11//12// %A = alloca i3213// store i32 0, i32* %A ; emits st.u3214//15// will be transformed to16//17// %A = alloca i3218// %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.u3221//22// And we will rely on NVPTXInferAddressSpaces to combine the last two23// 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 34using namespace llvm;35 36namespace {37class NVPTXLowerAlloca : public FunctionPass {38 bool runOnFunction(Function &F) override;39 40public:41 static char ID; // Pass identification, replacement for typeid42 NVPTXLowerAlloca() : FunctionPass(ID) {}43 StringRef getPassName() const override {44 return "convert address space of alloca'ed memory to local";45 }46};47} // namespace48 49char NVPTXLowerAlloca::ID = 1;50 51INITIALIZE_PASS(NVPTXLowerAlloca, "nvptx-lower-alloca", "Lower Alloca", false,52 false)53 54// =============================================================================55// Main function for this pass.56// =============================================================================57bool NVPTXLowerAlloca::runOnFunction(Function &F) {58 if (skipFunction(F))59 return false;60 61 bool Changed = false;62 for (auto &BB : F)63 for (auto &I : BB) {64 if (auto allocaInst = dyn_cast<AllocaInst>(&I)) {65 Changed = true;66 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.");74 75 Instruction *AllocaInLocalAS = allocaInst;76 auto ETy = allocaInst->getAllocatedType();77 78 // We need to make sure that LLVM has info that alloca needs to go to79 // ADDRESS_SPACE_LOCAL for InferAddressSpace pass.80 //81 // For allocas in ADDRESS_SPACE_LOCAL, we add addrspacecast to82 // ADDRESS_SPACE_LOCAL and back to ADDRESS_SPACE_GENERIC, so that83 // the alloca's users still use a generic pointer to operate on.84 //85 // For allocas already in ADDRESS_SPACE_LOCAL, we just need86 // 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 }94 95 auto AllocaInGenericAS = new AddrSpaceCastInst(96 AllocaInLocalAS,97 PointerType::get(ETy->getContext(), ADDRESS_SPACE_GENERIC), "");98 AllocaInGenericAS->insertAfter(AllocaInLocalAS->getIterator());99 100 for (Use &AllocaUse : llvm::make_early_inc_range(allocaInst->uses())) {101 // Check Load, Store, GEP, and BitCast Uses on alloca and make them102 // use the converted generic address, in order to expose non-generic103 // addrspacecast to NVPTXInferAddressSpaces. For other types104 // of instructions this is unnecessary and may introduce redundant105 // 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;132}133 134FunctionPass *llvm::createNVPTXLowerAllocaPass() {135 return new NVPTXLowerAlloca();136}137