[ConstantFold] Fold binary arithmetic on scalable vector splats.

It's a nice simplification, and it confuses instcombine if we don't do
it.

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

Author: efriedma-quic

llvm/lib/IR/ConstantFold.cpp

@@ -1408,12 +1408,7 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
         return ConstantFP::get(C1->getContext(), C3V);
       }
     }
-  } else if (IsScalableVector) {
-    // Do not iterate on scalable vector. The number of elements is unknown at
-    // compile-time.
-    // FIXME: this branch can potentially be removed
-    return nullptr;
-  } else if (auto *VTy = dyn_cast<FixedVectorType>(C1->getType())) {
+  } else if (auto *VTy = dyn_cast<VectorType>(C1->getType())) {
     // Fast path for splatted constants.
     if (Constant *C2Splat = C2->getSplatValue()) {
       if (Instruction::isIntDivRem(Opcode) && C2Splat->isNullValue())
@@ -1425,22 +1420,24 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
       }
     }
 
-    // Fold each element and create a vector constant from those constants.
-    SmallVector<Constant*, 16> Result;
-    Type *Ty = IntegerType::get(VTy->getContext(), 32);
-    for (unsigned i = 0, e = VTy->getNumElements(); i != e; ++i) {
-      Constant *ExtractIdx = ConstantInt::get(Ty, i);
-      Constant *LHS = ConstantExpr::getExtractElement(C1, ExtractIdx);
-      Constant *RHS = ConstantExpr::getExtractElement(C2, ExtractIdx);
+    if (auto *FVTy = dyn_cast<FixedVectorType>(VTy)) {
+      // Fold each element and create a vector constant from those constants.
+      SmallVector<Constant*, 16> Result;
+      Type *Ty = IntegerType::get(FVTy->getContext(), 32);
+      for (unsigned i = 0, e = FVTy->getNumElements(); i != e; ++i) {
+        Constant *ExtractIdx = ConstantInt::get(Ty, i);
+        Constant *LHS = ConstantExpr::getExtractElement(C1, ExtractIdx);
+        Constant *RHS = ConstantExpr::getExtractElement(C2, ExtractIdx);
 
-      // If any element of a divisor vector is zero, the whole op is undef.
-      if (Instruction::isIntDivRem(Opcode) && RHS->isNullValue())
-        return UndefValue::get(VTy);
+        // If any element of a divisor vector is zero, the whole op is undef.
+        if (Instruction::isIntDivRem(Opcode) && RHS->isNullValue())
+          return UndefValue::get(VTy);
 
-      Result.push_back(ConstantExpr::get(Opcode, LHS, RHS));
-    }
+        Result.push_back(ConstantExpr::get(Opcode, LHS, RHS));
+      }
 
-    return ConstantVector::get(Result);
+      return ConstantVector::get(Result);
+    }
   }
 
   if (ConstantExpr *CE1 = dyn_cast<ConstantExpr>(C1)) {

llvm/test/Transforms/InstSimplify/ConstProp/vscale.ll

@@ -41,6 +41,14 @@ define <vscale x 4 x i32> @sub() {
   ret <vscale x 4 x i32> %r
 }
 
+define <vscale x 4 x i32> @sub_splat() {
+; CHECK-LABEL: @sub_splat(
+; CHECK-NEXT:    ret <vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 -16, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer)
+;
+  %r = sub <vscale x 4 x i32> zeroinitializer, shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 16, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer)
+  ret <vscale x 4 x i32> %r
+}
+
 define <vscale x 4 x float> @fsub() {
 ; CHECK-LABEL: @fsub(
 ; CHECK-NEXT:    ret <vscale x 4 x float> undef
@@ -73,6 +81,14 @@ define <vscale x 4 x i32> @udiv() {
   ret <vscale x 4 x i32> %r
 }
 
+define <vscale x 4 x i32> @udiv_splat_zero() {
+; CHECK-LABEL: @udiv_splat_zero(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = udiv <vscale x 4 x i32> zeroinitializer, zeroinitializer
+  ret <vscale x 4 x i32> %r
+}
+
 define <vscale x 4 x i32> @sdiv() {
 ; CHECK-LABEL: @sdiv(
 ; CHECK-NEXT:    ret <vscale x 4 x i32> undef