blob: 8ca35c0061ff79362c27088df0d1fd0faf18808e [file] [log] [blame]
Jingyue Wu13755602016-03-20 20:59:20 +00001//===-- NVPTXInferAddressSpace.cpp - ---------------------*- C++ -*-===//
2//
3// The LLVM Compiler Infrastructure
4//
5// This file is distributed under the University of Illinois Open Source
6// License. See LICENSE.TXT for details.
7//
8//===----------------------------------------------------------------------===//
9//
10// CUDA C/C++ includes memory space designation as variable type qualifers (such
11// as __global__ and __shared__). Knowing the space of a memory access allows
12// CUDA compilers to emit faster PTX loads and stores. For example, a load from
13// shared memory can be translated to `ld.shared` which is roughly 10% faster
14// than a generic `ld` on an NVIDIA Tesla K40c.
15//
16// Unfortunately, type qualifiers only apply to variable declarations, so CUDA
17// compilers must infer the memory space of an address expression from
18// type-qualified variables.
19//
20// LLVM IR uses non-zero (so-called) specific address spaces to represent memory
21// spaces (e.g. addrspace(3) means shared memory). The Clang frontend
22// places only type-qualified variables in specific address spaces, and then
23// conservatively `addrspacecast`s each type-qualified variable to addrspace(0)
24// (so-called the generic address space) for other instructions to use.
25//
26// For example, the Clang translates the following CUDA code
27// __shared__ float a[10];
28// float v = a[i];
29// to
30// %0 = addrspacecast [10 x float] addrspace(3)* @a to [10 x float]*
31// %1 = gep [10 x float], [10 x float]* %0, i64 0, i64 %i
32// %v = load float, float* %1 ; emits ld.f32
33// @a is in addrspace(3) since it's type-qualified, but its use from %1 is
34// redirected to %0 (the generic version of @a).
35//
36// The optimization implemented in this file propagates specific address spaces
37// from type-qualified variable declarations to its users. For example, it
38// optimizes the above IR to
39// %1 = gep [10 x float] addrspace(3)* @a, i64 0, i64 %i
40// %v = load float addrspace(3)* %1 ; emits ld.shared.f32
41// propagating the addrspace(3) from @a to %1. As the result, the NVPTX
42// codegen is able to emit ld.shared.f32 for %v.
43//
44// Address space inference works in two steps. First, it uses a data-flow
45// analysis to infer as many generic pointers as possible to point to only one
46// specific address space. In the above example, it can prove that %1 only
47// points to addrspace(3). This algorithm was published in
48// CUDA: Compiling and optimizing for a GPU platform
49// Chakrabarti, Grover, Aarts, Kong, Kudlur, Lin, Marathe, Murphy, Wang
50// ICCS 2012
51//
52// Then, address space inference replaces all refinable generic pointers with
53// equivalent specific pointers.
54//
55// The major challenge of implementing this optimization is handling PHINodes,
56// which may create loops in the data flow graph. This brings two complications.
57//
58// First, the data flow analysis in Step 1 needs to be circular. For example,
59// %generic.input = addrspacecast float addrspace(3)* %input to float*
60// loop:
61// %y = phi [ %generic.input, %y2 ]
62// %y2 = getelementptr %y, 1
63// %v = load %y2
64// br ..., label %loop, ...
65// proving %y specific requires proving both %generic.input and %y2 specific,
66// but proving %y2 specific circles back to %y. To address this complication,
67// the data flow analysis operates on a lattice:
68// uninitialized > specific address spaces > generic.
69// All address expressions (our implementation only considers phi, bitcast,
70// addrspacecast, and getelementptr) start with the uninitialized address space.
71// The monotone transfer function moves the address space of a pointer down a
72// lattice path from uninitialized to specific and then to generic. A join
73// operation of two different specific address spaces pushes the expression down
74// to the generic address space. The analysis completes once it reaches a fixed
75// point.
76//
77// Second, IR rewriting in Step 2 also needs to be circular. For example,
78// converting %y to addrspace(3) requires the compiler to know the converted
79// %y2, but converting %y2 needs the converted %y. To address this complication,
80// we break these cycles using "undef" placeholders. When converting an
81// instruction `I` to a new address space, if its operand `Op` is not converted
82// yet, we let `I` temporarily use `undef` and fix all the uses of undef later.
83// For instance, our algorithm first converts %y to
84// %y' = phi float addrspace(3)* [ %input, undef ]
85// Then, it converts %y2 to
86// %y2' = getelementptr %y', 1
87// Finally, it fixes the undef in %y' so that
88// %y' = phi float addrspace(3)* [ %input, %y2' ]
89//
Jingyue Wu13755602016-03-20 20:59:20 +000090//===----------------------------------------------------------------------===//
91
Matt Arsenault850657a2017-01-31 01:10:58 +000092#include "llvm/Transforms/Scalar.h"
Jingyue Wu13755602016-03-20 20:59:20 +000093#include "llvm/ADT/DenseSet.h"
94#include "llvm/ADT/Optional.h"
95#include "llvm/ADT/SetVector.h"
Matt Arsenault42b64782017-01-30 23:02:12 +000096#include "llvm/Analysis/TargetTransformInfo.h"
Jingyue Wu13755602016-03-20 20:59:20 +000097#include "llvm/IR/Function.h"
98#include "llvm/IR/InstIterator.h"
99#include "llvm/IR/Instructions.h"
100#include "llvm/IR/Operator.h"
Jingyue Wu13755602016-03-20 20:59:20 +0000101#include "llvm/Support/Debug.h"
102#include "llvm/Support/raw_ostream.h"
103#include "llvm/Transforms/Utils/Local.h"
104#include "llvm/Transforms/Utils/ValueMapper.h"
105
Matt Arsenault850657a2017-01-31 01:10:58 +0000106#define DEBUG_TYPE "infer-address-spaces"
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000107
Jingyue Wu13755602016-03-20 20:59:20 +0000108using namespace llvm;
109
110namespace {
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000111static const unsigned UnknownAddressSpace = ~0u;
Jingyue Wu13755602016-03-20 20:59:20 +0000112
113using ValueToAddrSpaceMapTy = DenseMap<const Value *, unsigned>;
114
Matt Arsenault850657a2017-01-31 01:10:58 +0000115/// \brief InferAddressSpaces
116class InferAddressSpaces: public FunctionPass {
Matt Arsenault42b64782017-01-30 23:02:12 +0000117 /// Target specific address space which uses of should be replaced if
118 /// possible.
119 unsigned FlatAddrSpace;
120
Jingyue Wu13755602016-03-20 20:59:20 +0000121public:
122 static char ID;
123
Matt Arsenault850657a2017-01-31 01:10:58 +0000124 InferAddressSpaces() : FunctionPass(ID) {}
Jingyue Wu13755602016-03-20 20:59:20 +0000125
Matt Arsenault32b96002017-01-27 17:30:39 +0000126 void getAnalysisUsage(AnalysisUsage &AU) const override {
127 AU.setPreservesCFG();
Matt Arsenault42b64782017-01-30 23:02:12 +0000128 AU.addRequired<TargetTransformInfoWrapperPass>();
Matt Arsenault32b96002017-01-27 17:30:39 +0000129 }
130
Jingyue Wu13755602016-03-20 20:59:20 +0000131 bool runOnFunction(Function &F) override;
132
133private:
134 // Returns the new address space of V if updated; otherwise, returns None.
135 Optional<unsigned>
136 updateAddressSpace(const Value &V,
Matt Arsenault42b64782017-01-30 23:02:12 +0000137 const ValueToAddrSpaceMapTy &InferredAddrSpace) const;
Jingyue Wu13755602016-03-20 20:59:20 +0000138
139 // Tries to infer the specific address space of each address expression in
140 // Postorder.
141 void inferAddressSpaces(const std::vector<Value *> &Postorder,
Matt Arsenault42b64782017-01-30 23:02:12 +0000142 ValueToAddrSpaceMapTy *InferredAddrSpace) const;
Jingyue Wu13755602016-03-20 20:59:20 +0000143
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000144 // Changes the flat address expressions in function F to point to specific
Jingyue Wu13755602016-03-20 20:59:20 +0000145 // address spaces if InferredAddrSpace says so. Postorder is the postorder of
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000146 // all flat expressions in the use-def graph of function F.
Jingyue Wu13755602016-03-20 20:59:20 +0000147 bool
148 rewriteWithNewAddressSpaces(const std::vector<Value *> &Postorder,
149 const ValueToAddrSpaceMapTy &InferredAddrSpace,
Matt Arsenault42b64782017-01-30 23:02:12 +0000150 Function *F) const;
151
152 void appendsFlatAddressExpressionToPostorderStack(
153 Value *V, std::vector<std::pair<Value *, bool>> *PostorderStack,
154 DenseSet<Value *> *Visited) const;
155
156 std::vector<Value *> collectFlatAddressExpressions(Function &F) const;
157 Value *cloneValueWithNewAddressSpace(
158 Value *V, unsigned NewAddrSpace,
159 const ValueToValueMapTy &ValueWithNewAddrSpace,
160 SmallVectorImpl<const Use *> *UndefUsesToFix) const;
161 unsigned joinAddressSpaces(unsigned AS1, unsigned AS2) const;
Jingyue Wu13755602016-03-20 20:59:20 +0000162};
163} // end anonymous namespace
164
Matt Arsenault850657a2017-01-31 01:10:58 +0000165char InferAddressSpaces::ID = 0;
Jingyue Wu13755602016-03-20 20:59:20 +0000166
167namespace llvm {
Matt Arsenault850657a2017-01-31 01:10:58 +0000168void initializeInferAddressSpacesPass(PassRegistry &);
Jingyue Wu13755602016-03-20 20:59:20 +0000169}
Matt Arsenault850657a2017-01-31 01:10:58 +0000170
171INITIALIZE_PASS(InferAddressSpaces, DEBUG_TYPE, "Infer address spaces",
Jingyue Wu13755602016-03-20 20:59:20 +0000172 false, false)
173
174// Returns true if V is an address expression.
175// TODO: Currently, we consider only phi, bitcast, addrspacecast, and
176// getelementptr operators.
177static bool isAddressExpression(const Value &V) {
178 if (!isa<Operator>(V))
179 return false;
180
181 switch (cast<Operator>(V).getOpcode()) {
182 case Instruction::PHI:
183 case Instruction::BitCast:
184 case Instruction::AddrSpaceCast:
185 case Instruction::GetElementPtr:
186 return true;
187 default:
188 return false;
189 }
190}
191
192// Returns the pointer operands of V.
193//
194// Precondition: V is an address expression.
195static SmallVector<Value *, 2> getPointerOperands(const Value &V) {
196 assert(isAddressExpression(V));
197 const Operator& Op = cast<Operator>(V);
198 switch (Op.getOpcode()) {
199 case Instruction::PHI: {
200 auto IncomingValues = cast<PHINode>(Op).incoming_values();
201 return SmallVector<Value *, 2>(IncomingValues.begin(),
202 IncomingValues.end());
203 }
204 case Instruction::BitCast:
205 case Instruction::AddrSpaceCast:
206 case Instruction::GetElementPtr:
207 return {Op.getOperand(0)};
208 default:
209 llvm_unreachable("Unexpected instruction type.");
210 }
211}
212
Matt Arsenault42b64782017-01-30 23:02:12 +0000213// If V is an unvisited flat address expression, appends V to PostorderStack
Jingyue Wu13755602016-03-20 20:59:20 +0000214// and marks it as visited.
Matt Arsenault850657a2017-01-31 01:10:58 +0000215void InferAddressSpaces::appendsFlatAddressExpressionToPostorderStack(
216 Value *V, std::vector<std::pair<Value *, bool>> *PostorderStack,
217 DenseSet<Value *> *Visited) const {
Jingyue Wu13755602016-03-20 20:59:20 +0000218 assert(V->getType()->isPointerTy());
219 if (isAddressExpression(*V) &&
Matt Arsenault42b64782017-01-30 23:02:12 +0000220 V->getType()->getPointerAddressSpace() == FlatAddrSpace) {
Jingyue Wu13755602016-03-20 20:59:20 +0000221 if (Visited->insert(V).second)
222 PostorderStack->push_back(std::make_pair(V, false));
223 }
224}
225
Matt Arsenault42b64782017-01-30 23:02:12 +0000226// Returns all flat address expressions in function F. The elements are ordered
227// in postorder.
228std::vector<Value *>
Matt Arsenault850657a2017-01-31 01:10:58 +0000229InferAddressSpaces::collectFlatAddressExpressions(Function &F) const {
Jingyue Wu13755602016-03-20 20:59:20 +0000230 // This function implements a non-recursive postorder traversal of a partial
231 // use-def graph of function F.
232 std::vector<std::pair<Value*, bool>> PostorderStack;
233 // The set of visited expressions.
234 DenseSet<Value*> Visited;
235 // We only explore address expressions that are reachable from loads and
236 // stores for now because we aim at generating faster loads and stores.
237 for (Instruction &I : instructions(F)) {
238 if (isa<LoadInst>(I)) {
Matt Arsenault42b64782017-01-30 23:02:12 +0000239 appendsFlatAddressExpressionToPostorderStack(
Matt Arsenault850657a2017-01-31 01:10:58 +0000240 I.getOperand(0), &PostorderStack, &Visited);
Jingyue Wu13755602016-03-20 20:59:20 +0000241 } else if (isa<StoreInst>(I)) {
Matt Arsenault42b64782017-01-30 23:02:12 +0000242 appendsFlatAddressExpressionToPostorderStack(
Matt Arsenault850657a2017-01-31 01:10:58 +0000243 I.getOperand(1), &PostorderStack, &Visited);
Jingyue Wu13755602016-03-20 20:59:20 +0000244 }
245 }
246
247 std::vector<Value *> Postorder; // The resultant postorder.
248 while (!PostorderStack.empty()) {
249 // If the operands of the expression on the top are already explored,
250 // adds that expression to the resultant postorder.
251 if (PostorderStack.back().second) {
252 Postorder.push_back(PostorderStack.back().first);
253 PostorderStack.pop_back();
254 continue;
255 }
256 // Otherwise, adds its operands to the stack and explores them.
257 PostorderStack.back().second = true;
258 for (Value *PtrOperand : getPointerOperands(*PostorderStack.back().first)) {
Matt Arsenault42b64782017-01-30 23:02:12 +0000259 appendsFlatAddressExpressionToPostorderStack(
Matt Arsenault850657a2017-01-31 01:10:58 +0000260 PtrOperand, &PostorderStack, &Visited);
Jingyue Wu13755602016-03-20 20:59:20 +0000261 }
262 }
263 return Postorder;
264}
265
266// A helper function for cloneInstructionWithNewAddressSpace. Returns the clone
267// of OperandUse.get() in the new address space. If the clone is not ready yet,
268// returns an undef in the new address space as a placeholder.
269static Value *operandWithNewAddressSpaceOrCreateUndef(
Matt Arsenault850657a2017-01-31 01:10:58 +0000270 const Use &OperandUse, unsigned NewAddrSpace,
271 const ValueToValueMapTy &ValueWithNewAddrSpace,
272 SmallVectorImpl<const Use *> *UndefUsesToFix) {
Jingyue Wu13755602016-03-20 20:59:20 +0000273 Value *Operand = OperandUse.get();
274 if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand))
275 return NewOperand;
276
277 UndefUsesToFix->push_back(&OperandUse);
278 return UndefValue::get(
Matt Arsenault850657a2017-01-31 01:10:58 +0000279 Operand->getType()->getPointerElementType()->getPointerTo(NewAddrSpace));
Jingyue Wu13755602016-03-20 20:59:20 +0000280}
281
282// Returns a clone of `I` with its operands converted to those specified in
283// ValueWithNewAddrSpace. Due to potential cycles in the data flow graph, an
284// operand whose address space needs to be modified might not exist in
285// ValueWithNewAddrSpace. In that case, uses undef as a placeholder operand and
286// adds that operand use to UndefUsesToFix so that caller can fix them later.
287//
288// Note that we do not necessarily clone `I`, e.g., if it is an addrspacecast
289// from a pointer whose type already matches. Therefore, this function returns a
290// Value* instead of an Instruction*.
291static Value *cloneInstructionWithNewAddressSpace(
Matt Arsenault850657a2017-01-31 01:10:58 +0000292 Instruction *I, unsigned NewAddrSpace,
293 const ValueToValueMapTy &ValueWithNewAddrSpace,
294 SmallVectorImpl<const Use *> *UndefUsesToFix) {
Jingyue Wu13755602016-03-20 20:59:20 +0000295 Type *NewPtrType =
Matt Arsenault850657a2017-01-31 01:10:58 +0000296 I->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
Jingyue Wu13755602016-03-20 20:59:20 +0000297
298 if (I->getOpcode() == Instruction::AddrSpaceCast) {
299 Value *Src = I->getOperand(0);
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000300 // Because `I` is flat, the source address space must be specific.
Jingyue Wu13755602016-03-20 20:59:20 +0000301 // Therefore, the inferred address space must be the source space, according
302 // to our algorithm.
303 assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
304 if (Src->getType() != NewPtrType)
305 return new BitCastInst(Src, NewPtrType);
306 return Src;
307 }
308
309 // Computes the converted pointer operands.
310 SmallVector<Value *, 4> NewPointerOperands;
311 for (const Use &OperandUse : I->operands()) {
312 if (!OperandUse.get()->getType()->isPointerTy())
313 NewPointerOperands.push_back(nullptr);
314 else
315 NewPointerOperands.push_back(operandWithNewAddressSpaceOrCreateUndef(
Matt Arsenault850657a2017-01-31 01:10:58 +0000316 OperandUse, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix));
Jingyue Wu13755602016-03-20 20:59:20 +0000317 }
318
319 switch (I->getOpcode()) {
320 case Instruction::BitCast:
321 return new BitCastInst(NewPointerOperands[0], NewPtrType);
322 case Instruction::PHI: {
323 assert(I->getType()->isPointerTy());
324 PHINode *PHI = cast<PHINode>(I);
325 PHINode *NewPHI = PHINode::Create(NewPtrType, PHI->getNumIncomingValues());
326 for (unsigned Index = 0; Index < PHI->getNumIncomingValues(); ++Index) {
327 unsigned OperandNo = PHINode::getOperandNumForIncomingValue(Index);
328 NewPHI->addIncoming(NewPointerOperands[OperandNo],
329 PHI->getIncomingBlock(Index));
330 }
331 return NewPHI;
332 }
333 case Instruction::GetElementPtr: {
334 GetElementPtrInst *GEP = cast<GetElementPtrInst>(I);
335 GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
Matt Arsenault850657a2017-01-31 01:10:58 +0000336 GEP->getSourceElementType(), NewPointerOperands[0],
337 SmallVector<Value *, 4>(GEP->idx_begin(), GEP->idx_end()));
Jingyue Wu13755602016-03-20 20:59:20 +0000338 NewGEP->setIsInBounds(GEP->isInBounds());
339 return NewGEP;
340 }
341 default:
342 llvm_unreachable("Unexpected opcode");
343 }
344}
345
346// Similar to cloneInstructionWithNewAddressSpace, returns a clone of the
347// constant expression `CE` with its operands replaced as specified in
348// ValueWithNewAddrSpace.
349static Value *cloneConstantExprWithNewAddressSpace(
Matt Arsenault850657a2017-01-31 01:10:58 +0000350 ConstantExpr *CE, unsigned NewAddrSpace,
351 const ValueToValueMapTy &ValueWithNewAddrSpace) {
Jingyue Wu13755602016-03-20 20:59:20 +0000352 Type *TargetType =
Matt Arsenault850657a2017-01-31 01:10:58 +0000353 CE->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
Jingyue Wu13755602016-03-20 20:59:20 +0000354
355 if (CE->getOpcode() == Instruction::AddrSpaceCast) {
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000356 // Because CE is flat, the source address space must be specific.
Jingyue Wu13755602016-03-20 20:59:20 +0000357 // Therefore, the inferred address space must be the source space according
358 // to our algorithm.
359 assert(CE->getOperand(0)->getType()->getPointerAddressSpace() ==
360 NewAddrSpace);
361 return ConstantExpr::getBitCast(CE->getOperand(0), TargetType);
362 }
363
364 // Computes the operands of the new constant expression.
365 SmallVector<Constant *, 4> NewOperands;
366 for (unsigned Index = 0; Index < CE->getNumOperands(); ++Index) {
367 Constant *Operand = CE->getOperand(Index);
368 // If the address space of `Operand` needs to be modified, the new operand
369 // with the new address space should already be in ValueWithNewAddrSpace
370 // because (1) the constant expressions we consider (i.e. addrspacecast,
371 // bitcast, and getelementptr) do not incur cycles in the data flow graph
372 // and (2) this function is called on constant expressions in postorder.
373 if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand)) {
374 NewOperands.push_back(cast<Constant>(NewOperand));
375 } else {
376 // Otherwise, reuses the old operand.
377 NewOperands.push_back(Operand);
378 }
379 }
380
381 if (CE->getOpcode() == Instruction::GetElementPtr) {
382 // Needs to specify the source type while constructing a getelementptr
383 // constant expression.
384 return CE->getWithOperands(
Matt Arsenault850657a2017-01-31 01:10:58 +0000385 NewOperands, TargetType, /*OnlyIfReduced=*/false,
386 NewOperands[0]->getType()->getPointerElementType());
Jingyue Wu13755602016-03-20 20:59:20 +0000387 }
388
389 return CE->getWithOperands(NewOperands, TargetType);
390}
391
392// Returns a clone of the value `V`, with its operands replaced as specified in
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000393// ValueWithNewAddrSpace. This function is called on every flat address
Jingyue Wu13755602016-03-20 20:59:20 +0000394// expression whose address space needs to be modified, in postorder.
395//
396// See cloneInstructionWithNewAddressSpace for the meaning of UndefUsesToFix.
Matt Arsenault850657a2017-01-31 01:10:58 +0000397Value *InferAddressSpaces::cloneValueWithNewAddressSpace(
Matt Arsenault42b64782017-01-30 23:02:12 +0000398 Value *V, unsigned NewAddrSpace,
399 const ValueToValueMapTy &ValueWithNewAddrSpace,
400 SmallVectorImpl<const Use *> *UndefUsesToFix) const {
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000401 // All values in Postorder are flat address expressions.
Jingyue Wu13755602016-03-20 20:59:20 +0000402 assert(isAddressExpression(*V) &&
Matt Arsenault42b64782017-01-30 23:02:12 +0000403 V->getType()->getPointerAddressSpace() == FlatAddrSpace);
Jingyue Wu13755602016-03-20 20:59:20 +0000404
405 if (Instruction *I = dyn_cast<Instruction>(V)) {
406 Value *NewV = cloneInstructionWithNewAddressSpace(
Matt Arsenault850657a2017-01-31 01:10:58 +0000407 I, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix);
Jingyue Wu13755602016-03-20 20:59:20 +0000408 if (Instruction *NewI = dyn_cast<Instruction>(NewV)) {
409 if (NewI->getParent() == nullptr) {
410 NewI->insertBefore(I);
411 NewI->takeName(I);
412 }
413 }
414 return NewV;
415 }
416
417 return cloneConstantExprWithNewAddressSpace(
Matt Arsenault850657a2017-01-31 01:10:58 +0000418 cast<ConstantExpr>(V), NewAddrSpace, ValueWithNewAddrSpace);
Jingyue Wu13755602016-03-20 20:59:20 +0000419}
420
421// Defines the join operation on the address space lattice (see the file header
422// comments).
Matt Arsenault850657a2017-01-31 01:10:58 +0000423unsigned InferAddressSpaces::joinAddressSpaces(unsigned AS1,
424 unsigned AS2) const {
Matt Arsenault42b64782017-01-30 23:02:12 +0000425 if (AS1 == FlatAddrSpace || AS2 == FlatAddrSpace)
426 return FlatAddrSpace;
Jingyue Wu13755602016-03-20 20:59:20 +0000427
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000428 if (AS1 == UnknownAddressSpace)
Jingyue Wu13755602016-03-20 20:59:20 +0000429 return AS2;
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000430 if (AS2 == UnknownAddressSpace)
Jingyue Wu13755602016-03-20 20:59:20 +0000431 return AS1;
432
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000433 // The join of two different specific address spaces is flat.
Matt Arsenault42b64782017-01-30 23:02:12 +0000434 return (AS1 == AS2) ? AS1 : FlatAddrSpace;
Jingyue Wu13755602016-03-20 20:59:20 +0000435}
436
Matt Arsenault850657a2017-01-31 01:10:58 +0000437bool InferAddressSpaces::runOnFunction(Function &F) {
Andrew Kaylor87b10dd2016-04-26 23:44:31 +0000438 if (skipFunction(F))
439 return false;
440
Matt Arsenault42b64782017-01-30 23:02:12 +0000441 const TargetTransformInfo &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
442 FlatAddrSpace = TTI.getFlatAddressSpace();
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000443 if (FlatAddrSpace == UnknownAddressSpace)
Matt Arsenault42b64782017-01-30 23:02:12 +0000444 return false;
445
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000446 // Collects all flat address expressions in postorder.
Matt Arsenault42b64782017-01-30 23:02:12 +0000447 std::vector<Value *> Postorder = collectFlatAddressExpressions(F);
Jingyue Wu13755602016-03-20 20:59:20 +0000448
449 // Runs a data-flow analysis to refine the address spaces of every expression
450 // in Postorder.
451 ValueToAddrSpaceMapTy InferredAddrSpace;
452 inferAddressSpaces(Postorder, &InferredAddrSpace);
453
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000454 // Changes the address spaces of the flat address expressions who are inferred
455 // to point to a specific address space.
Jingyue Wu13755602016-03-20 20:59:20 +0000456 return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace, &F);
457}
458
Matt Arsenault850657a2017-01-31 01:10:58 +0000459void InferAddressSpaces::inferAddressSpaces(
460 const std::vector<Value *> &Postorder,
461 ValueToAddrSpaceMapTy *InferredAddrSpace) const {
Jingyue Wu13755602016-03-20 20:59:20 +0000462 SetVector<Value *> Worklist(Postorder.begin(), Postorder.end());
463 // Initially, all expressions are in the uninitialized address space.
464 for (Value *V : Postorder)
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000465 (*InferredAddrSpace)[V] = UnknownAddressSpace;
Jingyue Wu13755602016-03-20 20:59:20 +0000466
467 while (!Worklist.empty()) {
468 Value* V = Worklist.pop_back_val();
469
470 // Tries to update the address space of the stack top according to the
471 // address spaces of its operands.
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000472 DEBUG(dbgs() << "Updating the address space of\n " << *V << '\n');
Jingyue Wu13755602016-03-20 20:59:20 +0000473 Optional<unsigned> NewAS = updateAddressSpace(*V, *InferredAddrSpace);
474 if (!NewAS.hasValue())
475 continue;
476 // If any updates are made, grabs its users to the worklist because
477 // their address spaces can also be possibly updated.
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000478 DEBUG(dbgs() << " to " << NewAS.getValue() << '\n');
Jingyue Wu13755602016-03-20 20:59:20 +0000479 (*InferredAddrSpace)[V] = NewAS.getValue();
480
481 for (Value *User : V->users()) {
482 // Skip if User is already in the worklist.
483 if (Worklist.count(User))
484 continue;
485
486 auto Pos = InferredAddrSpace->find(User);
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000487 // Our algorithm only updates the address spaces of flat address
Jingyue Wu13755602016-03-20 20:59:20 +0000488 // expressions, which are those in InferredAddrSpace.
489 if (Pos == InferredAddrSpace->end())
490 continue;
491
492 // Function updateAddressSpace moves the address space down a lattice
Matt Arsenault850657a2017-01-31 01:10:58 +0000493 // path. Therefore, nothing to do if User is already inferred as flat (the
494 // bottom element in the lattice).
Matt Arsenault42b64782017-01-30 23:02:12 +0000495 if (Pos->second == FlatAddrSpace)
Jingyue Wu13755602016-03-20 20:59:20 +0000496 continue;
497
498 Worklist.insert(User);
499 }
500 }
501}
502
Matt Arsenault850657a2017-01-31 01:10:58 +0000503Optional<unsigned> InferAddressSpaces::updateAddressSpace(
504 const Value &V, const ValueToAddrSpaceMapTy &InferredAddrSpace) const {
Jingyue Wu13755602016-03-20 20:59:20 +0000505 assert(InferredAddrSpace.count(&V));
506
507 // The new inferred address space equals the join of the address spaces
508 // of all its pointer operands.
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000509 unsigned NewAS = UnknownAddressSpace;
Jingyue Wu13755602016-03-20 20:59:20 +0000510 for (Value *PtrOperand : getPointerOperands(V)) {
511 unsigned OperandAS;
512 if (InferredAddrSpace.count(PtrOperand))
513 OperandAS = InferredAddrSpace.lookup(PtrOperand);
514 else
515 OperandAS = PtrOperand->getType()->getPointerAddressSpace();
516 NewAS = joinAddressSpaces(NewAS, OperandAS);
Matt Arsenault850657a2017-01-31 01:10:58 +0000517
518 // join(flat, *) = flat. So we can break if NewAS is already flat.
Matt Arsenault42b64782017-01-30 23:02:12 +0000519 if (NewAS == FlatAddrSpace)
Jingyue Wu13755602016-03-20 20:59:20 +0000520 break;
521 }
522
523 unsigned OldAS = InferredAddrSpace.lookup(&V);
Matt Arsenault42b64782017-01-30 23:02:12 +0000524 assert(OldAS != FlatAddrSpace);
Jingyue Wu13755602016-03-20 20:59:20 +0000525 if (OldAS == NewAS)
526 return None;
527 return NewAS;
528}
529
Matt Arsenault850657a2017-01-31 01:10:58 +0000530bool InferAddressSpaces::rewriteWithNewAddressSpaces(
531 const std::vector<Value *> &Postorder,
532 const ValueToAddrSpaceMapTy &InferredAddrSpace, Function *F) const {
Jingyue Wu13755602016-03-20 20:59:20 +0000533 // For each address expression to be modified, creates a clone of it with its
534 // pointer operands converted to the new address space. Since the pointer
535 // operands are converted, the clone is naturally in the new address space by
536 // construction.
537 ValueToValueMapTy ValueWithNewAddrSpace;
538 SmallVector<const Use *, 32> UndefUsesToFix;
539 for (Value* V : Postorder) {
540 unsigned NewAddrSpace = InferredAddrSpace.lookup(V);
541 if (V->getType()->getPointerAddressSpace() != NewAddrSpace) {
542 ValueWithNewAddrSpace[V] = cloneValueWithNewAddressSpace(
Matt Arsenault850657a2017-01-31 01:10:58 +0000543 V, NewAddrSpace, ValueWithNewAddrSpace, &UndefUsesToFix);
Jingyue Wu13755602016-03-20 20:59:20 +0000544 }
545 }
546
547 if (ValueWithNewAddrSpace.empty())
548 return false;
549
550 // Fixes all the undef uses generated by cloneInstructionWithNewAddressSpace.
551 for (const Use* UndefUse : UndefUsesToFix) {
552 User *V = UndefUse->getUser();
553 User *NewV = cast<User>(ValueWithNewAddrSpace.lookup(V));
554 unsigned OperandNo = UndefUse->getOperandNo();
555 assert(isa<UndefValue>(NewV->getOperand(OperandNo)));
556 NewV->setOperand(OperandNo, ValueWithNewAddrSpace.lookup(UndefUse->get()));
557 }
558
559 // Replaces the uses of the old address expressions with the new ones.
560 for (Value *V : Postorder) {
561 Value *NewV = ValueWithNewAddrSpace.lookup(V);
562 if (NewV == nullptr)
563 continue;
564
565 SmallVector<Use *, 4> Uses;
566 for (Use &U : V->uses())
567 Uses.push_back(&U);
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000568
569 DEBUG(dbgs() << "Replacing the uses of " << *V
570 << "\n with\n " << *NewV << '\n');
571
Jingyue Wu13755602016-03-20 20:59:20 +0000572 for (Use *U : Uses) {
573 if (isa<LoadInst>(U->getUser()) ||
Matt Arsenault9f432ec2017-01-30 23:27:11 +0000574 (isa<StoreInst>(U->getUser()) &&
575 U->getOperandNo() == StoreInst::getPointerOperandIndex())) {
Jingyue Wu13755602016-03-20 20:59:20 +0000576 // If V is used as the pointer operand of a load/store, sets the pointer
577 // operand to NewV. This replacement does not change the element type,
578 // so the resultant load/store is still valid.
579 U->set(NewV);
580 } else if (isa<Instruction>(U->getUser())) {
Matt Arsenault850657a2017-01-31 01:10:58 +0000581 // Otherwise, replaces the use with flat(NewV).
Jingyue Wu13755602016-03-20 20:59:20 +0000582 // TODO: Some optimization opportunities are missed. For example, in
583 // %0 = icmp eq float* %p, %q
584 // if both p and q are inferred to be shared, we can rewrite %0 as
585 // %0 = icmp eq float addrspace(3)* %new_p, %new_q
586 // instead of currently
Matt Arsenault850657a2017-01-31 01:10:58 +0000587 // %flat_p = addrspacecast float addrspace(3)* %new_p to float*
588 // %flat_q = addrspacecast float addrspace(3)* %new_q to float*
589 // %0 = icmp eq float* %flat_p, %flat_q
Jingyue Wu13755602016-03-20 20:59:20 +0000590 if (Instruction *I = dyn_cast<Instruction>(V)) {
591 BasicBlock::iterator InsertPos = std::next(I->getIterator());
592 while (isa<PHINode>(InsertPos))
593 ++InsertPos;
594 U->set(new AddrSpaceCastInst(NewV, V->getType(), "", &*InsertPos));
595 } else {
596 U->set(ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV),
597 V->getType()));
598 }
599 }
600 }
601 if (V->use_empty())
602 RecursivelyDeleteTriviallyDeadInstructions(V);
603 }
604
605 return true;
606}
607
Matt Arsenault850657a2017-01-31 01:10:58 +0000608FunctionPass *llvm::createInferAddressSpacesPass() {
609 return new InferAddressSpaces();
Jingyue Wu13755602016-03-20 20:59:20 +0000610}