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/ADT/DenseSet.h"
  93 #include "llvm/ADT/Optional.h"
  94 #include "llvm/ADT/SetVector.h"
  95 #include "llvm/Analysis/TargetTransformInfo.h"
  96 #include "llvm/IR/Function.h"
  97 #include "llvm/IR/InstIterator.h"
  98 #include "llvm/IR/Instructions.h"
  99 #include "llvm/IR/Operator.h"
 100 #include "llvm/Support/Debug.h"
 101 #include "llvm/Support/raw_ostream.h"
 102 #include "llvm/Transforms/Scalar.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(ArrayRef<WeakTrackingVH> 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(ArrayRef<WeakTrackingVH> 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<WeakTrackingVH> 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->isZero())
 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
 278   // Generic addressing expressions may be hidden in nested constant
 279   // expressions.
 280   if (ConstantExpr *CE = dyn_cast<ConstantExpr>(V)) {
 281     // TODO: Look in non-address parts, like icmp operands.
 282     if (isAddressExpression(*CE) && Visited.insert(CE).second)
 283       PostorderStack.push_back(std::make_pair(CE, false));
 284
 285     return;
 286   }
 287
 288   if (isAddressExpression(*V) &&
 289       V->getType()->getPointerAddressSpace() == FlatAddrSpace) {
 290     if (Visited.insert(V).second) {
 291       PostorderStack.push_back(std::make_pair(V, false));
 292
 293       Operator *Op = cast<Operator>(V);
 294       for (unsigned I = 0, E = Op->getNumOperands(); I != E; ++I) {
 295         if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Op->getOperand(I))) {
 296           if (isAddressExpression(*CE) && Visited.insert(CE).second)
 297             PostorderStack.emplace_back(CE, false);
 298         }
 299       }
 300     }
 301   }
 302 }
 303
 304 // Returns all flat address expressions in function F. The elements are ordered
 305 // ordered in postorder.
 306 std::vector<WeakTrackingVH>
 307 InferAddressSpaces::collectFlatAddressExpressions(Function &F) const {
 308   // This function implements a non-recursive postorder traversal of a partial
 309   // use-def graph of function F.
 310   std::vector<std::pair<Value *, bool>> PostorderStack;
 311   // The set of visited expressions.
 312   DenseSet<Value *> Visited;
 313
 314   auto PushPtrOperand = [&](Value *Ptr) {
 315     appendsFlatAddressExpressionToPostorderStack(Ptr, PostorderStack,
 316                                                  Visited);
 317   };
 318
 319   // Look at operations that may be interesting accelerate by moving to a known
 320   // address space. We aim at generating after loads and stores, but pure
 321   // addressing calculations may also be faster.
 322   for (Instruction &I : instructions(F)) {
 323     if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
 324       if (!GEP->getType()->isVectorTy())
 325         PushPtrOperand(GEP->getPointerOperand());
 326     } else if (auto *LI = dyn_cast<LoadInst>(&I))
 327       PushPtrOperand(LI->getPointerOperand());
 328     else if (auto *SI = dyn_cast<StoreInst>(&I))
 329       PushPtrOperand(SI->getPointerOperand());
 330     else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
 331       PushPtrOperand(RMW->getPointerOperand());
 332     else if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(&I))
 333       PushPtrOperand(CmpX->getPointerOperand());
 334     else if (auto *MI = dyn_cast<MemIntrinsic>(&I)) {
 335       // For memset/memcpy/memmove, any pointer operand can be replaced.
 336       PushPtrOperand(MI->getRawDest());
 337
 338       // Handle 2nd operand for memcpy/memmove.
 339       if (auto *MTI = dyn_cast<MemTransferInst>(MI))
 340         PushPtrOperand(MTI->getRawSource());
 341     } else if (auto *II = dyn_cast<IntrinsicInst>(&I))
 342       collectRewritableIntrinsicOperands(II, PostorderStack, Visited);
 343     else if (ICmpInst *Cmp = dyn_cast<ICmpInst>(&I)) {
 344       // FIXME: Handle vectors of pointers
 345       if (Cmp->getOperand(0)->getType()->isPointerTy()) {
 346         PushPtrOperand(Cmp->getOperand(0));
 347         PushPtrOperand(Cmp->getOperand(1));
 348       }
 349     } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(&I)) {
 350       if (!ASC->getType()->isVectorTy())
 351         PushPtrOperand(ASC->getPointerOperand());
 352     }
 353   }
 354
 355   std::vector<WeakTrackingVH> Postorder; // The resultant postorder.
 356   while (!PostorderStack.empty()) {
 357     Value *TopVal = PostorderStack.back().first;
 358     // If the operands of the expression on the top are already explored,
 359     // adds that expression to the resultant postorder.
 360     if (PostorderStack.back().second) {
 361       if (TopVal->getType()->getPointerAddressSpace() == FlatAddrSpace)
 362         Postorder.push_back(TopVal);
 363       PostorderStack.pop_back();
 364       continue;
 365     }
 366     // Otherwise, adds its operands to the stack and explores them.
 367     PostorderStack.back().second = true;
 368     for (Value *PtrOperand : getPointerOperands(*TopVal)) {
 369       appendsFlatAddressExpressionToPostorderStack(PtrOperand, PostorderStack,
 370                                                    Visited);
 371     }
 372   }
 373   return Postorder;
 374 }
 375
 376 // A helper function for cloneInstructionWithNewAddressSpace. Returns the clone
 377 // of OperandUse.get() in the new address space. If the clone is not ready yet,
 378 // returns an undef in the new address space as a placeholder.
 379 static Value *operandWithNewAddressSpaceOrCreateUndef(
 380     const Use &OperandUse, unsigned NewAddrSpace,
 381     const ValueToValueMapTy &ValueWithNewAddrSpace,
 382     SmallVectorImpl<const Use *> *UndefUsesToFix) {
 383   Value *Operand = OperandUse.get();
 384
 385   Type *NewPtrTy =
 386       Operand->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
 387
 388   if (Constant *C = dyn_cast<Constant>(Operand))
 389     return ConstantExpr::getAddrSpaceCast(C, NewPtrTy);
 390
 391   if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand))
 392     return NewOperand;
 393
 394   UndefUsesToFix->push_back(&OperandUse);
 395   return UndefValue::get(NewPtrTy);
 396 }
 397
 398 // Returns a clone of `I` with its operands converted to those specified in
 399 // ValueWithNewAddrSpace. Due to potential cycles in the data flow graph, an
 400 // operand whose address space needs to be modified might not exist in
 401 // ValueWithNewAddrSpace. In that case, uses undef as a placeholder operand and
 402 // adds that operand use to UndefUsesToFix so that caller can fix them later.
 403 //
 404 // Note that we do not necessarily clone `I`, e.g., if it is an addrspacecast
 405 // from a pointer whose type already matches. Therefore, this function returns a
 406 // Value* instead of an Instruction*.
 407 static Value *cloneInstructionWithNewAddressSpace(
 408     Instruction *I, unsigned NewAddrSpace,
 409     const ValueToValueMapTy &ValueWithNewAddrSpace,
 410     SmallVectorImpl<const Use *> *UndefUsesToFix) {
 411   Type *NewPtrType =
 412       I->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
 413
 414   if (I->getOpcode() == Instruction::AddrSpaceCast) {
 415     Value *Src = I->getOperand(0);
 416     // Because `I` is flat, the source address space must be specific.
 417     // Therefore, the inferred address space must be the source space, according
 418     // to our algorithm.
 419     assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
 420     if (Src->getType() != NewPtrType)
 421       return new BitCastInst(Src, NewPtrType);
 422     return Src;
 423   }
 424
 425   // Computes the converted pointer operands.
 426   SmallVector<Value *, 4> NewPointerOperands;
 427   for (const Use &OperandUse : I->operands()) {
 428     if (!OperandUse.get()->getType()->isPointerTy())
 429       NewPointerOperands.push_back(nullptr);
 430     else
 431       NewPointerOperands.push_back(operandWithNewAddressSpaceOrCreateUndef(
 432                                      OperandUse, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix));
 433   }
 434
 435   switch (I->getOpcode()) {
 436   case Instruction::BitCast:
 437     return new BitCastInst(NewPointerOperands[0], NewPtrType);
 438   case Instruction::PHI: {
 439     assert(I->getType()->isPointerTy());
 440     PHINode *PHI = cast<PHINode>(I);
 441     PHINode *NewPHI = PHINode::Create(NewPtrType, PHI->getNumIncomingValues());
 442     for (unsigned Index = 0; Index < PHI->getNumIncomingValues(); ++Index) {
 443       unsigned OperandNo = PHINode::getOperandNumForIncomingValue(Index);
 444       NewPHI->addIncoming(NewPointerOperands[OperandNo],
 445                           PHI->getIncomingBlock(Index));
 446     }
 447     return NewPHI;
 448   }
 449   case Instruction::GetElementPtr: {
 450     GetElementPtrInst *GEP = cast<GetElementPtrInst>(I);
 451     GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
 452         GEP->getSourceElementType(), NewPointerOperands[0],
 453         SmallVector<Value *, 4>(GEP->idx_begin(), GEP->idx_end()));
 454     NewGEP->setIsInBounds(GEP->isInBounds());
 455     return NewGEP;
 456   }
 457   case Instruction::Select: {
 458     assert(I->getType()->isPointerTy());
 459     return SelectInst::Create(I->getOperand(0), NewPointerOperands[1],
 460                               NewPointerOperands[2], "", nullptr, I);
 461   }
 462   default:
 463     llvm_unreachable("Unexpected opcode");
 464   }
 465 }
 466
 467 // Similar to cloneInstructionWithNewAddressSpace, returns a clone of the
 468 // constant expression `CE` with its operands replaced as specified in
 469 // ValueWithNewAddrSpace.
 470 static Value *cloneConstantExprWithNewAddressSpace(
 471   ConstantExpr *CE, unsigned NewAddrSpace,
 472   const ValueToValueMapTy &ValueWithNewAddrSpace) {
 473   Type *TargetType =
 474     CE->getType()->getPointerElementType()->getPointerTo(NewAddrSpace);
 475
 476   if (CE->getOpcode() == Instruction::AddrSpaceCast) {
 477     // Because CE is flat, the source address space must be specific.
 478     // Therefore, the inferred address space must be the source space according
 479     // to our algorithm.
 480     assert(CE->getOperand(0)->getType()->getPointerAddressSpace() ==
 481            NewAddrSpace);
 482     return ConstantExpr::getBitCast(CE->getOperand(0), TargetType);
 483   }
 484
 485   if (CE->getOpcode() == Instruction::BitCast) {
 486     if (Value *NewOperand = ValueWithNewAddrSpace.lookup(CE->getOperand(0)))
 487       return ConstantExpr::getBitCast(cast<Constant>(NewOperand), TargetType);
 488     return ConstantExpr::getAddrSpaceCast(CE, TargetType);
 489   }
 490
 491   if (CE->getOpcode() == Instruction::Select) {
 492     Constant *Src0 = CE->getOperand(1);
 493     Constant *Src1 = CE->getOperand(2);
 494     if (Src0->getType()->getPointerAddressSpace() ==
 495         Src1->getType()->getPointerAddressSpace()) {
 496
 497       return ConstantExpr::getSelect(
 498           CE->getOperand(0), ConstantExpr::getAddrSpaceCast(Src0, TargetType),
 499           ConstantExpr::getAddrSpaceCast(Src1, TargetType));
 500     }
 501   }
 502
 503   // Computes the operands of the new constant expression.
 504   bool IsNew = false;
 505   SmallVector<Constant *, 4> NewOperands;
 506   for (unsigned Index = 0; Index < CE->getNumOperands(); ++Index) {
 507     Constant *Operand = CE->getOperand(Index);
 508     // If the address space of `Operand` needs to be modified, the new operand
 509     // with the new address space should already be in ValueWithNewAddrSpace
 510     // because (1) the constant expressions we consider (i.e. addrspacecast,
 511     // bitcast, and getelementptr) do not incur cycles in the data flow graph
 512     // and (2) this function is called on constant expressions in postorder.
 513     if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand)) {
 514       IsNew = true;
 515       NewOperands.push_back(cast<Constant>(NewOperand));
 516     } else {
 517       // Otherwise, reuses the old operand.
 518       NewOperands.push_back(Operand);
 519     }
 520   }
 521
 522   // If !IsNew, we will replace the Value with itself. However, replaced values
 523   // are assumed to wrapped in a addrspace cast later so drop it now.
 524   if (!IsNew)
 525     return nullptr;
 526
 527   if (CE->getOpcode() == Instruction::GetElementPtr) {
 528     // Needs to specify the source type while constructing a getelementptr
 529     // constant expression.
 530     return CE->getWithOperands(
 531       NewOperands, TargetType, /*OnlyIfReduced=*/false,
 532       NewOperands[0]->getType()->getPointerElementType());
 533   }
 534
 535   return CE->getWithOperands(NewOperands, TargetType);
 536 }
 537
 538 // Returns a clone of the value `V`, with its operands replaced as specified in
 539 // ValueWithNewAddrSpace. This function is called on every flat address
 540 // expression whose address space needs to be modified, in postorder.
 541 //
 542 // See cloneInstructionWithNewAddressSpace for the meaning of UndefUsesToFix.
 543 Value *InferAddressSpaces::cloneValueWithNewAddressSpace(
 544   Value *V, unsigned NewAddrSpace,
 545   const ValueToValueMapTy &ValueWithNewAddrSpace,
 546   SmallVectorImpl<const Use *> *UndefUsesToFix) const {
 547   // All values in Postorder are flat address expressions.
 548   assert(isAddressExpression(*V) &&
 549          V->getType()->getPointerAddressSpace() == FlatAddrSpace);
 550
 551   if (Instruction *I = dyn_cast<Instruction>(V)) {
 552     Value *NewV = cloneInstructionWithNewAddressSpace(
 553       I, NewAddrSpace, ValueWithNewAddrSpace, UndefUsesToFix);
 554     if (Instruction *NewI = dyn_cast<Instruction>(NewV)) {
 555       if (NewI->getParent() == nullptr) {
 556         NewI->insertBefore(I);
 557         NewI->takeName(I);
 558       }
 559     }
 560     return NewV;
 561   }
 562
 563   return cloneConstantExprWithNewAddressSpace(
 564     cast<ConstantExpr>(V), NewAddrSpace, ValueWithNewAddrSpace);
 565 }
 566
 567 // Defines the join operation on the address space lattice (see the file header
 568 // comments).
 569 unsigned InferAddressSpaces::joinAddressSpaces(unsigned AS1,
 570                                                unsigned AS2) const {
 571   if (AS1 == FlatAddrSpace || AS2 == FlatAddrSpace)
 572     return FlatAddrSpace;
 573
 574   if (AS1 == UninitializedAddressSpace)
 575     return AS2;
 576   if (AS2 == UninitializedAddressSpace)
 577     return AS1;
 578
 579   // The join of two different specific address spaces is flat.
 580   return (AS1 == AS2) ? AS1 : FlatAddrSpace;
 581 }
 582
 583 bool InferAddressSpaces::runOnFunction(Function &F) {
 584   if (skipFunction(F))
 585     return false;
 586
 587   const TargetTransformInfo &TTI =
 588       getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 589   FlatAddrSpace = TTI.getFlatAddressSpace();
 590   if (FlatAddrSpace == UninitializedAddressSpace)
 591     return false;
 592
 593   // Collects all flat address expressions in postorder.
 594   std::vector<WeakTrackingVH> Postorder = collectFlatAddressExpressions(F);
 595
 596   // Runs a data-flow analysis to refine the address spaces of every expression
 597   // in Postorder.
 598   ValueToAddrSpaceMapTy InferredAddrSpace;
 599   inferAddressSpaces(Postorder, &InferredAddrSpace);
 600
 601   // Changes the address spaces of the flat address expressions who are inferred
 602   // to point to a specific address space.
 603   return rewriteWithNewAddressSpaces(Postorder, InferredAddrSpace, &F);
 604 }
 605
 606 // Constants need to be tracked through RAUW to handle cases with nested
 607 // constant expressions, so wrap values in WeakTrackingVH.
 608 void InferAddressSpaces::inferAddressSpaces(
 609     ArrayRef<WeakTrackingVH> Postorder,
 610     ValueToAddrSpaceMapTy *InferredAddrSpace) const {
 611   SetVector<Value *> Worklist(Postorder.begin(), Postorder.end());
 612   // Initially, all expressions are in the uninitialized address space.
 613   for (Value *V : Postorder)
 614     (*InferredAddrSpace)[V] = UninitializedAddressSpace;
 615
 616   while (!Worklist.empty()) {
 617     Value *V = Worklist.pop_back_val();
 618
 619     // Tries to update the address space of the stack top according to the
 620     // address spaces of its operands.
 621     DEBUG(dbgs() << "Updating the address space of\n  " << *V << '\n');
 622     Optional<unsigned> NewAS = updateAddressSpace(*V, *InferredAddrSpace);
 623     if (!NewAS.hasValue())
 624       continue;
 625     // If any updates are made, grabs its users to the worklist because
 626     // their address spaces can also be possibly updated.
 627     DEBUG(dbgs() << "  to " << NewAS.getValue() << '\n');
 628     (*InferredAddrSpace)[V] = NewAS.getValue();
 629
 630     for (Value *User : V->users()) {
 631       // Skip if User is already in the worklist.
 632       if (Worklist.count(User))
 633         continue;
 634
 635       auto Pos = InferredAddrSpace->find(User);
 636       // Our algorithm only updates the address spaces of flat address
 637       // expressions, which are those in InferredAddrSpace.
 638       if (Pos == InferredAddrSpace->end())
 639         continue;
 640
 641       // Function updateAddressSpace moves the address space down a lattice
 642       // path. Therefore, nothing to do if User is already inferred as flat (the
 643       // bottom element in the lattice).
 644       if (Pos->second == FlatAddrSpace)
 645         continue;
 646
 647       Worklist.insert(User);
 648     }
 649   }
 650 }
 651
 652 Optional<unsigned> InferAddressSpaces::updateAddressSpace(
 653     const Value &V, const ValueToAddrSpaceMapTy &InferredAddrSpace) const {
 654   assert(InferredAddrSpace.count(&V));
 655
 656   // The new inferred address space equals the join of the address spaces
 657   // of all its pointer operands.
 658   unsigned NewAS = UninitializedAddressSpace;
 659
 660   const Operator &Op = cast<Operator>(V);
 661   if (Op.getOpcode() == Instruction::Select) {
 662     Value *Src0 = Op.getOperand(1);
 663     Value *Src1 = Op.getOperand(2);
 664
 665     auto I = InferredAddrSpace.find(Src0);
 666     unsigned Src0AS = (I != InferredAddrSpace.end()) ?
 667       I->second : Src0->getType()->getPointerAddressSpace();
 668
 669     auto J = InferredAddrSpace.find(Src1);
 670     unsigned Src1AS = (J != InferredAddrSpace.end()) ?
 671       J->second : Src1->getType()->getPointerAddressSpace();
 672
 673     auto *C0 = dyn_cast<Constant>(Src0);
 674     auto *C1 = dyn_cast<Constant>(Src1);
 675
 676     // If one of the inputs is a constant, we may be able to do a constant
 677     // addrspacecast of it. Defer inferring the address space until the input
 678     // address space is known.
 679     if ((C1 && Src0AS == UninitializedAddressSpace) ||
 680         (C0 && Src1AS == UninitializedAddressSpace))
 681       return None;
 682
 683     if (C0 && isSafeToCastConstAddrSpace(C0, Src1AS))
 684       NewAS = Src1AS;
 685     else if (C1 && isSafeToCastConstAddrSpace(C1, Src0AS))
 686       NewAS = Src0AS;
 687     else
 688       NewAS = joinAddressSpaces(Src0AS, Src1AS);
 689   } else {
 690     for (Value *PtrOperand : getPointerOperands(V)) {
 691       auto I = InferredAddrSpace.find(PtrOperand);
 692       unsigned OperandAS = I != InferredAddrSpace.end() ?
 693         I->second : PtrOperand->getType()->getPointerAddressSpace();
 694
 695       // join(flat, *) = flat. So we can break if NewAS is already flat.
 696       NewAS = joinAddressSpaces(NewAS, OperandAS);
 697       if (NewAS == FlatAddrSpace)
 698         break;
 699     }
 700   }
 701
 702   unsigned OldAS = InferredAddrSpace.lookup(&V);
 703   assert(OldAS != FlatAddrSpace);
 704   if (OldAS == NewAS)
 705     return None;
 706   return NewAS;
 707 }
 708
 709 /// \p returns true if \p U is the pointer operand of a memory instruction with
 710 /// a single pointer operand that can have its address space changed by simply
 711 /// mutating the use to a new value.
 712 static bool isSimplePointerUseValidToReplace(Use &U) {
 713   User *Inst = U.getUser();
 714   unsigned OpNo = U.getOperandNo();
 715
 716   if (auto *LI = dyn_cast<LoadInst>(Inst))
 717     return OpNo == LoadInst::getPointerOperandIndex() && !LI->isVolatile();
 718
 719   if (auto *SI = dyn_cast<StoreInst>(Inst))
 720     return OpNo == StoreInst::getPointerOperandIndex() && !SI->isVolatile();
 721
 722   if (auto *RMW = dyn_cast<AtomicRMWInst>(Inst))
 723     return OpNo == AtomicRMWInst::getPointerOperandIndex() && !RMW->isVolatile();
 724
 725   if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(Inst)) {
 726     return OpNo == AtomicCmpXchgInst::getPointerOperandIndex() &&
 727            !CmpX->isVolatile();
 728   }
 729
 730   return false;
 731 }
 732
 733 /// Update memory intrinsic uses that require more complex processing than
 734 /// simple memory instructions. Thse require re-mangling and may have multiple
 735 /// pointer operands.
 736 static bool handleMemIntrinsicPtrUse(MemIntrinsic *MI, Value *OldV,
 737                                      Value *NewV) {
 738   IRBuilder<> B(MI);
 739   MDNode *TBAA = MI->getMetadata(LLVMContext::MD_tbaa);
 740   MDNode *ScopeMD = MI->getMetadata(LLVMContext::MD_alias_scope);
 741   MDNode *NoAliasMD = MI->getMetadata(LLVMContext::MD_noalias);
 742
 743   if (auto *MSI = dyn_cast<MemSetInst>(MI)) {
 744     B.CreateMemSet(NewV, MSI->getValue(),
 745                    MSI->getLength(), MSI->getAlignment(),
 746                    false, // isVolatile
 747                    TBAA, ScopeMD, NoAliasMD);
 748   } else if (auto *MTI = dyn_cast<MemTransferInst>(MI)) {
 749     Value *Src = MTI->getRawSource();
 750     Value *Dest = MTI->getRawDest();
 751
 752     // Be careful in case this is a self-to-self copy.
 753     if (Src == OldV)
 754       Src = NewV;
 755
 756     if (Dest == OldV)
 757       Dest = NewV;
 758
 759     if (isa<MemCpyInst>(MTI)) {
 760       MDNode *TBAAStruct = MTI->getMetadata(LLVMContext::MD_tbaa_struct);
 761       B.CreateMemCpy(Dest, Src, MTI->getLength(),
 762                      MTI->getAlignment(),
 763                      false, // isVolatile
 764                      TBAA, TBAAStruct, ScopeMD, NoAliasMD);
 765     } else {
 766       assert(isa<MemMoveInst>(MTI));
 767       B.CreateMemMove(Dest, Src, MTI->getLength(),
 768                       MTI->getAlignment(),
 769                       false, // isVolatile
 770                       TBAA, ScopeMD, NoAliasMD);
 771     }
 772   } else
 773     llvm_unreachable("unhandled MemIntrinsic");
 774
 775   MI->eraseFromParent();
 776   return true;
 777 }
 778
 779 // \p returns true if it is OK to change the address space of constant \p C with
 780 // a ConstantExpr addrspacecast.
 781 bool InferAddressSpaces::isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const {
 782   assert(NewAS != UninitializedAddressSpace);
 783
 784   unsigned SrcAS = C->getType()->getPointerAddressSpace();
 785   if (SrcAS == NewAS || isa<UndefValue>(C))
 786     return true;
 787
 788   // Prevent illegal casts between different non-flat address spaces.
 789   if (SrcAS != FlatAddrSpace && NewAS != FlatAddrSpace)
 790     return false;
 791
 792   if (isa<ConstantPointerNull>(C))
 793     return true;
 794
 795   if (auto *Op = dyn_cast<Operator>(C)) {
 796     // If we already have a constant addrspacecast, it should be safe to cast it
 797     // off.
 798     if (Op->getOpcode() == Instruction::AddrSpaceCast)
 799       return isSafeToCastConstAddrSpace(cast<Constant>(Op->getOperand(0)), NewAS);
 800
 801     if (Op->getOpcode() == Instruction::IntToPtr &&
 802         Op->getType()->getPointerAddressSpace() == FlatAddrSpace)
 803       return true;
 804   }
 805
 806   return false;
 807 }
 808
 809 static Value::use_iterator skipToNextUser(Value::use_iterator I,
 810                                           Value::use_iterator End) {
 811   User *CurUser = I->getUser();
 812   ++I;
 813
 814   while (I != End && I->getUser() == CurUser)
 815     ++I;
 816
 817   return I;
 818 }
 819
 820 bool InferAddressSpaces::rewriteWithNewAddressSpaces(
 821     ArrayRef<WeakTrackingVH> Postorder,
 822     const ValueToAddrSpaceMapTy &InferredAddrSpace, Function *F) const {
 823   // For each address expression to be modified, creates a clone of it with its
 824   // pointer operands converted to the new address space. Since the pointer
 825   // operands are converted, the clone is naturally in the new address space by
 826   // construction.
 827   ValueToValueMapTy ValueWithNewAddrSpace;
 828   SmallVector<const Use *, 32> UndefUsesToFix;
 829   for (Value* V : Postorder) {
 830     unsigned NewAddrSpace = InferredAddrSpace.lookup(V);
 831     if (V->getType()->getPointerAddressSpace() != NewAddrSpace) {
 832       ValueWithNewAddrSpace[V] = cloneValueWithNewAddressSpace(
 833         V, NewAddrSpace, ValueWithNewAddrSpace, &UndefUsesToFix);
 834     }
 835   }
 836
 837   if (ValueWithNewAddrSpace.empty())
 838     return false;
 839
 840   // Fixes all the undef uses generated by cloneInstructionWithNewAddressSpace.
 841   for (const Use *UndefUse : UndefUsesToFix) {
 842     User *V = UndefUse->getUser();
 843     User *NewV = cast<User>(ValueWithNewAddrSpace.lookup(V));
 844     unsigned OperandNo = UndefUse->getOperandNo();
 845     assert(isa<UndefValue>(NewV->getOperand(OperandNo)));
 846     NewV->setOperand(OperandNo, ValueWithNewAddrSpace.lookup(UndefUse->get()));
 847   }
 848
 849   SmallVector<Instruction *, 16> DeadInstructions;
 850
 851   // Replaces the uses of the old address expressions with the new ones.
 852   for (const WeakTrackingVH &WVH : Postorder) {
 853     assert(WVH && "value was unexpectedly deleted");
 854     Value *V = WVH;
 855     Value *NewV = ValueWithNewAddrSpace.lookup(V);
 856     if (NewV == nullptr)
 857       continue;
 858
 859     DEBUG(dbgs() << "Replacing the uses of " << *V
 860                  << "\n  with\n  " << *NewV << '\n');
 861
 862     if (Constant *C = dyn_cast<Constant>(V)) {
 863       Constant *Replace = ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV),
 864                                                          C->getType());
 865       if (C != Replace) {
 866         DEBUG(dbgs() << "Inserting replacement const cast: "
 867               << Replace << ": " << *Replace << '\n');
 868         C->replaceAllUsesWith(Replace);
 869         V = Replace;
 870       }
 871     }
 872
 873     Value::use_iterator I, E, Next;
 874     for (I = V->use_begin(), E = V->use_end(); I != E; ) {
 875       Use &U = *I;
 876
 877       // Some users may see the same pointer operand in multiple operands. Skip
 878       // to the next instruction.
 879       I = skipToNextUser(I, E);
 880
 881       if (isSimplePointerUseValidToReplace(U)) {
 882         // If V is used as the pointer operand of a compatible memory operation,
 883         // sets the pointer operand to NewV. This replacement does not change
 884         // the element type, so the resultant load/store is still valid.
 885         U.set(NewV);
 886         continue;
 887       }
 888
 889       User *CurUser = U.getUser();
 890       // Handle more complex cases like intrinsic that need to be remangled.
 891       if (auto *MI = dyn_cast<MemIntrinsic>(CurUser)) {
 892         if (!MI->isVolatile() && handleMemIntrinsicPtrUse(MI, V, NewV))
 893           continue;
 894       }
 895
 896       if (auto *II = dyn_cast<IntrinsicInst>(CurUser)) {
 897         if (rewriteIntrinsicOperands(II, V, NewV))
 898           continue;
 899       }
 900
 901       if (isa<Instruction>(CurUser)) {
 902         if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUser)) {
 903           // If we can infer that both pointers are in the same addrspace,
 904           // transform e.g.
 905           //   %cmp = icmp eq float* %p, %q
 906           // into
 907           //   %cmp = icmp eq float addrspace(3)* %new_p, %new_q
 908
 909           unsigned NewAS = NewV->getType()->getPointerAddressSpace();
 910           int SrcIdx = U.getOperandNo();
 911           int OtherIdx = (SrcIdx == 0) ? 1 : 0;
 912           Value *OtherSrc = Cmp->getOperand(OtherIdx);
 913
 914           if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(OtherSrc)) {
 915             if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
 916               Cmp->setOperand(OtherIdx, OtherNewV);
 917               Cmp->setOperand(SrcIdx, NewV);
 918               continue;
 919             }
 920           }
 921
 922           // Even if the type mismatches, we can cast the constant.
 923           if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
 924             if (isSafeToCastConstAddrSpace(KOtherSrc, NewAS)) {
 925               Cmp->setOperand(SrcIdx, NewV);
 926               Cmp->setOperand(OtherIdx,
 927                 ConstantExpr::getAddrSpaceCast(KOtherSrc, NewV->getType()));
 928               continue;
 929             }
 930           }
 931         }
 932
 933         if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(CurUser)) {
 934           unsigned NewAS = NewV->getType()->getPointerAddressSpace();
 935           if (ASC->getDestAddressSpace() == NewAS) {
 936             ASC->replaceAllUsesWith(NewV);
 937             DeadInstructions.push_back(ASC);
 938             continue;
 939           }
 940         }
 941
 942         // Otherwise, replaces the use with flat(NewV).
 943         if (Instruction *I = dyn_cast<Instruction>(V)) {
 944           BasicBlock::iterator InsertPos = std::next(I->getIterator());
 945           while (isa<PHINode>(InsertPos))
 946             ++InsertPos;
 947           U.set(new AddrSpaceCastInst(NewV, V->getType(), "", &*InsertPos));
 948         } else {
 949           U.set(ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV),
 950                                                V->getType()));
 951         }
 952       }
 953     }
 954
 955     if (V->use_empty()) {
 956       if (Instruction *I = dyn_cast<Instruction>(V))
 957         DeadInstructions.push_back(I);
 958     }
 959   }
 960
 961   for (Instruction *I : DeadInstructions)
 962     RecursivelyDeleteTriviallyDeadInstructions(I);
 963
 964   return true;
 965 }
 966
 967 FunctionPass *llvm::createInferAddressSpacesPass() {
 968   return new InferAddressSpaces();
 969 }