Fix crash with scalable shifts by handling scalable vectors in getSafeVectorConstantForBinop

Author: lukel97

llvm/include/llvm/Transforms/InstCombine/InstCombiner.h

@@ -279,7 +279,7 @@ class LLVM_LIBRARY_VISIBILITY InstCombiner {
   static Constant *
   getSafeVectorConstantForBinop(BinaryOperator::BinaryOps Opcode, Constant *In,
                                 bool IsRHSConstant) {
-    auto *InVTy = cast<FixedVectorType>(In->getType());
+    auto *InVTy = cast<VectorType>(In->getType());
 
     Type *EltTy = InVTy->getElementType();
     auto *SafeC = ConstantExpr::getBinOpIdentity(Opcode, EltTy, IsRHSConstant);
@@ -320,7 +320,13 @@ class LLVM_LIBRARY_VISIBILITY InstCombiner {
       }
     }
     assert(SafeC && "Must have safe constant for binop");
-    unsigned NumElts = InVTy->getNumElements();
+    if (isa<ScalableVectorType>(InVTy)) {
+      if (isa<UndefValue>(In->getSplatValue()))
+        return ConstantVector::getSplat(InVTy->getElementCount(), SafeC);
+      return In;
+    }
+
+    unsigned NumElts = cast<FixedVectorType>(InVTy)->getNumElements();
     SmallVector<Constant *, 16> Out(NumElts);
     for (unsigned i = 0; i != NumElts; ++i) {
       Constant *C = In->getAggregateElement(i);

llvm/test/Transforms/InstCombine/vec_shuffle.ll

@@ -938,6 +938,28 @@ define <2 x i32> @shl_splat_constant1(<2 x i32> %x) {
   ret <2 x i32> %r
 }
 
+define <vscale x 2 x i32> @shl_splat_constant0_scalable(<vscale x 2 x i32> %x) {
+; CHECK-LABEL: @shl_splat_constant0_scalable(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl <vscale x 2 x i32> splat (i32 5), [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 2 x i32> [[R]]
+;
+  %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+  %r = shl <vscale x 2 x i32> splat (i32 5), %splat
+  ret <vscale x 2 x i32> %r
+}
+
+define <vscale x 2 x i32> @shl_splat_constant1_scalable(<vscale x 2 x i32> %x) {
+; CHECK-LABEL: @shl_splat_constant1_scalable(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl <vscale x 2 x i32> [[X:%.*]], splat (i32 5)
+; CHECK-NEXT:    [[R:%.*]] = shufflevector <vscale x 2 x i32> [[TMP1]], <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <vscale x 2 x i32> [[R]]
+;
+  %splat = shufflevector <vscale x 2 x i32> %x, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+  %r = shl <vscale x 2 x i32> %splat, splat (i32 5)
+  ret <vscale x 2 x i32> %r
+}
+
 define <2 x i32> @ashr_splat_constant0(<2 x i32> %x) {
 ; CHECK-LABEL: @ashr_splat_constant0(
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> <i32 5, i32 poison>, [[X:%.*]]