[InstCombine] Fold (x << y) >> y -> x & (-1 >> y)

[lebedev.ri]

Summary

We already do it for splat constants, but not just values.
Also, undef cases are mostly non-functional.

https://bugs.llvm.org/show_bug.cgi?id=37603
https://rise4fun.com/Alive/cplX

[lebedev.ri]

Simplify matcher, NFC.

[spatel]

LGTM.

lib/Transforms/InstCombine/InstCombineShifts.cpp

741–744

As with D47981, we could consolidate this, but the constant version doesn't need the one-use check.

I think it's fine either way, but let's keep both patches consistent in their structure.

[lebedev.ri]

In D47980#1127542, @spatel wrote:

LGTM.

Thank you for the review!

lib/Transforms/InstCombine/InstCombineShifts.cpp

741–744

As with D47981, we could consolidate this, but the constant version doesn't need the one-use check.

I haven't really though about that..
Sounds like i should add multi-use tests with constants to the tests.

[spatel]

lib/Transforms/InstCombine/InstCombineShifts.cpp

741–744

Oops - misplaced this inline comment. It was meant for the block where ShlAmt == ShAmt.

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineShifts.cpp

741–744

As with D47981, we could consolidate this, but the constant version doesn't need the one-use check.

Hmm, this is actually broken. That is only true if either the shifts are the same constant,
or shl is nuw, else it produces one extra instruction.

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineShifts.cpp

741–744

I will add tests, but there is one test in test/Transforms/InstCombine/shift.ll that breaks
if i 'fix' this, so maybe this is intentional..

[lebedev.ri]

Ok, that didn't go as planned :)
Did not notice the AMDGPU test change (yay tests), so the bots broke.
Did not evaluate the effect of that change on the AMDGPU codegen yet, will do tomorrow.

[nhaehnle]

test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll

898–899

Indeed, the pattern matching in the backend does not recognize this and ends up generating worse code.

[lebedev.ri]

This comment has been deleted.

test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll

898–899

Yep, see D48005 for the backend tests.

[lebedev.ri]

@nhaehnle AMDGPU patches (D48007, D48010, D48012) are ready to land, please accept this differential :)

[rampitec]

In general for AMDGPU two shifts are better. Any shift immediate can be folded right into the shift instruction while a rather big mask produced by this change would require either extra 4 bytes in the encoding or even worse a move and a register.

What's the rational for the folding?

In addition as tests suggest we would expect the pattern to be folded into a bfe instruction but D48005 shows it is at best "bfm" (with an extra register to hold a mask) and "and". I.e. it basically shows a regression for our target. There probably would be no concern if the sequence is converted to a bfe as expected.

[lebedev.ri]

In D47980#1128464, @rampitec wrote:

In general for AMDGPU two shifts are better. Any shift immediate can be folded right into the shift instruction while a rather big mask produced by this change would require either extra 4 bytes in the encoding or even worse a move and a register.

What's the rational for the folding?

In addition as tests suggest we would expect the pattern to be folded into a bfe instruction but D48005 shows it is at best "bfm" (with an extra register to hold a mask) and "and". I.e. it basically shows a regression for our target. There probably would be no concern if the sequence is converted to a bfe as expected.

You did note that i have posted D48007, D48010, D48012 to fix the AMDGPU backend, right?
Any other concerns ontop of that?

[rampitec]

Ah! I see the reviews to fold it into a bfe. I have no concern then, LGTM.

[lebedev.ri]

In D47980#1128473, @rampitec wrote:

Ah! I see the reviews to fold it into a bfe. I have no concern then, LGTM.

Thank you for reassurance!
Though i'm still waiting for @nhaehnle to re-accept this.

[lebedev.ri]

Thanks!
Do note that there does not seem to be any folds for when the offset is not 0.
E.g. (x >> start) & (-1 >> (width - y)) might be foldable to (UBFE
$x, $start, $width)

[lebedev.ri]

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

741–744 ↗

