[SelectionDAG] Fix lowering of vector geps

[jroelofs]

Summary

This fixes an assertion failure that was being triggered in SelectionDAG::getZeroExtendInReg(), where it was trying to extend the <2xi32> to i64 (which should have been <2xi64>).

[dmgreen]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 void SelectionDAGBuilder::visitGetElementPtr(const User &I) {
   Value *Op0 = I.getOperand(0);
   // Note that the pointer operand may be a vector of pointers. Take the scalar
   // element which holds a pointer.
   unsigned AS = Op0->getType()->getScalarType()->getPointerAddressSpace();
   SDValue N = getValue(Op0);
   SDLoc dl = getCurSDLoc();
   auto &TLI = DAG.getTargetLoweringInfo();
-  MVT PtrTy = TLI.getPointerTy(DAG.getDataLayout(), AS);
-  MVT PtrMemTy = TLI.getPointerMemTy(DAG.getDataLayout(), AS);
 
   // Normalize Vector GEP - all scalar operands should be converted to the
   // splat vector.
   bool IsVectorGEP = I.getType()->isVectorTy();
   ElementCount VectorElementCount =
       IsVectorGEP ? cast<VectorType>(I.getType())->getElementCount()
                   : ElementCount(0, false);
 
if (StructType *StTy = GTI.getStructTypeOrNull()) {
         }
       }
 
       N = DAG.getNode(ISD::ADD, dl,
                       N.getValueType(), N, IdxN);
     }
   }
 
+  MVT PtrTy = TLI.getPointerTy(DAG.getDataLayout(), AS);
+  MVT PtrMemTy = TLI.getPointerMemTy(DAG.getDataLayout(), AS);
+  if (IsVectorGEP) {
+    PtrTy = MVT::getVectorVT(PtrTy, VectorElementCount);
+    PtrMemTy = MVT::getVectorVT(PtrMemTy, VectorElementCount);
+  }
+
   if (PtrMemTy != PtrTy && !cast<GEPOperator>(I).isInBounds())
     N = DAG.getPtrExtendInReg(N, dl, PtrMemTy);
 
   setValue(&I, N);
 }
 
 void SelectionDAGBuilder::visitAlloca(const AllocaInst &I) {
   // If this is a fixed sized alloca in the entry block of the function,

llvm/test/CodeGen/AArch64/vector-gep.ll

+; RUN: llc < %s -mtriple=arm64_32-apple-watchos2.0.0 --aarch64-neon-syntax=generic | FileCheck %s
+
+target datalayout = "e-m:o-p:32:32-i64:64-i128:128-n32:64-S128"
+target triple = "arm64_32-apple-watchos2.0.0"
+
+; CHECK-LABEL: lCPI0_0:
+; CHECK-NEXT:    .quad 36
+; CHECK-NEXT:    .quad 4804
+
+define <2 x i8*> @vector_gep(<2 x i8*> %0) {
+; CHECK-LABEL: vector_gep:
+; CHECK:         adrp x[[REG8:[123]?[0-9]]], lCPI0_0@PAGE
+; CHECK:         ldr q[[REG1:[0-9]+]], [x[[REG8]], lCPI0_0@PAGEOFF]
+; CHECK:         add v[[REG0:[0-9]+]].2d, v[[REG0]].2d, v[[REG1]].2d
+; CHECK:         movi v[[REG1]].2d, #0x000000ffffffff
+; CHECK:         and v[[REG0]].16b, v[[REG0]].16b, v[[REG1]].16b
+; CHECK:         ret
+entry:
+  %1 = getelementptr i8, <2 x i8*> %0, <2 x i32> <i32 36, i32 4804>
+  ret <2 x i8*> %1
+}