[InstCombine] Shift amount reassociation (PR42391)

[lebedev.ri]

Summary

Given pattern:
(x shiftopcode Q) shiftopcode K
we should rewrite it as
x shiftopcode (Q+K) iff (Q+K) u< bitwidth(x)
This is valid for any shift, but they must be identical.

• https://rise4fun.com/Alive/9E2
• exact on both lshr => exact https://rise4fun.com/Alive/plHk
• exact on both ashr => exact https://rise4fun.com/Alive/QDAA
• nuw on both shl => nuw https://rise4fun.com/Alive/5Uk
• nsw on both shl => nsw https://rise4fun.com/Alive/0plg

Should fix PR42391.

[nikic]

lib/Transforms/InstCombine/InstCombineShifts.cpp

37

Nice typo ;)

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineShifts.cpp

37

Sigh, i thought i only had this in next WIP patch, apparently i copied it from here :/

[lebedev.ri]

Rebased

[nikic]

lib/Transforms/InstCombine/InstCombineShifts.cpp

30

Should probably be const SimplifyQuery &SQ?

51

Would it be possible to preserve nowrap/exact flags that exist on both shifts?

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineShifts.cpp

51

• exact on both lshr => exact https://rise4fun.com/Alive/plHk
• exact on both ashr => exact https://rise4fun.com/Alive/QDAA
• nuw on both shl => nuw https://rise4fun.com/Alive/5Uk
• nsw on both shl => nsw https://rise4fun.com/Alive/0plg

So basically if the same flag is set on both original shifts -> set it on new shift.
Don't think we can do anything with non-matching flags on shl.

[lebedev.ri]

Preserve flag (nuw/nsw/exact) if both shifts had it.
PTAL.

[nikic]

LGTM

lib/Transforms/InstCombine/InstCombineShifts.cpp

37

Personally I'd prefer auto *Sh1 = cast<BinaryOperator>(Sh0->getOperand(0)) over the use of m_CombineAnd() here, which makes the match less straightforward.

41

Can replace cast<Instruction>(Sh0->getOperand(0)) with Sh1 here.

[lebedev.ri]

Thanks for the review!

Other than D63829, unless that is magically sufficient to fix all the
performance problems i'm observing, i suspect i will have some more
folds like this...

lib/Transforms/InstCombine/InstCombineShifts.cpp

37

It really differs on case-by-case basis. I think m_CombineAnd() may be better in general,
at least compared to backend lack of it.

41

Whoops, yes, i wanted to do that but forgot. Thank you.

lib/Transforms/InstCombine/InstCombineShifts.cpp

 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PatternMatch.h"
 using namespace llvm;
 using namespace PatternMatch;
 
 #define DEBUG_TYPE "instcombine"
 
