contrib/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

   1 //===-- 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 //
  90 //===----------------------------------------------------------------------===//
  91
  92 #include "llvm/Transforms/Scalar.h"
  93 #include "llvm/ADT/DenseSet.h"
  94 #include "llvm/ADT/Optional.h"
  95 #include "llvm/ADT/SetVector.h"
  96 #include "llvm/Analysis/TargetTransformInfo.h"
  97 #include "llvm/IR/Function.h"
  98 #include "llvm/IR/InstIterator.h"
  99 #include "llvm/IR/Instructions.h"
 100 #include "llvm/IR/Operator.h"
 101 #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
 106 #define DEBUG_TYPE "infer-address-spaces"
 107
 108 using namespace llvm;
 109
 110 namespace {
 111 static const unsigned UninitializedAddressSpace = ~0u;
 112
 113 using ValueToAddrSpaceMapTy = DenseMap<const Value *, unsigned>;
 114
 115 /// \brief InferAddressSpaces
 116 class InferAddressSpaces : public FunctionPass {
 117   /// Target specific address space which uses of should be replaced if
 118   /// possible.
 119   unsigned FlatAddrSpace;
 120
 121 public:
 122   static char ID;
 123
 124   InferAddressSpaces() : FunctionPass(ID) {}
 125
 126   void getAnalysisUsage(AnalysisUsage &AU) const override {
 127     AU.setPreservesCFG();
 128     AU.addRequired<TargetTransformInfoWrapperPass>();
 129   }
 130
 131   bool runOnFunction(Function &F) override;
 132
 133 private:
 134   // Returns the new address space of V if updated; otherwise, returns None.
 135   Optional<unsigned>
 136   updateAddressSpace(const Value &V,
 137                      const ValueToAddrSpaceMapTy &InferredAddrSpace) const;
 138
 139   // Tries to infer the specific address space of each address expression in
 140   // Postorder.
 141   void inferAddressSpaces(const std::vector<Value *> &Postorder,
 142                           ValueToAddrSpaceMapTy *InferredAddrSpace) const;
 143
 144   bool isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const;
 145
 146   // Changes the flat address expressions in function F to point to specific
 147   // address spaces if InferredAddrSpace says so. Postorder is the postorder of
 148   // all flat expressions in the use-def graph of function F.
 149   bool
 150   rewriteWithNewAddressSpaces(const std::vector<Value *> &Postorder,
 151                               const ValueToAddrSpaceMapTy &InferredAddrSpace,
 152                               Function *F) const;
 153
 154   void appendsFlatAddressExpressionToPostorderStack(
 155     Value *V, std::vector<std::pair<Value *, bool>> &PostorderStack,
 156     DenseSet<Value *> &Visited) const;
 157
 158   bool rewriteIntrinsicOperands(IntrinsicInst *II,
 159                                 Value *OldV, Value *NewV) const;
 160   void collectRewritableIntrinsicOperands(
 161     IntrinsicInst *II,
 162     std::vector<std::pair<Value *, bool>> &PostorderStack,
 163     DenseSet<Value *> &Visited) const;
 164
 165   std::vector<Value *> collectFlatAddressExpressions(Function &F) const;
 166
 167   Value *cloneValueWithNewAddressSpace(
 168     Value *V, unsigned NewAddrSpace,
 169     const ValueToValueMapTy &ValueWithNewAddrSpace,
 170     SmallVectorImpl<const Use *> *UndefUsesToFix) const;
 171   unsigned joinAddressSpaces(unsigned AS1, unsigned AS2) const;
 172 };
 173 } // end anonymous namespace
 174
 175 char InferAddressSpaces::ID = 0;
 176
 177 namespace llvm {
 178 void initializeInferAddressSpacesPass(PassRegistry &);
 179 }
 180
 181 INITIALIZE_PASS(InferAddressSpaces, DEBUG_TYPE, "Infer address spaces",
 182                 false, false)
 183
 184 // Returns true if V is an address expression.
 185 // TODO: Currently, we consider only phi, bitcast, addrspacecast, and
 186 // getelementptr operators.
 187 static bool isAddressExpression(const Value &V) {
 188   if (!isa<Operator>(V))
 189     return false;
 190
 191   switch (cast<Operator>(V).getOpcode()) {
 192   case Instruction::PHI:
 193   case Instruction::BitCast:
 194   case Instruction::AddrSpaceCast:
 195   case Instruction::GetElementPtr:
 196   case Instruction::Select:
 197     return true;
 198   default:
 199     return false;
 200   }
 201 }
 202
 203 // Returns the pointer operands of V.
 204 //
 205 // Precondition: V is an address expression.
 206 static SmallVector<Value *, 2> getPointerOperands(const Value &V) {
 207   const Operator &Op = cast<Operator>(V);
 208   switch (Op.getOpcode()) {
 209   case Instruction::PHI: {
 210     auto IncomingValues = cast<PHINode>(Op).incoming_values();
 211     return SmallVector<Value *, 2>(IncomingValues.begin(),
 212                                    IncomingValues.end());
 213   }
 214   case Instruction::BitCast:
 215   case Instruction::AddrSpaceCast:
 216   case Instruction::GetElementPtr:
 217     return {Op.getOperand(0)};
 218   case Instruction::Select:
 219     return {Op.getOperand(1), Op.getOperand(2)};
 220   default:
 221     llvm_unreachable("Unexpected instruction type.");
 222   }
 223 }
 224
 225 // TODO: Move logic to TTI?
 226 bool InferAddressSpaces::rewriteIntrinsicOperands(IntrinsicInst *II,
 227                                                   Value *OldV,
 228                                                   Value *NewV) const {
 229   Module *M = II->getParent()->getParent()->getParent();
 230
 231   switch (II->getIntrinsicID()) {
 232   case Intrinsic::amdgcn_atomic_inc:
 233   case Intrinsic::amdgcn_atomic_dec:{
 234     const ConstantInt *IsVolatile = dyn_cast<ConstantInt>(II->getArgOperand(4));
 235     if (!IsVolatile || !IsVolatile->isNullValue())
 236       return false;
 237
 238     LLVM_FALLTHROUGH;
 239   }
 240   case Intrinsic::objectsize: {
 241     Type *DestTy = II->getType();
 242     Type *SrcTy = NewV->getType();
 243     Function *NewDecl =
 244         Intrinsic::getDeclaration(M, II->getIntrinsicID(), {DestTy, SrcTy});
 245     II->setArgOperand(0, NewV);
 246     II->setCalledFunction(NewDecl);
 247     return true;
 248   }
 249   default:
 250     return false;
 251   }
 252 }
 253
 254 // TODO: Move logic to TTI?
 255 void InferAddressSpaces::collectRewritableIntrinsicOperands(
 256     IntrinsicInst *II, std::vector<std::pair<Value *, bool>> &PostorderStack,
 257     DenseSet<Value *> &Visited) const {
 258   switch (II->getIntrinsicID()) {
 259   case Intrinsic::objectsize:
 260   case Intrinsic::amdgcn_atomic_inc:
 261   case Intrinsic::amdgcn_atomic_dec:
 262     appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(0),
 263                                                  PostorderStack, Visited);
 264     break;
 265   default:
 266     break;
 267   }
 268 }
 269
 270 // Returns all flat address expressions in function F. The elements are
 271 // If V is an unvisited flat address expression, appends V to PostorderStack
 272 // and marks it as visited.
 273 void InferAddressSpaces::appendsFlatAddressExpressionToPostorderStack(
 274     Value *V, std::vector<std::pair<Value *, bool>> &PostorderStack,
 275     DenseSet<Value *> &Visited) const {
 276   assert(V->getType()->isPointerTy());
 277   if (isAddressExpression(*V) &&
 278       V->getType()->getPointerAddressSpace() == FlatAddrSpace) {
 279     if (Visited.insert(V).second)
 280       PostorderStack.push_back(std::make_pair(V, false));
 281   }
 282 }
 283
 284 // Returns all flat address expressions in function F. The elements are ordered
 285 // ordered in postorder.
 286 std::vector<Value *>
 287 InferAddressSpaces::collectFlatAddressExpressions(Function &F) const {
 288   // This function implements a non-recursive postorder traversal of a partial
 289   // use-def graph of function F.
 290   std::vector<std::pair<Value *, bool>> PostorderStack;
 291   // The set of visited expressions.
 292   DenseSet<Value *> Visited;
 293
 294   auto PushPtrOperand = [&](Value *Ptr) {
 295     appendsFlatAddressExpressionToPostorderStack(Ptr, PostorderStack,
 296                                                  Visited);
 297   };
 298
 299   // Look at operations that may be interesting accelerate by moving to a known
 300   // address space. We aim at generating after loads and stores, but pure
 301   // addressing calculations may also be faster.
 302   for (Instruction &I : instructions(F)) {
 303     if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
 304       if (!GEP->getType()->isVectorTy())
 305         PushPtrOperand(GEP->getPointerOperand());
 306     } else if (auto *LI = dyn_cast<LoadInst>(&I))
 307       PushPtrOperand(LI->getPointerOperand());
 308     else if (auto *SI = dyn_cast<StoreInst>(&I))
 309       PushPtrOperand(SI->getPointerOperand());
 310     else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
 311       PushPtrOperand(RMW->getPointerOperand());
 312     else if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(&I))
 313       PushPtrOperand(CmpX->getPointerOperand());
 314     else if (auto *MI = dyn_cast<MemIntrinsic>(&I)) {
 315       // For memset/memcpy/memmove, any pointer operand can be replaced.
 316       PushPtrOperand(MI->getRawDest());
 317
 318       // Handle 2nd operand for memcpy/memmove.
 319       if (auto *MTI = dyn_cast<MemTransferInst>(MI))
 320         PushPtrOperand(MTI->getRawSource());
 321     } else if (auto *II = dyn_cast<IntrinsicInst>(&I))
 322       collectRewritableIntrinsicOperands(II, PostorderStack, Visited);
 323     else if (ICmpInst *Cmp = dyn_cast<ICmpInst>(&I)) {
 324       // FIXME: Handle vectors of pointers
 325       if (Cmp->getOperand(0)->getType()->isPointerTy()) {
 326         PushPtrOperand(Cmp->getOperand(0));
 327         PushPtrOperand(Cmp->getOperand(1));
 328       }
 329     }
 330   }
 331
 332   std::vector<Value *> Postorder; // The resultant postorder.
 333   while (!PostorderStack.empty()) {
 334     // If the operands of the expression on the top are already explored,
 335     // adds that expression to the resultant postorder.
 336     if (PostorderStack.back().second) {
 337       Postorder.push_back(PostorderStack.back().first);
 338       PostorderStack.pop_back();
 339       continue;
 340     }
 341     // Otherwise, adds its operands to the stack and explores them.
 342     PostorderStack.back().second = true;
 343     for (Value *PtrOperand : getPointerOperands(*PostorderStack.back().first)) {
 344       appendsFlatAddressExpressionToPostorderStack(PtrOperand, PostorderStack,
 345                                                    Visited);
 346     }
 347   }
 348   return Postorder;
 349 }
 350
 351 // A helper function for cloneInstructionWithNewAddressSpace. Returns the clone
 352 // of OperandUse.get() in the new address space. If the clone is not ready yet,
 353 // returns an undef in the new address space as a placeholder.
 354 static Value *operandWithNewAddressSpaceOrCreateUndef(
 355     const Use &OperandUse, unsigned NewAddrSpace,
 356     const ValueToValueMapTy &ValueWithNewAddrSpace,
 357     SmallVectorImpl<const Use *> *UndefUsesToFix) {
 358   Value *Operand = OperandUse.get();
 359
 360   Type *NewPtrTy =
 361       Operand->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
 362
 363   if (Constant *C = dyn_cast<Constant>(Operand))
 364     return ConstantExpr::getAddrSpaceCast(C, NewPtrTy);
 365
 366   if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand))
 367     return NewOperand;
 368
 369   UndefUsesToFix->push_back(&OperandUse);
 370   return UndefValue::get(NewPtrTy);
 371 }
 372
 373 // Returns a clone of `I` with its operands converted to those specified in
 374 // ValueWithNewAddrSpace. Due to potential cycles in the data flow graph, an
 375 // operand whose address space needs to be modified might not exist in
 376 // ValueWithNewAddrSpace. In that case, uses undef as a placeholder operand and
 377 // adds that operand use to UndefUsesToFix so that caller can fix them later.
 378 //
 379 // Note that we do not necessarily clone `I`, e.g., if it is an addrspacecast
 380 // from a pointer whose type already matches. Therefore, this function returns a
 381 // Value* instead of an Instruction*.
 382 static Value *cloneInstructionWithNewAddressSpace(
 383     Instruction *I, unsigned NewAddrSpace,
 384     const ValueToValueMapTy &ValueWithNewAddrSpace,
 385     SmallVectorImpl<const Use *> *UndefUsesToFix) {
 386   Type *NewPtrType =
 387       I->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
 388
 389   if (I->getOpcode() == Instruction::AddrSpaceCast) {
 390     Value *Src = I->getOperand(0);
 391     // Because `I` is flat, the source address space must be specific.
 392     // Therefore, the inferred address space must be the source space, according
 393     // to our algorithm.
 394     assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
 395     if (Src->getType() != NewPtrType)
 396       return new BitCastInst(Src, NewPtrType);
 397     return Src;
 398   }
 399
 400   // Computes the converted pointer operands.
 401   SmallVector<Value *, 4> NewPointerOperands;
 402   for (const Use &OperandUse : I->operands()) {
 403     if (!OperandUse.get()->getType()->isPointerTy())
 404       NewPointerOperands.push_back(nullptr);
 405     else
 406       NewPointerOperands.push_back(operandWithNewAddressSpaceOrCreateUndef(
 407                                      OperandUse, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix));
 408   }
 409
 410   switch (I->getOpcode()) {
 411   case Instruction::BitCast:
 412     return new BitCastInst(NewPointerOperands[0], NewPtrType);
 413   case Instruction::PHI: {
 414     assert(I->getType()->isPointerTy());
 415     PHINode *PHI = cast<PHINode>(I);
 416     PHINode *NewPHI = PHINode::Create(NewPtrType, PHI->getNumIncomingValues());
 417     for (unsigned Index = 0; Index < PHI->getNumIncomingValues(); ++Index) {
 418       unsigned OperandNo = PHINode::getOperandNumForIncomingValue(Index);
 419       NewPHI->addIncoming(NewPointerOperands[OperandNo],
 420                           PHI->getIncomingBlock(Index));
 421     }
 422     return NewPHI;
 423   }
 424   case Instruction::GetElementPtr: {
 425     GetElementPtrInst *GEP = cast<GetElementPtrInst>(I);
 426     GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
 427         GEP->getSourceElementType(), NewPointerOperands[0],
 428         SmallVector<Value *, 4>(GEP->idx_begin(), GEP->idx_end()));
 429     NewGEP->setIsInBounds(GEP->isInBounds());
 430     return NewGEP;
 431   }
 432   case Instruction::Select: {
 433     assert(I->getType()->isPointerTy());
 434     return SelectInst::Create(I->getOperand(0), NewPointerOperands[1],
 435                               NewPointerOperands[2], "", nullptr, I);
 436   }
 437   default:
 438     llvm_unreachable("Unexpected opcode");
 439   }
 440 }
 441
 442 // Similar to cloneInstructionWithNewAddressSpace, returns a clone of the
 443 // constant expression `CE` with its operands replaced as specified in
 444 // ValueWithNewAddrSpace.
 445 static Value *cloneConstantExprWithNewAddressSpace(
 446   ConstantExpr *CE, unsigned NewAddrSpace,
 447   const ValueToValueMapTy &ValueWithNewAddrSpace) {
 448   Type *TargetType =
 449     CE->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
 450
 451   if (CE->getOpcode() == Instruction::AddrSpaceCast) {
 452     // Because CE is flat, the source address space must be specific.
 453     // Therefore, the inferred address space must be the source space according
 454     // to our algorithm.
 455     assert(CE->getOperand(0)->getType()->getPointerAddressSpace() ==
 456            NewAddrSpace);
 457     return ConstantExpr::getBitCast(CE->getOperand(0), TargetType);
 458   }
 459
 460   if (CE->getOpcode() == Instruction::BitCast) {
 461     if (Value *NewOperand = ValueWithNewAddrSpace.lookup(CE->getOperand(0)))
 462       return ConstantExpr::getBitCast(cast<Constant>(NewOperand), TargetType);
 463     return ConstantExpr::getAddrSpaceCast(CE, TargetType);
 464   }
 465
 466   if (CE->getOpcode() == Instruction::Select) {
 467     Constant *Src0 = CE->getOperand(1);
 468     Constant *Src1 = CE->getOperand(2);
 469     if (Src0->getType()->getPointerAddressSpace() ==
 470         Src1->getType()->getPointerAddressSpace()) {
 471
 472       return ConstantExpr::getSelect(
 473           CE->getOperand(0), ConstantExpr::getAddrSpaceCast(Src0, TargetType),
 474           ConstantExpr::getAddrSpaceCast(Src1, TargetType));
 475     }
 476   }
 477
 478   // Computes the operands of the new constant expression.
 479   SmallVector<Constant *, 4> NewOperands;
 480   for (unsigned Index = 0; Index < CE->getNumOperands(); ++Index) {
 481     Constant *Operand = CE->getOperand(Index);
 482     // If the address space of `Operand` needs to be modified, the new operand
 483     // with the new address space should already be in ValueWithNewAddrSpace
 484     // because (1) the constant expressions we consider (i.e. addrspacecast,
 485     // bitcast, and getelementptr) do not incur cycles in the data flow graph
 486     // and (2) this function is called on constant expressions in postorder.
 487     if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand)) {
 488       NewOperands.push_back(cast<Constant>(NewOperand));
 489     } else {
 490       // Otherwise, reuses the old operand.
 491       NewOperands.push_back(Operand);
 492     }
 493   }
 494
 495   if (CE->getOpcode() == Instruction::GetElementPtr) {
 496     // Needs to specify the source type while constructing a getelementptr
 497     // constant expression.
 498     return CE->getWithOperands(
 499       NewOperands, TargetType, /*OnlyIfReduced=*/false,
 500       NewOperands[0]->getType()->getPointerElementType());
 501   }
 502
 503   return CE->getWithOperands(NewOperands, TargetType);
 504 }
 505
 506 // Returns a clone of the value `V`, with its operands replaced as specified in
 507 // ValueWithNewAddrSpace. This function is called on every flat address
 508 // expression whose address space needs to be modified, in postorder.
 509 //
 510 // See cloneInstructionWithNewAddressSpace for the meaning of UndefUsesToFix.
 511 Value *InferAddressSpaces::cloneValueWithNewAddressSpace(
 512   Value *V, unsigned NewAddrSpace,
 513   const ValueToValueMapTy &ValueWithNewAddrSpace,
 514   SmallVectorImpl<const Use *> *UndefUsesToFix) const {
 515   // All values in Postorder are flat address expressions.
 516   assert(isAddressExpression(*V) &&
 517          V->getType()->getPointerAddressSpace() == FlatAddrSpace);
 518
 519   if (Instruction *I = dyn_cast<Instruction>(V)) {
 520     Value *NewV = cloneInstructionWithNewAddressSpace(
 521       I, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix);
 522     if (Instruction *NewI = dyn_cast<Instruction>(NewV)) {
 523       if (NewI->getParent() == nullptr) {
 524         NewI->insertBefore(I);
 525         NewI->takeName(I);
 526       }
 527     }
 528     return NewV;
 529   }
 530
 531   return cloneConstantExprWithNewAddressSpace(
 532     cast<ConstantExpr>(V), NewAddrSpace, ValueWithNewAddrSpace);
 533 }
 534
 535 // Defines the join operation on the address space lattice (see the file header
 536 // comments).
 537 unsigned InferAddressSpaces::joinAddressSpaces(unsigned AS1,
 538                                                unsigned AS2) const {
 539   if (AS1 == FlatAddrSpace || AS2 == FlatAddrSpace)
 540     return FlatAddrSpace;
 541
 542   if (AS1 == UninitializedAddressSpace)
 543     return AS2;
 544   if (AS2 == UninitializedAddressSpace)
 545     return AS1;
 546
 547   // The join of two different specific address spaces is flat.
 548   return (AS1 == AS2) ? AS1 : FlatAddrSpace;
 549 }
 550
 551 bool InferAddressSpaces::runOnFunction(Function &F) {
 552   if (skipFunction(F))
 553     return false;
 554
 555   const TargetTransformInfo &TTI =
 556       getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 557   FlatAddrSpace = TTI.getFlatAddressSpace();
 558   if (FlatAddrSpace == UninitializedAddressSpace)
 559     return false;
 560
 561   // Collects all flat address expressions in postorder.
 562   std::vector<Value *> Postorder = collectFlatAddressExpressions(F);
 563
 564   // Runs a data-flow analysis to refine the address spaces of every expression
 565   // in Postorder.
 566   ValueToAddrSpaceMapTy InferredAddrSpace;
 567   inferAddressSpaces(Postorder, &InferredAddrSpace);
 568
 569   // Changes the address spaces of the flat address expressions who are inferred
 570   // to point to a specific address space.
 571   return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace, &F);
 572 }
 573
 574 void InferAddressSpaces::inferAddressSpaces(
 575     const std::vector<Value *> &Postorder,
 576     ValueToAddrSpaceMapTy *InferredAddrSpace) const {
 577   SetVector<Value *> Worklist(Postorder.begin(), Postorder.end());
 578   // Initially, all expressions are in the uninitialized address space.
 579   for (Value *V : Postorder)
 580     (*InferredAddrSpace)[V] = UninitializedAddressSpace;
 581
 582   while (!Worklist.empty()) {
 583     Value *V = Worklist.pop_back_val();
 584
 585     // Tries to update the address space of the stack top according to the
 586     // address spaces of its operands.
 587     DEBUG(dbgs() << "Updating the address space of\n  " << *V << '\n');
 588     Optional<unsigned> NewAS = updateAddressSpace(*V, *InferredAddrSpace);
 589     if (!NewAS.hasValue())
 590       continue;
 591     // If any updates are made, grabs its users to the worklist because
 592     // their address spaces can also be possibly updated.
 593     DEBUG(dbgs() << "  to " << NewAS.getValue() << '\n');
 594     (*InferredAddrSpace)[V] = NewAS.getValue();
 595
 596     for (Value *User : V->users()) {
 597       // Skip if User is already in the worklist.
 598       if (Worklist.count(User))
 599         continue;
 600
 601       auto Pos = InferredAddrSpace->find(User);
 602       // Our algorithm only updates the address spaces of flat address
 603       // expressions, which are those in InferredAddrSpace.
 604       if (Pos == InferredAddrSpace->end())
 605         continue;
 606
 607       // Function updateAddressSpace moves the address space down a lattice
 608       // path. Therefore, nothing to do if User is already inferred as flat (the
 609       // bottom element in the lattice).
 610       if (Pos->second == FlatAddrSpace)
 611         continue;
 612
 613       Worklist.insert(User);
 614     }
 615   }
 616 }
 617
 618 Optional<unsigned> InferAddressSpaces::updateAddressSpace(
 619     const Value &V, const ValueToAddrSpaceMapTy &InferredAddrSpace) const {
 620   assert(InferredAddrSpace.count(&V));
 621
 622   // The new inferred address space equals the join of the address spaces
 623   // of all its pointer operands.
 624   unsigned NewAS = UninitializedAddressSpace;
 625
 626   const Operator &Op = cast<Operator>(V);
 627   if (Op.getOpcode() == Instruction::Select) {
 628     Value *Src0 = Op.getOperand(1);
 629     Value *Src1 = Op.getOperand(2);
 630
 631     auto I = InferredAddrSpace.find(Src0);
 632     unsigned Src0AS = (I != InferredAddrSpace.end()) ?
 633       I->second : Src0->getType()->getPointerAddressSpace();
 634
 635     auto J = InferredAddrSpace.find(Src1);
 636     unsigned Src1AS = (J != InferredAddrSpace.end()) ?
 637       J->second : Src1->getType()->getPointerAddressSpace();
 638
 639     auto *C0 = dyn_cast<Constant>(Src0);
 640     auto *C1 = dyn_cast<Constant>(Src1);
 641
 642     // If one of the inputs is a constant, we may be able to do a constant
 643     // addrspacecast of it. Defer inferring the address space until the input
 644     // address space is known.
 645     if ((C1 && Src0AS == UninitializedAddressSpace) ||
 646         (C0 && Src1AS == UninitializedAddressSpace))
 647       return None;
 648
 649     if (C0 && isSafeToCastConstAddrSpace(C0, Src1AS))
 650       NewAS = Src1AS;
 651     else if (C1 && isSafeToCastConstAddrSpace(C1, Src0AS))
 652       NewAS = Src0AS;
 653     else
 654       NewAS = joinAddressSpaces(Src0AS, Src1AS);
 655   } else {
 656     for (Value *PtrOperand : getPointerOperands(V)) {
 657       auto I = InferredAddrSpace.find(PtrOperand);
 658       unsigned OperandAS = I != InferredAddrSpace.end() ?
 659         I->second : PtrOperand->getType()->getPointerAddressSpace();
 660
 661       // join(flat, *) = flat. So we can break if NewAS is already flat.
 662       NewAS = joinAddressSpaces(NewAS, OperandAS);
 663       if (NewAS == FlatAddrSpace)
 664         break;
 665     }
 666   }
 667
 668   unsigned OldAS = InferredAddrSpace.lookup(&V);
 669   assert(OldAS != FlatAddrSpace);
 670   if (OldAS == NewAS)
 671     return None;
 672   return NewAS;
 673 }
 674
 675 /// \p returns true if \p U is the pointer operand of a memory instruction with
 676 /// a single pointer operand that can have its address space changed by simply
 677 /// mutating the use to a new value.
 678 static bool isSimplePointerUseValidToReplace(Use &U) {
 679   User *Inst = U.getUser();
 680   unsigned OpNo = U.getOperandNo();
 681
 682   if (auto *LI = dyn_cast<LoadInst>(Inst))
 683     return OpNo == LoadInst::getPointerOperandIndex() && !LI->isVolatile();
 684
 685   if (auto *SI = dyn_cast<StoreInst>(Inst))
 686     return OpNo == StoreInst::getPointerOperandIndex() && !SI->isVolatile();
 687
 688   if (auto *RMW = dyn_cast<AtomicRMWInst>(Inst))
 689     return OpNo == AtomicRMWInst::getPointerOperandIndex() && !RMW->isVolatile();
 690
 691   if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Inst)) {
 692     return OpNo == AtomicCmpXchgInst::getPointerOperandIndex() &&
 693            !CmpX->isVolatile();
 694   }
 695
 696   return false;
 697 }
 698
 699 /// Update memory intrinsic uses that require more complex processing than
 700 /// simple memory instructions. Thse require re-mangling and may have multiple
 701 /// pointer operands.
 702 static bool handleMemIntrinsicPtrUse(MemIntrinsic *MI, Value *OldV,
 703                                      Value *NewV) {
 704   IRBuilder<> B(MI);
 705   MDNode *TBAA = MI->getMetadata(LLVMContext::MD_tbaa);
 706   MDNode *ScopeMD = MI->getMetadata(LLVMContext::MD_alias_scope);
 707   MDNode *NoAliasMD = MI->getMetadata(LLVMContext::MD_noalias);
 708
 709   if (auto *MSI = dyn_cast<MemSetInst>(MI)) {
 710     B.CreateMemSet(NewV, MSI->getValue(),
 711                    MSI->getLength(), MSI->getAlignment(),
 712                    false, // isVolatile
 713                    TBAA, ScopeMD, NoAliasMD);
 714   } else if (auto *MTI = dyn_cast<MemTransferInst>(MI)) {
 715     Value *Src = MTI->getRawSource();
 716     Value *Dest = MTI->getRawDest();
 717
 718     // Be careful in case this is a self-to-self copy.
 719     if (Src == OldV)
 720       Src = NewV;
 721
 722     if (Dest == OldV)
 723       Dest = NewV;
 724
 725     if (isa<MemCpyInst>(MTI)) {
 726       MDNode *TBAAStruct = MTI->getMetadata(LLVMContext::MD_tbaa_struct);
 727       B.CreateMemCpy(Dest, Src, MTI->getLength(),
 728                      MTI->getAlignment(),
 729                      false, // isVolatile
 730                      TBAA, TBAAStruct, ScopeMD, NoAliasMD);
 731     } else {
 732       assert(isa<MemMoveInst>(MTI));
 733       B.CreateMemMove(Dest, Src, MTI->getLength(),
 734                       MTI->getAlignment(),
 735                       false, // isVolatile
 736                       TBAA, ScopeMD, NoAliasMD);
 737     }
 738   } else
 739     llvm_unreachable("unhandled MemIntrinsic");
 740
 741   MI->eraseFromParent();
 742   return true;
 743 }
 744
 745 // \p returns true if it is OK to change the address space of constant \p C with
 746 // a ConstantExpr addrspacecast.
 747 bool InferAddressSpaces::isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const {
 748   assert(NewAS != UninitializedAddressSpace);
 749
 750   unsigned SrcAS = C->getType()->getPointerAddressSpace();
 751   if (SrcAS == NewAS || isa<UndefValue>(C))
 752     return true;
 753
 754   // Prevent illegal casts between different non-flat address spaces.
 755   if (SrcAS != FlatAddrSpace && NewAS != FlatAddrSpace)
 756     return false;
 757
 758   if (isa<ConstantPointerNull>(C))
 759     return true;
 760
 761   if (auto *Op = dyn_cast<Operator>(C)) {
 762     // If we already have a constant addrspacecast, it should be safe to cast it
 763     // off.
 764     if (Op->getOpcode() == Instruction::AddrSpaceCast)
 765       return isSafeToCastConstAddrSpace(cast<Constant>(Op->getOperand(0)), NewAS);
 766
 767     if (Op->getOpcode() == Instruction::IntToPtr &&
 768         Op->getType()->getPointerAddressSpace() == FlatAddrSpace)
 769       return true;
 770   }
 771
 772   return false;
 773 }
 774
 775 static Value::use_iterator skipToNextUser(Value::use_iterator I,
 776                                           Value::use_iterator End) {
 777   User *CurUser = I->getUser();
 778   ++I;
 779
 780   while (I != End && I->getUser() == CurUser)
 781     ++I;
 782
 783   return I;
 784 }
 785
 786 bool InferAddressSpaces::rewriteWithNewAddressSpaces(
 787   const std::vector<Value *> &Postorder,
 788   const ValueToAddrSpaceMapTy &InferredAddrSpace, Function *F) const {
 789   // For each address expression to be modified, creates a clone of it with its
 790   // pointer operands converted to the new address space. Since the pointer
 791   // operands are converted, the clone is naturally in the new address space by
 792   // construction.
 793   ValueToValueMapTy ValueWithNewAddrSpace;
 794   SmallVector<const Use *, 32> UndefUsesToFix;
 795   for (Value* V : Postorder) {
 796     unsigned NewAddrSpace = InferredAddrSpace.lookup(V);
 797     if (V->getType()->getPointerAddressSpace() != NewAddrSpace) {
 798       ValueWithNewAddrSpace[V] = cloneValueWithNewAddressSpace(
 799         V, NewAddrSpace, ValueWithNewAddrSpace, &UndefUsesToFix);
 800     }
 801   }
 802
 803   if (ValueWithNewAddrSpace.empty())
 804     return false;
 805
 806   // Fixes all the undef uses generated by cloneInstructionWithNewAddressSpace.
 807   for (const Use *UndefUse : UndefUsesToFix) {
 808     User *V = UndefUse->getUser();
 809     User *NewV = cast<User>(ValueWithNewAddrSpace.lookup(V));
 810     unsigned OperandNo = UndefUse->getOperandNo();
 811     assert(isa<UndefValue>(NewV->getOperand(OperandNo)));
 812     NewV->setOperand(OperandNo, ValueWithNewAddrSpace.lookup(UndefUse->get()));
 813   }
 814
 815   // Replaces the uses of the old address expressions with the new ones.
 816   for (Value *V : Postorder) {
 817     Value *NewV = ValueWithNewAddrSpace.lookup(V);
 818     if (NewV == nullptr)
 819       continue;
 820
 821     DEBUG(dbgs() << "Replacing the uses of " << *V
 822                  << "\n  with\n  " << *NewV << '\n');
 823
 824     Value::use_iterator I, E, Next;
 825     for (I = V->use_begin(), E = V->use_end(); I != E; ) {
 826       Use &U = *I;
 827
 828       // Some users may see the same pointer operand in multiple operands. Skip
 829       // to the next instruction.
 830       I = skipToNextUser(I, E);
 831
 832       if (isSimplePointerUseValidToReplace(U)) {
 833         // If V is used as the pointer operand of a compatible memory operation,
 834         // sets the pointer operand to NewV. This replacement does not change
 835         // the element type, so the resultant load/store is still valid.
 836         U.set(NewV);
 837         continue;
 838       }
 839
 840       User *CurUser = U.getUser();
 841       // Handle more complex cases like intrinsic that need to be remangled.
 842       if (auto *MI = dyn_cast<MemIntrinsic>(CurUser)) {
 843         if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, V, NewV))
 844           continue;
 845       }
 846
 847       if (auto *II = dyn_cast<IntrinsicInst>(CurUser)) {
 848         if (rewriteIntrinsicOperands(II, V, NewV))
 849           continue;
 850       }
 851
 852       if (isa<Instruction>(CurUser)) {
 853         if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUser)) {
 854           // If we can infer that both pointers are in the same addrspace,
 855           // transform e.g.
 856           //   %cmp = icmp eq float* %p, %q
 857           // into
 858           //   %cmp = icmp eq float addrspace(3)* %new_p, %new_q
 859
 860           unsigned NewAS = NewV->getType()->getPointerAddressSpace();
 861           int SrcIdx = U.getOperandNo();
 862           int OtherIdx = (SrcIdx == 0) ? 1 : 0;
 863           Value *OtherSrc = Cmp->getOperand(OtherIdx);
 864
 865           if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(OtherSrc)) {
 866             if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
 867               Cmp->setOperand(OtherIdx, OtherNewV);
 868               Cmp->setOperand(SrcIdx, NewV);
 869               continue;
 870             }
 871           }
 872
 873           // Even if the type mismatches, we can cast the constant.
 874           if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
 875             if (isSafeToCastConstAddrSpace(KOtherSrc, NewAS)) {
 876               Cmp->setOperand(SrcIdx, NewV);
 877               Cmp->setOperand(OtherIdx,
 878                 ConstantExpr::getAddrSpaceCast(KOtherSrc, NewV->getType()));
 879               continue;
 880             }
 881           }
 882         }
 883
 884         // Otherwise, replaces the use with flat(NewV).
 885         if (Instruction *I = dyn_cast<Instruction>(V)) {
 886           BasicBlock::iterator InsertPos = std::next(I->getIterator());
 887           while (isa<PHINode>(InsertPos))
 888             ++InsertPos;
 889           U.set(new AddrSpaceCastInst(NewV, V->getType(), "", &*InsertPos));
 890         } else {
 891           U.set(ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV),
 892                                                V->getType()));
 893         }
 894       }
 895     }
 896
 897     if (V->use_empty())
 898       RecursivelyDeleteTriviallyDeadInstructions(V);
 899   }
 900
 901   return true;
 902 }
 903
 904 FunctionPass *llvm::createInferAddressSpacesPass() {
 905   return new InferAddressSpaces();
 906 }