[ConstantFold] Fold binary arithmetic on scalable vector splats.

[efriedma]

Summary

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

[david-arm]

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

44

Is it worth also adding a test case for the integer divrem case where C2=0 that we now support where we return undef?

[efriedma]

Add an extra test so the divide path is exercised.

[david-arm]

LGTM! Thanks for adding the test!

llvm/lib/IR/ConstantFold.cpp

if (ConstantFP *CFP2 = dyn_cast<ConstantFP>(C2)) {
       case Instruction::FDiv:
         (void)C3V.divide(C2V, APFloat::rmNearestTiesToEven);
         return ConstantFP::get(C1->getContext(), C3V);
       case Instruction::FRem:
         (void)C3V.mod(C2V);
         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())
         return UndefValue::get(VTy);
       if (Constant *C1Splat = C1->getSplatValue()) {
         return ConstantVector::getSplat(
             VTy->getElementCount(),
             ConstantExpr::get(Opcode, C1Splat, C2Splat));
       }
     }
 
+    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(VTy->getContext(), 32);
-    for (unsigned i = 0, e = VTy->getNumElements(); i != e; ++i) {
+      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);
 
         Result.push_back(ConstantExpr::get(Opcode, LHS, RHS));
       }
 
       return ConstantVector::get(Result);
     }
+  }
 
   if (ConstantExpr *CE1 = dyn_cast<ConstantExpr>(C1)) {
     // There are many possible foldings we could do here.  We should probably
     // at least fold add of a pointer with an integer into the appropriate
     // getelementptr.  This will improve alias analysis a bit.
 
     // Given ((a + b) + c), if (b + c) folds to something interesting, return
     // (a + (b + c)).

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

 define <vscale x 4 x i32> @sub() {
 ; CHECK-LABEL: @sub(
 ; CHECK-NEXT:    ret <vscale x 4 x i32> undef
 ;
   %r = sub <vscale x 4 x i32> undef, undef
   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
 ;
   %r = fsub <vscale x 4 x float> undef, undef
   ret <vscale x 4 x float> %r
 }
 

 define <vscale x 4 x i32> @udiv() {
 ; CHECK-LABEL: @udiv(
 ; CHECK-NEXT:    ret <vscale x 4 x i32> undef
 ;
   %r = udiv <vscale x 4 x i32> undef, undef
   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
 ;
   %r = sdiv <vscale x 4 x i32> undef, undef
   ret <vscale x 4 x i32> %r
 }