[IR] Prefer scalar type for struct indexes in GEP constant expressions.

[efriedma]

Summary

Using a scalar constant instead of a vector constant has two advantages: one, it's simpler to understand, and two, it doesn't require complex pattern matching to work with scalable vectors.

Also includes a small fix to DataLayout to allow the scalable vector testcase to work correctly.

[ctetreau]

llvm/lib/IR/Constants.cpp

2178–2180

can this use auto?

2178–2180

what is the change to iterators from indices buying us?

llvm/test/Analysis/ConstantFolding/vectorgep-crash.ll

27

the old thing should still work right? And presumably it's testing some crash that happened, perhaps we should leave it alone?

[efriedma]

llvm/lib/IR/Constants.cpp

2178–2180

I think this can use auto, yes.

The point of using the iterators is that they provide the isStruct() and isSequential() methods.

llvm/test/Analysis/ConstantFolding/vectorgep-crash.ll

27

I'm only changing the CHECK line, not the input IR

[ctetreau]

llvm/lib/IR/Constants.cpp

2178–2180

Could these not be replaced by isa checks? isSequential seems especially useless since it checks that the thing is a Type *

[efriedma]

llvm/lib/IR/Constants.cpp

2178–2180

At a high level, generic_gep_type_iterator takes a GEP base type and some indexes, and steps through the indexed types. The computed types can be queried with isStruct()/isSequential()/etc. This is information which we can't get directly from the indexes themselves.

At a low level, I think you're misreading the implementation of generic_gep_type_iterator; look at the documentation for PointerUnion.

[sdesmalen]

llvm/lib/IR/Constants.cpp

2190–2191

What is the reason for not always using a scalar index value?

[efriedma]

llvm/lib/IR/Constants.cpp

2190–2191

In my thinking, there are two reasons:

1. I wanted to be conservative; arbitrary constant expressions are legal in this position anyway, so the harm from code failing to understand a vector splat is much smaller.
2. We don't want to accidentally change the result type of a GEP if every operand is a splat.

[sdesmalen]

llvm/lib/IR/Constants.cpp

2190–2191

Okay so if I understand that correctly, we can do this change safely for struct types without risking to change the result type. because in e.g.:

%gep = getelementptr %struct, %struct* null, i32 0, <4 x i32> <i32 1, i32 1, i32 1, i32 1>

the i32 0 will always be expanded into a vector:

%gep = getelementptr %struct, %struct* null, <4 x i32> zeroinitializer, i32 1

And so there is always at least one index that will remain a vector (splat).

[efriedma]

llvm/lib/IR/Constants.cpp

2190–2191

Yes, exactly.

[sdesmalen]

LGTM!

[ctetreau]

I guess my only real remaining issue is to use auto for the type of the iterators. Otherwise this looks good to me.

llvm/lib/IR/Constants.cpp

