[llvm] goldsteinn/shuf non zero (PR #87702)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-analysis

@llvm/pr-subscribers-llvm-transforms

Author: None (goldsteinn)

<details>
<summary>Changes</summary>

- **[ValueTracking] Add tests for `shufflevector` in `isKnownNonZero`**
- **[ValueTracking] Add support for `shufflevector` in `isKnownNonZero`**


---
Full diff: https://github.com/llvm/llvm-project/pull/87702.diff


2 Files Affected:

- (modified) llvm/lib/Analysis/ValueTracking.cpp (+15) 
- (modified) llvm/test/Transforms/InstSimplify/known-non-zero.ll (+116) 


``````````diff
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 5ad4da43bca7db..cb25d118cf344e 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2765,6 +2765,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 b647f11af4461d..a1f2570ec0af33 100644
--- a/llvm/test/Transforms/InstSimplify/known-non-zero.ll
+++ b/llvm/test/Transforms/InstSimplify/known-non-zero.ll
@@ -166,3 +166,119 @@ A:
 B:
   ret i1 0
 }
+
+define <4 x i1> @shuf_nonzero_both(<4 x i8> %xx, <4 x i8> %yy) {
+; CHECK-LABEL: @shuf_nonzero_both(
+; 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>
+
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 4, i32 7, i32 2>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_both_fail(<4 x i8> %xx, <4 x i8> %yy) {
+; CHECK-LABEL: @shuf_nonzero_both_fail(
+; 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 0>
+; 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]]
+;
+  %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 0>
+
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 4, i32 7, i32 2>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_both_fail2(<4 x i8> %xx, <4 x i8> %yy) {
+; CHECK-LABEL: @shuf_nonzero_both_fail2(
+; CHECK-NEXT:    [[X:%.*]] = add nuw <4 x i8> [[XX:%.*]], <i8 1, i8 1, i8 1, i8 1>
+; CHECK-NEXT:    [[Y:%.*]] = add <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]]
+;
+  %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 1>
+  %y = add <4 x i8> %yy, <i8 1, i8 1, i8 1, i8 1>
+
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> %y, <4 x i32> <i32 0, i32 4, i32 7, i32 2>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_lhs(<4 x i8> %xx) {
+; CHECK-LABEL: @shuf_nonzero_lhs(
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
+;
+  %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 1>
+
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> poison, <4 x i32> <i32 2, i32 2, i32 0, i32 1>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_lhs2(<4 x i8> %xx) {
+; CHECK-LABEL: @shuf_nonzero_lhs2(
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
+;
+  %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 0>
+
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> poison, <4 x i32> <i32 2, i32 0, i32 1, i32 1>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_lhs2_fail(<4 x i8> %xx) {
+; CHECK-LABEL: @shuf_nonzero_lhs2_fail(
+; 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 3, i32 1>
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
+; CHECK-NEXT:    ret <4 x i1> [[R]]
+;
+  %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 0>
+
+  %shuf = shufflevector <4 x i8> %x, <4 x i8> poison, <4 x i32> <i32 2, i32 0, i32 3, i32 1>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_rhs(<4 x i8> %xx) {
+; CHECK-LABEL: @shuf_nonzero_rhs(
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
+;
+  %x = add nuw <4 x i8> %xx, <i8 1, i8 1, i8 1, i8 1>
+
+  %shuf = shufflevector <4 x i8> poison, <4 x i8> %x, <4 x i32> <i32 6, i32 7, i32 5, i32 4>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_rhs2(<4 x i8> %xx) {
+; CHECK-LABEL: @shuf_nonzero_rhs2(
+; CHECK-NEXT:    ret <4 x i1> zeroinitializer
+;
+  %x = add nuw <4 x i8> %xx, <i8 0, i8 0, i8 1, i8 1>
+
+  %shuf = shufflevector <4 x i8> poison, <4 x i8> %x, <4 x i32> <i32 6, i32 7, i32 7, i32 6>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}
+
+define <4 x i1> @shuf_nonzero_rhs2_fail(<4 x i8> %xx) {
+; CHECK-LABEL: @shuf_nonzero_rhs2_fail(
+; 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 5, i32 6>
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <4 x i8> [[SHUF]], zeroinitializer
+; CHECK-NEXT:    ret <4 x i1> [[R]]
+;
+  %x = add nuw <4 x i8> %xx, <i8 0, i8 0, i8 1, i8 1>
+
+  %shuf = shufflevector <4 x i8> poison, <4 x i8> %x, <4 x i32> <i32 6, i32 7, i32 5, i32 6>
+  %r = icmp eq <4 x i8> %shuf, zeroinitializer
+  ret <4 x i1> %r
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/87702