[InstCombine] convert mask and shift of power-of-2 to cmp+select

[spatel]

Summary

When the mask is a power-of-2 constant and op0 is a shifted-power-of-2 constant, test if the shift amount equals the offset bit index:

(ShiftC << X) & C --> X == (log2(C) - log2(ShiftC)) ? C : 0
(ShiftC >> X) & C --> X == (log2(ShiftC) - log2(C)) ? C : 0

This is an alternate to D127610 with a more general pattern. We match only shift+and instead of the trailing xor, so we see a few more tests diffs. I think we discussed this initially in D126617.

Here are proofs for shifts in both directions:
https://alive2.llvm.org/ce/z/CFrLs4

The test diffs look equal or better for IR, and this makes the patterns more uniform in IR. The backend would likely invert this in both cases if that is profitable, and there's already code in place to do that.

[bcl5980]

I agree with that the transform is more general. And my question is how to invert the transform on backend if shift + and is faster.
For now SelectionDag has the transform select Cond, Pow2, 0 --> (zext Cond) << log2(Pow2)

// select Cond, Pow2, 0 --> (zext Cond) << log2(Pow2)
if (C1Val.isPowerOf2() && C2Val.isZero()) {
  if (VT != MVT::i1)
    Cond = DAG.getNode(ISD::ZERO_EXTEND, DL, VT, Cond);
  SDValue ShAmtC =
      DAG.getShiftAmountConstant(C1Val.exactLogBase2(), VT, DL);
  return DAG.getNode(ISD::SHL, DL, VT, Cond, ShAmtC);
}

And if we want to do this
X == (log2(C) - log2(ShiftC)) ? C : 0 --> (ShiftC << X) & C
X == (log2(ShiftC) - log2(C)) ? C : 0 --> (ShiftC >> X) & C
We need to assume X < BitWidth https://alive2.llvm.org/ce/z/-GjVsu

[spatel]

In D127801#3584097, @bcl5980 wrote:

We need to assume X < BitWidth https://alive2.llvm.org/ce/z/-GjVsu

Ah, yes good point. The expansion back to the shift form requires an extra safety check for an over-shift in IR. Here's another version of the proof to show that:
https://alive2.llvm.org/ce/z/E_KCpi

And this is a hard-coded example of what the backend does currently when it creates a shift:
https://alive2.llvm.org/ce/z/6EfMnj

But in the backend, the target can specify the exact behavior for an over-shift - it does not have to be undefined (create poison). If a target returns 0 for an over-shift, then no masking or extra cmp+select may be needed.

Do you have an example for a target where the code looks worse with the cmp+select IR?

I tried some tests with x86 and AArch, and they do not look worse to me with the cmp+select IR. Here is an example:

define i32 @src(i32 %x) {
  %shl = shl i32 2, %x
  %r = and i32 %shl, 64
  ret i32 %r
}

define i32 @tgt(i32 %x) {
  %i = icmp eq i32 %x, 5
  %r = select i1 %i, i32 64, i32 0
  ret i32 %r
}

% llc -o - sh2.ll -mtriple=x86_64 
src: 
	movl	%edi, %ecx
	movl	$2, %eax
	shll	%cl, %eax
	andl	$64, %eax

tgt:  
	xorl	%eax, %eax
	cmpl	$5, %edi
	sete	%al
	shll	$6, %eax

% llc -o - sh2.ll -mtriple=aarch64
src:                    
	mov	w8, #2
	lsl	w8, w8, w0
	and	w0, w8, #0x40

tgt:        
	cmp	w0, #5
	cset	w8, eq
	lsl	w0, w8, #6

[bcl5980]

Actually even if the instruction number are the same, we still prefer shift+and because cmp port is less than shift for most modern CPU I think.

Do you have an example for a target where the code looks worse with the cmp+select IR?

And I try to create 2 cases: https://alive2.llvm.org/ce/z/exvD7T

I agree that they have very little difference generally.
I also agree that we shouldn't consider too much CPU micro arch for a IR transform.
Maybe we can add metadata MD_range for x to keep the information then backend can get enough information to invert the transform.

[spatel]

In D127801#3585983, @bcl5980 wrote:

Actually even if the instruction number are the same, we still prefer shift+and because cmp port is less than shift for most modern CPU I think.

Yes, some targets may have better throughput, but I don't think there is enough difference to change the decision in IR.

As a counter-example, consider the horrible result for the shift+and pattern with x86 vectors (a target that may not have good support for variable shifts):
https://alive2.llvm.org/ce/z/ciMo-S

So we just have to decide if this canonicalization is worthwhile to improve IR. The backend has to do some work either way to produce optimal code for all targets.

And I try to create 2 cases: https://alive2.llvm.org/ce/z/exvD7T

