[ConstantFold][SVE] Fix constant folding for scalable vector compare instruction.

[huihuiz]

Summary

Do not iterate on scalable vector. Also do not return constant scalable vector
from ConstantInt::get().
Fix result type by using getElementCount() instead of getNumElements().

[huihuiz]

current upstream crash at : 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 test.ll

define <vscale x 4 x i1> @icmp_undef() {
  %r = icmp eq <vscale x 4 x i32> undef, undef
  ret <vscale x 4 x i1> %r
}

run : ./bin/opt -S -constprop t.ll -o -

[efriedma]

We should fix ConstantInt::get and Constant::getAllOnesValue to work with scalable types; it doesn't make sense to avoid them given we can generate an appropriate constant (splat using insertelement/shufflevector).

[huihuiz]

Thanks Eli for the feedback!

Supporting scalable vector for ConstantInt::get() and Constant::getAllOnesValue() in D74386

[efriedma]

LGTM

llvm/lib/IR/ConstantFold.cpp

   return ICmpInst::BAD_ICMP_PREDICATE;
 }
 
 Constant *llvm::ConstantFoldCompareInstruction(unsigned short pred,
                                                Constant *C1, Constant *C2) {
   Type *ResultTy;
   if (VectorType *VT = dyn_cast<VectorType>(C1->getType()))
     ResultTy = VectorType::get(Type::getInt1Ty(C1->getContext()),
-                               VT->getNumElements());
+                               VT->getElementCount());
   else
     ResultTy = Type::getInt1Ty(C1->getContext());
 
+  bool IsScalableVector =
+      C1->getType()->isVectorTy() && C1->getType()->getVectorIsScalable();
+
   // Fold FCMP_FALSE/FCMP_TRUE unconditionally.
   if (pred == FCmpInst::FCMP_FALSE)
     return Constant::getNullValue(ResultTy);
 
