[ConstantFolding, InstSimplify] Handle more vector GEPs

[mssimpso]

Summary

This patch addresses some additional cases where the compiler crashes upon encountering vector GEPs. There are a few identity-like cases (e.g., getelementptr %p, 0 -> %p) present in InstSimplify and ConstantFolding and a case in ConstantFolding that tries to mark GEPs inbounds. I'm happy to split this patch up if preferable. This should fix PR36116.

Reference: https://bugs.llvm.org/show_bug.cgi?id=36116

[haicheng]

LGTM

lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyGEPInst(Type *SrcTy, ArrayRef<Value *> Ops,
   else if (VectorType *VT = dyn_cast<VectorType>(Ops[1]->getType()))
     GEPTy = VectorType::get(GEPTy, VT->getNumElements());
 
   if (isa<UndefValue>(Ops[0]))
     return UndefValue::get(GEPTy);
 
   if (Ops.size() == 2) {
     // getelementptr P, 0 -> P.
-    if (match(Ops[1], m_Zero()))
+    if (match(Ops[1], m_Zero()) && Ops[0]->getType() == GEPTy)
       return Ops[0];
 
     Type *Ty = SrcTy;
     if (Ty->isSized()) {
       Value *P;
       uint64_t C;
       uint64_t TyAllocSize = Q.DL.getTypeAllocSize(Ty);
       // getelementptr P, N -> P if P points to a type of zero size.
-      if (TyAllocSize == 0)
+      if (TyAllocSize == 0 && Ops[0]->getType() == GEPTy)
         return Ops[0];
 
       // The following transforms are only safe if the ptrtoint cast
       // doesn't truncate the pointers.
       if (Ops[1]->getType()->getScalarSizeInBits() ==
           Q.DL.getIndexSizeInBits(AS)) {
         auto PtrToIntOrZero = [GEPTy](Value *P) -> Value * {
           if (match(P, m_Zero()))

lib/IR/ConstantFold.cpp

static bool isInBoundsIndices(ArrayRef<IndexTy> Idxs) {
   // No indices means nothing that could be out of bounds.
   if (Idxs.empty()) return true;
 
   // If the first index is zero, it's in bounds.
   if (cast<Constant>(Idxs[0])->isNullValue()) return true;
 
   // If the first index is one and all the rest are zero, it's in bounds,
   // by the one-past-the-end rule.
-  if (!cast<ConstantInt>(Idxs[0])->isOne())
+  if (auto *CI = dyn_cast<ConstantInt>(Idxs[0])) {
+    if (!CI->isOne())
       return false;
+  } else {
+    auto *CV = cast<ConstantDataVector>(Idxs[0]);
+    CI = dyn_cast_or_null<ConstantInt>(CV->getSplatValue());
+    if (!CI || !CI->isOne())
+      return false;
+  }
+
   for (unsigned i = 1, e = Idxs.size(); i != e; ++i)
     if (!cast<Constant>(Idxs[i])->isNullValue())
       return false;
   return true;
 }
 
 /// Test whether a given ConstantInt is in-range for a SequentialType.
 static bool isIndexInRangeOfArrayType(uint64_t NumElements,

 }
 
 Constant *llvm::ConstantFoldGetElementPtr(Type *PointeeTy, Constant *C,
                                           bool InBounds,
                                           Optional<unsigned> InRangeIndex,
                                           ArrayRef<Value *> Idxs) {
   if (Idxs.empty()) return C;
 
-  if (isa<UndefValue>(C)) {
-    Type *GEPTy = GetElementPtrInst::getGEPReturnType(
-        C, makeArrayRef((Value * const *)Idxs.data(), Idxs.size()));
+      C, makeArrayRef((Value *const *)Idxs.data(), Idxs.size()));
+
+  if (isa<UndefValue>(C))
     return UndefValue::get(GEPTy);
-  }
 
   Constant *Idx0 = cast<Constant>(Idxs[0]);
   if (Idxs.size() == 1 && (Idx0->isNullValue() || isa<UndefValue>(Idx0)))
-    return C;
+    return GEPTy->isVectorTy() && !C->getType()->isVectorTy()
+               ? ConstantVector::getSplat(
+                     cast<VectorType>(GEPTy)->getNumElements(), C)
+               : C;
 
   if (C->isNullValue()) {
     bool isNull = true;
     for (unsigned i = 0, e = Idxs.size(); i != e; ++i)
       if (!isa<UndefValue>(Idxs[i]) &&
           !cast<Constant>(Idxs[i])->isNullValue()) {
         isNull = false;
         break;

test/Analysis/ConstantFolding/vectorgep-crash.ll

 
 ; CHECK-LABEL: @test2
 ; CHECK: ret <16 x i32*> getelementptr ([65 x %struct.A], [65 x %struct.A]* @G, <16 x i64> zeroinitializer, <16 x i64> <i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7, i64 8, i64 9, 
 define <16 x i32*> @test2() {
 vector.body:
   %VectorGep = getelementptr [65 x %struct.A], [65 x %struct.A]* @G, <16 x i32> zeroinitializer, <16 x i64> <i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7, i64 8, i64 9, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, <16 x i32> zeroinitializer
   ret <16 x i32*> %VectorGep
 }
+
+@g = external global i8, align 1
+
+define <2 x i8*> @constant_zero_index() {
+; CHECK-LABEL: @constant_zero_index(
+; CHECK-NEXT:    ret <2 x i8*> <i8* @g, i8* @g>
+;
+  %gep = getelementptr i8, i8* @g, <2 x i64> zeroinitializer
+  ret <2 x i8*> %gep
+}
+
+define <2 x i8*> @constant_undef_index() {
+; CHECK-LABEL: @constant_undef_index(
+; CHECK-NEXT:    ret <2 x i8*> <i8* @g, i8* @g>
+;
+  %gep = getelementptr i8, i8* @g, <2 x i64> undef
+  ret <2 x i8*> %gep
+}
+
+define <2 x i8*> @constant_inbounds() {
+; CHECK-LABEL: @constant_inbounds(
+; CHECK-NEXT:    ret <2 x i8*> getelementptr inbounds (i8, i8* @g, <2 x i64> <i64 1, i64 1>)
+;
+  %gep = getelementptr i8, i8* @g, <2 x i64> <i64 1, i64 1>
+  ret <2 x i8*> %gep
+}

test/Transforms/InstSimplify/vector_gep.ll

 
 ; PR32697
 ; CHECK-LABEL: laalaa(
 ; CHECK-NEXT: ret <4 x i8*> undef
 define <4 x i8*> @laalaa() {
   %patatino = getelementptr i8, <4 x i8 *> undef, i64 undef
   ret <4 x i8*> %patatino
 }
+
+define <2 x i8*> @zero_index(i8* %p) {
+; CHECK-LABEL: @zero_index(
+; CHECK-NEXT:    %gep = getelementptr i8, i8* %p, <2 x i64> zeroinitializer
+; CHECK-NEXT:    ret <2 x i8*> %gep
+;
+  %gep = getelementptr i8, i8* %p, <2 x i64> zeroinitializer
+  ret <2 x i8*> %gep
+}
+
+define <2 x {}*> @unsized({}* %p) {
+; CHECK-LABEL: @unsized(
+; CHECK-NEXT:    %gep = getelementptr {}, {}* %p, <2 x i64> undef
+; CHECK-NEXT:    ret <2 x {}*> %gep
+;
+  %gep = getelementptr {}, {}* %p, <2 x i64> undef
+  ret <2 x {}*> %gep
+}