[SimplifyInst] Use correct type for GEPs with vector indices.

[fhahn]

Summary

The current code does not properly handle vector indices unless they are
the first index.

At the moment LangRef gives the impression that the vector index must be
the one and only index (https://llvm.org/docs/LangRef.html#getelementptr-instruction).

But vector indices can appear at any position and according to the
verifier there may be multiple vector indices. If that's the case, the
number of elements must match.

This patch updates SimplifyGEPInst to properly handle those additional
cases.

[nikic]

This looks okay, but I wonder whether we shouldn't be adjusting the getIndexedType() API instead. We have similar code elsewhere, for example in https://github.com/llvm/llvm-project/blob/9c5ebf0358960adf28931569a0c801b56c8008d9/llvm/lib/IR/Constants.cpp#L2415-L2430.

[nikic]

It probably makes sense to land this first in any case, so LGTM.

[fhahn]

In D99961#2671524, @nikic wrote:

This looks okay, but I wonder whether we shouldn't be adjusting the getIndexedType() API instead. We have similar code elsewhere, for example in https://github.com/llvm/llvm-project/blob/9c5ebf0358960adf28931569a0c801b56c8008d9/llvm/lib/IR/Constants.cpp#L2415-L2430.

SGTM, I'll look into this as follow-up

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyGEPInst(Type *SrcTy, ArrayRef<Value *> Ops,
 
   // getelementptr P -> P.
   if (Ops.size() == 1)
     return Ops[0];
 
   // Compute the (pointer) type returned by the GEP instruction.
   Type *LastType = GetElementPtrInst::getIndexedType(SrcTy, Ops.slice(1));
   Type *GEPTy = PointerType::get(LastType, AS);
-  if (VectorType *VT = dyn_cast<VectorType>(Ops[0]->getType()))
-    GEPTy = VectorType::get(GEPTy, VT->getElementCount());
-  else if (VectorType *VT = dyn_cast<VectorType>(Ops[1]->getType()))
+  for (Value *Op : Ops) {
+    // If one of the operands is a vector, the result type is a vector of
+    // pointers. All vector operands must have the same number of elements.
+    if (VectorType *VT = dyn_cast<VectorType>(Op->getType())) {
       GEPTy = VectorType::get(GEPTy, VT->getElementCount());
+      break;
+    }
+  }
 
   // getelementptr poison, idx -> poison
   // getelementptr baseptr, poison -> poison
   if (any_of(Ops, [](const auto *V) { return isa<PoisonValue>(V); }))
     return PoisonValue::get(GEPTy);
 
   if (Q.isUndefValue(Ops[0]))
     return UndefValue::get(GEPTy);

llvm/test/Transforms/InstSimplify/gep.ll

;
   %c2 = getelementptr inbounds i32, i32* %c1, i64 %offset
   %ptrtoint1 = ptrtoint i32* %c1 to i64
   %ptrtoint2 = ptrtoint i32* %c2 to i64
   %sub = sub i64 %ptrtoint2, %ptrtoint1
   %ashr = ashr exact i64 %sub, 2
   %gep = getelementptr inbounds i32, i32* %c1, i64 %ashr
   ret i32* %gep
 }
+
+define <8 x i32*> @gep_vector_index_op2_poison([144 x i32]* %ptr) {
+; CHECK-LABEL: @gep_vector_index_op2_poison(
+; CHECK-NEXT:    ret <8 x i32*> poison
+;
+  %res = getelementptr inbounds [144 x i32], [144 x i32]* %ptr, i64 0, <8 x i64> poison
+  ret <8 x i32*> %res
+}
+
+%t.1 = type { i32, [144 x i32] }
+
+define <8 x i32*> @gep_vector_index_op3_poison(%t.1* %ptr) {
+; CHECK-LABEL: @gep_vector_index_op3_poison(
+; CHECK-NEXT:    ret <8 x i32*> poison
+;
+  %res = getelementptr inbounds %t.1, %t.1* %ptr, i64 0, i32 1, <8 x i64> poison
+  ret <8 x i32*> %res
+}
+
+%t.2 = type { i32, i32 }
+%t.3 = type { i32, [144 x %t.2 ] }
+
+define <8 x i32*> @gep_vector_index_op3_poison_constant_index_afterwards(%t.3* %ptr) {
+; CHECK-LABEL: @gep_vector_index_op3_poison_constant_index_afterwards(
+; CHECK-NEXT:    ret <8 x i32*> poison
+;
+  %res = getelementptr inbounds %t.3, %t.3* %ptr, i64 0, i32 1, <8 x i64> poison, i32 1
+  ret <8 x i32*> %res
+}