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

[huihuiz]

Summary

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

[huihuiz]

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

take this test.ll

define <vscale x 2 x double> @fneg(<vscale x 2 x double> %val) {
  %r = fneg <vscale x 2 x double> undef
  ret <vscale x 2 x double> %r
}

run: opt -S -constprop test.ll -o -

[efriedma]

LGTM

llvm/lib/IR/ConstantFold.cpp

   if (ArrayType *AT = dyn_cast<ArrayType>(Agg->getType()))
     return ConstantArray::get(AT, Result);
   return ConstantVector::get(Result);
 }
 
 Constant *llvm::ConstantFoldUnaryInstruction(unsigned Opcode, Constant *C) {
   assert(Instruction::isUnaryOp(Opcode) && "Non-unary instruction detected");
 
-  // Handle scalar UndefValue. Vectors are always evaluated per element.
-  bool HasScalarUndef = !C->getType()->isVectorTy() && isa<UndefValue>(C);
+  // Handle scalar UndefValue and scalable vector UndefValue. Fixed-length
+  // 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
     case Instruction::UnaryOpsEnd:
       llvm_unreachable("Invalid UnaryOp");
     }
   }
 
   // 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");
 
   if (ConstantFP *CFP = dyn_cast<ConstantFP>(C)) {
     const APFloat &CV = CFP->getValueAPF();
     switch (Opcode) {
     default:
       break;
     case Instruction::FNeg:
       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);
     for (unsigned i = 0, e = VTy->getNumElements(); i != e; ++i) {
       Constant *ExtractIdx = ConstantInt::get(Ty, i);
       Constant *Elt = ConstantExpr::getExtractElement(C, ExtractIdx);
 
       Result.push_back(ConstantExpr::get(Opcode, Elt));

llvm/test/Analysis/ConstantFolding/vscale.ll

 ; 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
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 define <vscale x 4 x i32> @add() {
 ; CHECK-LABEL: @add(
 ; CHECK-NEXT:    ret <vscale x 4 x i32> undef