Fix "Combine bit test + conditional or into simple math"

This fixes the optimization introduced in r179748 and reverted in r179750.

While the optimization was sound, it did not properly respect differences in
bit-width.

llvm-svn: 180777

2 files changed, 173 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index fa946e63ac6..a34b73ace49 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -350,6 +350,67 @@ static Value *SimplifyWithOpReplaced(Value *V, Value *Op, Value *RepOp,
   return 0;
 }
 
+/// foldSelectICmpAndOr - We want to turn:
+///   (select (icmp eq (and X, C1), 0), Y, (or Y, C2))
+/// into:
+///   (or (shl (and X, C1), C3), y)
+/// iff:
+///   C1 and C2 are both powers of 2
+/// where:
+///   C3 = Log(C2) - Log(C1)
+///
+/// This transform handles cases where:
+/// 1. The icmp predicate is inverted
+/// 2. The select operands are reversed
+/// 3. The magnitude of C2 and C1 are flipped
+static Value *foldSelectICmpAndOr(const SelectInst &SI, Value *TrueVal,
+                                  Value *FalseVal,
+                                  InstCombiner::BuilderTy *Builder) {
+  const ICmpInst *IC = dyn_cast<ICmpInst>(SI.getCondition());
+  if (!IC || !IC->isEquality())
+    return 0;
+
+  Value *CmpLHS = IC->getOperand(0);
+  Value *CmpRHS = IC->getOperand(1);
+
+  if (!match(CmpRHS, m_Zero()))
+    return 0;
+
+  Value *X;
+  const APInt *C1;
+  if (!match(CmpLHS, m_And(m_Value(X), m_Power2(C1))))
+    return 0;
+
+  const APInt *C2;
+  bool OrOnTrueVal = false;
+  bool OrOnFalseVal = match(FalseVal, m_Or(m_Specific(TrueVal), m_Power2(C2)));
+  if (!OrOnFalseVal)
+    OrOnTrueVal = match(TrueVal, m_Or(m_Specific(FalseVal), m_Power2(C2)));
+
+  if (!OrOnFalseVal && !OrOnTrueVal)
+    return 0;
+
+  Value *V = CmpLHS;
+  Value *Y = OrOnFalseVal ? TrueVal : FalseVal;
+
+  unsigned C1Log = C1->logBase2();
+  unsigned C2Log = C2->logBase2();
+  if (C2Log > C1Log) {
+    V = Builder->CreateZExtOrTrunc(V, Y->getType());
+    V = Builder->CreateShl(V, C2Log - C1Log);
+  } else if (C1Log > C2Log) {
+    V = Builder->CreateLShr(V, C1Log - C2Log);
+    V = Builder->CreateZExtOrTrunc(V, Y->getType());
+  }
+
+  ICmpInst::Predicate Pred = IC->getPredicate();
+  if ((Pred == ICmpInst::ICMP_NE && OrOnFalseVal) ||
+      (Pred == ICmpInst::ICMP_EQ && OrOnTrueVal))
+    V = Builder->CreateXor(V, *C2);
+
+  return Builder->CreateOr(V, Y);
+}
+
 /// visitSelectInstWithICmp - Visit a SelectInst that has an
 /// ICmpInst as its first operand.
 ///
@@ -521,6 +582,9 @@ Instruction *InstCombiner::visitSelectInstWithICmp(SelectInst &SI,
     }
   }
 
+  if (Value *V = foldSelectICmpAndOr(SI, TrueVal, FalseVal, Builder))
+    return ReplaceInstUsesWith(SI, V);
+
   return Changed ? &SI : 0;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index cc3aacdce3c..08487a3696d 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -863,3 +863,112 @@ while.body:
 ; CHECK: @test64
 ; CHECK-NOT: select
 }
