From 2fab04a4e45b6b00d97e024f8627aa7ab45e63c2 Mon Sep 17 00:00:00 2001
From: Huihui Zhang <huihuiz@quicinc.com>
Date: Thu, 12 Mar 2020 15:59:14 -0700
Subject: [PATCH] [ConstantFold][SVE] Fix constant folding for scalable vector
 compare instruction.

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().

Reviewers: sdesmalen, efriedma, apazos, huntergr, willlovett

Reviewed By: efriedma

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

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D73753
---
 lib/IR/ConstantFold.cpp                 |  7 +++-
 lib/IR/Constants.cpp                    |  4 +--
 test/Analysis/ConstantFolding/vscale.ll | 48 +++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/lib/IR/ConstantFold.cpp b/lib/IR/ConstantFold.cpp
index 0da027f56fd..dc78c5537be 100644
--- a/lib/IR/ConstantFold.cpp
+++ b/lib/IR/ConstantFold.cpp
@@ -1840,7 +1840,7 @@ Constant *llvm::ConstantFoldCompareInstruction(unsigned short pred,
   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());
 
@@ -1971,6 +1971,11 @@ Constant *llvm::ConstantFoldCompareInstruction(unsigned short pred,
                                         R==APFloat::cmpEqual);
     }
   } else if (C1->getType()->isVectorTy()) {
+    // Do not iterate on scalable vector. The number of elements is unknown at
+    // compile-time.
+    if (C1->getType()->getVectorIsScalable())
+      return nullptr;
+
     // If we can constant fold the comparison of each element, constant fold
     // the whole vector comparison.
     SmallVector<Constant*, 4> ResElts;
diff --git a/lib/IR/Constants.cpp b/lib/IR/Constants.cpp
index eb0e5894ae5..c609add1f12 100644
--- a/lib/IR/Constants.cpp
+++ b/lib/IR/Constants.cpp
@@ -2175,7 +2175,7 @@ Constant *ConstantExpr::getICmp(unsigned short pred, Constant *LHS,
 
   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);
@@ -2200,7 +2200,7 @@ Constant *ConstantExpr::getFCmp(unsigned short pred, Constant *LHS,
 
   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);
diff --git a/test/Analysis/ConstantFolding/vscale.ll b/test/Analysis/ConstantFolding/vscale.ll
index 1616c70f200..f00e74bd14a 100644
--- a/test/Analysis/ConstantFolding/vscale.ll
+++ b/test/Analysis/ConstantFolding/vscale.ll
@@ -235,3 +235,51 @@ define <vscale x 16 x i8> @call() {
   %r =  call <vscale x 16 x i8> @llvm.sadd.sat.nxv16i8(<vscale x 16 x i8> undef, <vscale x 16 x i8> undef)
   ret <vscale x 16 x i8> %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> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> undef, i1 true, i32 0), <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer)
+;
+  %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> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> undef, i1 true, i32 0), <vscale x 4 x i1> undef, <vscale x 4 x i32> zeroinitializer)
+;
+  %r = icmp ule <vscale x 4 x i32> undef, zeroinitializer
+  ret <vscale x 4 x i1> %r
+}