[SVE] Make ConstantFoldGetElementPtr work for scalable vectors of indices

[david-arm]

Summary

This patch fixes a compiler crash that was hit when trying to simplify
the following code:

getelementptr [2 x i64], [2 x i64]* null, i64 0, <vscale x 2 x i64> zeroinitializer

For the case where we have a null pointer value like above, we just
need to ensure we don't assume the indices are always fixed width.

[fpetrogalli]

llvm/test/Transforms/InstSimplify/gep.ll

173

Don't you want to test what IR is generated? Mostly for consistency with the other tests. It also helps understanding what is the expected output for GEP mixing fixed vectors, scalars and scalable vectors.

[david-arm]

llvm/test/Transforms/InstSimplify/gep.ll

173

Yes you're right. I missed that for some reason. I'll push up a new patch. Thanks for spotting it!

[david-arm]

• Added full CHECK lines for new test

[fpetrogalli]

LGTM, thanks!

llvm/lib/IR/ConstantFold.cpp

for (unsigned i = 0, e = Idxs.size(); i != e; ++i)
       }
     if (isNull) {
       PointerType *PtrTy = cast<PointerType>(C->getType()->getScalarType());
       Type *Ty = GetElementPtrInst::getIndexedType(PointeeTy, Idxs);
 
       assert(Ty && "Invalid indices for GEP!");
       Type *OrigGEPTy = PointerType::get(Ty, PtrTy->getAddressSpace());
       Type *GEPTy = PointerType::get(Ty, PtrTy->getAddressSpace());
-      if (VectorType *VT = dyn_cast<VectorType>(C->getType())) {
-        // FIXME: handle scalable vectors (use getElementCount())
-        GEPTy = FixedVectorType::get(
-            OrigGEPTy, cast<FixedVectorType>(VT)->getNumElements());
-      }
+      if (VectorType *VT = dyn_cast<VectorType>(C->getType()))
+        GEPTy = VectorType::get(OrigGEPTy, VT->getElementCount());
+
       // The GEP returns a vector of pointers when one of more of
       // its arguments is a vector.
       for (unsigned i = 0, e = Idxs.size(); i != e; ++i) {
         if (auto *VT = dyn_cast<VectorType>(Idxs[i]->getType())) {
-          // FIXME: handle scalable vectors
-          GEPTy = FixedVectorType::get(
-              OrigGEPTy, cast<FixedVectorType>(VT)->getNumElements());
+          assert((!isa<VectorType>(GEPTy) || isa<ScalableVectorType>(GEPTy) ==
+                                                 isa<ScalableVectorType>(VT)) &&
+                 "Mismatched GEPTy vector types");
+          GEPTy = VectorType::get(OrigGEPTy, VT->getElementCount());
           break;
         }
       }
 
       return Constant::getNullValue(GEPTy);
     }
   }
 

llvm/test/Transforms/InstSimplify/gep.ll

 ; 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
 }
 
+define <vscale x 2 x i64*> @ptr_idx_mix_scalar_scalable_vector() {
+; CHECK-LABEL: @ptr_idx_mix_scalar_scalable_vector(
+; CHECK-NEXT:    ret <vscale x 2 x i64*> zeroinitializer
+;
+  %v = getelementptr [2 x i64], [2 x i64]* null, i64 0, <vscale x 2 x i64> zeroinitializer
+  ret <vscale x 2 x i64*> %v
+}
+
 ; Check ConstantExpr::getGetElementPtr() using ElementCount for size queries - end.