+
+; CHECK: @select_icmp_eq_and_1_0_or_2
+; CHECK-NEXT: [[SHL:%[a-z0-9]+]] = shl i32 %x, 1
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 [[SHL]], 2
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[AND]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_eq_and_1_0_or_2(i32 %x, i32 %y) {
+  %and = and i32 %x, 1
+  %cmp = icmp eq i32 %and, 0
+  %or = or i32 %y, 2
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+; CHECK: @select_icmp_eq_and_32_0_or_8
+; CHECK-NEXT: [[LSHR:%[a-z0-9]+]] = lshr i32 %x, 2
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 [[LSHR]], 8
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[AND]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_eq_and_32_0_or_8(i32 %x, i32 %y) {
+  %and = and i32 %x, 32
+  %cmp = icmp eq i32 %and, 0
+  %or = or i32 %y, 8
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+; CHECK: @select_icmp_ne_0_and_4096_or_4096
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 %x, 4096
+; CHECK-NEXT: [[XOR:%[a-z0-9]+]] = xor i32 [[AND]], 4096
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[XOR]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_ne_0_and_4096_or_4096(i32 %x, i32 %y) {
+  %and = and i32 %x, 4096
+  %cmp = icmp ne i32 0, %and
+  %or = or i32 %y, 4096
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+; CHECK: @select_icmp_eq_and_4096_0_or_4096
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 %x, 4096
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[AND]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_eq_and_4096_0_or_4096(i32 %x, i32 %y) {
+  %and = and i32 %x, 4096
+  %cmp = icmp eq i32 %and, 0
+  %or = or i32 %y, 4096
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+; CHECK: @select_icmp_ne_0_and_4096_or_32
+; CHECK-NEXT: [[LSHR:%[a-z0-9]+]] = lshr i32 %x, 7
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 [[LSHR]], 32
+; CHECK-NEXT: [[XOR:%[a-z0-9]+]] = xor i32 [[AND]], 32
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[XOR]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_ne_0_and_4096_or_32(i32 %x, i32 %y) {
+  %and = and i32 %x, 4096
+  %cmp = icmp ne i32 0, %and
+  %or = or i32 %y, 32
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+; CHECK: @select_icmp_ne_0_and_32_or_4096
+; CHECK-NEXT: [[SHL:%[a-z0-9]+]] = shl i32 %x, 7
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 [[SHL]], 4096
+; CHECK-NEXT: [[XOR:%[a-z0-9]+]] = xor i32 [[AND]], 4096
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[XOR]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_ne_0_and_32_or_4096(i32 %x, i32 %y) {
+  %and = and i32 %x, 32
+  %cmp = icmp ne i32 0, %and
+  %or = or i32 %y, 4096
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+; CHECK: @select_icmp_ne_0_and_1073741824_or_8
+; CHECK-NEXT: [[LSHR:%[a-z0-9]+]] = lshr i32 %x, 27
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i32 [[LSHR]], 8
+; CHECK-NEXT: [[TRUNC:%[a-z0-9]+]] = trunc i32 [[AND]] to i8
+; CHECK-NEXT: [[XOR:%[a-z0-9]+]] = xor i8 [[TRUNC]], 8
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i8 [[XOR]], %y
+; CHECK-NEXT: ret i8 [[OR]]
+define i8 @select_icmp_ne_0_and_1073741824_or_8(i32 %x, i8 %y) {
+  %and = and i32 %x, 1073741824
+  %cmp = icmp ne i32 0, %and
+  %or = or i8 %y, 8
+  %select = select i1 %cmp, i8 %y, i8 %or
+  ret i8 %select
+}
+
+; CHECK: @select_icmp_ne_0_and_8_or_1073741824
+; CHECK-NEXT: [[AND:%[a-z0-9]+]] = and i8 %x, 8
+; CHECK-NEXT: [[ZEXT:%[a-z0-9]+]] = zext i8 [[AND]] to i32
+; CHECK-NEXT: [[SHL:%[a-z0-9]+]] = shl nuw nsw i32 [[ZEXT]], 27
+; CHECK-NEXT: [[XOR:%[a-z0-9]+]] = xor i32 [[SHL]], 1073741824
+; CHECK-NEXT: [[OR:%[a-z0-9]+]] = or i32 [[XOR]], %y
+; CHECK-NEXT: ret i32 [[OR]]
+define i32 @select_icmp_ne_0_and_8_or_1073741824(i8 %x, i32 %y) {
+  %and = and i8 %x, 8
+  %cmp = icmp ne i8 0, %and
+  %or = or i32 %y, 1073741824
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}