(On Diff #151480)

BTW it's interesting to note that all these masks are not fine-grained, isn't it?
Alive says https://rise4fun.com/Alive/Yes (lol)

Though in practice, from what i have seen from the tests, somehow the mask seems to be adjusted later.

lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op0, m_Shl(m_Value(X), m_APInt(ShOp1)))) {
       if (ShlAmt < ShAmt) {
         Constant *ShiftDiff = ConstantInt::get(Ty, ShAmt - ShlAmt);
         if (cast<BinaryOperator>(Op0)->hasNoUnsignedWrap()) {
           // (X <<nuw C1) >>u C2 --> X >>u (C2 - C1)
           auto *NewLShr = BinaryOperator::CreateLShr(X, ShiftDiff);
           NewLShr->setIsExact(I.isExact());
           return NewLShr;
         }
         // (X << C1) >>u C2  --> (X >>u (C2 - C1)) & (-1 >> C2)
         Value *NewLShr = Builder.CreateLShr(X, ShiftDiff, "", I.isExact());
         APInt Mask(APInt::getLowBitsSet(BitWidth, BitWidth - ShAmt));
         return BinaryOperator::CreateAnd(NewLShr, ConstantInt::get(Ty, Mask));
       }
       if (ShlAmt > ShAmt) {
         Constant *ShiftDiff = ConstantInt::get(Ty, ShlAmt - ShAmt);
         if (cast<BinaryOperator>(Op0)->hasNoUnsignedWrap()) {
           // (X <<nuw C1) >>u C2 --> X <<nuw (C1 - C2)
           auto *NewShl = BinaryOperator::CreateShl(X, ShiftDiff);
           NewShl->setHasNoUnsignedWrap(true);
           return NewShl;
if (match(Op1, m_APInt(ShAmtAPInt))) {
 
     // If the shifted-out value is known-zero, then this is an exact shift.
     if (!I.isExact() &&
         MaskedValueIsZero(Op0, APInt::getLowBitsSet(BitWidth, ShAmt), 0, &I)) {
       I.setIsExact();
       return &I;
     }
   }
+
+  // Transform  (x << y) >> y  to  x & (-1 >> y)
+  Value *X;
+  if (match(Op0, m_OneUse(m_Shl(m_Value(X), m_Specific(Op1))))) {
+    Constant *AllOnes = ConstantInt::getAllOnesValue(Ty);
+    Value *Mask = Builder.CreateLShr(AllOnes, Op1);
+    return BinaryOperator::CreateAnd(Mask, X);
+  }
+
   return nullptr;
 }
 
 Instruction *InstCombiner::visitAShr(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   if (Value *V =
           SimplifyAShrInst(Op0, Op1, I.isExact(), SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);

test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll

 ; CHECK-NEXT: %bfe = call i32 @llvm.amdgcn.ubfe.i32(i32 %src, i32 1, i32 %width)
 define i32 @ubfe_offset_33(i32 %src, i32 %width) {
   %bfe = call i32 @llvm.amdgcn.ubfe.i32(i32 %src, i32 33, i32 %width)
   ret i32 %bfe
 }
 
 ; CHECK-LABEL: @ubfe_offset_0(
 ; CHECK-NEXT: %1 = sub i32 32, %width
-; CHECK-NEXT: %2 = shl i32 %src, %1
-; CHECK-NEXT: %bfe = lshr i32 %2, %1
+; CHECK-NEXT: %2 = lshr i32 -1, %1
+; CHECK-NEXT: %bfe = and i32 %2, %src
 ; CHECK-NEXT: ret i32 %bfe
 define i32 @ubfe_offset_0(i32 %src, i32 %width) {
   %bfe = call i32 @llvm.amdgcn.ubfe.i32(i32 %src, i32 0, i32 %width)
   ret i32 %bfe
 }
 
 ; CHECK-LABEL: @ubfe_offset_32(
 ; CHECK-NEXT: %1 = sub i32 32, %width
-; CHECK-NEXT: %2 = shl i32 %src, %1
-; CHECK-NEXT: %bfe = lshr i32 %2, %1
+; CHECK-NEXT: %2 = lshr i32 -1, %1
+; CHECK-NEXT: %bfe = and i32 %2, %src
 ; CHECK-NEXT: ret i32 %bfe
 define i32 @ubfe_offset_32(i32 %src, i32 %width) {
   %bfe = call i32 @llvm.amdgcn.ubfe.i32(i32 %src, i32 32, i32 %width)
   ret i32 %bfe
 }
 
 ; CHECK-LABEL: @ubfe_offset_31(
 ; CHECK-NEXT: %1 = sub i32 32, %width
-; CHECK-NEXT: %2 = shl i32 %src, %1
-; CHECK-NEXT: %bfe = lshr i32 %2, %1
+; CHECK-NEXT: %2 = lshr i32 -1, %1
+; CHECK-NEXT: %bfe = and i32 %2, %src
 ; CHECK-NEXT: ret i32 %bfe
 define i32 @ubfe_offset_31(i32 %src, i32 %width) {
   %bfe = call i32 @llvm.amdgcn.ubfe.i32(i32 %src, i32 32, i32 %width)
   ret i32 %bfe
 }
 
 ; CHECK-LABEL: @ubfe_offset_0_width_0(
 ; CHECK-NEXT: ret i32 0

 define i64 @ubfe_offset_33_width_4_i64(i64 %src) {
   %bfe = call i64 @llvm.amdgcn.ubfe.i64(i64 %src, i32 33, i32 4)
   ret i64 %bfe
 }
 
 ; CHECK-LABEL: @ubfe_offset_0_i64(
 ; CHECK-NEXT: %1 = sub i32 64, %width
 ; CHECK-NEXT: %2 = zext i32 %1 to i64
-; CHECK-NEXT: %3 = shl i64 %src, %2
-; CHECK-NEXT: %bfe = lshr i64 %3, %2
+; CHECK-NEXT: %3 = lshr i64 -1, %2
+; CHECK-NEXT: %bfe = and i64 %3, %src
 ; CHECK-NEXT: ret i64 %bfe
 define i64 @ubfe_offset_0_i64(i64 %src, i32 %width) {
   %bfe = call i64 @llvm.amdgcn.ubfe.i64(i64 %src, i32 0, i32 %width)
   ret i64 %bfe
 }
 
 ; CHECK-LABEL: @ubfe_offset_32_width_32_i64(
 ; CHECK-NEXT: %bfe = lshr i64 %src, 32

test/Transforms/InstCombine/canonicalize-shl-lshr-to-masking.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; https://bugs.llvm.org/show_bug.cgi?id=37603
 
 ; Pattern:
 ;   x << y >> y
 ; Should be transformed into:
 ;   x & (-1 >> y)
 
 ; ============================================================================ ;
 ; Basic positive tests
 ; ============================================================================ ;
 
 define i32 @positive_samevar(i32 %x, i32 %y) {
 ; CHECK-LABEL: @positive_samevar(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = lshr i32 [[TMP0]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 -1, [[Y:%.*]]
+; CHECK-NEXT:    [[RET:%.*]] = and i32 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %tmp0 = shl i32 %x, %y
   %ret = lshr i32 %tmp0, %y
   ret i32 %ret
 }
 
 define i32 @positive_sameconst(i32 %x) {

 }
 
 ; ============================================================================ ;
 ; Vector
 ; ============================================================================ ;
 
 define <2 x i32> @positive_samevar_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @positive_samevar_vec(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl <2 x i32> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = lshr <2 x i32> [[TMP0]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> <i32 -1, i32 -1>, [[Y:%.*]]
+; CHECK-NEXT:    [[RET:%.*]] = and <2 x i32> [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[RET]]
 ;
   %tmp0 = shl <2 x i32> %x, %y
   %ret = lshr <2 x i32> %tmp0, %y
   ret <2 x i32> %ret
 }
 
 ; ============================================================================ ;

 ;
   %tmp0 = shl <3 x i32> %x, <i32 5, i32 5, i32 5>
   %ret = lshr <3 x i32> %tmp0, <i32 5, i32 undef, i32 5>
   ret <3 x i32> %ret
 }
 
 define <3 x i32> @positive_sameconst_vec_undef2(<3 x i32> %x) {
 ; CHECK-LABEL: @positive_sameconst_vec_undef2(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl <3 x i32> [[X:%.*]], <i32 5, i32 undef, i32 5>
-; CHECK-NEXT:    [[RET:%.*]] = lshr <3 x i32> [[TMP0]], <i32 5, i32 undef, i32 5>
+; CHECK-NEXT:    [[RET:%.*]] = and <3 x i32> [[X:%.*]], <i32 134217727, i32 undef, i32 134217727>
 ; CHECK-NEXT:    ret <3 x i32> [[RET]]
 ;
   %tmp0 = shl <3 x i32> %x, <i32 5, i32 undef, i32 5>
   %ret = lshr <3 x i32> %tmp0, <i32 5, i32 undef, i32 5>
   ret <3 x i32> %ret
 }
 
 define <2 x i32> @positive_biggerShl_vec(<2 x i32> %x) {