[llvm] goldsteinn/insert non zero (PR #87703)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: None (goldsteinn)

<details>
<summary>Changes</summary>

- **[ValueTracking] Add tests for `insertelement` in `isKnownNonZero`; NFC**
- **[ValueTracking] Add support for `insertelement` in `isKnownNonZero`**


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


2 Files Affected:

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


``````````diff
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 5ad4da43bca7db..fce22e3af12027 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -2751,6 +2751,29 @@ static bool isKnownNonZeroFromOperator(const Operator *I,
       return isKnownNonZero(U.get(), DemandedElts, NewDepth, RecQ);
     });
   }
+  case Instruction::InsertElement: {
+    if (isa<ScalableVectorType>(I->getType()))
+      break;
+
+    const Value *Vec = I->getOperand(0);
+    const Value *Elt = I->getOperand(1);
+    auto *CIdx = dyn_cast<ConstantInt>(I->getOperand(2));
+
+    unsigned NumElts = DemandedElts.getBitWidth();
+    APInt DemandedVecElts = DemandedElts;
+    bool SkipElt = false;
+    // If we know the index we are inserting too, clear it from Vec check.
+    if (CIdx && CIdx->getValue().ult(NumElts)) {
+      DemandedVecElts.clearBit(CIdx->getZExtValue());
+      SkipElt = !DemandedElts[CIdx->getZExtValue()];
+    }
+
+    // Result is zero if Elt is non-zero and rest of the demanded elts in Vec
+    // are non-zero.
+    return (SkipElt || isKnownNonZero(Elt, Depth, Q)) &&
+           (DemandedVecElts.isZero() ||
+            isKnownNonZero(Vec, DemandedVecElts, Depth, Q));
+  }
   case Instruction::ExtractElement:
     if (const auto *EEI = dyn_cast<ExtractElementInst>(I)) {
       const Value *Vec = EEI->getVectorOperand();
diff --git a/llvm/test/Transforms/InstSimplify/known-non-zero.ll b/llvm/test/Transforms/InstSimplify/known-non-zero.ll
index b647f11af4461d..a54309be3e4163 100644
--- a/llvm/test/Transforms/InstSimplify/known-non-zero.ll
+++ b/llvm/test/Transforms/InstSimplify/known-non-zero.ll
@@ -166,3 +166,87 @@ A:
 B:
   ret i1 0
 }
+
+define <2 x i1> @insert_nonzero0(<2 x i8> %xx, i8 %yy) {
+; CHECK-LABEL: @insert_nonzero0(
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
+;
+  %x = add nuw <2 x i8> %xx, <i8 1, i8 0>
+  %y = add nuw i8 %yy, 1
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 1
+  %r = icmp eq <2 x i8> %ins, zeroinitializer
+  ret <2 x i1> %r
+}
+
+define <2 x i1> @insert_nonzero1(<2 x i8> %xx, i8 %yy) {
+; CHECK-LABEL: @insert_nonzero1(
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
+;
+  %x = add nuw <2 x i8> %xx, <i8 0, i8 1>
+  %y = add nuw i8 %yy, 1
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 0
+  %r = icmp eq <2 x i8> %ins, zeroinitializer
+  ret <2 x i1> %r
+}
+
+define <2 x i1> @insert_nonzero_fail(<2 x i8> %xx, i8 %yy) {
+; CHECK-LABEL: @insert_nonzero_fail(
+; CHECK-NEXT:    [[X:%.*]] = add nuw <2 x i8> [[XX:%.*]], <i8 1, i8 0>
+; CHECK-NEXT:    [[Y:%.*]] = add nuw i8 [[YY:%.*]], 1
+; CHECK-NEXT:    [[INS:%.*]] = insertelement <2 x i8> [[X]], i8 [[Y]], i32 0
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i8> [[INS]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[R]]
+;
+  %x = add nuw <2 x i8> %xx, <i8 1, i8 0>
+  %y = add nuw i8 %yy, 1
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 0
+  %r = icmp eq <2 x i8> %ins, zeroinitializer
+  ret <2 x i1> %r
+}
+
+define <2 x i1> @insert_nonzero_fail2(<2 x i8> %xx, i8 %yy) {
+; CHECK-LABEL: @insert_nonzero_fail2(
+; CHECK-NEXT:    [[X:%.*]] = add nuw <2 x i8> [[XX:%.*]], <i8 0, i8 1>
+; CHECK-NEXT:    [[Y:%.*]] = add i8 [[YY:%.*]], 1
+; CHECK-NEXT:    [[INS:%.*]] = insertelement <2 x i8> [[X]], i8 [[Y]], i32 0
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i8> [[INS]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[R]]
+;
+  %x = add nuw <2 x i8> %xx, <i8 0, i8 1>
+  %y = add i8 %yy, 1
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 0
+  %r = icmp eq <2 x i8> %ins, zeroinitializer
+  ret <2 x i1> %r
+}
+
+define <2 x i1> @insert_nonzero_any_idx(<2 x i8> %xx, i8 %yy, i32 %idx) {
+; CHECK-LABEL: @insert_nonzero_any_idx(
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
+;
+  %x = add nuw <2 x i8> %xx, <i8 1, i8 1>
+  %y = add nuw i8 %yy, 1
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 %idx
+  %r = icmp eq <2 x i8> %ins, zeroinitializer
+  ret <2 x i1> %r
+}
+
+define <2 x i1> @insert_nonzero_any_idx_fail(<2 x i8> %xx, i8 %yy, i32 %idx) {
+; CHECK-LABEL: @insert_nonzero_any_idx_fail(
+; CHECK-NEXT:    [[X:%.*]] = add nuw <2 x i8> [[XX:%.*]], <i8 1, i8 0>
+; CHECK-NEXT:    [[Y:%.*]] = add nuw i8 [[YY:%.*]], 1
+; CHECK-NEXT:    [[INS:%.*]] = insertelement <2 x i8> [[X]], i8 [[Y]], i32 [[IDX:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i8> [[INS]], zeroinitializer
+; CHECK-NEXT:    ret <2 x i1> [[R]]
+;
+  %x = add nuw <2 x i8> %xx, <i8 1, i8 0>
+  %y = add nuw i8 %yy, 1
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 %idx
+  %r = icmp eq <2 x i8> %ins, zeroinitializer
+  ret <2 x i1> %r
+}

``````````

</details>


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