[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.

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