[ConstantFold][SVE] Fix constant folding for scalable vector binary operations.

[huihuiz]

Summary

Scalable vector should not be evaluated element by element.
Add support to handle scalable vector UndefValue.

[huihuiz]

Similar issue , current upstream crash at void llvm::Value::doRAUW(llvm::Value *, llvm::Value::ReplaceMetadataUses): Assertion `New->getType() == getType() && "replaceAllUses of value with new value of different type!"' failed.

opt -S -constprop add.ll -o -

add.ll

define <vscale x 8 x i16> @add() {
  %r = add <vscale x 8 x i16> undef, undef
  ret <vscale x 8 x i16> %r
}

I will also be moving other sve tests under test/Analysis/ConstantFolding/ to vscale.ll

[efriedma]

I'd prefer to handle all the issues in ConstantFoldBinaryInstruction as one review, I think.

[huihuiz]

Hang on a bit.
Let me see if I can find more issues in ConstantFoldBinaryInstruction.

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

1019

nit: maybe move the subexpression

C1->getType()->isVectorTy() && C1->getType()->getVectorIsScalable()

to a separate IsScalableVector variable, so that this expression is easier to read?

[huihuiz]

Apologies for the delay, was on vacation, also got caught up with a work project in the past two weeks.

There is more we can do for scalable vector for ConstantFoldBinaryInstruction.
E.g., Constant fold binary op with constant splats, refer to https://reviews.llvm.org/D71712

<binop> (<vty> splat(<ty> <C1>)),
                (<vty> splat(<ty> <C2>))
==>  <vty> splat (<ty> <binop> (C1, C2))

Enable this will require changing API of ConstantVector::getSplat() to take in ElementCount for number of elements.
Probably better to do in a separate patch.

For this patch, we might want to focus on not iterating over scalable vectors.

[efriedma]

LGTM

llvm/test/Analysis/ConstantFolding/vscale.ll

2

You can drop the first line, since you edited the test by hand.

llvm/lib/IR/ConstantFold.cpp

Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
   Constant *Identity = ConstantExpr::getBinOpIdentity(Opcode, C1->getType());
   if (Identity) {
     if (C1 == Identity)
       return C2;
     if (C2 == Identity)
       return C1;
   }
 
-  // Handle scalar UndefValue. Vectors are always evaluated per element.
-  bool HasScalarUndef = !C1->getType()->isVectorTy() &&
-                        (isa<UndefValue>(C1) || isa<UndefValue>(C2));
-  if (HasScalarUndef) {
+  // Handle scalar UndefValue and scalable vector UndefValue. Fixed-length
+  // vectors are always evaluated per element.
+  bool IsScalableVector =
+      C1->getType()->isVectorTy() && C1->getType()->getVectorIsScalable();
+  bool HasScalarUndefOrScalableVectorUndef =
+      (!C1->getType()->isVectorTy() || IsScalableVector) &&
+  if (HasScalarUndefOrScalableVectorUndef) {
     switch (static_cast<Instruction::BinaryOps>(Opcode)) {
     case Instruction::Xor:
       if (isa<UndefValue>(C1) && isa<UndefValue>(C2))
         // Handle undef ^ undef -> 0 special case. This is a common
         // idiom (misuse).
         return Constant::getNullValue(C1->getType());
       LLVM_FALLTHROUGH;
     case Instruction::Add:
case Instruction::FRem:
       // with a NaN operand will propagate NaN.
       return ConstantFP::getNaN(C1->getType());
     case Instruction::BinaryOpsEnd:
       llvm_unreachable("Invalid BinaryOp");
     }
   }
 
   // Neither constant should be UndefValue, unless these are vector constants.
-  assert(!HasScalarUndef && "Unexpected UndefValue");
+  assert((!HasScalarUndefOrScalableVectorUndef) && "Unexpected UndefValue");
 
   // Handle simplifications when the RHS is a constant int.
   if (ConstantInt *CI2 = dyn_cast<ConstantInt>(C2)) {
     switch (Opcode) {
     case Instruction::Add:
       if (CI2->isZero()) return C1;                             // X + 0 == X
       break;
     case Instruction::Sub:
if (ConstantFP *CFP2 = dyn_cast<ConstantFP>(C2)) {
         (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 (VectorType *VTy = dyn_cast<VectorType>(C1->getType())) {
+    // Do not iterate on scalable vector. The number of elements is unknown at
+    // compile-time.
+    if (IsScalableVector)
+      return nullptr;
+
     // 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);
 

llvm/test/Analysis/ConstantFolding/vscale.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -constprop -S | FileCheck %s
+
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Binary Operations
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define <vscale x 4 x i32> @add() {
+; CHECK-LABEL: @add(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = add <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x float> @fadd() {
+; CHECK-LABEL: @fadd(
+; CHECK-NEXT:    ret <vscale x 4 x float> undef
+;
+  %r = fadd <vscale x 4 x float> undef, undef
+  ret <vscale x 4 x float> %r
+}
+
+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 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> @mul() {
+; CHECK-LABEL: @mul(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = mul <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x float> @fmul() {
+; CHECK-LABEL: @fmul(
+; CHECK-NEXT:    ret <vscale x 4 x float> undef
+;
+  %r = fmul <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> @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
+}
+
+define <vscale x 4 x float> @fdiv() {
+; CHECK-LABEL: @fdiv(
+; CHECK-NEXT:    ret <vscale x 4 x float> undef
+;
+  %r = fdiv <vscale x 4 x float> undef, undef
+  ret <vscale x 4 x float> %r
+}
+
+define <vscale x 4 x i32> @urem() {
+; CHECK-LABEL: @urem(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = urem <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @srem() {
+; CHECK-LABEL: @srem(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = srem <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x float> @frem() {
+; CHECK-LABEL: @frem(
+; CHECK-NEXT:    ret <vscale x 4 x float> undef
+;
+  %r = frem <vscale x 4 x float> undef, undef
+  ret <vscale x 4 x float> %r
+}
+
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Bitwise Binary Operations
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define <vscale x 4 x i32> @shl() {
+; CHECK-LABEL: @shl(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = shl <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @lshr() {
+; CHECK-LABEL: @lshr(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = lshr <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @ashr() {
+; CHECK-LABEL: @ashr(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = ashr <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @and() {
+; CHECK-LABEL: @and(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = and <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @or() {
+; CHECK-LABEL: @or(
+; CHECK-NEXT:    ret <vscale x 4 x i32> undef
+;
+  %r = or <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}
+
+define <vscale x 4 x i32> @xor() {
+; CHECK-LABEL: @xor(
+; CHECK-NEXT:    ret <vscale x 4 x i32> zeroinitializer
+;
+  %r = xor <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i32> %r
+}