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

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

Reviewers: sdesmalen, efriedma, apazos, huntergr, willlovett

Reviewed By: efriedma

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

Tags: #llvm

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

llvm/lib/IR/ConstantFold.cpp

@@ -953,10 +953,14 @@ Constant *llvm::ConstantFoldInsertValueInstruction(Constant *Agg,
 Constant *llvm::ConstantFoldUnaryInstruction(unsigned Opcode, Constant *C) {
   assert(Instruction::isUnaryOp(Opcode) && "Non-unary instruction detected");
 
-  // Handle scalar UndefValue. Vectors are always evaluated per element.
+  // Handle scalar UndefValue and scalable vector UndefValue. Fixed-length
-  bool HasScalarUndef = !C->getType()->isVectorTy() && isa<UndefValue>(C);
+  // vectors are always evaluated per element.
+  bool IsScalableVector =
+      C->getType()->isVectorTy() && C->getType()->getVectorIsScalable();
+  bool HasScalarUndefOrScalableVectorUndef =
+      (!C->getType()->isVectorTy() || IsScalableVector) && isa<UndefValue>(C);
 
-  if (HasScalarUndef) {
+  if (HasScalarUndefOrScalableVectorUndef) {
     switch (static_cast<Instruction::UnaryOps>(Opcode)) {
     case Instruction::FNeg:
       return C; // -undef -> undef
@@ -966,7 +970,7 @@ Constant *llvm::ConstantFoldUnaryInstruction(unsigned Opcode, Constant *C) {
   }
 
   // Constant should not be UndefValue, unless these are vector constants.
-  assert(!HasScalarUndef && "Unexpected UndefValue");
+  assert(!HasScalarUndefOrScalableVectorUndef && "Unexpected UndefValue");
   // We only have FP UnaryOps right now.
   assert(!isa<ConstantInt>(C) && "Unexpected Integer UnaryOp");
 
@@ -979,6 +983,11 @@ Constant *llvm::ConstantFoldUnaryInstruction(unsigned Opcode, Constant *C) {
       return ConstantFP::get(C->getContext(), neg(CV));
     }
   } else if (VectorType *VTy = dyn_cast<VectorType>(C->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);

llvm/test/Analysis/ConstantFolding/vscale.ll

@@ -1,6 +1,17 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -constprop -S | FileCheck %s
 
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; Unary Operations
+;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+
+define <vscale x 2 x double> @fneg(<vscale x 2 x double> %val) {
+; CHECK-LABEL: @fneg(
+; CHECK-NEXT:    ret <vscale x 2 x double> undef
+;
+  %r = fneg <vscale x 2 x double> undef
+  ret <vscale x 2 x double> %r
+}
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; Binary Operations