[llvm] [InstCombine] Fold `(X & Mask) == 0 ? TC : FC -> TC +/- (X & Mask)` (PR #100437)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

<details>
<summary>Changes</summary>

Alive2: https://alive2.llvm.org/ce/z/d9wV7N


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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+35-24) 
- (modified) llvm/test/Transforms/InstCombine/select-icmp-and.ll (+27) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index aaf4ece3249a2..b094d7830307a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -155,42 +155,54 @@ static Value *foldSelectICmpAnd(SelectInst &Sel, ICmpInst *Cmp,
   } else {
     return nullptr;
   }
+  if (Pred == ICmpInst::ICMP_NE)
+    std::swap(SelTC, SelFC);
 
   // In general, when both constants are non-zero, we would need an offset to
   // replace the select. This would require more instructions than we started
   // with. But there's one special-case that we handle here because it can
   // simplify/reduce the instructions.
-  APInt TC = *SelTC;
-  APInt FC = *SelFC;
+  const APInt &TC = *SelTC;
+  const APInt &FC = *SelFC;
   if (!TC.isZero() && !FC.isZero()) {
+    if (TC.getBitWidth() != AndMask.getBitWidth())
+      return nullptr;
+    // If we have to create an 'and', then we must kill the cmp to not
+    // increase the instruction count.
+    if (CreateAnd && !Cmp->hasOneUse())
+      return nullptr;
+
     // If the select constants differ by exactly one bit and that's the same
     // bit that is masked and checked by the select condition, the select can
     // be replaced by bitwise logic to set/clear one bit of the constant result.
-    if (TC.getBitWidth() != AndMask.getBitWidth() || (TC ^ FC) != AndMask)
-      return nullptr;
-    if (CreateAnd) {
-      // If we have to create an 'and', then we must kill the cmp to not
-      // increase the instruction count.
-      if (!Cmp->hasOneUse())
-        return nullptr;
-      V = Builder.CreateAnd(V, ConstantInt::get(SelType, AndMask));
-    }
-    bool ExtraBitInTC = TC.ugt(FC);
-    if (Pred == ICmpInst::ICMP_EQ) {
-      // If the masked bit in V is clear, clear or set the bit in the result:
-      // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) ^ TC
-      // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) | TC
+    // If the masked bit in V is clear, clear or set the bit in the result:
+    // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) ^ TC
+    // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) | TC
+    if ((TC ^ FC) == AndMask) {
+      if (CreateAnd)
+        V = Builder.CreateAnd(V, ConstantInt::get(SelType, AndMask));
+      bool ExtraBitInTC = TC.ugt(FC);
       Constant *C = ConstantInt::get(SelType, TC);
       return ExtraBitInTC ? Builder.CreateXor(V, C) : Builder.CreateOr(V, C);
     }
-    if (Pred == ICmpInst::ICMP_NE) {
-      // If the masked bit in V is set, set or clear the bit in the result:
-      // (V & AndMaskC) != 0 ? TC : FC --> (V & AndMaskC) | FC
-      // (V & AndMaskC) != 0 ? TC : FC --> (V & AndMaskC) ^ FC
-      Constant *C = ConstantInt::get(SelType, FC);
-      return ExtraBitInTC ? Builder.CreateOr(V, C) : Builder.CreateXor(V, C);
+
+    // (V & AndMaskC) == 0 ? TC : FC --> (V & AndMaskC) + TC
+    if (TC + AndMask == FC) {
+      if (CreateAnd)
+        V = Builder.CreateAnd(V, ConstantInt::get(SelType, AndMask));
+      Constant *C = ConstantInt::get(SelType, TC);
+      return Builder.CreateAdd(V, C);
     }
-    llvm_unreachable("Only expecting equality predicates");
+
+    // (V & AndMaskC) == 0 ? TC : FC --> TC - (V & AndMaskC)
+    if (TC - AndMask == FC) {
+      if (CreateAnd)
+        V = Builder.CreateAnd(V, ConstantInt::get(SelType, AndMask));
+      Constant *C = ConstantInt::get(SelType, TC);
+      return Builder.CreateSub(C, V);
+    }
+
+    return nullptr;
   }
 
   // Make sure one of the select arms is a power-of-2.
@@ -203,7 +215,6 @@ static Value *foldSelectICmpAnd(SelectInst &Sel, ICmpInst *Cmp,
   unsigned ValZeros = ValC.logBase2();
   unsigned AndZeros = AndMask.logBase2();
   bool ShouldNotVal = !TC.isZero();
-  ShouldNotVal ^= Pred == ICmpInst::ICMP_NE;
 
   // If we would need to create an 'and' + 'shift' + 'xor' to replace a 'select'
   // + 'icmp', then this transformation would result in more instructions and
diff --git a/llvm/test/Transforms/InstCombine/select-icmp-and.ll b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
index 8bedf699dc922..b847773034aef 100644
--- a/llvm/test/Transforms/InstCombine/select-icmp-and.ll
+++ b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
@@ -629,3 +629,30 @@ define i8 @set_to_clear_decomposebittest_extra_use(i8 %x) {
   ret i8 %t3
 }
 
+define i32 @select_bittest_to_add(i32 %x) {
+; CHECK-LABEL: @select_bittest_to_add(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[X:%.*]], 1
+; CHECK-NEXT:    [[RET:%.*]] = add nuw nsw i32 [[AND]], 3
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+entry:
+  %and = and i32 %x, 1
+  %cmp = icmp eq i32 %and, 0
+  %ret = select i1 %cmp, i32 3, i32 4
+  ret i32 %ret
+}
+
+define i32 @select_bittest_to_sub(i32 %x) {
+; CHECK-LABEL: @select_bittest_to_sub(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[X:%.*]], 1
+; CHECK-NEXT:    [[RET:%.*]] = sub nuw nsw i32 4, [[AND]]
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+entry:
+  %and = and i32 %x, 1
+  %cmp = icmp eq i32 %and, 0
+  %ret = select i1 %cmp, i32 4, i32 3
+  ret i32 %ret
+}

``````````

</details>


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