[llvm] 964df09 - [ValueTracking] Support non-constant idx for `computeKnownBits` of `insertelement`

[Noah Goldstein via llvm-commits]

Author: Noah Goldstein
Date: 2024-04-09T01:01:41-05:00
New Revision: 964df099e1f8afcb9d052f61e065da82b19cc81b

URL: https://github.com/llvm/llvm-project/commit/964df099e1f8afcb9d052f61e065da82b19cc81b
DIFF: https://github.com/llvm/llvm-project/commit/964df099e1f8afcb9d052f61e065da82b19cc81b.diff

LOG: [ValueTracking] Support non-constant idx for `computeKnownBits` of `insertelement`

Its same logic as before, we just need to intersect what we know about
the new Elt and the entire pre-existing Vec.

Closes #87707

Added: 
    

Modified: 
    llvm/lib/Analysis/ValueTracking.cpp
    llvm/test/Transforms/InstCombine/insertelement.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 6602db424fc260..ca48cfe7738154 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -1727,26 +1727,25 @@ static void computeKnownBitsFromOperator(const Operator *I,
     const Value *Vec = I->getOperand(0);
     const Value *Elt = I->getOperand(1);
     auto *CIdx = dyn_cast<ConstantInt>(I->getOperand(2));
-    // Early out if the index is non-constant or out-of-range.
     unsigned NumElts = DemandedElts.getBitWidth();
-    if (!CIdx || CIdx->getValue().uge(NumElts)) {
-      Known.resetAll();
-      return;
+    APInt DemandedVecElts = DemandedElts;
+    bool NeedsElt = true;
+    // If we know the index we are inserting too, clear it from Vec check.
+    if (CIdx && CIdx->getValue().ult(NumElts)) {
+      DemandedVecElts.clearBit(CIdx->getZExtValue());
+      NeedsElt = DemandedElts[CIdx->getZExtValue()];
     }
+
     Known.One.setAllBits();
     Known.Zero.setAllBits();
-    unsigned EltIdx = CIdx->getZExtValue();
-    // Do we demand the inserted element?
-    if (DemandedElts[EltIdx]) {
+    if (NeedsElt) {
       computeKnownBits(Elt, Known, Depth + 1, Q);
       // If we don't know any bits, early out.
       if (Known.isUnknown())
         break;
     }
-    // We don't need the base vector element that has been inserted.
-    APInt DemandedVecElts = DemandedElts;
-    DemandedVecElts.clearBit(EltIdx);
-    if (!!DemandedVecElts) {
+
+    if (!DemandedVecElts.isZero()) {
       computeKnownBits(Vec, DemandedVecElts, Known2, Depth + 1, Q);
       Known = Known.intersectWith(Known2);
     }

diff  --git a/llvm/test/Transforms/InstCombine/insertelement.ll b/llvm/test/Transforms/InstCombine/insertelement.ll
index 178e2e25474a3a..c8df2db6e70caa 100644
--- a/llvm/test/Transforms/InstCombine/insertelement.ll
+++ b/llvm/test/Transforms/InstCombine/insertelement.ll
@@ -17,8 +17,7 @@ define <4 x i32> @insert_unknown_idx(<4 x i32> %x, i32 %idx) {
 ; CHECK-LABEL: @insert_unknown_idx(
 ; CHECK-NEXT:    [[V1:%.*]] = and <4 x i32> [[X:%.*]], <i32 7, i32 7, i32 7, i32 7>
 ; CHECK-NEXT:    [[V2:%.*]] = insertelement <4 x i32> [[V1]], i32 6, i32 [[IDX:%.*]]
-; CHECK-NEXT:    [[V3:%.*]] = and <4 x i32> [[V2]], <i32 7, i32 7, i32 7, i32 7>
-; CHECK-NEXT:    ret <4 x i32> [[V3]]
+; CHECK-NEXT:    ret <4 x i32> [[V2]]
 ;
   %v1 = and <4 x i32> %x, <i32 7, i32 7, i32 7, i32 7>
   %v2 = insertelement <4 x i32> %v1, i32 6, i32 %idx
@@ -28,11 +27,7 @@ define <4 x i32> @insert_unknown_idx(<4 x i32> %x, i32 %idx) {
 
 define <2 x i8> @insert_known_any_idx(<2 x i8> %xx, i8 %yy, i32 %idx) {
 ; CHECK-LABEL: @insert_known_any_idx(
-; CHECK-NEXT:    [[X:%.*]] = or <2 x i8> [[XX:%.*]], <i8 16, i8 16>
-; CHECK-NEXT:    [[Y:%.*]] = or i8 [[YY:%.*]], 16
-; CHECK-NEXT:    [[INS:%.*]] = insertelement <2 x i8> [[X]], i8 [[Y]], i32 [[IDX:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = and <2 x i8> [[INS]], <i8 16, i8 16>
-; CHECK-NEXT:    ret <2 x i8> [[R]]
+; CHECK-NEXT:    ret <2 x i8> <i8 16, i8 16>
 ;
   %x = or <2 x i8> %xx, <i8 16, i8 16>
   %y = or i8 %yy, 16