[ValueTracking] Add support for insertelement in isKnownNonZero

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

Closes llvm#87703

Author: goldsteinn

llvm/lib/Analysis/ValueTracking.cpp

@@ -2763,6 +2763,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();

llvm/test/Transforms/InstSimplify/known-non-zero.ll

@@ -296,11 +296,7 @@ define <4 x i1> @shuf_nonzero_rhs2_fail(<4 x i8> %xx) {
 
 define <2 x i1> @insert_nonzero0(<2 x i8> %xx, i8 %yy) {
 ; CHECK-LABEL: @insert_nonzero0(
-; 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 1
-; CHECK-NEXT:    [[R:%.*]] = icmp eq <2 x i8> [[INS]], zeroinitializer
-; CHECK-NEXT:    ret <2 x i1> [[R]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %x = add nuw <2 x i8> %xx, <i8 1, i8 0>
   %y = add nuw i8 %yy, 1
@@ -312,11 +308,7 @@ define <2 x i1> @insert_nonzero0(<2 x i8> %xx, i8 %yy) {
 
 define <2 x i1> @insert_nonzero1(<2 x i8> %xx, i8 %yy) {
 ; CHECK-LABEL: @insert_nonzero1(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <2 x i8> [[XX:%.*]], <i8 0, i8 1>
-; 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]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %x = add nuw <2 x i8> %xx, <i8 0, i8 1>
   %y = add nuw i8 %yy, 1
@@ -360,11 +352,7 @@ define <2 x i1> @insert_nonzero_fail2(<2 x i8> %xx, i8 %yy) {
 
 define <2 x i1> @insert_nonzero_any_idx(<2 x i8> %xx, i8 %yy, i32 %idx) {
 ; CHECK-LABEL: @insert_nonzero_any_idx(
-; CHECK-NEXT:    [[X:%.*]] = add nuw <2 x i8> [[XX:%.*]], <i8 1, i8 1>
-; 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]]
+; CHECK-NEXT:    ret <2 x i1> zeroinitializer
 ;
   %x = add nuw <2 x i8> %xx, <i8 1, i8 1>
   %y = add nuw i8 %yy, 1