[llvm] 9c545a1 - [ValueTracking] Add support for `insertelement` in `isKnownNonZero`

[Noah Goldstein via llvm-commits]

Author: Noah Goldstein
Date: 2024-04-10T13:13:43-05:00
New Revision: 9c545a14c09051b011358854655c1f466d656e79

URL: https://github.com/llvm/llvm-project/commit/9c545a14c09051b011358854655c1f466d656e79
DIFF: https://github.com/llvm/llvm-project/commit/9c545a14c09051b011358854655c1f466d656e79.diff

LOG: [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 #87703

Added: 
    

Modified: 
    llvm/lib/Analysis/ValueTracking.cpp
    llvm/test/Transforms/InstSimplify/known-non-zero.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index b3029f440ca2a4..9f16eaf9e09905 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/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();

diff  --git a/llvm/test/Transforms/InstSimplify/known-non-zero.ll b/llvm/test/Transforms/InstSimplify/known-non-zero.ll
index 1417e86ee678c8..d9b8f5eed32390 100644
--- a/llvm/test/Transforms/InstSimplify/known-non-zero.ll
+++ b/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