Thanks - I understand your point. Although for RISCV64, I don't think we should see the sign extensions in a normal example from source:
https://godbolt.org/z/KWEPPTsx1

Maybe we can add metadata MD_range for x to keep the information then backend can get enough information to invert the transform.

I haven't looked at metadata in a long time, so I don't know if it is practical. Range metadata is limited on where it can be added (for example, not on function parameters?), and then we have to make sure it is still available for the backend to use.

Since we made similar transforms before this one, this seemed like a good canonicalization to me, but if you think there's too much risk for the backend, we can abandon. If we do that, then we'd likely want to recreate this transform in SDAG, so we can still improve examples like the x86 code I showed in the previous link.

[bcl5980]

Thanks for the x86 counter example that I haven't thought about.
My point is backend can do this transform also. But it is hard to invert it if we canonicalize IR shift+and to icmp+select.
I think we need someone else help to review the discussion @RKSimon @nikic @craig.topper.

[nikic]

LG from my side. I don't think the differences in backend codegen are bad enough that we would even bother with a reverse transform at this point. We can reevaluate matching a more complex pattern if it becomes an issue.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1917

op0 -> Op0

[spatel]

Thanks - let's give this a try.
It seems like it can help IR analysis, and there are potential wins for codegen too. If we find perf regressions that are not easy to reverse, then we can revert.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

if (match(Op0, m_OneUse(m_Xor(m_Value(X), m_Value(Y)))) ||
       if (!isa<Constant>(Y) && MaskedValueIsZero(Y, NotAndMask, 0, &I)) {
         // Not masking anything out for the RHS, move mask to LHS.
         // and ({x}or X, Y), C --> {x}or (and X, C), Y
         Value *NewLHS = Builder.CreateAnd(X, Op1, X->getName() + ".masked");
         return BinaryOperator::Create(BinOp, NewLHS, Y);
       }
     }
 
+    // When the mask is a power-of-2 constant and op0 is a shifted-power-of-2
+    // constant, test if the shift amount equals the offset bit index:
+    // (ShiftC << X) & C --> X == (log2(C) - log2(ShiftC)) ? C : 0
+    // (ShiftC >> X) & C --> X == (log2(ShiftC) - log2(C)) ? C : 0
+    if (C->isPowerOf2() &&
+        match(Op0, m_OneUse(m_LogicalShift(m_Power2(ShiftC), m_Value(X))))) {
+      int Log2ShiftC = ShiftC->exactLogBase2();
+      int Log2C = C->exactLogBase2();
+      bool IsShiftLeft =
+         cast<BinaryOperator>(Op0)->getOpcode() == Instruction::Shl;
+      int BitNum = IsShiftLeft ? Log2C - Log2ShiftC : Log2ShiftC - Log2C;
+      assert(BitNum >= 0 && "Expected demanded bits to handle impossible mask");
+      Value *Cmp = Builder.CreateICmpEQ(X, ConstantInt::get(Ty, BitNum));
+      return SelectInst::Create(Cmp, ConstantInt::get(Ty, *C),
+                                ConstantInt::getNullValue(Ty));
+    }
+
     Constant *C1, *C2;
     const APInt *C3 = C;
     Value *X;
     if (C3->isPowerOf2()) {
       Constant *Log2C3 = ConstantInt::get(Ty, C3->countTrailingZeros());
       if (match(Op0, m_OneUse(m_LShr(m_Shl(m_ImmConstant(C1), m_Value(X)),
                                      m_ImmConstant(C2)))) &&
           match(C1, m_Power2())) {

llvm/test/Transforms/InstCombine/and.ll

;
   %r = and <3 x i16> %lshr, <i16 undef, i16 8, i16 8>
   ret <3 x i16> %r
 }
 
 ; LShrC < CTTZ(ShlC) < LShrC + CTTZ(AndC)
 
 define i16 @shl_lshr_pow2_const_case2(i16 %x) {
 ; CHECK-LABEL: @shl_lshr_pow2_const_case2(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i16 2, [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[TMP1]], 8
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i16 [[X:%.*]], 2
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP1]], i16 8, i16 0
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %shl = shl i16 16, %x
   %lshr = lshr i16 %shl, 3
   %r = and i16 %lshr, 8
   ret i16 %r
 }
 
