[llvm] [InstCombine]: Fold OR bit tests (PR #178366)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (arrowten)

<details>
<summary>Changes</summary>

Fold: OR of bit tests with shifted constants.
((X & (C0 << Y)) != 0) || ((X & (C1 << Y)) != 0) --> (X & ((C0 | C1) << Y)) != 0

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+55) 
- (added) llvm/test/Transforms/InstCombine/fold-or-bit-tests.ll (+21) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index fb30a4545cffe..d2761c3e09e2f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -4422,6 +4422,61 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
     if (Instruction *Result = foldSelectInstWithICmp(SI, ICI))
       return Result;
 
+  // Fold: OR of bit tests with shifted constants.
+  //
+  // ((X & (C0 << Y)) != 0) || ((X & (C1 << Y)) != 0)
+  //   --> (X & ((C0 | C1) << Y)) != 0
+  auto *CondICmp = dyn_cast<ICmpInst>(CondVal);
+  auto *FalseICmp = dyn_cast<ICmpInst>(FalseVal);
+
+  if (CondICmp && FalseICmp && match(TrueVal, m_One()) &&
+      CondICmp->getPredicate() == ICmpInst::ICMP_NE &&
+      FalseICmp->getPredicate() == ICmpInst::ICMP_NE &&
+      match(CondICmp->getOperand(1), m_Zero()) &&
+      match(FalseICmp->getOperand(1), m_Zero())) {
+    Value *X, *Mask0, *Mask1;
+
+    if (match(CondICmp->getOperand(0), m_And(m_Value(X), m_Value(Mask0))) &&
+        match(FalseICmp->getOperand(0), m_And(m_Specific(X), m_Value(Mask1)))) {
+      auto Extract = [&](Value *V, Value *&Shift, const APInt *&Base) -> bool {
+        if (auto *T = dyn_cast<TruncInst>(V))
+          V = T->getOperand(0);
+
+        Value *LHS, *RHS;
+
+        if (!match(V, m_Shl(m_Value(LHS), m_Value(RHS))))
+          return false;
+
+        if (!match(LHS, m_APInt(Base)))
+          return false;
+
+        Shift = RHS;
+
+        return true;
+      };
+
+      Value *Shift0, *Shift1;
+      const APInt *Base0, *Base1;
+
+      if (Extract(Mask0, Shift0, Base0) &&
+          Extract(Mask1, Shift1, Base1) && (Shift0 == Shift1)) {
+        APInt NewBase = *Base0 | *Base1;
+        Value *NewShl = Builder.CreateShl(
+          ConstantInt::get(Shift0->getType(), NewBase), Shift0);
+        Value *NewMask = NewShl;
+
+        if (Mask0->getType() != NewShl->getType())
+          NewMask = Builder.CreateTrunc(NewShl, Mask0->getType());
+
+        Value *NewAnd = Builder.CreateAnd(X, NewMask);
+        Instruction *NewCmp = cast<Instruction>(Builder.CreateICmpNE(
+          NewAnd, ConstantInt::getNullValue(NewAnd->getType())));
+
+        return replaceInstUsesWith(SI, NewCmp);
+      }
+    }
+  }
+
   if (Value *V = foldSelectBitTest(SI, CondVal, TrueVal, FalseVal, Builder, SQ))
     return replaceInstUsesWith(SI, V);
 
diff --git a/llvm/test/Transforms/InstCombine/fold-or-bit-tests.ll b/llvm/test/Transforms/InstCombine/fold-or-bit-tests.ll
new file mode 100644
index 0000000000000..afea70e3a6df5
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fold-or-bit-tests.ll
@@ -0,0 +1,21 @@
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+define i1 @src(i64 %arg0, ptr %arg1) {
+; CHECK: %[[SH:.*]] = shl i32 5, %[[SHIFT:.*]]
+; CHECK: %[[TR:.*]] = trunc i32 %[[SH]] to i8
+; CHECK: %[[AND:.*]] = and i8 %{{.*}}, %[[TR]]
+; CHECK: icmp ne i8 %[[AND]], 0
+
+  %v0 = load i8, ptr %arg1, align 1
+  %v1 = trunc nuw nsw i64 %arg0 to i32
+  %v2 = shl nuw nsw i32 1, %v1
+  %v3 = trunc nuw nsw i32 %v2 to i8
+  %v4 = and i8 %v0, %v3
+  %v5 = shl nuw nsw i32 4, %v1
+  %v6 = trunc nuw nsw i32 %v5 to i8
+  %v7 = and i8 %v0, %v6
+  %v8 = icmp ne i8 %v4, 0
+  %v9 = icmp ne i8 %v7, 0
+  %v10 = select i1 %v8, i1 true, i1 %v9
+  ret i1 %v10
+}

``````````

</details>


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