[ValueTracking] Add support for `shufflevector` in `isKnownNonZero`

Shuffles don't modify the data, so if all elements that end up in the
destination are non-zero the result is non-zero.

Closes #87702

2 files changed, 20 insertions, 21 deletions

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 5ef1969..b3029f4 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2777,6 +2777,21 @@ static bool isKnownNonZeroFromOperator(const Operator *I,
       }
     }
     break;
+  case Instruction::ShuffleVector: {
+    auto *Shuf = dyn_cast<ShuffleVectorInst>(I);
+    if (!Shuf)
+      break;
+    APInt DemandedLHS, DemandedRHS;
+    // For undef elements, we don't know anything about the common state of
+    // the shuffle result.
+    if (!getShuffleDemandedElts(Shuf, DemandedElts, DemandedLHS, DemandedRHS))
+      break;
+    // If demanded elements for both vecs are non-zero, the shuffle is non-zero.
+    return (DemandedRHS.isZero() ||
+            isKnownNonZero(Shuf->getOperand(1), DemandedRHS, Depth, Q)) &&
+           (DemandedLHS.isZero() ||
+            isKnownNonZero(Shuf->getOperand(0), DemandedLHS, Depth, Q));
+  }
   case Instruction::Freeze:
     return isKnownNonZero(I->getOperand(0), Depth, Q) &&
            isGuaranteedNotToBePoison(I->getOperand(0), Q.AC, Q.CxtI, Q.DT,
diff --git a/llvm/test/Transforms/InstSimplify/known-non-zero.ll b/llvm/test/Transforms/InstSimplify/known-non-zero.ll
index 9486708..c443dc6 100644
--- a/llvm/test/Transforms/InstSimplify/known-non-zero.ll
+++ b/llvm/test/Transforms/InstSimplify/known-non-zero.ll
@@ -180,11 +180,7 @@ define i1 @smax_non_zero(i8 %xx, i8 %y) {
 
 define <4 x i1> @shuf_nonzero_both(<4 x i8> %xx, <4 x i8> %yy) {
 ; CHECK-LABEL: @shuf_nonzero_both(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <4 x i8> [[XX:%.*]], <i8 1, i8 1, i8 1, i8 1>
-; CHECK-NEXT:    [[Y:%.*]] = add nuw <4 x i8> [[YY:%.*]], <i8 1, i8 1, i8 1, i8 1>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> [[Y]], <4 x i32> <i32 0, i32 4, i32 7, i32 2>
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
-; CHECK-NEXT:    ret <4 x i1> [[R]]
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
 ;
   %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 1>
   %y = add nuw <4 x i8> %yy, <i8 1, i8 1, i8 1, i8 1>
@@ -228,10 +224,7 @@ define <4 x i1> @shuf_nonzero_both_fail2(<4 x i8> %xx, <4 x i8> %yy) {
 
 define <4 x i1> @shuf_nonzero_lhs(<4 x i8> %xx) {
 ; CHECK-LABEL: @shuf_nonzero_lhs(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <4 x i8> [[XX:%.*]], <i8 1, i8 1, i8 1, i8 1>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> poison, <4 x i32> <i32 2, i32 2, i32 0, i32 1>
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
-; CHECK-NEXT:    ret <4 x i1> [[R]]
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
 ;
   %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 1>
 
@@ -242,10 +235,7 @@ define <4 x i1> @shuf_nonzero_lhs(<4 x i8> %xx) {
 
 define <4 x i1> @shuf_nonzero_lhs2(<4 x i8> %xx) {
 ; CHECK-LABEL: @shuf_nonzero_lhs2(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <4 x i8> [[XX:%.*]], <i8 1, i8 1, i8 1, i8 0>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i8> [[X]], <4 x i8> poison, <4 x i32> <i32 2, i32 0, i32 1, i32 1>
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
-; CHECK-NEXT:    ret <4 x i1> [[R]]
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
 ;
   %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 0>
 
@@ -270,10 +260,7 @@ define <4 x i1> @shuf_nonzero_lhs2_fail(<4 x i8> %xx) {
 
 define <4 x i1> @shuf_nonzero_rhs(<4 x i8> %xx) {
 ; CHECK-LABEL: @shuf_nonzero_rhs(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <4 x i8> [[XX:%.*]], <i8 1, i8 1, i8 1, i8 1>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i8> poison, <4 x i8> [[X]], <4 x i32> <i32 6, i32 7, i32 5, i32 4>
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
-; CHECK-NEXT:    ret <4 x i1> [[R]]
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
 ;
   %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 1>
 
@@ -284,10 +271,7 @@ define <4 x i1> @shuf_nonzero_rhs(<4 x i8> %xx) {
 
 define <4 x i1> @shuf_nonzero_rhs2(<4 x i8> %xx) {
 ; CHECK-LABEL: @shuf_nonzero_rhs2(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <4 x i8> [[XX:%.*]], <i8 0, i8 0, i8 1, i8 1>
-; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i8> poison, <4 x i8> [[X]], <4 x i32> <i32 6, i32 7, i32 7, i32 6>
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
-; CHECK-NEXT:    ret <4 x i1> [[R]]
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
 ;
   %x = add nuw <4 x i8> %xx, <i8 0, i8 0, i8 1, i8 1>