[InstCombine] Teach foldSelectICmpAndOr to recognize (select (icmp slt (trunc (X)), 0), Y, (or Y, C2))

[craig.topper]

Summary

InstCombine likes to turn (icmp eq (and X, C1), 0) into (icmp slt (trunc (X)), 0) sometimes. This breaks foldSelectICmpAndOr's ability to recognize (select (icmp eq (and X, C1), 0), Y, (or Y, C2))->(or (shl (and X, C1), C3), y).

This patch tries to recover this. I had to flip around some of the early out checks so that I could create a new And instruction during the compare processing without it possibly never getting used.

[spatel]

The enhancement looks right, but I'm worried that the existing transform can increase the instruction count if any of the pieces in the pattern have more than one use. Are we missing m_OneUse() in some/all of the matchers?

lib/Transforms/InstCombine/InstCombineSelect.cpp

365 ↗

(On Diff #102451)

Formatting.

[craig.topper]

I think you're right. The 'icmp' and the 'or' pieces aren't checked for single use. In the existing code we do reuse the 'and'. In my new code i did make sure the trunc only had one use so that putting back an 'and' doesn't increase anything. So in the existing code we potentially create an 'or' and possibly a 'shift' and only delete a 'select'. And the really bad case we create an 'xor', an 'or', and a 'shift' while only removing a 'select'. This is kinda weird to handle because we don't always create a shift and don't always create an xor.

[spatel]

In D34184#781467, @craig.topper wrote:

I think you're right. The 'icmp' and the 'or' pieces aren't checked for single use. In the existing code we do reuse the 'and'. In my new code i did make sure the trunc only had one use so that putting back an 'and' doesn't increase anything. So in the existing code we potentially create an 'or' and possibly a 'shift' and only delete a 'select'. And the really bad case we create an 'xor', an 'or', and a 'shift' while only removing a 'select'. This is kinda weird to handle because we don't always create a shift and don't always create an xor.

Yeah, it's a mess. I think the safe thing would be:

1. Add tests with extra uses to show how this can go bad.
2. Check one-use on everything to avoid that.
3. Make the enhancement proposed in this patch.
4. Separate out the various cases and remove one-use if possible (I think this would be a rare enhancement, so probably not too important...but I don't have any actual data).

[craig.topper]

Extra use tests committed in r305509

[craig.topper]

Rebase on top of the one use fixes to the existing code.

[spatel]

LGTM. See inline for a nit.

lib/Transforms/InstCombine/InstCombineSelect.cpp

384–388 ↗

(On Diff #103483)

That "line" still looks awful. :)
Give the OneBitSet variable a name and its own line(s), so clang-format can do better?

llvm/trunk/lib/Transforms/InstCombine/InstCombineSelect.cpp

 /// 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() || !SI.getType()->isIntegerTy())
+  if (!IC || !SI.getType()->isIntegerTy())
     return nullptr;
 
   Value *CmpLHS = IC->getOperand(0);
   Value *CmpRHS = IC->getOperand(1);
 
+  Value *V;
+  unsigned C1Log;
+  bool IsEqualZero;
+  bool NeedAnd = false;
+  if (IC->isEquality()) {
     if (!match(CmpRHS, m_Zero()))
       return nullptr;
 
-  Value *X;
     const APInt *C1;
-  if (!match(CmpLHS, m_And(m_Value(X), m_Power2(C1))))
+    if (!match(CmpLHS, m_And(m_Value(), m_Power2(C1))))
+      return nullptr;
+
+    V = CmpLHS;
+    C1Log = C1->logBase2();
+    IsEqualZero = IC->getPredicate() == ICmpInst::ICMP_EQ;
+  } else if (IC->getPredicate() == ICmpInst::ICMP_SLT ||
+             IC->getPredicate() == ICmpInst::ICMP_SGT) {
+    // We also need to recognize (icmp slt (trunc (X)), 0) and
+    // (icmp sgt (trunc (X)), -1).
+    IsEqualZero = IC->getPredicate() == ICmpInst::ICMP_SGT;
+    if ((IsEqualZero && !match(CmpRHS, m_AllOnes())) ||
+        (!IsEqualZero && !match(CmpRHS, m_Zero())))
+      return nullptr;
+
+    if (!match(CmpLHS, m_OneUse(m_Trunc(m_Value(V)))))
       return nullptr;
 
+    C1Log = CmpLHS->getType()->getScalarSizeInBits() - 1;
+    NeedAnd = true;
+  } else {
+    return nullptr;
+  }
+
   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 nullptr;
 
-  Value *V = CmpLHS;
   Value *Y = OrOnFalseVal ? TrueVal : FalseVal;
 
-  unsigned C1Log = C1->logBase2();
   unsigned C2Log = C2->logBase2();
 
-  ICmpInst::Predicate Pred = IC->getPredicate();
-  bool NeedXor = (Pred == ICmpInst::ICMP_NE && OrOnFalseVal) ||
-                 (Pred == ICmpInst::ICMP_EQ && OrOnTrueVal);
+  bool NeedXor = (!IsEqualZero && OrOnFalseVal) || (IsEqualZero && OrOnTrueVal);
   bool NeedShift = C1Log != C2Log;
   bool NeedZExtTrunc = Y->getType()->getIntegerBitWidth() !=
                        V->getType()->getIntegerBitWidth();
 
   // Make sure we don't create more instructions than we save.
   Value *Or = OrOnFalseVal ? FalseVal : TrueVal;
   if ((NeedShift + NeedXor + NeedZExtTrunc) >
       (IC->hasOneUse() + Or->hasOneUse()))
     return nullptr;
 
+  if (NeedAnd) {
+    // Insert the AND instruction on the input to the truncate.
+    APInt C1 = APInt::getOneBitSet(V->getType()->getScalarSizeInBits(), C1Log);
+    V = Builder->CreateAnd(V, ConstantInt::get(V->getType(), C1));
+  }
+
   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());
   } else
     V = Builder->CreateZExtOrTrunc(V, Y->getType());

llvm/trunk/test/Transforms/InstCombine/select-with-bitwise-ops.ll

;
   %1 = and i16 %x, 4
   %2 = icmp ne i16 %1, 0
   %3 = select i1 %2, i32 40, i32 42
   ret i32 %3
 }
 
 define i32 @test68(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test68(
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[TMP1]], -1
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[Y:%.*]], 2
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i32 [[Y]], i32 [[OR]]
-; CHECK-NEXT:    ret i32 [[SELECT]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 6
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 2
+; CHECK-NEXT:    [[TMP3:%.*]] = or i32 [[TMP2]], [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP3]]
 ;
   %and = and i32 %x, 128
   %cmp = icmp eq i32 %and, 0
   %or = or i32 %y, 2
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }
 
 define i32 @test69(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test69(
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i8
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i8 [[TMP1]], 0
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[Y:%.*]], 2
-; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i32 [[Y]], i32 [[OR]]
-; CHECK-NEXT:    ret i32 [[SELECT]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[X:%.*]], 6
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 2
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i32 [[TMP2]], 2
+; CHECK-NEXT:    [[TMP4:%.*]] = or i32 [[TMP3]], [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[TMP4]]
 ;
   %and = and i32 %x, 128
   %cmp = icmp ne i32 %and, 0
   %or = or i32 %y, 2
   %select = select i1 %cmp, i32 %y, i32 %or
   ret i32 %select
 }