-  if (pred == FCmpInst::FCMP_TRUE)
+  if (!IsScalableVector && pred == FCmpInst::FCMP_TRUE)
     return Constant::getAllOnesValue(ResultTy);
 
   // Handle some degenerate cases first
   if (isa<UndefValue>(C1) || isa<UndefValue>(C2)) {
     CmpInst::Predicate Predicate = CmpInst::Predicate(pred);
     bool isIntegerPredicate = ICmpInst::isIntPredicate(Predicate);
     // For EQ and NE, we can always pick a value for the undef to make the
     // predicate pass or fail, so we can return undef.
     // Also, if both operands are undef, we can return undef for int comparison.
     if (ICmpInst::isEquality(Predicate) || (isIntegerPredicate && C1 == C2))
       return UndefValue::get(ResultTy);
 
-    // Otherwise, for integer compare, pick the same value as the non-undef
-    // operand, and fold it to true or false.
-    if (isIntegerPredicate)
-      return ConstantInt::get(ResultTy, CmpInst::isTrueWhenEqual(Predicate));
+    if (!IsScalableVector) {
+      // Otherwise, for integer compare, pick the same value as the non-undef
+      // operand, and fold it to true or false.
+      if (isIntegerPredicate)
+        return ConstantInt::get(ResultTy, CmpInst::isTrueWhenEqual(Predicate));
 
-    // Choosing NaN for the undef will always make unordered comparison succeed
-    // and ordered comparison fails.
+      // Choosing NaN for the undef will always make unordered comparison
+      // succeed and ordered comparison fails.
       return ConstantInt::get(ResultTy, CmpInst::isUnordered(Predicate));
+    }
   }
 
   // icmp eq/ne(null,GV) -> false/true
   if (C1->isNullValue()) {
     if (const GlobalValue *GV = dyn_cast<GlobalValue>(C2))
       // Don't try to evaluate aliases.  External weak GV can be null.
       if (!isa<GlobalAlias>(GV) && !GV->hasExternalWeakLinkage() &&
           !NullPointerIsDefined(nullptr /* F */,

                                         R==APFloat::cmpEqual);
     case FCmpInst::FCMP_UGE:
       return ConstantInt::get(ResultTy, R!=APFloat::cmpLessThan);
     case FCmpInst::FCMP_OGE:
       return ConstantInt::get(ResultTy, R==APFloat::cmpGreaterThan ||
                                         R==APFloat::cmpEqual);
     }
   } else if (C1->getType()->isVectorTy()) {
+    // Do not iterate on scalable vector. The number of elements is unknown at
+    // compile-time.
+    if (IsScalableVector)
+      return nullptr;
+
     // If we can constant fold the comparison of each element, constant fold
     // the whole vector comparison.
     SmallVector<Constant*, 4> ResElts;
     Type *Ty = IntegerType::get(C1->getContext(), 32);
     // Compare the elements, producing an i1 result or constant expr.
     for (unsigned i = 0, e = C1->getType()->getVectorNumElements(); i != e;++i){
       Constant *C1E =
         ConstantExpr::getExtractElement(C1, ConstantInt::get(Ty, i));

llvm/lib/IR/Constants.cpp

 
   // Look up the constant in the table first to ensure uniqueness
   Constant *ArgVec[] = { LHS, RHS };
   // Get the key type with both the opcode and predicate
   const ConstantExprKeyType Key(Instruction::ICmp, ArgVec, pred);
 
   Type *ResultTy = Type::getInt1Ty(LHS->getContext());
   if (VectorType *VT = dyn_cast<VectorType>(LHS->getType()))
-    ResultTy = VectorType::get(ResultTy, VT->getNumElements());
+    ResultTy = VectorType::get(ResultTy, VT->getElementCount());
 
   LLVMContextImpl *pImpl = LHS->getType()->getContext().pImpl;
   return pImpl->ExprConstants.getOrCreate(ResultTy, Key);
 }
 
 Constant *ConstantExpr::getFCmp(unsigned short pred, Constant *LHS,
                                 Constant *RHS, bool OnlyIfReduced) {
   assert(LHS->getType() == RHS->getType());
   assert(CmpInst::isFPPredicate((CmpInst::Predicate)pred) &&
          "Invalid FCmp Predicate");
 
   if (Constant *FC = ConstantFoldCompareInstruction(pred, LHS, RHS))
     return FC;          // Fold a few common cases...
 
   if (OnlyIfReduced)
     return nullptr;
 
   // Look up the constant in the table first to ensure uniqueness
   Constant *ArgVec[] = { LHS, RHS };
   // Get the key type with both the opcode and predicate
   const ConstantExprKeyType Key(Instruction::FCmp, ArgVec, pred);
 
   Type *ResultTy = Type::getInt1Ty(LHS->getContext());
   if (VectorType *VT = dyn_cast<VectorType>(LHS->getType()))
-    ResultTy = VectorType::get(ResultTy, VT->getNumElements());
+    ResultTy = VectorType::get(ResultTy, VT->getElementCount());
 
   LLVMContextImpl *pImpl = LHS->getType()->getContext().pImpl;
   return pImpl->ExprConstants.getOrCreate(ResultTy, Key);
 }
 
 Constant *ConstantExpr::getExtractElement(Constant *Val, Constant *Idx,
                                           Type *OnlyIfReducedTy) {
   assert(Val->getType()->isVectorTy() &&

llvm/test/Analysis/ConstantFolding/vscale.ll

 
 define <vscale x 4 x i32> @select() {
 ; CHECK-LABEL: @select(
 ; CHECK-NEXT:    ret <vscale x 4 x i32> undef
 ;
   %r = select <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> undef
   ret <vscale x 4 x i32> %r
 }
+
+define <vscale x 4 x i1> @icmp_undef() {
+; CHECK-LABEL: @icmp_undef(
+; CHECK-NEXT:    ret <vscale x 4 x i1> undef
+;
+  %r = icmp eq <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i1> %r
+}
+
+define <vscale x 4 x i1> @icmp_zero() {
+; CHECK-LABEL: @icmp_zero(
+; CHECK-NEXT:    ret <vscale x 4 x i1> icmp eq (<vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> zeroinitializer)
+;
+  %r = icmp eq <vscale x 4 x i32> zeroinitializer, zeroinitializer
+  ret <vscale x 4 x i1> %r
+}
+
+define <vscale x 4 x i1> @fcmp_true() {
+; CHECK-LABEL: @fcmp_true(
+; CHECK-NEXT:    ret <vscale x 4 x i1> fcmp true (<vscale x 4 x float> undef, <vscale x 4 x float> undef)
+;
+  %r = fcmp true <vscale x 4 x float> undef, undef
+  ret <vscale x 4 x i1> %r
+}
+
+define <vscale x 4 x i1> @fcmp_false() {
+; CHECK-LABEL: @fcmp_false(
+; CHECK-NEXT:    ret <vscale x 4 x i1> zeroinitializer
+;
+  %r = fcmp false <vscale x 4 x float> undef, undef
+  ret <vscale x 4 x i1> %r
+}
+
+define <vscale x 4 x i1> @fcmp_undef() {
+; CHECK-LABEL: @fcmp_undef(
+; CHECK-NEXT:    ret <vscale x 4 x i1> undef
+;
+  %r = icmp ne <vscale x 4 x i32> undef, undef
+  ret <vscale x 4 x i1> %r
+}
+
+define <vscale x 4 x i1> @fcmp_not_equality() {
+; CHECK-LABEL: @fcmp_not_equality(
+; CHECK-NEXT:    ret <vscale x 4 x i1> icmp ule (<vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer)
+;
+  %r = icmp ule <vscale x 4 x i32> undef, zeroinitializer
+  ret <vscale x 4 x i1> %r
+}