[InferAddressSpaces] Handle vector of pointers type

llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

 
 INITIALIZE_PASS_BEGIN(InferAddressSpaces, DEBUG_TYPE, "Infer address spaces",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(InferAddressSpaces, DEBUG_TYPE, "Infer address spaces",
                     false, false)
 
+static unsigned getPtrOrVecOfPtrsAddressSpace(Type *Ty) {
+  if (Ty->isVectorTy()) {
+    Ty = cast<VectorType>(Ty)->getElementType();
+  }
+  assert(Ty->isPointerTy());
+  return Ty->getPointerAddressSpace();
+}
+
+static bool isPtrOrVecOfPtrsType(Type *Ty) {
+  if (Ty->isVectorTy()) {
+    Ty = cast<VectorType>(Ty)->getElementType();
+  }
+  return Ty->isPointerTy();
+}
+
+static Type *getPtrOrVecOfPtrsWithNewAS(Type *Ty, unsigned NewAddrSpace) {
+  if (!Ty->isVectorTy()) {
+    assert(Ty->isPointerTy());
+    return PointerType::getWithSamePointeeType(cast<PointerType>(Ty),
+                                               NewAddrSpace);
+  }
+
+  Type *PT = cast<VectorType>(Ty)->getElementType();
+  assert(PT->isPointerTy());
+
+  Type *NPT =
+      PointerType::getWithSamePointeeType(cast<PointerType>(PT), NewAddrSpace);
+  return VectorType::get(NPT, cast<VectorType>(Ty)->getElementCount());
+}
+
+static bool hasSameElementOfPtrOrVecPtrs(Type *Ty1, Type *Ty2) {
+  assert(isPtrOrVecOfPtrsType(Ty1) && isPtrOrVecOfPtrsType(Ty2));
+  assert(Ty1->isVectorTy() == Ty2->isVectorTy());
+  if (Ty1->isVectorTy()) {
+    Ty1 = cast<VectorType>(Ty1)->getElementType();
+    Ty2 = cast<VectorType>(Ty2)->getElementType();
+  }
+
+  assert(Ty1->isPointerTy() && Ty2->isPointerTy());
+  return cast<PointerType>(Ty1)->hasSameElementTypeAs(cast<PointerType>(Ty2));
+}
+
 // Check whether that's no-op pointer bicast using a pair of
 // `ptrtoint`/`inttoptr` due to the missing no-op pointer bitcast over
 // different address spaces.
 static bool isNoopPtrIntCastPair(const Operator *I2P, const DataLayout &DL,
                                  const TargetTransformInfo *TTI) {
   assert(I2P->getOpcode() == Instruction::IntToPtr);
   auto *P2I = dyn_cast<Operator>(I2P->getOperand(0));
   if (!P2I || P2I->getOpcode() != Instruction::PtrToInt)
     return false;
   // Check it's really safe to treat that pair of `ptrtoint`/`inttoptr` as a
   // no-op cast. Besides checking both of them are no-op casts, as the
   // reinterpreted pointer may be used in other pointer arithmetic, we also
   // need to double-check that through the target-specific hook. That ensures
   // the underlying target also agrees that's a no-op address space cast and
   // pointer bits are preserved.
   // The current IR spec doesn't have clear rules on address space casts,
   // especially a clear definition for pointer bits in non-default address
   // spaces. It would be undefined if that pointer is dereferenced after an
   // invalid reinterpret cast. Also, due to the unclearness for the meaning of
   // bits in non-default address spaces in the current spec, the pointer
   // arithmetic may also be undefined after invalid pointer reinterpret cast.
   // However, as we confirm through the target hooks that it's a no-op
   // addrspacecast, it doesn't matter since the bits should be the same.
-  unsigned P2IOp0AS = P2I->getOperand(0)->getType()->getPointerAddressSpace();
-  unsigned I2PAS = I2P->getType()->getPointerAddressSpace();
+  unsigned P2IOp0AS =
+      getPtrOrVecOfPtrsAddressSpace(P2I->getOperand(0)->getType());
+  unsigned I2PAS = getPtrOrVecOfPtrsAddressSpace(I2P->getType());
   return CastInst::isNoopCast(Instruction::CastOps(I2P->getOpcode()),
                               I2P->getOperand(0)->getType(), I2P->getType(),
                               DL) &&
          CastInst::isNoopCast(Instruction::CastOps(P2I->getOpcode()),
                               P2I->getOperand(0)->getType(), P2I->getType(),
                               DL) &&
          (P2IOp0AS == I2PAS || TTI->isNoopAddrSpaceCast(P2IOp0AS, I2PAS));
 }
 
 // Returns true if V is an address expression.
 // TODO: Currently, we consider only phi, bitcast, addrspacecast, and
 // getelementptr operators.
 static bool isAddressExpression(const Value &V, const DataLayout &DL,
                                 const TargetTransformInfo *TTI) {
   const Operator *Op = dyn_cast<Operator>(&V);
   if (!Op)
     return false;
 
   switch (Op->getOpcode()) {
   case Instruction::PHI:
-    assert(Op->getType()->isPointerTy());
+    assert(isPtrOrVecOfPtrsType(Op->getType()));
     return true;
   case Instruction::BitCast:
   case Instruction::AddrSpaceCast:
   case Instruction::GetElementPtr:
     return true;
   case Instruction::Select:
-    return Op->getType()->isPointerTy();
+    return isPtrOrVecOfPtrsType(Op->getType());
   case Instruction::Call: {
     const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&V);
     return II && II->getIntrinsicID() == Intrinsic::ptrmask;
   }
   case Instruction::IntToPtr:
     return isNoopPtrIntCastPair(Op, DL, TTI);
   default:
     // That value is an address expression if it has an assumed address space.

 }
 
 // Returns all flat address expressions in function F. The elements are
 // If V is an unvisited flat address expression, appends V to PostorderStack
 // and marks it as visited.
 void InferAddressSpacesImpl::appendsFlatAddressExpressionToPostorderStack(
     Value *V, PostorderStackTy &PostorderStack,
     DenseSet<Value *> &Visited) const {
-  assert(V->getType()->isPointerTy());
+  assert(isPtrOrVecOfPtrsType(V->getType()));
 
   // Generic addressing expressions may be hidden in nested constant
   // expressions.
   if (ConstantExpr *CE = dyn_cast<ConstantExpr>(V)) {
     // TODO: Look in non-address parts, like icmp operands.
     if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
       PostorderStack.emplace_back(CE, false);
 
     return;
   }
 
-  if (V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
+  if (getPtrOrVecOfPtrsAddressSpace(V->getType()) == FlatAddrSpace &&
       isAddressExpression(*V, *DL, TTI)) {
     if (Visited.insert(V).second) {
       PostorderStack.emplace_back(V, false);
 
       Operator *Op = cast<Operator>(V);
       for (unsigned I = 0, E = Op->getNumOperands(); I != E; ++I) {
         if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Op->getOperand(I))) {
           if (isAddressExpression(*CE, *DL, TTI) && Visited.insert(CE).second)
appendsFlatAddressExpressionToPostorderStack(Ptr, PostorderStack,
                                                  Visited);
   };
 
   // Look at operations that may be interesting accelerate by moving to a known
   // address space. We aim at generating after loads and stores, but pure
   // addressing calculations may also be faster.
   for (Instruction &I : instructions(F)) {
     if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
-      if (!GEP->getType()->isVectorTy())
-        PushPtrOperand(GEP->getPointerOperand());
     } else if (auto *LI = dyn_cast<LoadInst>(&I))
       PushPtrOperand(LI->getPointerOperand());
     else if (auto *SI = dyn_cast<StoreInst>(&I))
       PushPtrOperand(SI->getPointerOperand());
     else if (auto *RMW = dyn_cast<AtomicRMWInst>(&I))
       PushPtrOperand(RMW->getPointerOperand());
     else if (auto *CmpX = dyn_cast<AtomicCmpXchgInst>(&I))
       PushPtrOperand(CmpX->getPointerOperand());
     else if (auto *MI = dyn_cast<MemIntrinsic>(&I)) {
       // For memset/memcpy/memmove, any pointer operand can be replaced.
       PushPtrOperand(MI->getRawDest());
 
       // Handle 2nd operand for memcpy/memmove.
       if (auto *MTI = dyn_cast<MemTransferInst>(MI))
         PushPtrOperand(MTI->getRawSource());
     } else if (auto *II = dyn_cast<IntrinsicInst>(&I))
       collectRewritableIntrinsicOperands(II, PostorderStack, Visited);
     else if (ICmpInst *Cmp = dyn_cast<ICmpInst>(&I)) {
       // FIXME: Handle vectors of pointers
-      if (Cmp->getOperand(0)->getType()->isPointerTy()) {
+      if (getPtrOrVecOfPtrsAddressSpace(Cmp->getOperand(0)->getType())) {
         PushPtrOperand(Cmp->getOperand(0));
         PushPtrOperand(Cmp->getOperand(1));
       }
     } else if (auto *ASC = dyn_cast<AddrSpaceCastInst>(&I)) {
-      if (!ASC->getType()->isVectorTy())
-        PushPtrOperand(ASC->getPointerOperand());
     } else if (auto *I2P = dyn_cast<IntToPtrInst>(&I)) {
       if (isNoopPtrIntCastPair(cast<Operator>(I2P), *DL, TTI))
         PushPtrOperand(
             cast<Operator>(I2P->getOperand(0))->getOperand(0));
     }
   }
 
   std::vector<WeakTrackingVH> Postorder; // The resultant postorder.
   while (!PostorderStack.empty()) {
     Value *TopVal = PostorderStack.back().getPointer();
     // If the operands of the expression on the top are already explored,
     // adds that expression to the resultant postorder.
     if (PostorderStack.back().getInt()) {
-      if (TopVal->getType()->getPointerAddressSpace() == FlatAddrSpace)
+      if (getPtrOrVecOfPtrsAddressSpace(TopVal->getType()) == FlatAddrSpace)
         Postorder.push_back(TopVal);
       PostorderStack.pop_back();
       continue;
     }
     // Otherwise, adds its operands to the stack and explores them.
     PostorderStack.back().setInt(true);
     // Skip values with an assumed address space.
     if (TTI->getAssumedAddrSpace(TopVal) == UninitializedAddressSpace) {

 // returns an undef in the new address space as a placeholder.
 static Value *operandWithNewAddressSpaceOrCreateUndef(
     const Use &OperandUse, unsigned NewAddrSpace,
     const ValueToValueMapTy &ValueWithNewAddrSpace,
     const PredicatedAddrSpaceMapTy &PredicatedAS,
     SmallVectorImpl<const Use *> *UndefUsesToFix) {
   Value *Operand = OperandUse.get();
 
-  Type *NewPtrTy = PointerType::getWithSamePointeeType(
-      cast<PointerType>(Operand->getType()), NewAddrSpace);
+  Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(Operand->getType(), NewAddrSpace);
 
   if (Constant *C = dyn_cast<Constant>(Operand))
     return ConstantExpr::getAddrSpaceCast(C, NewPtrTy);
 
   if (Value *NewOperand = ValueWithNewAddrSpace.lookup(Operand))
     return NewOperand;
 
   Instruction *Inst = cast<Instruction>(OperandUse.getUser());
   auto I = PredicatedAS.find(std::make_pair(Inst, Operand));
   if (I != PredicatedAS.end()) {
     // Insert an addrspacecast on that operand before the user.
     unsigned NewAS = I->second;
-    Type *NewPtrTy = PointerType::getWithSamePointeeType(
-        cast<PointerType>(Operand->getType()), NewAS);
+    Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(Operand->getType(), NewAS);
     auto *NewI = new AddrSpaceCastInst(Operand, NewPtrTy);
     NewI->insertBefore(Inst);
     NewI->setDebugLoc(Inst->getDebugLoc());
     return NewI;
   }
 
   UndefUsesToFix->push_back(&OperandUse);
   return UndefValue::get(NewPtrTy);

 //
 // This may also return nullptr in the case the instruction could not be
 // rewritten.
 Value *InferAddressSpacesImpl::cloneInstructionWithNewAddressSpace(
     Instruction *I, unsigned NewAddrSpace,
     const ValueToValueMapTy &ValueWithNewAddrSpace,
     const PredicatedAddrSpaceMapTy &PredicatedAS,
     SmallVectorImpl<const Use *> *UndefUsesToFix) const {
-  Type *NewPtrType = PointerType::getWithSamePointeeType(
-      cast<PointerType>(I->getType()), NewAddrSpace);
+  Type *NewPtrType = getPtrOrVecOfPtrsWithNewAS(I->getType(), NewAddrSpace);
 
   if (I->getOpcode() == Instruction::AddrSpaceCast) {
     Value *Src = I->getOperand(0);
     // Because `I` is flat, the source address space must be specific.
     // Therefore, the inferred address space must be the source space, according
     // to our algorithm.
-    assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
+    assert(getPtrOrVecOfPtrsAddressSpace(Src->getType()) == NewAddrSpace);
     if (Src->getType() != NewPtrType)
       return new BitCastInst(Src, NewPtrType);
     return Src;
   }
 
   if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     // Technically the intrinsic ID is a pointer typed argument, so specially
     // handle calls early.
if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
 
     return nullptr;
   }
 
   unsigned AS = TTI->getAssumedAddrSpace(I);
   if (AS != UninitializedAddressSpace) {
     // For the assumed address space, insert an `addrspacecast` to make that
     // explicit.
-    Type *NewPtrTy = PointerType::getWithSamePointeeType(
-        cast<PointerType>(I->getType()), AS);
+    Type *NewPtrTy = getPtrOrVecOfPtrsWithNewAS(I->getType(), AS);
     auto *NewI = new AddrSpaceCastInst(I, NewPtrTy);
     NewI->insertAfter(I);
     return NewI;
   }
 
   // Computes the converted pointer operands.
   SmallVector<Value *, 4> NewPointerOperands;
   for (const Use &OperandUse : I->operands()) {
-    if (!OperandUse.get()->getType()->isPointerTy())
+    if (!isPtrOrVecOfPtrsType(OperandUse.get()->getType()))
       NewPointerOperands.push_back(nullptr);
     else
       NewPointerOperands.push_back(operandWithNewAddressSpaceOrCreateUndef(
           OperandUse, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS,
           UndefUsesToFix));
   }
 
   switch (I->getOpcode()) {
   case Instruction::BitCast:
     return new BitCastInst(NewPointerOperands[0], NewPtrType);
   case Instruction::PHI: {
-    assert(I->getType()->isPointerTy());
+    assert(isPtrOrVecOfPtrsType(I->getType()));
     PHINode *PHI = cast<PHINode>(I);
     PHINode *NewPHI = PHINode::Create(NewPtrType, PHI->getNumIncomingValues());
     for (unsigned Index = 0; Index < PHI->getNumIncomingValues(); ++Index) {
       unsigned OperandNo = PHINode::getOperandNumForIncomingValue(Index);
       NewPHI->addIncoming(NewPointerOperands[OperandNo],
                           PHI->getIncomingBlock(Index));
     }
     return NewPHI;
   }
   case Instruction::GetElementPtr: {
     GetElementPtrInst *GEP = cast<GetElementPtrInst>(I);
     GetElementPtrInst *NewGEP = GetElementPtrInst::Create(
         GEP->getSourceElementType(), NewPointerOperands[0],
         SmallVector<Value *, 4>(GEP->indices()));
     NewGEP->setIsInBounds(GEP->isInBounds());
     return NewGEP;
   }
   case Instruction::Select:
-    assert(I->getType()->isPointerTy());
+    assert(isPtrOrVecOfPtrsType(I->getType()));
     return SelectInst::Create(I->getOperand(0), NewPointerOperands[1],
                               NewPointerOperands[2], "", nullptr, I);
   case Instruction::IntToPtr: {
     assert(isNoopPtrIntCastPair(cast<Operator>(I), *DL, TTI));
     Value *Src = cast<Operator>(I->getOperand(0))->getOperand(0);
     if (Src->getType() == NewPtrType)
       return Src;
 

 
 // Similar to cloneInstructionWithNewAddressSpace, returns a clone of the
 // constant expression `CE` with its operands replaced as specified in
 // ValueWithNewAddrSpace.
 static Value *cloneConstantExprWithNewAddressSpace(
     ConstantExpr *CE, unsigned NewAddrSpace,
     const ValueToValueMapTy &ValueWithNewAddrSpace, const DataLayout *DL,
     const TargetTransformInfo *TTI) {
-  Type *TargetType = CE->getType()->isPointerTy()
-                         ? PointerType::getWithSamePointeeType(
-                               cast<PointerType>(CE->getType()), NewAddrSpace)
-                         : CE->getType();
+  Type *TargetType =
+      isPtrOrVecOfPtrsType(CE->getType())
+          ? getPtrOrVecOfPtrsWithNewAS(CE->getType(), NewAddrSpace)
 
   if (CE->getOpcode() == Instruction::AddrSpaceCast) {
     // Because CE is flat, the source address space must be specific.
     // Therefore, the inferred address space must be the source space according
     // to our algorithm.
-    assert(CE->getOperand(0)->getType()->getPointerAddressSpace() ==
+    assert(getPtrOrVecOfPtrsAddressSpace(CE->getOperand(0)->getType()) ==
            NewAddrSpace);
     return ConstantExpr::getBitCast(CE->getOperand(0), TargetType);
   }
 
   if (CE->getOpcode() == Instruction::BitCast) {
     if (Value *NewOperand = ValueWithNewAddrSpace.lookup(CE->getOperand(0)))
       return ConstantExpr::getBitCast(cast<Constant>(NewOperand), TargetType);
     return ConstantExpr::getAddrSpaceCast(CE, TargetType);
   }
 
   if (CE->getOpcode() == Instruction::Select) {
     Constant *Src0 = CE->getOperand(1);
     Constant *Src1 = CE->getOperand(2);
-    if (Src0->getType()->getPointerAddressSpace() ==
-        Src1->getType()->getPointerAddressSpace()) {
+    if (getPtrOrVecOfPtrsAddressSpace(Src0->getType()) ==
+        getPtrOrVecOfPtrsAddressSpace(Src1->getType())) {
 
       return ConstantExpr::getSelect(
           CE->getOperand(0), ConstantExpr::getAddrSpaceCast(Src0, TargetType),
           ConstantExpr::getAddrSpaceCast(Src1, TargetType));
     }
   }
 
   if (CE->getOpcode() == Instruction::IntToPtr) {
     assert(isNoopPtrIntCastPair(cast<Operator>(CE), *DL, TTI));
     Constant *Src = cast<ConstantExpr>(CE->getOperand(0))->getOperand(0);
-    assert(Src->getType()->getPointerAddressSpace() == NewAddrSpace);
+    assert(getPtrOrVecOfPtrsAddressSpace(Src->getType()) == NewAddrSpace);
     return ConstantExpr::getBitCast(Src, TargetType);
   }
 
   // Computes the operands of the new constant expression.
   bool IsNew = false;
   SmallVector<Constant *, 4> NewOperands;
   for (unsigned Index = 0; Index < CE->getNumOperands(); ++Index) {
     Constant *Operand = CE->getOperand(Index);

 //
 // See cloneInstructionWithNewAddressSpace for the meaning of UndefUsesToFix.
 Value *InferAddressSpacesImpl::cloneValueWithNewAddressSpace(
     Value *V, unsigned NewAddrSpace,
     const ValueToValueMapTy &ValueWithNewAddrSpace,
     const PredicatedAddrSpaceMapTy &PredicatedAS,
     SmallVectorImpl<const Use *> *UndefUsesToFix) const {
   // All values in Postorder are flat address expressions.
-  assert(V->getType()->getPointerAddressSpace() == FlatAddrSpace &&
+  assert(isPtrOrVecOfPtrsType(V->getType()) == FlatAddrSpace &&
          isAddressExpression(*V, *DL, TTI));
 
   if (Instruction *I = dyn_cast<Instruction>(V)) {
     Value *NewV = cloneInstructionWithNewAddressSpace(
         I, NewAddrSpace, ValueWithNewAddrSpace, PredicatedAS, UndefUsesToFix);
     if (Instruction *NewI = dyn_cast_or_null<Instruction>(NewV)) {
       if (NewI->getParent() == nullptr) {
         NewI->insertBefore(I);
bool InferAddressSpacesImpl::updateAddressSpace(
   unsigned NewAS = UninitializedAddressSpace;
 
   const Operator &Op = cast<Operator>(V);
   if (Op.getOpcode() == Instruction::Select) {
     Value *Src0 = Op.getOperand(1);
     Value *Src1 = Op.getOperand(2);
 
     auto I = InferredAddrSpace.find(Src0);
-    unsigned Src0AS = (I != InferredAddrSpace.end()) ?
-      I->second : Src0->getType()->getPointerAddressSpace();
+    unsigned Src0AS = (I != InferredAddrSpace.end())
+                          ? I->second
+                          : getPtrOrVecOfPtrsAddressSpace(Src0->getType());
 
     auto J = InferredAddrSpace.find(Src1);
-    unsigned Src1AS = (J != InferredAddrSpace.end()) ?
-      J->second : Src1->getType()->getPointerAddressSpace();
+    unsigned Src1AS = (J != InferredAddrSpace.end())
+                          ? J->second
+                          : getPtrOrVecOfPtrsAddressSpace(Src1->getType());
 
     auto *C0 = dyn_cast<Constant>(Src0);
     auto *C1 = dyn_cast<Constant>(Src1);
 
     // If one of the inputs is a constant, we may be able to do a constant
     // addrspacecast of it. Defer inferring the address space until the input
     // address space is known.
     if ((C1 && Src0AS == UninitializedAddressSpace) ||
if (AS != UninitializedAddressSpace) {
       // Use the assumed address space directly.
       NewAS = AS;
     } else {
       // Otherwise, infer the address space from its pointer operands.
       for (Value *PtrOperand : getPointerOperands(V, *DL, TTI)) {
         auto I = InferredAddrSpace.find(PtrOperand);
         unsigned OperandAS;
         if (I == InferredAddrSpace.end()) {
-          OperandAS = PtrOperand->getType()->getPointerAddressSpace();
+          OperandAS = getPtrOrVecOfPtrsAddressSpace(PtrOperand->getType());
           if (OperandAS == FlatAddrSpace) {
             // Check AC for assumption dominating V.
             unsigned AS = getPredicatedAddrSpace(V, PtrOperand);
             if (AS != UninitializedAddressSpace) {
               LLVM_DEBUG(dbgs()
                          << "  deduce operand AS from the predicate addrspace "
                          << AS << '\n');
               OperandAS = AS;

 }
 
 // \p returns true if it is OK to change the address space of constant \p C with
 // a ConstantExpr addrspacecast.
 bool InferAddressSpacesImpl::isSafeToCastConstAddrSpace(Constant *C,
                                                         unsigned NewAS) const {
   assert(NewAS != UninitializedAddressSpace);
 
-  unsigned SrcAS = C->getType()->getPointerAddressSpace();
+  unsigned SrcAS = getPtrOrVecOfPtrsAddressSpace(C->getType());
   if (SrcAS == NewAS || isa<UndefValue>(C))
     return true;
 
   // Prevent illegal casts between different non-flat address spaces.
   if (SrcAS != FlatAddrSpace && NewAS != FlatAddrSpace)
     return false;
 
   if (isa<ConstantPointerNull>(C))
     return true;
 
   if (auto *Op = dyn_cast<Operator>(C)) {
     // If we already have a constant addrspacecast, it should be safe to cast it
     // off.
     if (Op->getOpcode() == Instruction::AddrSpaceCast)
       return isSafeToCastConstAddrSpace(cast<Constant>(Op->getOperand(0)), NewAS);
 
     if (Op->getOpcode() == Instruction::IntToPtr &&
-        Op->getType()->getPointerAddressSpace() == FlatAddrSpace)
+        getPtrOrVecOfPtrsAddressSpace(Op->getType()) == FlatAddrSpace)
       return true;
   }
 
   return false;
 }
 
 static Value::use_iterator skipToNextUser(Value::use_iterator I,
                                           Value::use_iterator End) {
bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
   for (Value* V : Postorder) {
     unsigned NewAddrSpace = InferredAddrSpace.lookup(V);
 
     // In some degenerate cases (e.g. invalid IR in unreachable code), we may
     // not even infer the value to have its original address space.
     if (NewAddrSpace == UninitializedAddressSpace)
       continue;
 
-    if (V->getType()->getPointerAddressSpace() != NewAddrSpace) {
+    if (getPtrOrVecOfPtrsAddressSpace(V->getType()) != NewAddrSpace) {
       Value *New =
           cloneValueWithNewAddressSpace(V, NewAddrSpace, ValueWithNewAddrSpace,
                                         PredicatedAS, &UndefUsesToFix);
       if (New)
         ValueWithNewAddrSpace[V] = New;
     }
   }
 
for (const WeakTrackingVH &WVH : Postorder) {
     for (I = V->use_begin(), E = V->use_end(); I != E; ) {
       Use &U = *I;
 
       // Some users may see the same pointer operand in multiple operands. Skip
       // to the next instruction.
       I = skipToNextUser(I, E);
 
       if (isSimplePointerUseValidToReplace(
-              *TTI, U, V->getType()->getPointerAddressSpace())) {
+              *TTI, U, getPtrOrVecOfPtrsAddressSpace(V->getType()))) {
         // If V is used as the pointer operand of a compatible memory operation,
         // sets the pointer operand to NewV. This replacement does not change
         // the element type, so the resultant load/store is still valid.
         U.set(NewV);
         continue;
       }
 
       User *CurUser = U.getUser();
for (I = V->use_begin(), E = V->use_end(); I != E; ) {
       if (isa<Instruction>(CurUser)) {
         if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUser)) {
           // If we can infer that both pointers are in the same addrspace,
           // transform e.g.
           //   %cmp = icmp eq float* %p, %q
           // into
           //   %cmp = icmp eq float addrspace(3)* %new_p, %new_q
 
-          unsigned NewAS = NewV->getType()->getPointerAddressSpace();
+          unsigned NewAS = getPtrOrVecOfPtrsAddressSpace(NewV->getType());
           int SrcIdx = U.getOperandNo();
           int OtherIdx = (SrcIdx == 0) ? 1 : 0;
           Value *OtherSrc = Cmp->getOperand(OtherIdx);
 
           if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(OtherSrc)) {
-            if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
+            if (getPtrOrVecOfPtrsAddressSpace(OtherNewV->getType()) == NewAS) {
               Cmp->setOperand(OtherIdx, OtherNewV);
               Cmp->setOperand(SrcIdx, NewV);
               continue;
             }
           }
 
           // Even if the type mismatches, we can cast the constant.
           if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
             if (isSafeToCastConstAddrSpace(KOtherSrc, NewAS)) {
               Cmp->setOperand(SrcIdx, NewV);
               Cmp->setOperand(OtherIdx,
                 ConstantExpr::getAddrSpaceCast(KOtherSrc, NewV->getType()));
               continue;
             }
           }
         }
 
         if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(CurUser)) {
-          unsigned NewAS = NewV->getType()->getPointerAddressSpace();
-          if (ASC->getDestAddressSpace() == NewAS) {
-            if (!cast<PointerType>(ASC->getType())
-                    ->hasSameElementTypeAs(
-                        cast<PointerType>(NewV->getType()))) {
+          unsigned NewAS = getPtrOrVecOfPtrsAddressSpace(NewV->getType());
+          if (getPtrOrVecOfPtrsAddressSpace(ASC->getType()) == NewAS) {
+            if (!hasSameElementOfPtrOrVecPtrs(ASC->getType(),
+                                              NewV->getType())) {
               BasicBlock::iterator InsertPos;
               if (Instruction *NewVInst = dyn_cast<Instruction>(NewV))
                 InsertPos = std::next(NewVInst->getIterator());
               else if (Instruction *VInst = dyn_cast<Instruction>(V))
                 InsertPos = std::next(VInst->getIterator());
               else
                 InsertPos = ASC->getIterator();