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

[goldsteinn]

• [ValueTracking] Add tests for non-constant idx in computeKnownBits of insertelement; NFC
• [ValueTracking] Support non-constant idx for computeKnownBits of insertelement

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: None (goldsteinn)

Changes

• [ValueTracking] Add tests for non-constant idx in computeKnownBits of insertelement; NFC
• [ValueTracking] Support non-constant idx for computeKnownBits of insertelement

---

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

2 Files Affected:

• (modified) llvm/lib/Analysis/ValueTracking.cpp (+10-11)
• (modified) llvm/test/Transforms/InstCombine/insertelement.ll (+47-2)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 5ad4da43bca7db..15949142a8c77b 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 976c495465ce47..c8df2db6e70caa 100644
--- a/llvm/test/Transforms/InstCombine/insertelement.ll
+++ b/llvm/test/Transforms/InstCombine/insertelement.ll
@@ -17,11 +17,56 @@ 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
   %v3 = and <4 x i32> %v2, <i32 7, i32 7, i32 7, i32 7>
   ret <4 x i32> %v3
 }
+
+define <2 x i8> @insert_known_any_idx(<2 x i8> %xx, i8 %yy, i32 %idx) {
+; CHECK-LABEL: @insert_known_any_idx(
+; 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
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 %idx
+  %r = and <2 x i8> %ins, <i8 16, i8 16>
+  ret <2 x i8> %r
+}
+
+define <2 x i8> @insert_known_any_idx_fail1(<2 x i8> %xx, i8 %yy, i32 %idx) {
+; CHECK-LABEL: @insert_known_any_idx_fail1(
+; CHECK-NEXT:    [[X:%.*]] = or <2 x i8> [[XX:%.*]], <i8 17, i8 33>
+; 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]]
+;
+  %x = or <2 x i8> %xx, <i8 17, i8 33>
+  %y = or i8 %yy, 16
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 %idx
+  %r = and <2 x i8> %ins, <i8 16, i8 16>
+  ret <2 x i8> %r
+}
+
+
+define <2 x i8> @insert_known_any_idx_fail2(<2 x i8> %xx, i8 %yy, i32 %idx) {
+; CHECK-LABEL: @insert_known_any_idx_fail2(
+; CHECK-NEXT:    [[X:%.*]] = or <2 x i8> [[XX:%.*]], <i8 17, i8 31>
+; CHECK-NEXT:    [[Y:%.*]] = or i8 [[YY:%.*]], 15
+; 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]]
+;
+  %x = or <2 x i8> %xx, <i8 17, i8 31>
+  %y = or i8 %yy, 15
+
+  %ins = insertelement <2 x i8> %x, i8 %y, i32 %idx
+  %r = and <2 x i8> %ins, <i8 16, i8 16>
+  ret <2 x i8> %r
+}
+

[nikic]

LGTM