Constant *ConstantExpr::getGetElementPtr(Type *Ty, Constant *C,
 
   if (OnlyIfReducedTy == ReqTy)
     return nullptr;
 
   // Look up the constant in the table first to ensure uniqueness
   std::vector<Constant*> ArgVec;
   ArgVec.reserve(1 + Idxs.size());
   ArgVec.push_back(C);
-  for (unsigned i = 0, e = Idxs.size(); i != e; ++i) {
+  auto GTI = gep_type_begin(Ty, Idxs), GTE = gep_type_end(Ty, Idxs);
+  for (; GTI != GTE; ++GTI) {
+    auto *Idx = cast<Constant>(GTI.getOperand());
     assert(
-        (!isa<VectorType>(Idxs[i]->getType()) ||
-         cast<VectorType>(Idxs[i]->getType())->getElementCount() == EltCount) &&
+        (!isa<VectorType>(Idx->getType()) ||
+         cast<VectorType>(Idx->getType())->getElementCount() == EltCount) &&
         "getelementptr index type missmatch");
 
-    Constant *Idx = cast<Constant>(Idxs[i]);
-    if (EltCount.Min != 0 && !Idxs[i]->getType()->isVectorTy())
+    if (GTI.isStruct() && Idx->getType()->isVectorTy()) {
+      Idx = Idx->getSplatValue();
+    } else if (GTI.isSequential() && EltCount.Min != 0 &&
+               !Idx->getType()->isVectorTy()) {
       Idx = ConstantVector::getSplat(EltCount, Idx);
+    }
     ArgVec.push_back(Idx);
   }
 
   unsigned SubClassOptionalData = InBounds ? GEPOperator::IsInBounds : 0;
   if (InRangeIndex && *InRangeIndex < 63)
     SubClassOptionalData |= (*InRangeIndex + 1) << 1;
   const ConstantExprKeyType Key(Instruction::GetElementPtr, ArgVec, 0,
                                 SubClassOptionalData, None, None, Ty);

llvm/lib/IR/DataLayout.cpp

 }
 
 Type *DataLayout::getIndexType(Type *Ty) const {
   assert(Ty->isPtrOrPtrVectorTy() &&
          "Expected a pointer or pointer vector type.");
   unsigned NumBits = getIndexTypeSizeInBits(Ty);
   IntegerType *IntTy = IntegerType::get(Ty->getContext(), NumBits);
   if (VectorType *VecTy = dyn_cast<VectorType>(Ty))
-    return FixedVectorType::get(IntTy, VecTy->getNumElements());
+    return VectorType::get(IntTy, VecTy);
   return IntTy;
 }
 
 int64_t DataLayout::getIndexedOffsetInType(Type *ElemTy,
                                            ArrayRef<Value *> Indices) const {
   int64_t Result = 0;
 
   generic_gep_type_iterator<Value* const*>

llvm/test/Analysis/ConstantFolding/vectorgep-crash.ll

 }
 
 %struct.A = type { i32, %struct.B* }
 %struct.B = type { i64, %struct.C* }
 %struct.C = type { i64 }
 
 @G = internal global [65 x %struct.A] zeroinitializer, align 16
 ; CHECK-LABEL: @test
-; 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, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, <16 x i32> zeroinitializer)
+; 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, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, i32 0)
 define <16 x i32*> @test() {
 vector.body:
   %VectorGep = 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, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, <16 x i32> zeroinitializer
   ret <16 x i32*> %VectorGep
 }
 
 ; 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, 

llvm/test/Transforms/InstSimplify/gep.ll

 ;
   %gep = getelementptr i32, <4 x i32*> zeroinitializer, <4 x i64> <i64 1, i64 1, i64 1, i64 1>
   ret <4 x i32*> %gep
 }
 
 %struct = type { double, float }
 define <4 x float*> @vector_idx_mix_scalar_vector() {
 ; CHECK-LABEL: @vector_idx_mix_scalar_vector(
-; CHECK-NEXT:    ret <4 x float*> getelementptr (%struct, <4 x %struct*> zeroinitializer, <4 x i64> zeroinitializer, <4 x i32> <i32 1, i32 1, i32 1, i32 1>)
+; CHECK-NEXT:    ret <4 x float*> getelementptr (%struct, <4 x %struct*> zeroinitializer, <4 x i64> zeroinitializer, i32 1)
 ;
   %gep = getelementptr %struct, <4 x %struct*> zeroinitializer, i32 0, <4 x i32> <i32 1, i32 1, i32 1, i32 1>
   ret <4 x float*> %gep
 }
 
 ; Constant scalable
 
 define <vscale x 4 x i32*> @scalable_idx_scalar() {
 ; CHECK-LABEL: @scalable_idx_scalar(
 ; CHECK-NEXT:    ret <vscale x 4 x i32*> getelementptr (i32, <vscale x 4 x i32*> zeroinitializer, <vscale x 4 x i64> shufflevector (<vscale x 4 x i64> insertelement (<vscale x 4 x i64> undef, i64 1, i32 0), <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer))
 ;
   %gep = getelementptr i32, <vscale x 4 x i32*> zeroinitializer, i64 1
   ret <vscale x 4 x i32*> %gep
 }
 
+define <vscale x 4 x float*> @scalable_vector_idx_mix_scalar_vector() {
+; CHECK-LABEL: @scalable_vector_idx_mix_scalar_vector(
+; CHECK-NEXT:    ret <vscale x 4 x float*> getelementptr (%struct, <vscale x 4 x %struct*> zeroinitializer, <vscale x 4 x i64> zeroinitializer, i32 1)
+;
+  %gep = getelementptr %struct, <vscale x 4 x %struct*> zeroinitializer, i32 0, i32 1
+  ret <vscale x 4 x float*> %gep
+}
+
 ; Check ConstantExpr::getGetElementPtr() using ElementCount for size queries - end.