From 7686d4951712468e354350b6d26aa70d052c8db4 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Fri, 8 Dec 2023 17:02:48 +0100
Subject: [ValueTracking] Handle returned attribute with mismatched type

The returned attribute can be used when it is possible to
"losslessly bitcast" between the argument and return type,
including between two vector types.

computeKnownBits() would crash in this case, isKnownNonZero()
would potentially produce a miscompile.

Fixes https://github.com/llvm/llvm-project/issues/74722.
---
 llvm/lib/Analysis/ValueTracking.cpp           |  8 +++++---
 llvm/test/Transforms/InstSimplify/returned.ll | 25 +++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index eb2ee04..5445746 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -1465,8 +1465,10 @@ static void computeKnownBitsFromOperator(const Operator *I,
             Q.IIQ.getMetadata(cast<Instruction>(I), LLVMContext::MD_range))
       computeKnownBitsFromRangeMetadata(*MD, Known);
     if (const Value *RV = cast<CallBase>(I)->getReturnedArgOperand()) {
-      computeKnownBits(RV, Known2, Depth + 1, Q);
-      Known = Known.unionWith(Known2);
+      if (RV->getType() == I->getType()) {
+        computeKnownBits(RV, Known2, Depth + 1, Q);
+        Known = Known.unionWith(Known2);
+      }
     }
     if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
       switch (II->getIntrinsicID()) {
@@ -2712,7 +2714,7 @@ static bool isKnownNonZeroFromOperator(const Operator *I,
       if (const auto *RP = getArgumentAliasingToReturnedPointer(Call, true))
         return isKnownNonZero(RP, Depth, Q);
     } else if (const Value *RV = cast<CallBase>(I)->getReturnedArgOperand()) {
-      if (isKnownNonZero(RV, Depth, Q))
+      if (RV->getType() == I->getType() && isKnownNonZero(RV, Depth, Q))
         return true;
     }
 
diff --git a/llvm/test/Transforms/InstSimplify/returned.ll b/llvm/test/Transforms/InstSimplify/returned.ll
index 94a98ac6..2da1052 100644
--- a/llvm/test/Transforms/InstSimplify/returned.ll
+++ b/llvm/test/Transforms/InstSimplify/returned.ll
@@ -25,6 +25,31 @@ define i1 @gep3() {
   ret i1 %equal
 }
 
+define <8 x i1> @returned_vec_arg_casted(<2 x i32> %a) {
+; CHECK-LABEL: @returned_vec_arg_casted(
+; CHECK-NEXT:    [[X:%.*]] = call <8 x i8> @passthru_8i8v_from_2i32v(<2 x i32> [[A:%.*]])
+; CHECK-NEXT:    [[C:%.*]] = icmp slt <8 x i8> [[X]], zeroinitializer
+; CHECK-NEXT:    ret <8 x i1> [[C]]
+;
+  %x = call <8 x i8> @passthru_8i8v_from_2i32v(<2 x i32> %a)
+  %C = icmp slt <8 x i8> %x, zeroinitializer
+  ret <8 x i1> %C
+}
+
+define <8 x i1> @returned_vec_arg_casted2(<2 x i32> %a) {
+; CHECK-LABEL: @returned_vec_arg_casted2(
+; CHECK-NEXT:    [[OR:%.*]] = or <2 x i32> [[A:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[X:%.*]] = call <8 x i8> @passthru_8i8v_from_2i32v(<2 x i32> [[OR]])
+; CHECK-NEXT:    [[C:%.*]] = icmp ne <8 x i8> [[X]], zeroinitializer
+; CHECK-NEXT:    ret <8 x i1> [[C]]
+;
+  %or = or <2 x i32> %a, <i32 1, i32 1>
+  %x = call <8 x i8> @passthru_8i8v_from_2i32v(<2 x i32> %or)
+  %C = icmp ne <8 x i8> %x, zeroinitializer
+  ret <8 x i1> %C
+}
+
+declare <8 x i8> @passthru_8i8v_from_2i32v(<2 x i32> returned)
 declare ptr @func1(ptr returned) nounwind readnone willreturn
 declare ptr @func2(ptr returned) nounwind readnone willreturn
 
-- 
cgit v1.1