[SVE] Make ConstantFoldGetElementPtr work for scalable vectors of indices

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.

Differential Revision: https://reviews.llvm.org/D82183

diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index 81eaaaf..f3c3e9a 100644 (file)
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -2298,18 +2298,17 @@ Constant *llvm::ConstantFoldGetElementPtr(Type *PointeeTy, Constant *C,
       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;
         }
       }

diff --git a/llvm/test/Transforms/InstSimplify/gep.ll b/llvm/test/Transforms/InstSimplify/gep.ll
index ea5e72e..5044b50 100644 (file)
--- a/llvm/test/Transforms/InstSimplify/gep.ll
+++ b/llvm/test/Transforms/InstSimplify/gep.ll
@@ -176,4 +176,12 @@ define <vscale x 4 x float*> @scalable_vector_idx_mix_scalar_vector() {
   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.