[ConstantFold] Fix a crash when folding a GEP that has vector index

[haicheng]

Summary

LLVM crashes with following IRs if compiling it with opt -instsimplify. See PR34880

@block = global [64 x [8192 x i8]] zeroinitializer, align 1

define <2 x i8*> @foo() {
  %1 = getelementptr inbounds [64 x [8192 x i8]], [64 x [8192 x i8]]* @block, i64 0, <2 x i64> <i64 0, i64 1>, i64 8192
  ret <2 x i8*> %1
}

ConstantFolding tries to factor out the last index (8192) into the second last dimension. The code assumes the current index (8192) and the previous index (<i64 0, i64 1>) are both integer typed, but it now only checks if the current index is a ConstantInt. This patch exits early if the previous index is not a ConstantInt.

[majnemer]

lib/IR/ConstantFold.cpp

2213

Shouldn't this be !isa<Constant> ?

[haicheng]

lib/IR/ConstantFold.cpp

2213

Using !isa<Constant> still crashes the test case. In the test case, when Idxs[i-1] is <i64 0, i64 1> which is a Constant but not a ConstantInt, code below will cast its type to IntegerType and then llvm crashes.

[haicheng]

lib/IR/ConstantFold.cpp

2259

The crash happens here. Type::getIntegerBitWidth() casts underlying type to IntegerType and crashes when the underlying type is a vector.

2265

I think this line also does not expect a vector PrevIdx.

2272

I think this line also does not expect a vector PrevIdx.

[haicheng]

Kindly Ping.

This patch just adds a simply check to bail out early to prevent LLVM from crashing.

[mssimpso]

Haicheng,

