From 899fe2cf318f6e3c7a183125ff17ef53944b2989 Mon Sep 17 00:00:00 2001
From: goldsteinn <35538541+goldsteinn@users.noreply.github.com>
Date: Thu, 4 Jul 2024 15:14:39 +0800
Subject: [PATCH] [CVP][LVI] Fix incorrect scalar type when getting constant
 folded vec (#97682)

Fixes #97674

After #97428 added support for vectors, our constant ranges can now be
from splat vectors so when they reduce to a singe constant value, we
need to return the original type as opposed to just an int.
---
 llvm/lib/Analysis/LazyValueInfo.cpp           |  5 +++--
 .../CorrelatedValuePropagation/vectors.ll     | 21 +++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index b30e6a6a367c5..4209ee4a77953 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -1370,6 +1370,7 @@ LazyValueInfoImpl::getEdgeValueLocal(Value *Val, BasicBlock *BBFrom,
 
       // If V is the condition of the branch itself, then we know exactly what
       // it is.
+      // NB: The condition on a `br` can't be a vector type.
       if (Condition == Val)
         return ValueLatticeElement::get(ConstantInt::get(
                               Type::getInt1Ty(Val->getContext()), isTrueDest));
@@ -1723,7 +1724,7 @@ Constant *LazyValueInfo::getConstant(Value *V, Instruction *CxtI) {
   if (Result.isConstantRange()) {
     const ConstantRange &CR = Result.getConstantRange();
     if (const APInt *SingleVal = CR.getSingleElement())
-      return ConstantInt::get(V->getContext(), *SingleVal);
+      return ConstantInt::get(V->getType(), *SingleVal);
   }
   return nullptr;
 }
@@ -1758,7 +1759,7 @@ Constant *LazyValueInfo::getConstantOnEdge(Value *V, BasicBlock *FromBB,
   if (Result.isConstantRange()) {
     const ConstantRange &CR = Result.getConstantRange();
     if (const APInt *SingleVal = CR.getSingleElement())
-      return ConstantInt::get(V->getContext(), *SingleVal);
+      return ConstantInt::get(V->getType(), *SingleVal);
   }
   return nullptr;
 }
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll b/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
index 0024b0a5c75c9..caaed628ed43e 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/vectors.ll
@@ -220,3 +220,24 @@ define <2 x i16> @and_with_poison(<2 x i8> %a) {
   %res = and <2 x i16> %zext, <i16 u0xff, i16 poison>
   ret <2 x i16> %res
 }
+
+
+
+define <4 x i64> @issue_97674_getConstantOnEdge(i1 %cond) {
+entry:
+  br i1 %cond, label %if.then, label %if.end
+
+if.then:
+  %folds = add <4 x i64> zeroinitializer, <i64 1, i64 1, i64 1, i64 1>
+  br label %if.end
+
+if.end:
+  %r = phi <4 x i64> [ %folds, %if.then ], [ zeroinitializer, %entry ]
+  ret <4 x i64> %r
+}
+    
+define <4 x i64> @issue_97674_getConstant() {
+entry:
+  %folds = add <4 x i64> zeroinitializer, zeroinitializer
+  ret <4 x i64> %folds
+}