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

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

Reviewers: sdesmalen, huntergr, spatel, lebedev.ri, apazos, efriedma, willlovett

Reviewed By: efriedma

Subscribers: tschuett, hiraditya, rkruppe, psnobl, llvm-commits

Tags: #llvm

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

2 files changed, 169 insertions, 5 deletions

diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index 6e24f03..3d8ad88 100644
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -1013,10 +1013,14 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
       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) &&
+      (isa<UndefValue>(C1) || isa<UndefValue>(C2));
+  if (HasScalarUndefOrScalableVectorUndef) {
     switch (static_cast<Instruction::BinaryOps>(Opcode)) {
     case Instruction::Xor:
       if (isa<UndefValue>(C1) && isa<UndefValue>(C2))
@@ -1119,7 +1123,7 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
   }
 
   // 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)) {
@@ -1330,6 +1334,11 @@ Constant *llvm::ConstantFoldBinaryInstruction(unsigned Opcode, Constant *C1,
       }
     }
   } 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);
diff --git a/llvm/test/Analysis/ConstantFolding/vscale.ll b/llvm/test/Analysis/ConstantFolding/vscale.ll
new file mode 100644
index 0000000..2c6e8f1
--- /dev/null
+++ b/llvm/test/Analysis/ConstantFolding/vscale.ll
@@ -0,0 +1,155 @@
+; 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
+}