How hard do you think it would be to update this transformation to reason about vector GEPs rather than bailing out? (I've left some inline comments if you want to give it a try). Having said that, I'm fine with that patch as is. But you should at least add a comment explaining why we're giving up.

lib/IR/ConstantFold.cpp

2259

getScalarScalarSizeInBits should work for both scalars and vectors.

2265

isIntOrIntVectorTy should work for both scalars and vectors.

2272

I believe the add should work for vectors as well, but if PrevIdx is a vector, I think you would have to broadcast Div to a vector.

[haicheng]

Thank you, Matt.

As you pointed out, supporting vector previous index is not difficult. I choosed to bail out at first because

1. We already bail out when the current index is a vector.

2. I initially planned to fix the crash in this patch to unblock the release and add the vector index support for both current and previous index.

3. I think this piece of algorithm work better if we iterate from the last index to the first. So, we don't need to recursively call this function in the following case.

@block = global [8 x [64 x [8192 x i8]]] zeroinitializer, align 1

define i8* @vectorindex() {
  %1 = getelementptr inbounds [8 x [64 x [8192 x i8]]], [8 x [64 x [8192 x i8]]]* @block, i64 0, i64 0, i64 63, i64 8192
  ret i8* %1
}

[mssimpso]

In D38677#908574, @haicheng wrote:

Thank you, Matt.

As you pointed out, supporting vector previous index is not difficult. I choosed to bail out at first because

1. We already bail out when the current index is a vector.

2. I initially planned to fix the crash in this patch to unblock the release and add the vector index support for both current and previous index.

3. I think this piece of algorithm work better if we iterate from the last index to the first. So, we don't need to recursively call this function in the following case.

@block = global [8 x [64 x [8192 x i8]]] zeroinitializer, align 1

define i8* @vectorindex() {
  %1 = getelementptr inbounds [8 x [64 x [8192 x i8]]], [8 x [64 x [8192 x i8]]]* @block, i64 0, i64 0, i64 63, i64 8192
  ret i8* %1
}

OK, that sounds good to me.

[haicheng]

Add a FIXME.

Thank you, Matt.

lib/IR/ConstantFold.cpp

Constant *llvm::ConstantFoldGetElementPtr(Type *PointeeTy, Constant *C,
   for (unsigned i = 1, e = Idxs.size(); i != e;
        Prev = Ty, Ty = cast<CompositeType>(Ty)->getTypeAtIndex(Idxs[i]), ++i) {
     auto *CI = dyn_cast<ConstantInt>(Idxs[i]);
     if (!CI) {
       // We don't know if it's in range or not.
       Unknown = true;
       continue;
     }
+    if (!isa<ConstantInt>(Idxs[i - 1]))
+      continue;
     if (InRangeIndex && i == *InRangeIndex + 1) {
       // If an index is marked inrange, we cannot apply this canonicalization to
       // the following index, as that will cause the inrange index to point to
       // the wrong element.
       continue;
     }
     if (isa<StructType>(Ty)) {
       // The verify makes sure that GEPs into a struct are in range.
for (unsigned i = 1, e = Idxs.size(); i != e;
     ConstantInt *Factor = ConstantInt::get(CI->getType(), NumElements);
     NewIdxs[i] = ConstantExpr::getSRem(CI, Factor);
 
     Constant *PrevIdx = NewIdxs[i-1] ? NewIdxs[i-1] :
                            cast<Constant>(Idxs[i - 1]);
     Constant *Div = ConstantExpr::getSDiv(CI, Factor);
 
     unsigned CommonExtendedWidth =
         std::max(PrevIdx->getType()->getIntegerBitWidth(),
                  Div->getType()->getIntegerBitWidth());
     CommonExtendedWidth = std::max(CommonExtendedWidth, 64U);
 
     // Before adding, extend both operands to i64 to avoid
     // overflow trouble.
     if (!PrevIdx->getType()->isIntegerTy(CommonExtendedWidth))
       PrevIdx = ConstantExpr::getSExt(
           PrevIdx, Type::getIntNTy(Div->getContext(), CommonExtendedWidth));
     if (!Div->getType()->isIntegerTy(CommonExtendedWidth))
       Div = ConstantExpr::getSExt(
           Div, Type::getIntNTy(Div->getContext(), CommonExtendedWidth));
 
     NewIdxs[i - 1] = ConstantExpr::getAdd(PrevIdx, Div);
   }
 
   // If we did any factoring, start over with the adjusted indices.
   if (!NewIdxs.empty()) {
     for (unsigned i = 0, e = Idxs.size(); i != e; ++i)
       if (!NewIdxs[i]) NewIdxs[i] = cast<Constant>(Idxs[i]);
     return ConstantExpr::getGetElementPtr(PointeeTy, C, NewIdxs, InBounds,
                                           InRangeIndex);

test/Transforms/InstCombine/gep-vector.ll

 ; RUN: opt -instcombine %s -S | FileCheck %s
 
 ; CHECK-LABEL: patatino
 ; CHECK-NEXT: ret <8 x i64*> undef
 define <8 x i64*> @patatino() {
   %el = getelementptr i64, <8 x i64*> undef, <8 x i64> undef
   ret <8 x i64*> %el
 }
 
 ; CHECK-LABEL: patatino2
 ; CHECK-NEXT: ret <8 x i64*> undef
 define <8 x i64*> @patatino2() {
   %el = getelementptr inbounds i64, i64* undef, <8 x i64> undef
   ret <8 x i64*> %el
 }
+
+@block = global [64 x [8192 x i8]] zeroinitializer, align 1
+
+; CHECK-LABEL:vectorindex 
+; CHECK-NEXT: ret <2 x i8*> getelementptr inbounds ([64 x [8192 x i8]], [64 x [8192 x i8]]* @block, <2 x i64> zeroinitializer, <2 x i64> <i64 0, i64 1>, <2 x i64> <i64 8192, i64 8192>) 
+define <2 x i8*> @vectorindex() {
+  %1 = getelementptr inbounds [64 x [8192 x i8]], [64 x [8192 x i8]]* @block, i64 0, <2 x i64> <i64 0, i64 1>, i64 8192
+  ret <2 x i8*> %1
+}