1//===- InferAddressSpace.cpp - --------------------------------------------===//
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// CUDA C/C++ includes memory space designation as variable type qualifers (such
10// as __global__ and __shared__). Knowing the space of a memory access allows
11// CUDA compilers to emit faster PTX loads and stores. For example, a load from
12// shared memory can be translated to `ld.shared` which is roughly 10% faster
13// than a generic `ld` on an NVIDIA Tesla K40c.
14//
15// Unfortunately, type qualifiers only apply to variable declarations, so CUDA
16// compilers must infer the memory space of an address expression from
17// type-qualified variables.
18//
19// LLVM IR uses non-zero (so-called) specific address spaces to represent memory
20// spaces (e.g. addrspace(3) means shared memory). The Clang frontend
21// places only type-qualified variables in specific address spaces, and then
22// conservatively `addrspacecast`s each type-qualified variable to addrspace(0)
23// (so-called the generic address space) for other instructions to use.
24//
25// For example, the Clang translates the following CUDA code
26// __shared__ float a[10];
27// float v = a[i];
28// to
29// %0 = addrspacecast [10 x float] addrspace(3)* @a to [10 x float]*
30// %1 = gep [10 x float], [10 x float]* %0, i64 0, i64 %i
31// %v = load float, float* %1 ; emits ld.f32
32// @a is in addrspace(3) since it's type-qualified, but its use from %1 is
33// redirected to %0 (the generic version of @a).
34//
35// The optimization implemented in this file propagates specific address spaces
36// from type-qualified variable declarations to its users. For example, it
37// optimizes the above IR to
38// %1 = gep [10 x float] addrspace(3)* @a, i64 0, i64 %i
39// %v = load float addrspace(3)* %1 ; emits ld.shared.f32
40// propagating the addrspace(3) from @a to %1. As the result, the NVPTX
41// codegen is able to emit ld.shared.f32 for %v.
42//
43// Address space inference works in two steps. First, it uses a data-flow
44// analysis to infer as many generic pointers as possible to point to only one
45// specific address space. In the above example, it can prove that %1 only
46// points to addrspace(3). This algorithm was published in
47// CUDA: Compiling and optimizing for a GPU platform
48// Chakrabarti, Grover, Aarts, Kong, Kudlur, Lin, Marathe, Murphy, Wang
49// ICCS 2012
50//
51// Then, address space inference replaces all refinable generic pointers with
52// equivalent specific pointers.
53//
54// The major challenge of implementing this optimization is handling PHINodes,
55// which may create loops in the data flow graph. This brings two complications.
56//
57// First, the data flow analysis in Step 1 needs to be circular. For example,
58// %generic.input = addrspacecast float addrspace(3)* %input to float*
59// loop:
60// %y = phi [ %generic.input, %y2 ]
61// %y2 = getelementptr %y, 1
62// %v = load %y2
63// br ..., label %loop, ...
64// proving %y specific requires proving both %generic.input and %y2 specific,
65// but proving %y2 specific circles back to %y. To address this complication,
66// the data flow analysis operates on a lattice:
67// uninitialized > specific address spaces > generic.
68// All address expressions (our implementation only considers phi, bitcast,
69// addrspacecast, and getelementptr) start with the uninitialized address space.
70// The monotone transfer function moves the address space of a pointer down a
71// lattice path from uninitialized to specific and then to generic. A join
72// operation of two different specific address spaces pushes the expression down
73// to the generic address space. The analysis completes once it reaches a fixed
74// point.
75//
76// Second, IR rewriting in Step 2 also needs to be circular. For example,
77// converting %y to addrspace(3) requires the compiler to know the converted
78// %y2, but converting %y2 needs the converted %y. To address this complication,
79// we break these cycles using "poison" placeholders. When converting an
80// instruction `I` to a new address space, if its operand `Op` is not converted
81// yet, we let `I` temporarily use `poison` and fix all the uses later.
82// For instance, our algorithm first converts %y to
83// %y' = phi float addrspace(3)* [ %input, poison ]
84// Then, it converts %y2 to
85// %y2' = getelementptr %y', 1
86// Finally, it fixes the poison in %y' so that
87// %y' = phi float addrspace(3)* [ %input, %y2' ]
88//
89//===----------------------------------------------------------------------===//
90
91#include "llvm/Transforms/Scalar/InferAddressSpaces.h"
92#include "llvm/ADT/ArrayRef.h"
93#include "llvm/ADT/DenseMap.h"
94#include "llvm/ADT/DenseSet.h"
95#include "llvm/ADT/SetVector.h"
96#include "llvm/ADT/SmallVector.h"
97#include "llvm/Analysis/AssumptionCache.h"
98#include "llvm/Analysis/TargetTransformInfo.h"
99#include "llvm/Analysis/ValueTracking.h"
100#include "llvm/IR/Argument.h"
101#include "llvm/IR/BasicBlock.h"
102#include "llvm/IR/Constant.h"
103#include "llvm/IR/Constants.h"
104#include "llvm/IR/Dominators.h"
105#include "llvm/IR/Function.h"
106#include "llvm/IR/IRBuilder.h"
107#include "llvm/IR/InstIterator.h"
108#include "llvm/IR/Instruction.h"
109#include "llvm/IR/Instructions.h"
110#include "llvm/IR/IntrinsicInst.h"
111#include "llvm/IR/Intrinsics.h"
112#include "llvm/IR/LLVMContext.h"
113#include "llvm/IR/Operator.h"
114#include "llvm/IR/PassManager.h"
115#include "llvm/IR/PatternMatch.h"
116#include "llvm/IR/Type.h"
117#include "llvm/IR/Use.h"
118#include "llvm/IR/User.h"
119#include "llvm/IR/Value.h"
120#include "llvm/IR/ValueHandle.h"
121#include "llvm/InitializePasses.h"
122#include "llvm/Pass.h"
123#include "llvm/Support/Casting.h"
124#include "llvm/Support/CommandLine.h"
125#include "llvm/Support/Debug.h"
126#include "llvm/Support/ErrorHandling.h"
127#include "llvm/Support/KnownBits.h"
128#include "llvm/Support/raw_ostream.h"
129#include "llvm/Transforms/Scalar.h"
130#include "llvm/Transforms/Utils/Local.h"
131#include "llvm/Transforms/Utils/ValueMapper.h"
132#include <cassert>
133#include <iterator>
134#include <limits>
135#include <optional>
136#include <utility>
137#include <vector>
138
139#define DEBUG_TYPE "infer-address-spaces"
140
141using namespace llvm;
142using namespace llvm::PatternMatch;
143
144static cl::opt<bool> AssumeDefaultIsFlatAddressSpace(
145 "assume-default-is-flat-addrspace", cl::init(Val: false), cl::ReallyHidden,
146 cl::desc("The default address space is assumed as the flat address space. "
147 "This is mainly for test purpose."));
148
149static const unsigned UninitializedAddressSpace =
150 std::numeric_limits<unsigned>::max();
151
152namespace {
153
154using ValueToAddrSpaceMapTy = DenseMap<const Value *, unsigned>;
155// Different from ValueToAddrSpaceMapTy, where a new addrspace is inferred on
156// the *def* of a value, PredicatedAddrSpaceMapTy is map where a new
157// addrspace is inferred on the *use* of a pointer. This map is introduced to
158// infer addrspace from the addrspace predicate assumption built from assume
159// intrinsic. In that scenario, only specific uses (under valid assumption
160// context) could be inferred with a new addrspace.
161using PredicatedAddrSpaceMapTy =
162 DenseMap<std::pair<const Value *, const Value *>, unsigned>;
163using PostorderStackTy = llvm::SmallVector<PointerIntPair<Value *, 1, bool>, 4>;
164
165class InferAddressSpaces : public FunctionPass {
166 unsigned FlatAddrSpace = 0;
167
168public:
169 static char ID;
170
171 InferAddressSpaces()
172 : FunctionPass(ID), FlatAddrSpace(UninitializedAddressSpace) {
173 initializeInferAddressSpacesPass(*PassRegistry::getPassRegistry());
174 }
175 InferAddressSpaces(unsigned AS) : FunctionPass(ID), FlatAddrSpace(AS) {
176 initializeInferAddressSpacesPass(*PassRegistry::getPassRegistry());
177 }
178
179 void getAnalysisUsage(AnalysisUsage &AU) const override {
180 AU.setPreservesCFG();
181 AU.addRequired<AssumptionCacheTracker>();
182 AU.addRequired<TargetTransformInfoWrapperPass>();
183 }
184
185 bool runOnFunction(Function &F) override;
186};
187
188class InferAddressSpacesImpl {
189 AssumptionCache &AC;
190 Function *F = nullptr;
191 const DominatorTree *DT = nullptr;
192 const TargetTransformInfo *TTI = nullptr;
193 const DataLayout *DL = nullptr;
194
195 /// Target specific address space which uses of should be replaced if
196 /// possible.
197 unsigned FlatAddrSpace = 0;
198 DenseMap<const Value *, Value *> PtrIntCastPairs;
199
200 // Tries to find if the inttoptr instruction is derived from an pointer have
201 // specific address space, and is safe to propagate the address space to the
202 // new pointer that inttoptr produces.
203 Value *getIntToPtrPointerOperand(const Operator *I2P) const;
204 // Tries to find if the inttoptr instruction is derived from an pointer have
205 // specific address space, and is safe to propagate the address space to the
206 // new pointer that inttoptr produces. If the old pointer is found, cache the
207 // <OldPtr, inttoptr> pairs to a map.
208 void collectIntToPtrPointerOperand();
209 // Check if an old pointer is found ahead of time. The safety has been checked
210 // when collecting the inttoptr original pointer and the result is cached in
211 // PtrIntCastPairs.
212 bool isSafeToCastIntToPtrAddrSpace(const Operator *I2P) const {
213 return PtrIntCastPairs.contains(Val: I2P);
214 }
215 bool isAddressExpression(const Value &V, const DataLayout &DL,
216 const TargetTransformInfo *TTI) const;
217 Value *cloneConstantExprWithNewAddressSpace(
218 ConstantExpr *CE, unsigned NewAddrSpace,
219 const ValueToValueMapTy &ValueWithNewAddrSpace, const DataLayout *DL,
220 const TargetTransformInfo *TTI) const;
221
222 SmallVector<Value *, 2>
223 getPointerOperands(const Value &V, const DataLayout &DL,
224 const TargetTransformInfo *TTI) const;
225
226 // Try to update the address space of V. If V is updated, returns true and
227 // false otherwise.
228 bool updateAddressSpace(const Value &V,
229 ValueToAddrSpaceMapTy &InferredAddrSpace,
230 PredicatedAddrSpaceMapTy &PredicatedAS) const;
231
232 // Adds the users of V whose address space may still change to Worklist.
233 void enqueueUsers(Value &V, const ValueToAddrSpaceMapTy &InferredAddrSpace,
234 SetVector<Value *> &Worklist) const;
235
236 // Propagates address spaces out of Worklist until nothing changes.
237 void runToFixPoint(SetVector<Value *> &Worklist,
238 ValueToAddrSpaceMapTy &InferredAddrSpace,
239 PredicatedAddrSpaceMapTy &PredicatedAS) const;
240
241 // Tries to infer the specific address space of each address expression in
242 // Postorder.
243 void inferAddressSpaces(ArrayRef<WeakTrackingVH> Postorder,
244 ValueToAddrSpaceMapTy &InferredAddrSpace,
245 PredicatedAddrSpaceMapTy &PredicatedAS) const;
246
247 bool isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const;
248
249 Value *clonePtrMaskWithNewAddressSpace(
250 IntrinsicInst *I, unsigned NewAddrSpace,
251 const ValueToValueMapTy &ValueWithNewAddrSpace,
252 const PredicatedAddrSpaceMapTy &PredicatedAS,
253 SmallVectorImpl<const Use *> *PoisonUsesToFix) const;
254
255 Value *cloneInstructionWithNewAddressSpace(
256 Instruction *I, unsigned NewAddrSpace,
257 const ValueToValueMapTy &ValueWithNewAddrSpace,
258 const PredicatedAddrSpaceMapTy &PredicatedAS,
259 SmallVectorImpl<const Use *> *PoisonUsesToFix) const;
260
261 void performPointerReplacement(
262 Value *V, Value *NewV, Use &U, ValueToValueMapTy &ValueWithNewAddrSpace,
263 SmallVectorImpl<Instruction *> &DeadInstructions) const;
264
265 // Changes the flat address expressions in function F to point to specific
266 // address spaces if InferredAddrSpace says so. Postorder is the postorder of
267 // all flat expressions in the use-def graph of function F.
268 bool rewriteWithNewAddressSpaces(
269 ArrayRef<WeakTrackingVH> Postorder,
270 const ValueToAddrSpaceMapTy &InferredAddrSpace,
271 const PredicatedAddrSpaceMapTy &PredicatedAS) const;
272
273 void appendsFlatAddressExpressionToPostorderStack(
274 Value *V, PostorderStackTy &PostorderStack,
275 DenseSet<Value *> &Visited) const;
276
277 bool rewriteIntrinsicOperands(IntrinsicInst *II, Value *OldV,
278 Value *NewV) const;
279 void collectRewritableIntrinsicOperands(IntrinsicInst *II,
280 PostorderStackTy &PostorderStack,
281 DenseSet<Value *> &Visited) const;
282
283 std::vector<WeakTrackingVH> collectFlatAddressExpressions(Function &F) const;
284
285 Value *cloneValueWithNewAddressSpace(
286 Value *V, unsigned NewAddrSpace,
287 const ValueToValueMapTy &ValueWithNewAddrSpace,
288 const PredicatedAddrSpaceMapTy &PredicatedAS,
289 SmallVectorImpl<const Use *> *PoisonUsesToFix) const;
290 unsigned joinAddressSpaces(unsigned AS1, unsigned AS2) const;
291
292 unsigned getPredicatedAddrSpace(const Value &PtrV,
293 const Value *UserCtx) const;
294
295public:
296 InferAddressSpacesImpl(AssumptionCache &AC, const DominatorTree *DT,
297 const TargetTransformInfo *TTI, unsigned FlatAddrSpace)
298 : AC(AC), DT(DT), TTI(TTI), FlatAddrSpace(FlatAddrSpace) {}
299 bool run(Function &F);
300};
301
302} // end anonymous namespace
303
304char InferAddressSpaces::ID = 0;
305
306INITIALIZE_PASS_BEGIN(InferAddressSpaces, DEBUG_TYPE, "Infer address spaces",
307 false, false)
308INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
309INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
310INITIALIZE_PASS_END(InferAddressSpaces, DEBUG_TYPE, "Infer address spaces",
311 false, false)
312
313static Type *getPtrOrVecOfPtrsWithNewAS(Type *Ty, unsigned NewAddrSpace) {
314 assert(Ty->isPtrOrPtrVectorTy());
315 PointerType *NPT = PointerType::get(C&: Ty->getContext(), AddressSpace: NewAddrSpace);
316 return Ty->getWithNewType(EltTy: NPT);
317}
318
319// Check whether that's no-op pointer bitcast using a pair of
320// `ptrtoint`/`inttoptr` due to the missing no-op pointer bitcast over
321// different address spaces.
322static bool isNoopPtrIntCastPair(const Operator *I2P, const DataLayout &DL,
323 const TargetTransformInfo *TTI) {
324 assert(I2P->getOpcode() == Instruction::IntToPtr);
325 auto *P2I = dyn_cast<Operator>(Val: I2P->getOperand(i: 0));
326 if (!P2I || P2I->getOpcode() != Instruction::PtrToInt)
327 return false;
328 // Check it's really safe to treat that pair of `ptrtoint`/`inttoptr` as a
329 // no-op cast. Besides checking both of them are no-op casts, as the
330 // reinterpreted pointer may be used in other pointer arithmetic, we also
331 // need to double-check that through the target-specific hook. That ensures
332 // the underlying target also agrees that's a no-op address space cast and
333 // pointer bits are preserved.
334 // The current IR spec doesn't have clear rules on address space casts,
335 // especially a clear definition for pointer bits in non-default address
336 // spaces. It would be undefined if that pointer is dereferenced after an
337 // invalid reinterpret cast. Also, due to the unclearness for the meaning of
338 // bits in non-default address spaces in the current spec, the pointer
339 // arithmetic may also be undefined after invalid pointer reinterpret cast.
340 // However, as we confirm through the target hooks that it's a no-op
341 // addrspacecast, it doesn't matter since the bits should be the same.
342 unsigned P2IOp0AS = P2I->getOperand(i: 0)->getType()->getPointerAddressSpace();
343 unsigned I2PAS = I2P->getType()->getPointerAddressSpace();
344 return CastInst::isNoopCast(Opcode: Instruction::CastOps(I2P->getOpcode()),
345 SrcTy: I2P->getOperand(i: 0)->getType(), DstTy: I2P->getType(),
346 DL) &&
347 CastInst::isNoopCast(Opcode: Instruction::CastOps(P2I->getOpcode()),
348 SrcTy: P2I->getOperand(i: 0)->getType(), DstTy: P2I->getType(),
349 DL) &&
350 (P2IOp0AS == I2PAS || TTI->isNoopAddrSpaceCast(FromAS: P2IOp0AS, ToAS: I2PAS));
351}
352
353// Returns true if V is an address expression.
354// TODO: Currently, we only consider:
355// - arguments
356// - phi, bitcast, addrspacecast, and getelementptr operators
357bool InferAddressSpacesImpl::isAddressExpression(
358 const Value &V, const DataLayout &DL,
359 const TargetTransformInfo *TTI) const {
360
361 if (const Argument *Arg = dyn_cast<Argument>(Val: &V))
362 return Arg->getType()->isPointerTy() &&
363 TTI->getAssumedAddrSpace(V: &V) != UninitializedAddressSpace;
364
365 const Operator *Op = dyn_cast<Operator>(Val: &V);
366 if (!Op)
367 return false;
368
369 switch (Op->getOpcode()) {
370 case Instruction::PHI:
371 assert(Op->getType()->isPtrOrPtrVectorTy());
372 return true;
373 case Instruction::BitCast:
374 case Instruction::AddrSpaceCast:
375 case Instruction::GetElementPtr:
376 return true;
377 case Instruction::Select:
378 return Op->getType()->isPtrOrPtrVectorTy();
379 case Instruction::Call: {
380 const IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: &V);
381 return II && II->getIntrinsicID() == Intrinsic::ptrmask;
382 }
383 case Instruction::IntToPtr:
384 return isNoopPtrIntCastPair(I2P: Op, DL, TTI) ||
385 isSafeToCastIntToPtrAddrSpace(I2P: Op);
386 default:
387 // That value is an address expression if it has an assumed address space.
388 return TTI->getAssumedAddrSpace(V: &V) != UninitializedAddressSpace;
389 }
390}
391
392// Returns the pointer operands of V.
393//
394// Precondition: V is an address expression.
395SmallVector<Value *, 2> InferAddressSpacesImpl::getPointerOperands(
396 const Value &V, const DataLayout &DL,
397 const TargetTransformInfo *TTI) const {
398 if (isa<Argument>(Val: &V))
399 return {};
400
401 const Operator &Op = cast<Operator>(Val: V);
402 switch (Op.getOpcode()) {
403 case Instruction::PHI: {
404 auto IncomingValues = cast<PHINode>(Val: Op).incoming_values();
405 return {IncomingValues.begin(), IncomingValues.end()};
406 }
407 case Instruction::BitCast:
408 case Instruction::AddrSpaceCast:
409 case Instruction::GetElementPtr:
410 return {Op.getOperand(i: 0)};
411 case Instruction::Select:
412 return {Op.getOperand(i: 1), Op.getOperand(i: 2)};
413 case Instruction::Call: {
414 const IntrinsicInst &II = cast<IntrinsicInst>(Val: Op);
415 assert(II.getIntrinsicID() == Intrinsic::ptrmask &&
416 "unexpected intrinsic call");
417 return {II.getArgOperand(i: 0)};
418 }
419 case Instruction::IntToPtr: {
420 if (isNoopPtrIntCastPair(I2P: &Op, DL, TTI)) {
421 auto *P2I = cast<Operator>(Val: Op.getOperand(i: 0));
422 return {P2I->getOperand(i: 0)};
423 }
424 assert(isSafeToCastIntToPtrAddrSpace(&Op));
425 return {getIntToPtrPointerOperand(I2P: &Op)};
426 }
427 default:
428 llvm_unreachable("Unexpected instruction type.");
429 }
430}
431
432// Return mask. The 1 in mask indicate the bit is changed.
433// This helper function is to compute the max know changed bits for ptr1 and
434// ptr2 after the operation `ptr2 = ptr1 Op Mask`.
435static APInt computeMaxChangedPtrBits(const Operator *Op, const Value *Mask,
436 const DataLayout &DL, AssumptionCache *AC,
437 const DominatorTree *DT) {
438 KnownBits Known = computeKnownBits(V: Mask, DL, AC, CxtI: nullptr, DT);
439 switch (Op->getOpcode()) {
440 case Instruction::Xor:
441 case Instruction::Or:
442 return ~Known.Zero;
443 case Instruction::And:
444 return ~Known.One;
445 default:
446 return APInt::getAllOnes(numBits: Known.getBitWidth());
447 }
448}
449
450Value *
451InferAddressSpacesImpl::getIntToPtrPointerOperand(const Operator *I2P) const {
452 assert(I2P->getOpcode() == Instruction::IntToPtr);
453 if (I2P->getType()->isVectorTy())
454 return nullptr;
455
456 // If I2P has been accessed and has the corresponding old pointer value, just
457 // return true.
458 if (auto *OldPtr = PtrIntCastPairs.lookup(Val: I2P))
459 return OldPtr;
460
461 Value *LogicalOp = I2P->getOperand(i: 0);
462 Value *OldPtr, *Mask;
463 if (!match(V: LogicalOp,
464 P: m_c_BitwiseLogic(L: m_PtrToInt(Op: m_Value(V&: OldPtr)), R: m_Value(V&: Mask))))
465 return nullptr;
466
467 Operator *AsCast = dyn_cast<AddrSpaceCastOperator>(Val: OldPtr);
468 if (!AsCast)
469 return nullptr;
470
471 unsigned SrcAS = I2P->getType()->getPointerAddressSpace();
472 unsigned DstAS = AsCast->getOperand(i: 0)->getType()->getPointerAddressSpace();
473 APInt PreservedPtrMask = TTI->getAddrSpaceCastPreservedPtrMask(SrcAS, DstAS);
474 if (PreservedPtrMask.isZero())
475 return nullptr;
476 APInt ChangedPtrBits =
477 computeMaxChangedPtrBits(Op: cast<Operator>(Val: LogicalOp), Mask, DL: *DL, AC: &AC, DT);
478 // Check if the address bits change is within the preserved mask. If the bits
479 // change is not preserved, it is not safe to perform address space cast.
480 // The following pattern is not safe to cast address space.
481 // %1 = ptrtoint ptr addrspace(3) %sp to i32
482 // %2 = zext i32 %1 to i64
483 // %gp = inttoptr i64 %2 to ptr
484 assert(ChangedPtrBits.getBitWidth() == PreservedPtrMask.getBitWidth());
485 if (ChangedPtrBits.isSubsetOf(RHS: PreservedPtrMask))
486 return OldPtr;
487
488 return nullptr;
489}
490
491void InferAddressSpacesImpl::collectIntToPtrPointerOperand() {
492 // Only collect inttoptr instruction.
493 // TODO: We need to collect inttoptr constant expression as well.
494 for (Instruction &I : instructions(F)) {
495 if (!dyn_cast<IntToPtrInst>(Val: &I))
496 continue;
497 if (auto *OldPtr = getIntToPtrPointerOperand(I2P: cast<Operator>(Val: &I)))
498 PtrIntCastPairs.insert(KV: {&I, OldPtr});
499 }
500}
501
502bool InferAddressSpacesImpl::rewriteIntrinsicOperands(IntrinsicInst *II,
503 Value *OldV,
504 Value *NewV) const {
505 Module *M = II->getParent()->getParent()->getParent();
506 Intrinsic::ID IID = II->getIntrinsicID();
507 switch (IID) {
508 case Intrinsic::objectsize:
509 case Intrinsic::masked_load: {
510 Type *DestTy = II->getType();
511 Type *SrcTy = NewV->getType();
512 Function *NewDecl =
513 Intrinsic::getOrInsertDeclaration(M, id: IID, OverloadTys: {DestTy, SrcTy});
514 II->setArgOperand(i: 0, v: NewV);
515 II->setCalledFunction(NewDecl);
516 return true;
517 }
518 case Intrinsic::ptrmask:
519 // This is handled as an address expression, not as a use memory operation.
520 return false;
521 case Intrinsic::masked_gather: {
522 Type *RetTy = II->getType();
523 Type *NewPtrTy = NewV->getType();
524 Function *NewDecl =
525 Intrinsic::getOrInsertDeclaration(M, id: IID, OverloadTys: {RetTy, NewPtrTy});
526 II->setArgOperand(i: 0, v: NewV);
527 II->setCalledFunction(NewDecl);
528 return true;
529 }
530 case Intrinsic::masked_store:
531 case Intrinsic::masked_scatter: {
532 Type *ValueTy = II->getOperand(i_nocapture: 0)->getType();
533 Type *NewPtrTy = NewV->getType();
534 Function *NewDecl = Intrinsic::getOrInsertDeclaration(
535 M, id: II->getIntrinsicID(), OverloadTys: {ValueTy, NewPtrTy});
536 II->setArgOperand(i: 1, v: NewV);
537 II->setCalledFunction(NewDecl);
538 return true;
539 }
540 case Intrinsic::prefetch:
541 case Intrinsic::is_constant: {
542 Function *NewDecl = Intrinsic::getOrInsertDeclaration(
543 M, id: II->getIntrinsicID(), OverloadTys: {NewV->getType()});
544 II->setArgOperand(i: 0, v: NewV);
545 II->setCalledFunction(NewDecl);
546 return true;
547 }
548 case Intrinsic::fake_use: {
549 II->replaceUsesOfWith(From: OldV, To: NewV);
550 return true;
551 }
552 case Intrinsic::lifetime_start:
553 case Intrinsic::lifetime_end: {
554 // Always force lifetime markers to work directly on the alloca.
555 NewV = NewV->stripPointerCasts();
556 Function *NewDecl = Intrinsic::getOrInsertDeclaration(
557 M, id: II->getIntrinsicID(), OverloadTys: {NewV->getType()});
558 II->setArgOperand(i: 0, v: NewV);
559 II->setCalledFunction(NewDecl);
560 return true;
561 }
562 default: {
563 Value *Rewrite = TTI->rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
564 if (!Rewrite)
565 return false;
566 if (Rewrite != II)
567 II->replaceAllUsesWith(V: Rewrite);
568 return true;
569 }
570 }
571}
572
573void InferAddressSpacesImpl::collectRewritableIntrinsicOperands(
574 IntrinsicInst *II, PostorderStackTy &PostorderStack,
575 DenseSet<Value *> &Visited) const {
576 auto IID = II->getIntrinsicID();
577 switch (IID) {
578 case Intrinsic::ptrmask:
579 case Intrinsic::objectsize:
580 appendsFlatAddressExpressionToPostorderStack(V: II->getArgOperand(i: 0),
581 PostorderStack, Visited);
582 break;
583 case Intrinsic::is_constant: {
584 Value *Ptr = II->getArgOperand(i: 0);
585 if (Ptr->getType()->isPtrOrPtrVectorTy()) {
586 appendsFlatAddressExpressionToPostorderStack(V: Ptr, PostorderStack,
587 Visited);
588 }
589
590 break;
591 }
592 case Intrinsic::masked_load:
593 case Intrinsic::masked_gather:
594 case Intrinsic::prefetch:
595 appendsFlatAddressExpressionToPostorderStack(V: II->getArgOperand(i: 0),
596 PostorderStack, Visited);
597 break;
598 case Intrinsic::masked_store:
599 case Intrinsic::masked_scatter:
600 appendsFlatAddressExpressionToPostorderStack(V: II->getArgOperand(i: 1),
601 PostorderStack, Visited);
602 break;
603 case Intrinsic::fake_use: {
604 for (Value *Op : II->operands()) {
605 if (Op->getType()->isPtrOrPtrVectorTy()) {
606 appendsFlatAddressExpressionToPostorderStack(V: Op, PostorderStack,
607 Visited);
608 }
609 }
610
611 break;
612 }
613 case Intrinsic::lifetime_start:
614 case Intrinsic::lifetime_end: {
615 appendsFlatAddressExpressionToPostorderStack(V: II->getArgOperand(i: 0),
616 PostorderStack, Visited);
617 break;
618 }
619 default:
620 SmallVector<int, 2> OpIndexes;
621 if (TTI->collectFlatAddressOperands(OpIndexes, IID)) {
622 for (int Idx : OpIndexes) {
623 appendsFlatAddressExpressionToPostorderStack(V: II->getArgOperand(i: Idx),
624 PostorderStack, Visited);
625 }
626 }
627 break;
628 }
629}
630
631// Returns all flat address expressions in function F. The elements are
632// If V is an unvisited flat address expression, appends V to PostorderStack
633// and marks it as visited.
634void InferAddressSpacesImpl::appendsFlatAddressExpressionToPostorderStack(
635 Value *V, PostorderStackTy &PostorderStack,
636 DenseSet<Value *> &Visited) const {
637 assert(V->getType()->isPtrOrPtrVectorTy());
638
639 // Generic addressing expressions may be hidden in nested constant
640 // expressions.
641 if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Val: V)) {
642 // TODO: Look in non-address parts, like icmp operands.
643 if (isAddressExpression(V: *CE, DL: *DL, TTI) && Visited.insert(V: CE).second)
644 PostorderStack.emplace_back(Args&: CE, Args: false);
645
646 return;
647 }
648
649 if (V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
650 isAddressExpression(V: *V, DL: *DL, TTI)) {
651 if (Visited.insert(V).second) {
652 PostorderStack.emplace_back(Args&: V, Args: false);
653
654 if (auto *Op = dyn_cast<Operator>(Val: V))
655 for (auto &O : Op->operands())
656 if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Val&: O))
657 if (isAddressExpression(V: *CE, DL: *DL, TTI) && Visited.insert(V: CE).second)
658 PostorderStack.emplace_back(Args&: CE, Args: false);
659 }
660 }
661}
662
663// Returns all flat address expressions in function F. The elements are ordered
664// in postorder.
665std::vector<WeakTrackingVH>
666InferAddressSpacesImpl::collectFlatAddressExpressions(Function &F) const {
667 // This function implements a non-recursive postorder traversal of a partial
668 // use-def graph of function F.
669 PostorderStackTy PostorderStack;
670 // The set of visited expressions.
671 DenseSet<Value *> Visited;
672
673 auto PushPtrOperand = [&](Value *Ptr) {
674 appendsFlatAddressExpressionToPostorderStack(V: Ptr, PostorderStack, Visited);
675 };
676
677 // Look at operations that may be interesting accelerate by moving to a known
678 // address space. We aim at generating after loads and stores, but pure
679 // addressing calculations may also be faster.
680 for (Instruction &I : instructions(F)) {
681 if (auto *GEP = dyn_cast<GetElementPtrInst>(Val: &I)) {
682 PushPtrOperand(GEP->getPointerOperand());
683 } else if (auto *LI = dyn_cast<LoadInst>(Val: &I))
684 PushPtrOperand(LI->getPointerOperand());
685 else if (auto *SI = dyn_cast<StoreInst>(Val: &I))
686 PushPtrOperand(SI->getPointerOperand());
687 else if (auto *RMW = dyn_cast<AtomicRMWInst>(Val: &I))
688 PushPtrOperand(RMW->getPointerOperand());
689 else if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Val: &I))
690 PushPtrOperand(CmpX->getPointerOperand());
691 else if (auto *MI = dyn_cast<MemIntrinsic>(Val: &I)) {
692 // For memset/memcpy/memmove, any pointer operand can be replaced.
693 PushPtrOperand(MI->getRawDest());
694
695 // Handle 2nd operand for memcpy/memmove.
696 if (auto *MTI = dyn_cast<MemTransferInst>(Val: MI))
697 PushPtrOperand(MTI->getRawSource());
698 } else if (auto *II = dyn_cast<IntrinsicInst>(Val: &I))
699 collectRewritableIntrinsicOperands(II, PostorderStack, Visited);
700 else if (ICmpInst *Cmp = dyn_cast<ICmpInst>(Val: &I)) {
701 if (Cmp->getOperand(i_nocapture: 0)->getType()->isPtrOrPtrVectorTy()) {
702 PushPtrOperand(Cmp->getOperand(i_nocapture: 0));
703 PushPtrOperand(Cmp->getOperand(i_nocapture: 1));
704 }
705 } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(Val: &I)) {
706 PushPtrOperand(ASC->getPointerOperand());
707 } else if (auto *I2P = dyn_cast<IntToPtrInst>(Val: &I)) {
708 if (isNoopPtrIntCastPair(I2P: cast<Operator>(Val: I2P), DL: *DL, TTI))
709 PushPtrOperand(cast<Operator>(Val: I2P->getOperand(i_nocapture: 0))->getOperand(i: 0));
710 else if (isSafeToCastIntToPtrAddrSpace(I2P: cast<Operator>(Val: I2P)))
711 PushPtrOperand(getIntToPtrPointerOperand(I2P: cast<Operator>(Val: I2P)));
712 } else if (auto *RI = dyn_cast<ReturnInst>(Val: &I)) {
713 if (auto *RV = RI->getReturnValue();
714 RV && RV->getType()->isPtrOrPtrVectorTy())
715 PushPtrOperand(RV);
716 }
717 }
718
719 std::vector<WeakTrackingVH> Postorder; // The resultant postorder.
720 while (!PostorderStack.empty()) {
721 Value *TopVal = PostorderStack.back().getPointer();
722 // If the operands of the expression on the top are already explored,
723 // adds that expression to the resultant postorder.
724 if (PostorderStack.back().getInt()) {
725 if (TopVal->getType()->getPointerAddressSpace() == FlatAddrSpace)
726 Postorder.push_back(x: TopVal);
727 PostorderStack.pop_back();
728 continue;
729 }
730 // Otherwise, adds its operands to the stack and explores them.
731 PostorderStack.back().setInt(true);
732 // Skip values with an assumed address space.
733 if (TTI->getAssumedAddrSpace(V: TopVal) == UninitializedAddressSpace) {
734 for (Value *PtrOperand : getPointerOperands(V: *TopVal, DL: *DL, TTI)) {
735 appendsFlatAddressExpressionToPostorderStack(V: PtrOperand, PostorderStack,
736 Visited);
737 }
738 }
739 }
740 return Postorder;
741}
742
743// Inserts an addrspacecast for a phi node operand, handling the proper
744// insertion position based on the operand type.
745static Value *phiNodeOperandWithNewAddressSpace(AddrSpaceCastInst *NewI,
746 Value *Operand) {
747 auto InsertBefore = [NewI](auto It) {
748 NewI->insertBefore(It);
749 NewI->setDebugLoc(It->getDebugLoc());
750 return NewI;
751 };
752
753 if (auto *Arg = dyn_cast<Argument>(Val: Operand)) {
754 // For arguments, insert the cast at the beginning of entry block.
755 // Consider inserting at the dominating block for better placement.
756 Function *F = Arg->getParent();
757 auto InsertI = F->getEntryBlock().getFirstNonPHIIt();
758 return InsertBefore(InsertI);
759 }
760
761 // No check for Constant here, as constants are already handled.
762 assert(isa<Instruction>(Operand));
763
764 Instruction *OpInst = cast<Instruction>(Val: Operand);
765 if (LLVM_UNLIKELY(OpInst->getOpcode() == Instruction::PHI)) {
766 // If the operand is defined by another PHI node, insert after the first
767 // non-PHI instruction at the corresponding basic block.
768 auto InsertI = OpInst->getParent()->getFirstNonPHIIt();
769 return InsertBefore(InsertI);
770 }
771
772 // Otherwise, insert immediately after the operand definition.
773 NewI->insertAfter(InsertPos: OpInst->getIterator());
774 NewI->setDebugLoc(OpInst->getDebugLoc());
775 return NewI;
776}
777
778// A helper function for cloneInstructionWithNewAddressSpace. Returns the clone
779// of OperandUse.get() in the new address space. If the clone is not ready yet,
780// returns poison in the new address space as a placeholder.
781static Value *operandWithNewAddressSpaceOrCreatePoison(
782 const Use &OperandUse, unsigned NewAddrSpace,
783 const ValueToValueMapTy &ValueWithNewAddrSpace,
784 const PredicatedAddrSpaceMapTy &PredicatedAS,
785 SmallVectorImpl<const Use *> *PoisonUsesToFix) {
786 Value *Operand = OperandUse.get();
787
788 Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(Ty: Operand->getType(), NewAddrSpace);
789
790 if (Constant *C = dyn_cast<Constant>(Val: Operand))
791 return ConstantExpr::getAddrSpaceCast(C, Ty: NewPtrTy);
792
793 if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Val: Operand))
794 return NewOperand;
795
796 Instruction *Inst = cast<Instruction>(Val: OperandUse.getUser());
797 auto I = PredicatedAS.find(Val: std::make_pair(x&: Inst, y&: Operand));
798 if (I != PredicatedAS.end()) {
799 // Insert an addrspacecast on that operand before the user.
800 unsigned NewAS = I->second;
801 Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(Ty: Operand->getType(), NewAddrSpace: NewAS);
802 auto *NewI = new AddrSpaceCastInst(Operand, NewPtrTy);
803
804 if (LLVM_UNLIKELY(Inst->getOpcode() == Instruction::PHI))
805 return phiNodeOperandWithNewAddressSpace(NewI, Operand);
806
807 NewI->insertBefore(InsertPos: Inst->getIterator());
808 NewI->setDebugLoc(Inst->getDebugLoc());
809 return NewI;
810 }
811
812 PoisonUsesToFix->push_back(Elt: &OperandUse);
813 return PoisonValue::get(T: NewPtrTy);
814}
815
816// A helper function for cloneInstructionWithNewAddressSpace. Handles the
817// conversion of a ptrmask intrinsic instruction.
818Value *InferAddressSpacesImpl::clonePtrMaskWithNewAddressSpace(
819 IntrinsicInst *I, unsigned NewAddrSpace,
820 const ValueToValueMapTy &ValueWithNewAddrSpace,
821 const PredicatedAddrSpaceMapTy &PredicatedAS,
822 SmallVectorImpl<const Use *> *PoisonUsesToFix) const {
823 const Use &PtrOpUse = I->getArgOperandUse(i: 0);
824 unsigned OldAddrSpace = PtrOpUse->getType()->getPointerAddressSpace();
825 Value *MaskOp = I->getArgOperand(i: 1);
826 Type *MaskTy = MaskOp->getType();
827
828 KnownBits OldPtrBits{DL->getPointerSizeInBits(AS: OldAddrSpace)};
829 KnownBits NewPtrBits{DL->getPointerSizeInBits(AS: NewAddrSpace)};
830 if (!TTI->isNoopAddrSpaceCast(FromAS: OldAddrSpace, ToAS: NewAddrSpace)) {
831 std::tie(args&: OldPtrBits, args&: NewPtrBits) =
832 TTI->computeKnownBitsAddrSpaceCast(ToAS: NewAddrSpace, PtrOp: *PtrOpUse.get());
833 }
834
835 // If the pointers in both addrspaces have a bitwise representation and if the
836 // representation of the new pointer is smaller (fewer bits) than the old one,
837 // check if the mask is applicable to the ptr in the new addrspace. Any
838 // masking only clearing the low bits will also apply in the new addrspace
839 // Note: checking if the mask clears high bits is not sufficient as those
840 // might have already been 0 in the old ptr.
841 if (OldPtrBits.getBitWidth() > NewPtrBits.getBitWidth()) {
842 KnownBits MaskBits =
843 computeKnownBits(V: MaskOp, DL: *DL, /*AssumptionCache=*/AC: nullptr, CxtI: I);
844 // Set all unknown bits of the old ptr to 1, so that we are conservative in
845 // checking which bits are cleared by the mask.
846 OldPtrBits.One |= ~OldPtrBits.Zero;
847 // Check which bits are cleared by the mask in the old ptr.
848 KnownBits ClearedBits = KnownBits::sub(LHS: OldPtrBits, RHS: OldPtrBits & MaskBits);
849
850 // If the mask isn't applicable to the new ptr, leave the ptrmask as-is and
851 // insert an addrspacecast after it.
852 if (ClearedBits.countMaxActiveBits() > NewPtrBits.countMaxActiveBits()) {
853 std::optional<BasicBlock::iterator> InsertPoint =
854 I->getInsertionPointAfterDef();
855 assert(InsertPoint && "insertion after ptrmask should be possible");
856 Type *NewPtrType = getPtrOrVecOfPtrsWithNewAS(Ty: I->getType(), NewAddrSpace);
857 Instruction *AddrSpaceCast =
858 new AddrSpaceCastInst(I, NewPtrType, "", *InsertPoint);
859 AddrSpaceCast->setDebugLoc(I->getDebugLoc());
860 return AddrSpaceCast;
861 }
862 }
863
864 IRBuilder<> B(I);
865 if (NewPtrBits.getBitWidth() < MaskTy->getScalarSizeInBits()) {
866 MaskTy = MaskTy->getWithNewBitWidth(NewBitWidth: NewPtrBits.getBitWidth());
867 MaskOp = B.CreateTrunc(V: MaskOp, DestTy: MaskTy);
868 }
869 Value *NewPtr = operandWithNewAddressSpaceOrCreatePoison(
870 OperandUse: PtrOpUse, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS,
871 PoisonUsesToFix);
872 return B.CreateIntrinsic(ID: Intrinsic::ptrmask, OverloadTypes: {NewPtr->getType(), MaskTy},
873 Args: {NewPtr, MaskOp});
874}
875
876// Returns a clone of `I` with its operands converted to those specified in
877// ValueWithNewAddrSpace. Due to potential cycles in the data flow graph, an
878// operand whose address space needs to be modified might not exist in
879// ValueWithNewAddrSpace. In that case, uses poison as a placeholder operand and
880// adds that operand use to PoisonUsesToFix so that caller can fix them later.
881//
882// Note that we do not necessarily clone `I`, e.g., if it is an addrspacecast
883// from a pointer whose type already matches. Therefore, this function returns a
884// Value* instead of an Instruction*.
885Value *InferAddressSpacesImpl::cloneInstructionWithNewAddressSpace(
886 Instruction *I, unsigned NewAddrSpace,
887 const ValueToValueMapTy &ValueWithNewAddrSpace,
888 const PredicatedAddrSpaceMapTy &PredicatedAS,
889 SmallVectorImpl<const Use *> *PoisonUsesToFix) const {
890 Type *NewPtrType = getPtrOrVecOfPtrsWithNewAS(Ty: I->getType(), NewAddrSpace);
891
892 if (I->getOpcode() == Instruction::AddrSpaceCast) {
893 Value *Src = I->getOperand(i: 0);
894 // Because `I` is flat, the source address space must be specific.
895 // Therefore, the inferred address space must be the source space, according
896 // to our algorithm.
897 assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
898 return Src;
899 }
900
901 if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: I)) {
902 // Technically the intrinsic ID is a pointer typed argument, so specially
903 // handle calls early.
904 assert(II->getIntrinsicID() == Intrinsic::ptrmask);
905 return clonePtrMaskWithNewAddressSpace(
906 I: II, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS, PoisonUsesToFix);
907 }
908
909 unsigned AS = TTI->getAssumedAddrSpace(V: I);
910 if (AS != UninitializedAddressSpace) {
911 // For the assumed address space, insert an `addrspacecast` to make that
912 // explicit.
913 Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(Ty: I->getType(), NewAddrSpace: AS);
914 auto *NewI = new AddrSpaceCastInst(I, NewPtrTy);
915 NewI->insertAfter(InsertPos: I->getIterator());
916 NewI->setDebugLoc(I->getDebugLoc());
917 return NewI;
918 }
919
920 // Computes the converted pointer operands.
921 SmallVector<Value *, 4> NewPointerOperands;
922 for (const Use &OperandUse : I->operands()) {
923 if (!OperandUse.get()->getType()->isPtrOrPtrVectorTy())
924 NewPointerOperands.push_back(Elt: nullptr);
925 else
926 NewPointerOperands.push_back(Elt: operandWithNewAddressSpaceOrCreatePoison(
927 OperandUse, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS,
928 PoisonUsesToFix));
929 }
930
931 switch (I->getOpcode()) {
932 case Instruction::BitCast:
933 return new BitCastInst(NewPointerOperands[0], NewPtrType);
934 case Instruction::PHI: {
935 assert(I->getType()->isPtrOrPtrVectorTy());
936 PHINode *PHI = cast<PHINode>(Val: I);
937 PHINode *NewPHI = PHINode::Create(Ty: NewPtrType, NumReservedValues: PHI->getNumIncomingValues());
938 for (unsigned Index = 0; Index < PHI->getNumIncomingValues(); ++Index) {
939 unsigned OperandNo = PHINode::getOperandNumForIncomingValue(i: Index);
940 NewPHI->addIncoming(V: NewPointerOperands[OperandNo],
941 BB: PHI->getIncomingBlock(i: Index));
942 }
943 return NewPHI;
944 }
945 case Instruction::GetElementPtr: {
946 GetElementPtrInst *GEP = cast<GetElementPtrInst>(Val: I);
947 GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
948 PointeeType: GEP->getSourceElementType(), Ptr: NewPointerOperands[0],
949 IdxList: SmallVector<Value *, 4>(GEP->indices()));
950 NewGEP->setIsInBounds(GEP->isInBounds());
951 return NewGEP;
952 }
953 case Instruction::Select:
954 assert(I->getType()->isPtrOrPtrVectorTy());
955 return SelectInst::Create(C: I->getOperand(i: 0), S1: NewPointerOperands[1],
956 S2: NewPointerOperands[2], NameStr: "", InsertBefore: nullptr, MDFrom: I);
957 case Instruction::IntToPtr: {
958 if (isNoopPtrIntCastPair(I2P: cast<Operator>(Val: I), DL: *DL, TTI)) {
959 Value *Src = cast<Operator>(Val: I->getOperand(i: 0))->getOperand(i: 0);
960 if (Src->getType() == NewPtrType)
961 return Src;
962
963 // If we had a no-op inttoptr/ptrtoint pair, we may still have inferred a
964 // source address space from a generic pointer source need to insert a
965 // cast back.
966 return new AddrSpaceCastInst(Src, NewPtrType);
967 }
968 assert(isSafeToCastIntToPtrAddrSpace(cast<Operator>(I)));
969 AddrSpaceCastInst *AsCast = new AddrSpaceCastInst(I, NewPtrType);
970 AsCast->insertAfter(InsertPos: I);
971 return AsCast;
972 }
973 default:
974 llvm_unreachable("Unexpected opcode");
975 }
976}
977
978// Similar to cloneInstructionWithNewAddressSpace, returns a clone of the
979// constant expression `CE` with its operands replaced as specified in
980// ValueWithNewAddrSpace.
981Value *InferAddressSpacesImpl::cloneConstantExprWithNewAddressSpace(
982 ConstantExpr *CE, unsigned NewAddrSpace,
983 const ValueToValueMapTy &ValueWithNewAddrSpace, const DataLayout *DL,
984 const TargetTransformInfo *TTI) const {
985 Type *TargetType =
986 CE->getType()->isPtrOrPtrVectorTy()
987 ? getPtrOrVecOfPtrsWithNewAS(Ty: CE->getType(), NewAddrSpace)
988 : CE->getType();
989
990 if (CE->getOpcode() == Instruction::AddrSpaceCast) {
991 // Because CE is flat, the source address space must be specific.
992 // Therefore, the inferred address space must be the source space according
993 // to our algorithm.
994 assert(CE->getOperand(0)->getType()->getPointerAddressSpace() ==
995 NewAddrSpace);
996 return CE->getOperand(i_nocapture: 0);
997 }
998
999 if (CE->getOpcode() == Instruction::BitCast) {
1000 if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Val: CE->getOperand(i_nocapture: 0)))
1001 return ConstantExpr::getBitCast(C: cast<Constant>(Val: NewOperand), Ty: TargetType);
1002 return ConstantExpr::getAddrSpaceCast(C: CE, Ty: TargetType);
1003 }
1004
1005 if (CE->getOpcode() == Instruction::IntToPtr) {
1006 if (isNoopPtrIntCastPair(I2P: cast<Operator>(Val: CE), DL: *DL, TTI)) {
1007 Constant *Src = cast<ConstantExpr>(Val: CE->getOperand(i_nocapture: 0))->getOperand(i_nocapture: 0);
1008 assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
1009 return Src;
1010 }
1011 assert(isSafeToCastIntToPtrAddrSpace(cast<Operator>(CE)));
1012 return ConstantExpr::getAddrSpaceCast(C: CE, Ty: TargetType);
1013 }
1014
1015 // Computes the operands of the new constant expression.
1016 bool IsNew = false;
1017 SmallVector<Constant *, 4> NewOperands;
1018 for (unsigned Index = 0; Index < CE->getNumOperands(); ++Index) {
1019 Constant *Operand = CE->getOperand(i_nocapture: Index);
1020 // If the address space of `Operand` needs to be modified, the new operand
1021 // with the new address space should already be in ValueWithNewAddrSpace
1022 // because (1) the constant expressions we consider (i.e. addrspacecast,
1023 // bitcast, and getelementptr) do not incur cycles in the data flow graph
1024 // and (2) this function is called on constant expressions in postorder.
1025 if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Val: Operand)) {
1026 IsNew = true;
1027 NewOperands.push_back(Elt: cast<Constant>(Val: NewOperand));
1028 continue;
1029 }
1030 if (auto *CExpr = dyn_cast<ConstantExpr>(Val: Operand))
1031 if (Value *NewOperand = cloneConstantExprWithNewAddressSpace(
1032 CE: CExpr, NewAddrSpace, ValueWithNewAddrSpace, DL, TTI)) {
1033 IsNew = true;
1034 NewOperands.push_back(Elt: cast<Constant>(Val: NewOperand));
1035 continue;
1036 }
1037 // Otherwise, reuses the old operand.
1038 NewOperands.push_back(Elt: Operand);
1039 }
1040
1041 // If !IsNew, we will replace the Value with itself. However, replaced values
1042 // are assumed to wrapped in an addrspacecast cast later so drop it now.
1043 if (!IsNew)
1044 return nullptr;
1045
1046 if (CE->getOpcode() == Instruction::GetElementPtr) {
1047 // Needs to specify the source type while constructing a getelementptr
1048 // constant expression.
1049 return CE->getWithOperands(Ops: NewOperands, Ty: TargetType, /*OnlyIfReduced=*/false,
1050 SrcTy: cast<GEPOperator>(Val: CE)->getSourceElementType());
1051 }
1052
1053 return CE->getWithOperands(Ops: NewOperands, Ty: TargetType);
1054}
1055
1056// Returns a clone of the value `V`, with its operands replaced as specified in
1057// ValueWithNewAddrSpace. This function is called on every flat address
1058// expression whose address space needs to be modified, in postorder.
1059//
1060// See cloneInstructionWithNewAddressSpace for the meaning of PoisonUsesToFix.
1061Value *InferAddressSpacesImpl::cloneValueWithNewAddressSpace(
1062 Value *V, unsigned NewAddrSpace,
1063 const ValueToValueMapTy &ValueWithNewAddrSpace,
1064 const PredicatedAddrSpaceMapTy &PredicatedAS,
1065 SmallVectorImpl<const Use *> *PoisonUsesToFix) const {
1066 // All values in Postorder are flat address expressions.
1067 assert(V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
1068 isAddressExpression(*V, *DL, TTI));
1069
1070 if (auto *Arg = dyn_cast<Argument>(Val: V)) {
1071 // Arguments are address space casted in the function body, as we do not
1072 // want to change the function signature.
1073 Function *F = Arg->getParent();
1074 BasicBlock::iterator Insert = F->getEntryBlock().getFirstNonPHIIt();
1075
1076 Type *NewPtrTy = PointerType::get(C&: Arg->getContext(), AddressSpace: NewAddrSpace);
1077 auto *NewI = new AddrSpaceCastInst(Arg, NewPtrTy);
1078 NewI->insertBefore(InsertPos: Insert);
1079 return NewI;
1080 }
1081
1082 if (Instruction *I = dyn_cast<Instruction>(Val: V)) {
1083 Value *NewV = cloneInstructionWithNewAddressSpace(
1084 I, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS, PoisonUsesToFix);
1085 if (Instruction *NewI = dyn_cast_or_null<Instruction>(Val: NewV)) {
1086 if (NewI->getParent() == nullptr) {
1087 NewI->insertBefore(InsertPos: I->getIterator());
1088 NewI->takeName(V: I);
1089 NewI->setDebugLoc(I->getDebugLoc());
1090 }
1091 }
1092 return NewV;
1093 }
1094
1095 return cloneConstantExprWithNewAddressSpace(
1096 CE: cast<ConstantExpr>(Val: V), NewAddrSpace, ValueWithNewAddrSpace, DL, TTI);
1097}
1098
1099// Defines the join operation on the address space lattice (see the file header
1100// comments).
1101unsigned InferAddressSpacesImpl::joinAddressSpaces(unsigned AS1,
1102 unsigned AS2) const {
1103 if (AS1 == FlatAddrSpace || AS2 == FlatAddrSpace)
1104 return FlatAddrSpace;
1105
1106 if (AS1 == UninitializedAddressSpace)
1107 return AS2;
1108 if (AS2 == UninitializedAddressSpace)
1109 return AS1;
1110
1111 // The join of two different specific address spaces is flat.
1112 return (AS1 == AS2) ? AS1 : FlatAddrSpace;
1113}
1114
1115bool InferAddressSpacesImpl::run(Function &CurFn) {
1116 F = &CurFn;
1117 DL = &F->getDataLayout();
1118 PtrIntCastPairs.clear();
1119
1120 if (AssumeDefaultIsFlatAddressSpace)
1121 FlatAddrSpace = 0;
1122
1123 if (FlatAddrSpace == UninitializedAddressSpace) {
1124 FlatAddrSpace = TTI->getFlatAddressSpace();
1125 if (FlatAddrSpace == UninitializedAddressSpace)
1126 return false;
1127 }
1128
1129 collectIntToPtrPointerOperand();
1130 // Collects all flat address expressions in postorder.
1131 std::vector<WeakTrackingVH> Postorder = collectFlatAddressExpressions(F&: *F);
1132
1133 // Runs a data-flow analysis to refine the address spaces of every expression
1134 // in Postorder.
1135 ValueToAddrSpaceMapTy InferredAddrSpace;
1136 PredicatedAddrSpaceMapTy PredicatedAS;
1137 inferAddressSpaces(Postorder, InferredAddrSpace, PredicatedAS);
1138
1139 // Changes the address spaces of the flat address expressions who are inferred
1140 // to point to a specific address space.
1141 return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace,
1142 PredicatedAS);
1143}
1144
1145void InferAddressSpacesImpl::enqueueUsers(
1146 Value &V, const ValueToAddrSpaceMapTy &InferredAddrSpace,
1147 SetVector<Value *> &Worklist) const {
1148 for (Value *User : V.users()) {
1149 // Skip if User is already in the worklist.
1150 if (Worklist.count(key: User))
1151 continue;
1152
1153 ValueToAddrSpaceMapTy::const_iterator Pos = InferredAddrSpace.find(Val: User);
1154 // Our algorithm only updates the address spaces of flat address
1155 // expressions, which are those in InferredAddrSpace.
1156 if (Pos == InferredAddrSpace.end())
1157 continue;
1158
1159 // Function updateAddressSpace moves the address space down a lattice path.
1160 // Therefore, nothing to do if User is already inferred as flat (the bottom
1161 // element in the lattice).
1162 if (Pos->second == FlatAddrSpace)
1163 continue;
1164
1165 Worklist.insert(X: User);
1166 }
1167}
1168
1169void InferAddressSpacesImpl::runToFixPoint(
1170 SetVector<Value *> &Worklist, ValueToAddrSpaceMapTy &InferredAddrSpace,
1171 PredicatedAddrSpaceMapTy &PredicatedAS) const {
1172 while (!Worklist.empty()) {
1173 Value *V = Worklist.pop_back_val();
1174
1175 // Try to update the address space of the stack top according to the
1176 // address spaces of its operands.
1177 if (!updateAddressSpace(V: *V, InferredAddrSpace, PredicatedAS))
1178 continue;
1179
1180 enqueueUsers(V&: *V, InferredAddrSpace, Worklist);
1181 }
1182}
1183
1184// Constants need to be tracked through RAUW to handle cases with nested
1185// constant expressions, so wrap values in WeakTrackingVH.
1186void InferAddressSpacesImpl::inferAddressSpaces(
1187 ArrayRef<WeakTrackingVH> Postorder,
1188 ValueToAddrSpaceMapTy &InferredAddrSpace,
1189 PredicatedAddrSpaceMapTy &PredicatedAS) const {
1190 SetVector<Value *> Worklist(llvm::from_range, Postorder);
1191 // Initially, all expressions are in the uninitialized address space.
1192 for (Value *V : Postorder)
1193 InferredAddrSpace[V] = UninitializedAddressSpace;
1194
1195 runToFixPoint(Worklist, InferredAddrSpace, PredicatedAS);
1196
1197 // A value still uninitialized here is stuck in a cycle of uninitialized
1198 // values and carries no address space information. Lower it to flat so its
1199 // users join to flat, instead of being rewritten to reference an operand
1200 // that rewriteWithNewAddressSpaces() never converts.
1201 SmallVector<Value *, 4> Lowered;
1202 for (Value *V : Postorder) {
1203 ValueToAddrSpaceMapTy::iterator I = InferredAddrSpace.find(Val: V);
1204 if (I->second == UninitializedAddressSpace) {
1205 I->second = FlatAddrSpace;
1206 Lowered.push_back(Elt: V);
1207 }
1208 }
1209
1210 for (Value *V : Lowered)
1211 enqueueUsers(V&: *V, InferredAddrSpace, Worklist);
1212
1213 runToFixPoint(Worklist, InferredAddrSpace, PredicatedAS);
1214}
1215
1216unsigned
1217InferAddressSpacesImpl::getPredicatedAddrSpace(const Value &Ptr,
1218 const Value *UserCtx) const {
1219 const Instruction *UserCtxI = dyn_cast<Instruction>(Val: UserCtx);
1220 if (!UserCtxI)
1221 return UninitializedAddressSpace;
1222
1223 const Value *StrippedPtr = Ptr.stripInBoundsOffsets();
1224 for (auto &AssumeVH : AC.assumptionsFor(V: StrippedPtr)) {
1225 if (!AssumeVH)
1226 continue;
1227 CallInst *CI = cast<CallInst>(Val&: AssumeVH);
1228 if (!isValidAssumeForContext(I: CI, CxtI: UserCtxI, DT))
1229 continue;
1230
1231 const Value *Ptr;
1232 unsigned AS;
1233 std::tie(args&: Ptr, args&: AS) = TTI->getPredicatedAddrSpace(V: CI->getArgOperand(i: 0));
1234 if (Ptr)
1235 return AS;
1236 }
1237
1238 return UninitializedAddressSpace;
1239}
1240
1241bool InferAddressSpacesImpl::updateAddressSpace(
1242 const Value &V, ValueToAddrSpaceMapTy &InferredAddrSpace,
1243 PredicatedAddrSpaceMapTy &PredicatedAS) const {
1244 assert(InferredAddrSpace.count(&V));
1245
1246 LLVM_DEBUG(dbgs() << "Updating the address space of\n " << V << '\n');
1247
1248 // The new inferred address space equals the join of the address spaces
1249 // of all its pointer operands.
1250 unsigned NewAS = UninitializedAddressSpace;
1251
1252 // isAddressExpression should guarantee that V is an operator or an argument.
1253 assert(isa<Operator>(V) || isa<Argument>(V));
1254
1255 unsigned AS = TTI->getAssumedAddrSpace(V: &V);
1256 if (AS != UninitializedAddressSpace) {
1257 // Use the assumed address space directly.
1258 NewAS = AS;
1259 } else {
1260 // Otherwise, infer the address space from its pointer operands.
1261 SmallVector<Constant *, 2> ConstantPtrOps;
1262 SmallVector<Value *, 2> PtrOps = getPointerOperands(V, DL: *DL, TTI);
1263 for (Value *PtrOperand : PtrOps) {
1264 auto I = InferredAddrSpace.find(Val: PtrOperand);
1265 unsigned OperandAS;
1266 if (I == InferredAddrSpace.end()) {
1267 OperandAS = PtrOperand->getType()->getPointerAddressSpace();
1268 if (auto *C = dyn_cast<Constant>(Val: PtrOperand);
1269 C && OperandAS == FlatAddrSpace) {
1270 // Defer joining the address space of constant pointer operands.
1271 ConstantPtrOps.push_back(Elt: C);
1272 continue;
1273 }
1274 if (OperandAS == FlatAddrSpace) {
1275 // Check AC for assumption dominating V.
1276 unsigned AS = getPredicatedAddrSpace(Ptr: *PtrOperand, UserCtx: &V);
1277 if (AS != UninitializedAddressSpace) {
1278 LLVM_DEBUG(dbgs()
1279 << " deduce operand AS from the predicate addrspace "
1280 << AS << '\n');
1281 OperandAS = AS;
1282 // Record this use with the predicated AS.
1283 PredicatedAS[std::make_pair(x: &V, y&: PtrOperand)] = OperandAS;
1284 }
1285 }
1286 } else
1287 OperandAS = I->second;
1288
1289 // join(flat, *) = flat. So we can break if NewAS is already flat.
1290 NewAS = joinAddressSpaces(AS1: NewAS, AS2: OperandAS);
1291 if (NewAS == FlatAddrSpace)
1292 break;
1293 }
1294
1295 if (NewAS != FlatAddrSpace && NewAS != UninitializedAddressSpace) {
1296 if (any_of(Range&: ConstantPtrOps, P: [=](Constant *C) {
1297 return !isSafeToCastConstAddrSpace(C, NewAS);
1298 }))
1299 NewAS = FlatAddrSpace;
1300 }
1301
1302 // operator(flat const, flat const, ...) -> flat
1303 if (NewAS == UninitializedAddressSpace &&
1304 PtrOps.size() == ConstantPtrOps.size())
1305 NewAS = FlatAddrSpace;
1306 }
1307
1308 unsigned OldAS = InferredAddrSpace.lookup(Val: &V);
1309 assert(OldAS != FlatAddrSpace);
1310 if (OldAS == NewAS)
1311 return false;
1312
1313 // If any updates are made, grabs its users to the worklist because
1314 // their address spaces can also be possibly updated.
1315 LLVM_DEBUG(dbgs() << " to " << NewAS << '\n');
1316 InferredAddrSpace[&V] = NewAS;
1317 return true;
1318}
1319
1320/// Replace operand \p OpIdx in \p Inst, if the value is the same as \p OldVal
1321/// with \p NewVal.
1322static bool replaceOperandIfSame(Instruction *Inst, unsigned OpIdx,
1323 Value *OldVal, Value *NewVal) {
1324 Use &U = Inst->getOperandUse(i: OpIdx);
1325 if (U.get() == OldVal) {
1326 U.set(NewVal);
1327 return true;
1328 }
1329
1330 return false;
1331}
1332
1333template <typename InstrType>
1334static bool replaceSimplePointerUse(const TargetTransformInfo &TTI,
1335 InstrType *MemInstr, unsigned AddrSpace,
1336 Value *OldV, Value *NewV) {
1337 if (!MemInstr->isVolatile() || TTI.hasVolatileVariant(I: MemInstr, AddrSpace)) {
1338 return replaceOperandIfSame(MemInstr, InstrType::getPointerOperandIndex(),
1339 OldV, NewV);
1340 }
1341
1342 return false;
1343}
1344
1345/// If \p OldV is used as the pointer operand of a compatible memory operation
1346/// \p Inst, replaces the pointer operand with NewV.
1347///
1348/// This covers memory instructions with a single pointer operand that can have
1349/// its address space changed by simply mutating the use to a new value.
1350///
1351/// \p returns true the user replacement was made.
1352static bool replaceIfSimplePointerUse(const TargetTransformInfo &TTI,
1353 User *Inst, unsigned AddrSpace,
1354 Value *OldV, Value *NewV) {
1355 if (auto *LI = dyn_cast<LoadInst>(Val: Inst))
1356 return replaceSimplePointerUse(TTI, MemInstr: LI, AddrSpace, OldV, NewV);
1357
1358 if (auto *SI = dyn_cast<StoreInst>(Val: Inst))
1359 return replaceSimplePointerUse(TTI, MemInstr: SI, AddrSpace, OldV, NewV);
1360
1361 if (auto *RMW = dyn_cast<AtomicRMWInst>(Val: Inst))
1362 return replaceSimplePointerUse(TTI, MemInstr: RMW, AddrSpace, OldV, NewV);
1363
1364 if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Val: Inst))
1365 return replaceSimplePointerUse(TTI, MemInstr: CmpX, AddrSpace, OldV, NewV);
1366
1367 return false;
1368}
1369
1370/// Update memory intrinsic uses that require more complex processing than
1371/// simple memory instructions. These require re-mangling and may have multiple
1372/// pointer operands.
1373static bool handleMemIntrinsicPtrUse(MemIntrinsic *MI, Value *OldV,
1374 Value *NewV) {
1375 IRBuilder<> B(MI);
1376 if (auto *MSI = dyn_cast<MemSetInst>(Val: MI)) {
1377 B.CreateMemSet(Ptr: NewV, Val: MSI->getValue(), Size: MSI->getLength(), Align: MSI->getDestAlign(),
1378 isVolatile: false, // isVolatile
1379 AAInfo: MI->getAAMetadata());
1380 } else if (auto *MTI = dyn_cast<MemTransferInst>(Val: MI)) {
1381 Value *Src = MTI->getRawSource();
1382 Value *Dest = MTI->getRawDest();
1383
1384 // Be careful in case this is a self-to-self copy.
1385 if (Src == OldV)
1386 Src = NewV;
1387
1388 if (Dest == OldV)
1389 Dest = NewV;
1390
1391 if (auto *MCI = dyn_cast<MemCpyInst>(Val: MTI)) {
1392 if (MCI->isForceInlined())
1393 B.CreateMemCpyInline(Dst: Dest, DstAlign: MTI->getDestAlign(), Src,
1394 SrcAlign: MTI->getSourceAlign(), Size: MTI->getLength(),
1395 isVolatile: false, // isVolatile
1396 AAInfo: MI->getAAMetadata());
1397 else
1398 B.CreateMemCpy(Dst: Dest, DstAlign: MTI->getDestAlign(), Src, SrcAlign: MTI->getSourceAlign(),
1399 Size: MTI->getLength(),
1400 isVolatile: false, // isVolatile
1401 AAInfo: MI->getAAMetadata());
1402 } else {
1403 assert(isa<MemMoveInst>(MTI));
1404 B.CreateMemMove(Dst: Dest, DstAlign: MTI->getDestAlign(), Src, SrcAlign: MTI->getSourceAlign(),
1405 Size: MTI->getLength(),
1406 isVolatile: false, // isVolatile
1407 AAInfo: MI->getAAMetadata());
1408 }
1409 } else
1410 llvm_unreachable("unhandled MemIntrinsic");
1411
1412 MI->eraseFromParent();
1413 return true;
1414}
1415
1416// \p returns true if it is OK to change the address space of constant \p C with
1417// a ConstantExpr addrspacecast.
1418bool InferAddressSpacesImpl::isSafeToCastConstAddrSpace(Constant *C,
1419 unsigned NewAS) const {
1420 assert(NewAS != UninitializedAddressSpace);
1421
1422 unsigned SrcAS = C->getType()->getPointerAddressSpace();
1423 if (SrcAS == NewAS || isa<UndefValue>(Val: C))
1424 return true;
1425
1426 // Prevent illegal casts between different non-flat address spaces.
1427 if (SrcAS != FlatAddrSpace && NewAS != FlatAddrSpace)
1428 return false;
1429
1430 if (isa<ConstantPointerNull>(Val: C) || isa<ConstantAggregateZero>(Val: C))
1431 return true;
1432
1433 if (auto *Op = dyn_cast<Operator>(Val: C)) {
1434 // If we already have a constant addrspacecast, it should be safe to cast it
1435 // off.
1436 if (Op->getOpcode() == Instruction::AddrSpaceCast)
1437 return isSafeToCastConstAddrSpace(C: cast<Constant>(Val: Op->getOperand(i: 0)),
1438 NewAS);
1439
1440 if (Op->getOpcode() == Instruction::IntToPtr &&
1441 Op->getType()->getPointerAddressSpace() == FlatAddrSpace)
1442 return true;
1443 }
1444
1445 return false;
1446}
1447
1448static Value::use_iterator skipToNextUser(Value::use_iterator I,
1449 Value::use_iterator End) {
1450 User *CurUser = I->getUser();
1451 ++I;
1452
1453 while (I != End && I->getUser() == CurUser)
1454 ++I;
1455
1456 return I;
1457}
1458
1459void InferAddressSpacesImpl::performPointerReplacement(
1460 Value *V, Value *NewV, Use &U, ValueToValueMapTy &ValueWithNewAddrSpace,
1461 SmallVectorImpl<Instruction *> &DeadInstructions) const {
1462
1463 User *CurUser = U.getUser();
1464
1465 unsigned AddrSpace = V->getType()->getPointerAddressSpace();
1466 if (replaceIfSimplePointerUse(TTI: *TTI, Inst: CurUser, AddrSpace, OldV: V, NewV))
1467 return;
1468
1469 // Skip if the current user is the new value itself.
1470 if (CurUser == NewV)
1471 return;
1472
1473 auto *CurUserI = dyn_cast<Instruction>(Val: CurUser);
1474 if (!CurUserI || CurUserI->getFunction() != F)
1475 return;
1476
1477 // Handle more complex cases like intrinsic that need to be remangled.
1478 if (auto *MI = dyn_cast<MemIntrinsic>(Val: CurUser)) {
1479 if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, OldV: V, NewV))
1480 return;
1481 }
1482
1483 if (auto *II = dyn_cast<IntrinsicInst>(Val: CurUser)) {
1484 if (rewriteIntrinsicOperands(II, OldV: V, NewV))
1485 return;
1486 }
1487
1488 if (ICmpInst *Cmp = dyn_cast<ICmpInst>(Val: CurUserI)) {
1489 // If we can infer that both pointers are in the same addrspace,
1490 // transform e.g.
1491 // %cmp = icmp eq float* %p, %q
1492 // into
1493 // %cmp = icmp eq float addrspace(3)* %new_p, %new_q
1494
1495 unsigned NewAS = NewV->getType()->getPointerAddressSpace();
1496 int SrcIdx = U.getOperandNo();
1497 int OtherIdx = (SrcIdx == 0) ? 1 : 0;
1498 Value *OtherSrc = Cmp->getOperand(i_nocapture: OtherIdx);
1499
1500 if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(Val: OtherSrc)) {
1501 if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
1502 Cmp->setOperand(i_nocapture: OtherIdx, Val_nocapture: OtherNewV);
1503 Cmp->setOperand(i_nocapture: SrcIdx, Val_nocapture: NewV);
1504 return;
1505 }
1506 }
1507
1508 // Even if the type mismatches, we can cast the constant.
1509 if (auto *KOtherSrc = dyn_cast<Constant>(Val: OtherSrc)) {
1510 if (isSafeToCastConstAddrSpace(C: KOtherSrc, NewAS)) {
1511 Cmp->setOperand(i_nocapture: SrcIdx, Val_nocapture: NewV);
1512 Cmp->setOperand(i_nocapture: OtherIdx, Val_nocapture: ConstantExpr::getAddrSpaceCast(
1513 C: KOtherSrc, Ty: NewV->getType()));
1514 return;
1515 }
1516 }
1517 }
1518
1519 if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Val: CurUserI)) {
1520 unsigned NewAS = NewV->getType()->getPointerAddressSpace();
1521 if (ASC->getDestAddressSpace() == NewAS) {
1522 ASC->replaceAllUsesWith(V: NewV);
1523 DeadInstructions.push_back(Elt: ASC);
1524 return;
1525 }
1526 }
1527
1528 // Otherwise, replaces the use with flat(NewV).
1529 if (isa<Instruction>(Val: V) || isa<Instruction>(Val: NewV)) {
1530 // Don't create a copy of the original addrspacecast.
1531 if (U == V && isa<AddrSpaceCastInst>(Val: V))
1532 return;
1533
1534 // Insert the addrspacecast after NewV.
1535 BasicBlock::iterator InsertPos;
1536 if (Instruction *NewVInst = dyn_cast<Instruction>(Val: NewV))
1537 InsertPos = std::next(x: NewVInst->getIterator());
1538 else
1539 InsertPos = std::next(x: cast<Instruction>(Val: V)->getIterator());
1540
1541 while (isa<PHINode>(Val: InsertPos))
1542 ++InsertPos;
1543 // This instruction may contain multiple uses of V, update them all.
1544 CurUser->replaceUsesOfWith(
1545 From: V, To: new AddrSpaceCastInst(NewV, V->getType(), "", InsertPos));
1546 } else {
1547 CurUserI->replaceUsesOfWith(
1548 From: V, To: ConstantExpr::getAddrSpaceCast(C: cast<Constant>(Val: NewV), Ty: V->getType()));
1549 }
1550}
1551
1552bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
1553 ArrayRef<WeakTrackingVH> Postorder,
1554 const ValueToAddrSpaceMapTy &InferredAddrSpace,
1555 const PredicatedAddrSpaceMapTy &PredicatedAS) const {
1556 // For each address expression to be modified, creates a clone of it with its
1557 // pointer operands converted to the new address space. Since the pointer
1558 // operands are converted, the clone is naturally in the new address space by
1559 // construction.
1560 ValueToValueMapTy ValueWithNewAddrSpace;
1561 SmallVector<const Use *, 32> PoisonUsesToFix;
1562 for (Value *V : Postorder) {
1563 unsigned NewAddrSpace = InferredAddrSpace.lookup(Val: V);
1564
1565 // In some degenerate cases (e.g. invalid IR in unreachable code), we may
1566 // not even infer the value to have its original address space.
1567 if (NewAddrSpace == UninitializedAddressSpace)
1568 continue;
1569
1570 if (V->getType()->getPointerAddressSpace() != NewAddrSpace) {
1571 Value *New =
1572 cloneValueWithNewAddressSpace(V, NewAddrSpace, ValueWithNewAddrSpace,
1573 PredicatedAS, PoisonUsesToFix: &PoisonUsesToFix);
1574 if (New)
1575 ValueWithNewAddrSpace[V] = New;
1576 }
1577 }
1578
1579 if (ValueWithNewAddrSpace.empty())
1580 return false;
1581
1582 // Fixes all the poison uses generated by cloneInstructionWithNewAddressSpace.
1583 for (const Use *PoisonUse : PoisonUsesToFix) {
1584 User *V = PoisonUse->getUser();
1585 User *NewV = cast_or_null<User>(Val: ValueWithNewAddrSpace.lookup(Val: V));
1586 if (!NewV)
1587 continue;
1588
1589 unsigned OperandNo = PoisonUse->getOperandNo();
1590 assert(isa<PoisonValue>(NewV->getOperand(OperandNo)));
1591 WeakTrackingVH NewOp = ValueWithNewAddrSpace.lookup(Val: PoisonUse->get());
1592 assert(NewOp &&
1593 "poison replacements in ValueWithNewAddrSpace shouldn't be null");
1594 NewV->setOperand(i: OperandNo, Val: NewOp);
1595 }
1596
1597 SmallVector<Instruction *, 16> DeadInstructions;
1598 ValueToValueMapTy VMap;
1599 ValueMapper VMapper(VMap, RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
1600
1601 // Replaces the uses of the old address expressions with the new ones.
1602 for (const WeakTrackingVH &WVH : Postorder) {
1603 assert(WVH && "value was unexpectedly deleted");
1604 Value *V = WVH;
1605 Value *NewV = ValueWithNewAddrSpace.lookup(Val: V);
1606 if (NewV == nullptr)
1607 continue;
1608
1609 LLVM_DEBUG(dbgs() << "Replacing the uses of " << *V << "\n with\n "
1610 << *NewV << '\n');
1611
1612 if (Constant *C = dyn_cast<Constant>(Val: V)) {
1613 Constant *Replace =
1614 ConstantExpr::getAddrSpaceCast(C: cast<Constant>(Val: NewV), Ty: C->getType());
1615 if (C != Replace) {
1616 LLVM_DEBUG(dbgs() << "Inserting replacement const cast: " << Replace
1617 << ": " << *Replace << '\n');
1618 SmallVector<User *, 16> WorkList;
1619 for (User *U : make_early_inc_range(Range: C->users())) {
1620 if (auto *I = dyn_cast<Instruction>(Val: U)) {
1621 if (I->getFunction() == F)
1622 I->replaceUsesOfWith(From: C, To: Replace);
1623 } else {
1624 WorkList.append(in_start: U->user_begin(), in_end: U->user_end());
1625 }
1626 }
1627 if (!WorkList.empty()) {
1628 VMap[C] = Replace;
1629 DenseSet<User *> Visited{WorkList.begin(), WorkList.end()};
1630 while (!WorkList.empty()) {
1631 User *U = WorkList.pop_back_val();
1632 if (auto *I = dyn_cast<Instruction>(Val: U)) {
1633 if (I->getFunction() == F)
1634 VMapper.remapInstruction(I&: *I);
1635 continue;
1636 }
1637 for (User *U2 : U->users())
1638 if (Visited.insert(V: U2).second)
1639 WorkList.push_back(Elt: U2);
1640 }
1641 }
1642 V = Replace;
1643 }
1644 }
1645
1646 Value::use_iterator I, E, Next;
1647 for (I = V->use_begin(), E = V->use_end(); I != E;) {
1648 Use &U = *I;
1649
1650 // Some users may see the same pointer operand in multiple operands. Skip
1651 // to the next instruction.
1652 I = skipToNextUser(I, End: E);
1653
1654 performPointerReplacement(V, NewV, U, ValueWithNewAddrSpace,
1655 DeadInstructions);
1656 }
1657
1658 if (V->use_empty()) {
1659 if (Instruction *I = dyn_cast<Instruction>(Val: V))
1660 DeadInstructions.push_back(Elt: I);
1661 }
1662 }
1663
1664 for (Instruction *I : DeadInstructions)
1665 RecursivelyDeleteTriviallyDeadInstructions(V: I);
1666
1667 return true;
1668}
1669
1670bool InferAddressSpaces::runOnFunction(Function &F) {
1671 if (skipFunction(F))
1672 return false;
1673
1674 auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
1675 DominatorTree *DT = DTWP ? &DTWP->getDomTree() : nullptr;
1676 return InferAddressSpacesImpl(
1677 getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F), DT,
1678 &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F),
1679 FlatAddrSpace)
1680 .run(CurFn&: F);
1681}
1682
1683FunctionPass *llvm::createInferAddressSpacesPass(unsigned AddressSpace) {
1684 return new InferAddressSpaces(AddressSpace);
1685}
1686
1687InferAddressSpacesPass::InferAddressSpacesPass()
1688 : FlatAddrSpace(UninitializedAddressSpace) {}
1689InferAddressSpacesPass::InferAddressSpacesPass(unsigned AddressSpace)
1690 : FlatAddrSpace(AddressSpace) {}
1691
1692PreservedAnalyses InferAddressSpacesPass::run(Function &F,
1693 FunctionAnalysisManager &AM) {
1694 bool Changed =
1695 InferAddressSpacesImpl(AM.getResult<AssumptionAnalysis>(IR&: F),
1696 AM.getCachedResult<DominatorTreeAnalysis>(IR&: F),
1697 &AM.getResult<TargetIRAnalysis>(IR&: F), FlatAddrSpace)
1698 .run(CurFn&: F);
1699 if (Changed) {
1700 PreservedAnalyses PA;
1701 PA.preserveSet<CFGAnalyses>();
1702 return PA;
1703 }
1704 return PreservedAnalyses::all();
1705}
1706