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// Replace each generic alloca with an equivalent alloca in the local address
10// space, followed by an addrspacecast back to generic for its users. For
11// example,
12//
13// %A = alloca i32
14// store i32 0, ptr %A ; emits st.u32
15//
16// is transformed to
17//
18// %A = alloca i32, addrspace(5)
19// %A.generic = addrspacecast ptr addrspace(5) %A to ptr
20// store i32 0, ptr %A.generic
21//
22// This gives the alloca a local frame index, which stack lowering addresses
23// through the local frame pointer (%SPL). When NVPTXInferAddressSpaces runs
24// after this pass, it propagates the local address space into the users and
25// folds the cast away where possible (so the store above becomes st.local.u32).
26//
27//===----------------------------------------------------------------------===//
28
29#include "MCTargetDesc/NVPTXBaseInfo.h"
30#include "NVPTX.h"
31#include "llvm/ADT/SmallVector.h"
32#include "llvm/IR/DebugInfo.h"
33#include "llvm/IR/Function.h"
34#include "llvm/IR/Instructions.h"
35#include "llvm/IR/IntrinsicInst.h"
36#include "llvm/IR/Intrinsics.h"
37#include "llvm/IR/Type.h"
38#include "llvm/Pass.h"
39
40using namespace llvm;
41
42// =============================================================================
43// Main function for this pass.
44// =============================================================================
45static bool lowerAllocas(Function &F) {
46 // Mandatory lowering: later stack lowering relies on local allocas, so run
47 // even for optnone functions (optnone is intentionally not honored).
48 SmallVector<AllocaInst *, 8> GenericAllocas;
49 for (auto &BB : F)
50 for (auto &I : BB)
51 if (auto *AI = dyn_cast<AllocaInst>(Val: &I);
52 AI && AI->getAddressSpace() == ADDRESS_SPACE_GENERIC)
53 GenericAllocas.push_back(Elt: AI);
54
55 for (AllocaInst *AI : GenericAllocas) {
56 // Create an equivalent alloca in the local address space.
57 auto *LocalAlloca = new AllocaInst(AI->getAllocatedType(),
58 ADDRESS_SPACE_LOCAL, AI->getArraySize(),
59 AI->getAlign(), "", AI->getIterator());
60 LocalAlloca->copyMetadata(SrcInst: *AI);
61 LocalAlloca->setUsedWithInAlloca(AI->isUsedWithInAlloca());
62 LocalAlloca->setSwiftError(AI->isSwiftError());
63
64 // Debug records and lifetime markers have to reference the alloca itself,
65 // not a cast of it, so retarget them to the local alloca before rewriting
66 // the remaining users through a generic addrspacecast below:
67 // - the verifier requires an alloca operand for lifetime markers, and
68 // - pointing debug records at the alloca keeps the variable described by
69 // its (stable) stack slot rather than the cvta.local result.
70 SmallVector<DbgVariableRecord *, 2> DbgUsers;
71 findDbgUsers(V: AI, DbgVariableRecords&: DbgUsers);
72 for (DbgVariableRecord *DVR : DbgUsers)
73 DVR->replaceVariableLocationOp(OldValue: AI, NewValue: LocalAlloca);
74
75 for (Use &U : llvm::make_early_inc_range(Range: AI->uses())) {
76 auto *II = dyn_cast<IntrinsicInst>(Val: U.getUser());
77 if (!II || !isLifetimeIntrinsic(ID: II->getIntrinsicID()))
78 continue;
79 U.set(LocalAlloca);
80 Function *Decl = Intrinsic::getOrInsertDeclaration(
81 M: II->getModule(), id: II->getIntrinsicID(), OverloadTys: {LocalAlloca->getType()});
82 II->setCalledFunction(Decl);
83 }
84
85 // Everything else can go through a single generic addrspacecast.
86 // replaceAllUsesWith leaves the (already retargeted) lifetime markers and
87 // debug records untouched. NVPTXInferAddressSpaces folds the cast into the
88 // users that can operate on local memory directly.
89 auto *GenericPtr = new AddrSpaceCastInst(LocalAlloca, AI->getType(), "",
90 AI->getIterator());
91 GenericPtr->setDebugLoc(AI->getDebugLoc());
92 AI->replaceAllUsesWith(V: GenericPtr);
93 LocalAlloca->takeName(V: AI);
94 AI->eraseFromParent();
95 }
96
97 return !GenericAllocas.empty();
98}
99
100namespace {
101class NVPTXLowerAllocaLegacyPass : public FunctionPass {
102 bool runOnFunction(Function &F) override { return lowerAllocas(F); }
103
104public:
105 static char ID; // Pass identification, replacement for typeid
106 NVPTXLowerAllocaLegacyPass() : FunctionPass(ID) {}
107 StringRef getPassName() const override {
108 return "convert address space of alloca'ed memory to local";
109 }
110};
111} // namespace
112
113char NVPTXLowerAllocaLegacyPass::ID = 0;
114
115INITIALIZE_PASS(NVPTXLowerAllocaLegacyPass, "nvptx-lower-alloca",
116 "Lower Alloca", false, false)
117
118FunctionPass *llvm::createNVPTXLowerAllocaLegacyPass() {
119 return new NVPTXLowerAllocaLegacyPass();
120}
121
122PreservedAnalyses NVPTXLowerAllocaPass::run(Function &F,
123 FunctionAnalysisManager &FAM) {
124 if (!lowerAllocas(F))
125 return PreservedAnalyses::all();
126 return PreservedAnalyses::none().preserveSet<CFGAnalyses>();
127}
128