-; TODO: this pattern can be transform to icmp+select
-
 define i16 @shl_lshr_pow2_not_const_case2(i16 %x) {
 ; CHECK-LABEL: @shl_lshr_pow2_not_const_case2(
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i16 2, [[X:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i16 [[TMP1]], 8
-; CHECK-NEXT:    [[R:%.*]] = xor i16 [[AND]], 8
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i16 [[X:%.*]], 2
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP1]], i16 0, i16 8
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %shl = shl i16 16, %x
   %lshr = lshr i16 %shl, 3
   %and = and i16 %lshr, 8
   %r   = xor i16 %and, 8
   ret i16 %r
 }
;
   %shl = shl i16 3, %x
   %lshr = lshr i16 %shl, 6
   %r = and i16 %lshr, 7
   ret i16 %r
 }
 
 define i16 @lshr_lshr_pow2_const(i16 %x) {
 ; CHECK-LABEL: @lshr_lshr_pow2_const(
-; CHECK-NEXT:    [[LSHR2:%.*]] = lshr i16 32, [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[LSHR2]], 4
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i16 [[X:%.*]], 3
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP1]], i16 4, i16 0
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %lshr1 = lshr i16 2048, %x
   %lshr2 = lshr i16 %lshr1, 6
   %r = and i16 %lshr2, 4
   ret i16 %r
 }
 
;
   ret i16 %r
 }
 
 ; demanded bits path for lshr+shl+and
 ; Log2(LshrC) + ShlC < BitWidth
 
 define i16 @lshr_shl_pow2_const_case1(i16 %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case1(
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr i16 1024, [[X:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[TMP1]], 8
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i16 [[X:%.*]], 7
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP1]], i16 8, i16 0
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %lshr1 = lshr i16 256, %x
   %shl = shl i16 %lshr1, 2
   %r = and i16 %shl, 8
   ret i16 %r
 }
 
-; TODO: this pattern can be transform to icmp+select
-
 define i16 @lshr_shl_pow2_const_xor(i16 %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_xor(
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr i16 1024, [[X:%.*]]
-; CHECK-NEXT:    [[AND:%.*]] = and i16 [[TMP1]], 8
-; CHECK-NEXT:    [[R:%.*]] = xor i16 [[AND]], 8
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i16 [[X:%.*]], 7
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP1]], i16 0, i16 8
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %lshr1 = lshr i16 256, %x
   %shl = shl i16 %lshr1, 2
   %and = and i16 %shl, 8
   %r = xor i16 %and, 8
   ret i16 %r
 }

llvm/test/Transforms/InstCombine/icmp-and-shift.ll