+// Given pattern:
+//   (x shiftopcode Q) shiftopcode K
+// we should rewrite it as
+//   x shiftopcode (Q+K)  iff (Q+K) u< bitwidth(x)
+// This is valid for any shift, but they must be identical.
+static Instruction *
+reassociateShiftAmtsOfTwoSameDirectionShifts(BinaryOperator *Sh0,
+                                             const SimplifyQuery &SQ) {
+  // Look for:  (x shiftopcode ShAmt0) shiftopcode ShAmt1
+  Value *X, *ShAmt1, *Sh1Value, *ShAmt0;
+  if (!match(Sh0, m_Shift(m_CombineAnd(m_Shift(m_Value(X), m_Value(ShAmt1)),
+                                       m_Value(Sh1Value)),
+                          m_Value(ShAmt0))))
+    return nullptr;
+  auto *Sh1 = cast<BinaryOperator>(Sh1Value);
+
+  // The shift opcodes must be identical.
+  Instruction::BinaryOps ShiftOpcode = Sh0->getOpcode();
+  if (ShiftOpcode != cast<Instruction>(Sh0->getOperand(0))->getOpcode())
+    return nullptr;
+  // Can we fold (ShAmt0+ShAmt1) ?
+  Value *NewShAmt = SimplifyBinOp(Instruction::BinaryOps::Add, ShAmt0, ShAmt1,
+                                  SQ.getWithInstruction(Sh0));
+  if (!NewShAmt)
+    return nullptr; // Did not simplify.
+  // Is the new shift amount smaller than the bit width?
+  // FIXME: could also rely on ConstantRange.
+  unsigned BitWidth = X->getType()->getScalarSizeInBits();
+  if (!match(NewShAmt, m_SpecificInt_ULT(APInt(BitWidth, BitWidth))))
+    return nullptr;
+  // All good, we can do this fold.
+  BinaryOperator *NewShift = BinaryOperator::Create(ShiftOpcode, X, NewShAmt);
+  // If both of the original shifts had the same flag set, preserve the flag.
+  if (ShiftOpcode == Instruction::BinaryOps::Shl) {
+    NewShift->setHasNoUnsignedWrap(Sh0->hasNoUnsignedWrap() &&
+                                   Sh1->hasNoUnsignedWrap());
+    NewShift->setHasNoSignedWrap(Sh0->hasNoSignedWrap() &&
+                                 Sh1->hasNoSignedWrap());
+  } else {
+    NewShift->setIsExact(Sh0->isExact() && Sh1->isExact());
+  }
+  return NewShift;
+}
+
 Instruction *InstCombiner::commonShiftTransforms(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   assert(Op0->getType() == Op1->getType());
 
   // See if we can fold away this shift.
   if (SimplifyDemandedInstructionBits(I))
     return &I;
 
   // Try to fold constant and into select arguments.
   if (isa<Constant>(Op0))
     if (SelectInst *SI = dyn_cast<SelectInst>(Op1))
       if (Instruction *R = FoldOpIntoSelect(I, SI))
         return R;
 
   if (Constant *CUI = dyn_cast<Constant>(Op1))
     if (Instruction *Res = FoldShiftByConstant(Op0, CUI, I))
       return Res;
 
+  if (Instruction *NewShift =
+          reassociateShiftAmtsOfTwoSameDirectionShifts(&I, SQ))
+    return NewShift;
+
   // (C1 shift (A add C2)) -> (C1 shift C2) shift A)
   // iff A and C2 are both positive.
   Value *A;
   Constant *C;
   if (match(Op0, m_Constant()) && match(Op1, m_Add(m_Value(A), m_Constant(C))))
     if (isKnownNonNegative(A, DL, 0, &AC, &I, &DT) &&
         isKnownNonNegative(C, DL, 0, &AC, &I, &DT))
       return BinaryOperator::Create(

test/Transforms/InstCombine/shift-amount-reassociation.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -instcombine -S | FileCheck %s
 
 ; Given pattern:
 ;   (x shiftopcode Q) shiftopcode K
 ; we should rewrite it as
 ;   x shiftopcode (Q+K)  iff (Q+K) u< bitwidth(x)
 ; This is valid for any shift, but they must be identical.
 ; THIS FOLD DOES *NOT* REQUIRE ANY 'exact'/'nuw'/`nsw` FLAGS!
 
 ; Basic scalar test
 
 define i32 @t0(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t0(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = lshr exact i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = lshr i32 %x, %t0
   %t2 = add i32 %y, -2
   %t3 = lshr exact i32 %t1, %t2 ; while there, test that we don't propagate partial 'exact' flag
   ret i32 %t3
 }
 
 define <2 x i32> @t1_vec_splat(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @t1_vec_splat(
-; CHECK-NEXT:    [[T0:%.*]] = sub <2 x i32> <i32 32, i32 32>, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add <2 x i32> [[Y]], <i32 -2, i32 -2>
-; CHECK-NEXT:    [[T3:%.*]] = lshr <2 x i32> [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 30, i32 30>
 ; CHECK-NEXT:    ret <2 x i32> [[T3]]
 ;
   %t0 = sub <2 x i32> <i32 32, i32 32>, %y
   %t1 = lshr <2 x i32> %x, %t0
   %t2 = add <2 x i32> %y, <i32 -2, i32 -2>
   %t3 = lshr <2 x i32> %t1, %t2
   ret <2 x i32> %t3
 }
 
 define <2 x i32> @t2_vec_nonsplat(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @t2_vec_nonsplat(
-; CHECK-NEXT:    [[T0:%.*]] = sub <2 x i32> <i32 32, i32 30>, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr <2 x i32> [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add <2 x i32> [[Y]], <i32 -2, i32 0>
-; CHECK-NEXT:    [[T3:%.*]] = lshr <2 x i32> [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr <2 x i32> [[X:%.*]], <i32 30, i32 30>
 ; CHECK-NEXT:    ret <2 x i32> [[T3]]
 ;
   %t0 = sub <2 x i32> <i32 32, i32 30>, %y
   %t1 = lshr <2 x i32> %x, %t0
   %t2 = add <2 x i32> %y, <i32 -2, i32 0>
   %t3 = lshr <2 x i32> %t1, %t2
   ret <2 x i32> %t3
 }
 
 ; Basic vector tests
 
 define <3 x i32> @t3_vec_nonsplat_undef0(<3 x i32> %x, <3 x i32> %y) {
 ; CHECK-LABEL: @t3_vec_nonsplat_undef0(
-; CHECK-NEXT:    [[T0:%.*]] = sub <3 x i32> <i32 32, i32 undef, i32 32>, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr <3 x i32> [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add <3 x i32> [[Y]], <i32 -2, i32 -2, i32 -2>
-; CHECK-NEXT:    [[T3:%.*]] = lshr <3 x i32> [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr <3 x i32> [[X:%.*]], <i32 30, i32 undef, i32 30>
 ; CHECK-NEXT:    ret <3 x i32> [[T3]]
 ;
   %t0 = sub <3 x i32> <i32 32, i32 undef, i32 32>, %y
   %t1 = lshr <3 x i32> %x, %t0
   %t2 = add <3 x i32> %y, <i32 -2, i32 -2, i32 -2>
   %t3 = lshr <3 x i32> %t1, %t2
   ret <3 x i32> %t3
 }
 
 define <3 x i32> @t4_vec_nonsplat_undef1(<3 x i32> %x, <3 x i32> %y) {
 ; CHECK-LABEL: @t4_vec_nonsplat_undef1(
-; CHECK-NEXT:    [[T0:%.*]] = sub <3 x i32> <i32 32, i32 32, i32 32>, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr <3 x i32> [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add <3 x i32> [[Y]], <i32 -2, i32 undef, i32 -2>
-; CHECK-NEXT:    [[T3:%.*]] = lshr <3 x i32> [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr <3 x i32> [[X:%.*]], <i32 30, i32 undef, i32 30>
 ; CHECK-NEXT:    ret <3 x i32> [[T3]]
 ;
   %t0 = sub <3 x i32> <i32 32, i32 32, i32 32>, %y
   %t1 = lshr <3 x i32> %x, %t0
   %t2 = add <3 x i32> %y, <i32 -2, i32 undef, i32 -2>
   %t3 = lshr <3 x i32> %t1, %t2
   ret <3 x i32> %t3
 }
 
 define <3 x i32> @t5_vec_nonsplat_undef1(<3 x i32> %x, <3 x i32> %y) {
 ; CHECK-LABEL: @t5_vec_nonsplat_undef1(
-; CHECK-NEXT:    [[T0:%.*]] = sub <3 x i32> <i32 32, i32 undef, i32 32>, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr <3 x i32> [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add <3 x i32> [[Y]], <i32 -2, i32 undef, i32 -2>
-; CHECK-NEXT:    [[T3:%.*]] = lshr <3 x i32> [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr <3 x i32> [[X:%.*]], <i32 30, i32 undef, i32 30>
 ; CHECK-NEXT:    ret <3 x i32> [[T3]]
 ;
   %t0 = sub <3 x i32> <i32 32, i32 undef, i32 32>, %y
   %t1 = lshr <3 x i32> %x, %t0
   %t2 = add <3 x i32> %y, <i32 -2, i32 undef, i32 -2>
   %t3 = lshr <3 x i32> %t1, %t2
   ret <3 x i32> %t3
 }
 
 ; Some other shift opcodes
 define i32 @t6_shl(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t6_shl(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = shl nuw i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = shl nsw i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = shl i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = shl nuw i32 %x, %t0 ; while there, test that we don't propagate partial 'nuw' flag
   %t2 = add i32 %y, -2
   %t3 = shl nsw i32 %t1, %t2 ; while there, test that we don't propagate partial 'nsw' flag
   ret i32 %t3
 }
 define i32 @t7_ashr(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t7_ashr(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = ashr exact i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = ashr i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = ashr i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = ashr exact i32 %x, %t0 ; while there, test that we don't propagate partial 'exact' flag
   %t2 = add i32 %y, -2
   %t3 = ashr i32 %t1, %t2
   ret i32 %t3
 }
 
 ; If the same flag is present on both shifts, it can be kept.
 define i32 @t8_lshr_exact_flag_preservation(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t8_lshr_exact_flag_preservation(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = lshr exact i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = lshr exact i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = lshr exact i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = lshr exact i32 %x, %t0
   %t2 = add i32 %y, -2
   %t3 = lshr exact i32 %t1, %t2
   ret i32 %t3
 }
 define i32 @t9_ashr_exact_flag_preservation(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t9_ashr_exact_flag_preservation(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = ashr exact i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = ashr exact i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = ashr exact i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = ashr exact i32 %x, %t0
   %t2 = add i32 %y, -2
   %t3 = ashr exact i32 %t1, %t2
   ret i32 %t3
 }
 define i32 @t10_shl_nuw_flag_preservation(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t10_shl_nuw_flag_preservation(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = shl nuw i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = shl nuw nsw i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = shl nuw i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = shl nuw i32 %x, %t0
   %t2 = add i32 %y, -2
   %t3 = shl nuw nsw i32 %t1, %t2 ; only 'nuw' should be propagated.
   ret i32 %t3
 }
 define i32 @t11_shl_nsw_flag_preservation(i32 %x, i32 %y) {
 ; CHECK-LABEL: @t11_shl_nsw_flag_preservation(
-; CHECK-NEXT:    [[T0:%.*]] = sub i32 32, [[Y:%.*]]
-; CHECK-NEXT:    [[T1:%.*]] = shl nsw i32 [[X:%.*]], [[T0]]
-; CHECK-NEXT:    [[T2:%.*]] = add i32 [[Y]], -2
-; CHECK-NEXT:    [[T3:%.*]] = shl nuw nsw i32 [[T1]], [[T2]]
+; CHECK-NEXT:    [[T3:%.*]] = shl nsw i32 [[X:%.*]], 30
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t0 = sub i32 32, %y
   %t1 = shl nsw i32 %x, %t0
   %t2 = add i32 %y, -2
   %t3 = shl nsw nuw i32 %t1, %t2 ; only 'nuw' should be propagated.
   ret i32 %t3
 }