;
   %and  = and i32 %lshr, 1
   %cmp  = icmp eq i32 %and, 0
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 define <2 x i32> @icmp_ne_and1_lshr_pow2_vec(<2 x i32> %0) {
 ; CHECK-LABEL: @icmp_ne_and1_lshr_pow2_vec(
-; CHECK-NEXT:    [[AND:%.*]] = lshr <2 x i32> <i32 2, i32 2>, [[TMP0:%.*]]
-; CHECK-NEXT:    [[AND_LOBIT:%.*]] = and <2 x i32> [[AND]], <i32 1, i32 1>
-; CHECK-NEXT:    ret <2 x i32> [[AND_LOBIT]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq <2 x i32> [[TMP0:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[CONV:%.*]] = zext <2 x i1> [[TMP2]] to <2 x i32>
+; CHECK-NEXT:    ret <2 x i32> [[CONV]]
 ;
   %lshr = lshr <2 x i32> <i32 8, i32 8>, %0
   %and  = and <2 x i32> %lshr, <i32 4, i32 4>
   %cmp  = icmp ne <2 x i32> %and, <i32 0, i32 0>
   %conv = zext <2 x i1> %cmp to <2 x i32>
   ret <2 x i32> %conv
 }
 
 define i32 @icmp_eq_and_pow2_lshr_pow2(i32 %0) {
 ; CHECK-LABEL: @icmp_eq_and_pow2_lshr_pow2(
-; CHECK-NEXT:    [[AND:%.*]] = lshr i32 2, [[TMP0:%.*]]
-; CHECK-NEXT:    [[AND_LOBIT:%.*]] = and i32 [[AND]], 1
-; CHECK-NEXT:    [[TMP2:%.*]] = xor i32 [[AND_LOBIT]], 1
-; CHECK-NEXT:    ret i32 [[TMP2]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[TMP0:%.*]], 1
+; CHECK-NEXT:    [[CONV:%.*]] = zext i1 [[TMP2]] to i32
+; CHECK-NEXT:    ret i32 [[CONV]]
 ;
   %lshr = lshr i32 8, %0
   %and  = and i32 %lshr, 4
   %cmp  = icmp eq i32 %and, 0
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 define i32 @icmp_eq_and_pow2_lshr_pow2_case2(i32 %0) {
 ; CHECK-LABEL: @icmp_eq_and_pow2_lshr_pow2_case2(
 ; CHECK-NEXT:    ret i32 1
 ;
   %lshr = lshr i32 4, %0
   %and  = and i32 %lshr, 8
   %cmp  = icmp eq i32 %and, 0
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 define <2 x i32> @icmp_eq_and_pow2_lshr_pow2_vec(<2 x i32> %0) {
 ; CHECK-LABEL: @icmp_eq_and_pow2_lshr_pow2_vec(
-; CHECK-NEXT:    [[AND:%.*]] = lshr <2 x i32> <i32 2, i32 2>, [[TMP0:%.*]]
-; CHECK-NEXT:    [[AND_LOBIT:%.*]] = and <2 x i32> [[AND]], <i32 1, i32 1>
-; CHECK-NEXT:    [[TMP2:%.*]] = xor <2 x i32> [[AND_LOBIT]], <i32 1, i32 1>
-; CHECK-NEXT:    ret <2 x i32> [[TMP2]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne <2 x i32> [[TMP0:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[CONV:%.*]] = zext <2 x i1> [[TMP2]] to <2 x i32>
+; CHECK-NEXT:    ret <2 x i32> [[CONV]]
 ;
   %lshr = lshr <2 x i32> <i32 8, i32 8>, %0
   %and  = and <2 x i32> %lshr, <i32 4, i32 4>
   %cmp  = icmp eq <2 x i32> %and, <i32 0, i32 0>
   %conv = zext <2 x i1> %cmp to <2 x i32>
   ret <2 x i32> %conv
 }
 
 define i32 @icmp_ne_and_pow2_lshr_pow2(i32 %0) {
 ; CHECK-LABEL: @icmp_ne_and_pow2_lshr_pow2(
-; CHECK-NEXT:    [[AND:%.*]] = lshr i32 2, [[TMP0:%.*]]
-; CHECK-NEXT:    [[AND_LOBIT:%.*]] = and i32 [[AND]], 1
-; CHECK-NEXT:    [[TMP2:%.*]] = xor i32 [[AND_LOBIT]], 1
-; CHECK-NEXT:    ret i32 [[TMP2]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i32 [[TMP0:%.*]], 1
+; CHECK-NEXT:    [[CONV:%.*]] = zext i1 [[TMP2]] to i32
+; CHECK-NEXT:    ret i32 [[CONV]]
 ;
   %lshr = lshr i32 8, %0
   %and  = and i32 %lshr, 4
   %cmp  = icmp eq i32 %and, 0
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 define i32 @icmp_ne_and_pow2_lshr_pow2_case2(i32 %0) {
 ; CHECK-LABEL: @icmp_ne_and_pow2_lshr_pow2_case2(
 ; CHECK-NEXT:    ret i32 1
 ;
   %lshr = lshr i32 4, %0
   %and  = and i32 %lshr, 8
   %cmp  = icmp eq i32 %and, 0
   %conv = zext i1 %cmp to i32
   ret i32 %conv
 }
 
 define <2 x i32> @icmp_ne_and_pow2_lshr_pow2_vec(<2 x i32> %0) {
 ; CHECK-LABEL: @icmp_ne_and_pow2_lshr_pow2_vec(
-; CHECK-NEXT:    [[AND:%.*]] = lshr <2 x i32> <i32 2, i32 2>, [[TMP0:%.*]]
-; CHECK-NEXT:    [[AND_LOBIT:%.*]] = and <2 x i32> [[AND]], <i32 1, i32 1>
-; CHECK-NEXT:    ret <2 x i32> [[AND_LOBIT]]
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq <2 x i32> [[TMP0:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[CONV:%.*]] = zext <2 x i1> [[TMP2]] to <2 x i32>
+; CHECK-NEXT:    ret <2 x i32> [[CONV]]
 ;
   %lshr = lshr <2 x i32> <i32 8, i32 8>, %0
   %and  = and <2 x i32> %lshr, <i32 4, i32 4>
   %cmp  = icmp ne <2 x i32> %and, <i32 0, i32 0>
   %conv = zext <2 x i1> %cmp to <2 x i32>
   ret <2 x i32> %conv
 }
 

llvm/test/Transforms/InstCombine/lshr-and-signbit-icmpeq-zero.ll

;
   %lshr = lshr i32 12345, %y
   %and = and i32 %lshr, 2147483648
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
 define i1 @scalar_i32_lshr_and_negC_eq_X_is_constant2(i32 %y) {
 ; CHECK-LABEL: @scalar_i32_lshr_and_negC_eq_X_is_constant2(
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 -2147483648, [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = icmp sgt i32 [[LSHR]], -1
-; CHECK-NEXT:    ret i1 [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[Y:%.*]], 0
+; CHECK-NEXT:    ret i1 [[TMP1]]
 ;
   %lshr = lshr i32 2147483648, %y
   %and = and i32 %lshr, 2147483648
   %r = icmp eq i32 %and, 0
   ret i1 %r
 }
 
 ; Check 'slt' predicate