[InstCombine] canonicalize shift+logic+shift to reduce dependency chain

[spatel]

Summary

shift (logic (shift X, C0), Y), C1 --> logic (shift X, C0+C1), (shift Y, C1)

This is an IR translation of an existing SDAG transform added here:
rL370617

So we again have 9 possible patterns with a commuted IR variant of each pattern:
https://rise4fun.com/Alive/VlI
https://rise4fun.com/Alive/n1m
https://rise4fun.com/Alive/1Vn

Part of the motivation is to allow easier recognition and subsequent canonicalization of bswap patterns as discussed in PR43146:
https://bugs.llvm.org/show_bug.cgi?id=43146

We had to delay this transform because it used to allow the SLP vectorizer to create awful reductions out of simple load-combines. That problem was fixed with:
rL375025
(we'll bring back load combining in IR someday...)

The backend is also better equipped to deal with these patterns now using hooks like TLI.getShiftAmountThreshold().

The only remaining potential controversy is that the -reassociate pass tends to reverse this kind of pattern (to help GVN?). But since -reassociate doesn't do anything with these specific patterns, there is no conflict currently.

Finally, there's a new pass proposal at D67383 for general tree-height-reduction reassociation, and it could use a cost model to decide how to optimally rearrange these kinds of ops for a target. That patch appears to be stalled.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

304–305

Do you want to assert that I (and thus the inner 'shift') is actually a shift?

[spatel]

Patch updated:
Assert that the input argument is a shift; if that's not true, we're being called from some unexpected place.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

304–305

Sure - more asserts are good.

[lebedev.ri]

Okay, thanks, LG.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

321

This innermost shift only has to be one-use iff Y is not constant.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

321

Good point - I'll make it a TODO for this patch, and we can follow-up with an enhancement.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

auto *NewShift = BinaryOperator::Create(OuterShift->getOpcode(), X,
                                           OuterShift->getOperand(1));
   if (!NeedMask)
     return NewShift;
 
   Builder.Insert(NewShift);
   return BinaryOperator::Create(Instruction::And, NewShift, NewMask);
 }
 
+/// If we have a shift-by-constant of a bitwise logic op that itself has a
+/// shift-by-constant operand with identical opcode, we may be able to convert
+/// that into 2 independent shifts followed by the logic op. This eliminates a
+/// a use of an intermediate value (reduces dependency chain).
+static Instruction *foldShiftOfShiftedLogic(BinaryOperator &I,
+                                            InstCombiner::BuilderTy &Builder) {
+  assert(I.isShift() && "Expected a shift as input");
+  auto *LogicInst = dyn_cast<BinaryOperator>(I.getOperand(0));
+  if (!LogicInst || !LogicInst->isBitwiseLogicOp() || !LogicInst->hasOneUse())
+    return nullptr;
+
+  const APInt *C0, *C1;
+  if (!match(I.getOperand(1), m_APInt(C1)))
+    return nullptr;
+
+  Instruction::BinaryOps ShiftOpcode = I.getOpcode();
+  Type *Ty = I.getType();
+
+  // Find a matching one-use shift by constant. The fold is not valid if the sum
+  // of the shift values equals or exceeds bitwidth.
+  // TODO: Remove the one-use check if the other logic operand (Y) is constant.
+  Value *X, *Y;
+  auto matchFirstShift = [&](Value *V) {
+    return match(V, m_OneUse(m_Shift(m_Value(X), m_APInt(C0)))) &&
+           cast<BinaryOperator>(V)->getOpcode() == ShiftOpcode &&
+           (*C0 + *C1).ult(Ty->getScalarSizeInBits());
+  };
+
+  // Logic ops are commutative, so check each operand for a match.
+  if (matchFirstShift(LogicInst->getOperand(0)))
+    Y = LogicInst->getOperand(1);
+  else if (matchFirstShift(LogicInst->getOperand(1)))
+    Y = LogicInst->getOperand(0);
+  else
+    return nullptr;
+
+  // shift (logic (shift X, C0), Y), C1 -> logic (shift X, C0+C1), (shift Y, C1)
+  Constant *ShiftSumC = ConstantInt::get(Ty, *C0 + *C1);
+  Value *NewShift1 = Builder.CreateBinOp(ShiftOpcode, X, ShiftSumC);
+  Value *NewShift2 = Builder.CreateBinOp(ShiftOpcode, Y, I.getOperand(1));
+  return BinaryOperator::Create(LogicInst->getOpcode(), NewShift1, NewShift2);
+}
+
 Instruction *InstCombiner::commonShiftTransforms(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   assert(Op0->getType() == Op1->getType());
 
   // If the shift amount is a one-use `sext`, we can demote it to `zext`.
   Value *Y;
   if (match(Op1, m_OneUse(m_SExt(m_Value(Y))))) {
     Value *NewExt = Builder.CreateZExt(Y, I.getType(), Op1->getName());
if (Op1->hasOneUse() && match(Op1, m_SRem(m_Value(A), m_Power2(B)))) {
     // FIXME: Should this get moved into SimplifyDemandedBits by saying we don't
     // demand the sign bit (and many others) here??
     Value *Rem = Builder.CreateAnd(A, ConstantInt::get(I.getType(), *B - 1),
                                    Op1->getName());
     I.setOperand(1, Rem);
     return &I;
   }
 
+  if (Instruction *Logic = foldShiftOfShiftedLogic(I, Builder))
+    return Logic;
+
   return nullptr;
 }
 
 /// Return true if we can simplify two logical (either left or right) shifts
 /// that have constant shift amounts: OuterShift (InnerShift X, C1), C2.
 static bool canEvaluateShiftedShift(unsigned OuterShAmt, bool IsOuterShl,
                                     Instruction *InnerShift, InstCombiner &IC,
                                     Instruction *CxtI) {

llvm/test/Transforms/InstCombine/bswap.ll

 }
 
 ; Extra use should not prevent matching to bswap.
 ; swaphalf = (x << 16 | x >> 16)
 ; ((swaphalf << 8) & 0xff00ff00) | ((swaphalf >> 8) & 0x00ff00ff)
 
 define i32 @bswap32_shl_first_extra_use(i32 %x) {
 ; CHECK-LABEL: @bswap32_shl_first_extra_use(
-; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[X:%.*]], 16
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[X]], 16
-; CHECK-NEXT:    [[SWAPHALF:%.*]] = or i32 [[SHL]], [[SHR]]
-; CHECK-NEXT:    [[T:%.*]] = shl i32 [[SWAPHALF]], 8
+; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[X:%.*]], 16
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X]], 24
+; CHECK-NEXT:    [[TMP2:%.*]] = shl nuw nsw i32 [[SHR]], 8
+; CHECK-NEXT:    [[T:%.*]] = or i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    [[BSWAP:%.*]] = call i32 @llvm.bswap.i32(i32 [[X]])
 ; CHECK-NEXT:    call void @extra_use(i32 [[T]])
 ; CHECK-NEXT:    ret i32 [[BSWAP]]
 ;
   %shl = shl i32 %x, 16
   %shr = lshr i32 %x, 16
   %swaphalf = or i32 %shl, %shr
   %t = shl i32 %swaphalf, 8

llvm/test/Transforms/InstCombine/shift-logic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define i8 @shl_and(i8 %x, i8 %y) {
 ; CHECK-LABEL: @shl_and(
-; CHECK-NEXT:    [[SH0:%.*]] = shl i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[SH0]], [[Y:%.*]]
-; CHECK-NEXT:    [[SH1:%.*]] = shl i8 [[R]], 2
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i8 [[X:%.*]], 5
+; CHECK-NEXT:    [[TMP2:%.*]] = shl i8 [[Y:%.*]], 2
+; CHECK-NEXT:    [[SH1:%.*]] = and i8 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i8 [[SH1]]
 ;
   %sh0 = shl i8 %x, 3
   %r = and i8 %sh0, %y
   %sh1 = shl i8 %r, 2
   ret i8 %sh1
 }
 
 define i16 @shl_or(i16 %x, i16 %py) {
 ; CHECK-LABEL: @shl_or(
 ; CHECK-NEXT:    [[Y:%.*]] = srem i16 [[PY:%.*]], 42
-; CHECK-NEXT:    [[SH0:%.*]] = shl i16 [[X:%.*]], 5
-; CHECK-NEXT:    [[R:%.*]] = or i16 [[Y]], [[SH0]]
-; CHECK-NEXT:    [[SH1:%.*]] = shl i16 [[R]], 7
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i16 [[X:%.*]], 12
+; CHECK-NEXT:    [[TMP2:%.*]] = shl nsw i16 [[Y]], 7
+; CHECK-NEXT:    [[SH1:%.*]] = or i16 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i16 [[SH1]]
 ;
   %y = srem i16 %py, 42 ; thwart complexity-based canonicalization
   %sh0 = shl i16 %x, 5
   %r = or i16 %y, %sh0
   %sh1 = shl i16 %r, 7
   ret i16 %sh1
 }
 
 define i32 @shl_xor(i32 %x, i32 %y) {
 ; CHECK-LABEL: @shl_xor(
-; CHECK-NEXT:    [[SH0:%.*]] = shl i32 [[X:%.*]], 5
-; CHECK-NEXT:    [[R:%.*]] = xor i32 [[SH0]], [[Y:%.*]]
-; CHECK-NEXT:    [[SH1:%.*]] = shl i32 [[R]], 7
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 12
+; CHECK-NEXT:    [[TMP2:%.*]] = shl i32 [[Y:%.*]], 7
+; CHECK-NEXT:    [[SH1:%.*]] = xor i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[SH1]]
 ;
   %sh0 = shl i32 %x, 5
   %r = xor i32 %sh0, %y
   %sh1 = shl i32 %r, 7
   ret i32 %sh1
 }
 
 define i64 @lshr_and(i64 %x, i64 %py) {
 ; CHECK-LABEL: @lshr_and(
 ; CHECK-NEXT:    [[Y:%.*]] = srem i64 [[PY:%.*]], 42
-; CHECK-NEXT:    [[SH0:%.*]] = lshr i64 [[X:%.*]], 5
-; CHECK-NEXT:    [[R:%.*]] = and i64 [[Y]], [[SH0]]
-; CHECK-NEXT:    [[SH1:%.*]] = lshr i64 [[R]], 7
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[X:%.*]], 12
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i64 [[Y]], 7
+; CHECK-NEXT:    [[SH1:%.*]] = and i64 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i64 [[SH1]]
 ;
   %y = srem i64 %py, 42 ; thwart complexity-based canonicalization
   %sh0 = lshr i64 %x, 5
   %r = and i64 %y, %sh0
   %sh1 = lshr i64 %r, 7
   ret i64 %sh1
 }
 
 define <4 x i32> @lshr_or(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @lshr_or(
-; CHECK-NEXT:    [[SH0:%.*]] = lshr <4 x i32> [[X:%.*]], <i32 5, i32 5, i32 5, i32 5>
-; CHECK-NEXT:    [[R:%.*]] = or <4 x i32> [[SH0]], [[Y:%.*]]
-; CHECK-NEXT:    [[SH1:%.*]] = lshr <4 x i32> [[R]], <i32 7, i32 7, i32 7, i32 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <4 x i32> [[X:%.*]], <i32 12, i32 12, i32 12, i32 12>
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr <4 x i32> [[Y:%.*]], <i32 7, i32 7, i32 7, i32 7>
+; CHECK-NEXT:    [[SH1:%.*]] = or <4 x i32> [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret <4 x i32> [[SH1]]
 ;
   %sh0 = lshr <4 x i32> %x, <i32 5, i32 5, i32 5, i32 5>
   %r = or <4 x i32> %sh0, %y
   %sh1 = lshr <4 x i32> %r, <i32 7, i32 7, i32 7, i32 7>
   ret <4 x i32> %sh1
 }
 
 define <8 x i16> @lshr_xor(<8 x i16> %x, <8 x i16> %py) {
 ; CHECK-LABEL: @lshr_xor(
 ; CHECK-NEXT:    [[Y:%.*]] = srem <8 x i16> [[PY:%.*]], <i16 42, i16 42, i16 42, i16 42, i16 42, i16 42, i16 42, i16 42>
-; CHECK-NEXT:    [[SH0:%.*]] = lshr <8 x i16> [[X:%.*]], <i16 5, i16 5, i16 5, i16 5, i16 5, i16 5, i16 5, i16 5>
-; CHECK-NEXT:    [[R:%.*]] = xor <8 x i16> [[Y]], [[SH0]]
-; CHECK-NEXT:    [[SH1:%.*]] = lshr <8 x i16> [[R]], <i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <8 x i16> [[X:%.*]], <i16 12, i16 12, i16 12, i16 12, i16 12, i16 12, i16 12, i16 12>
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr <8 x i16> [[Y]], <i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7>
+; CHECK-NEXT:    [[SH1:%.*]] = xor <8 x i16> [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret <8 x i16> [[SH1]]
 ;
   %y = srem <8 x i16> %py, <i16 42, i16 42, i16 42, i16 42, i16 42, i16 42, i16 42, i16 -42> ; thwart complexity-based canonicalization
   %sh0 = lshr <8 x i16> %x, <i16 5, i16 5, i16 5, i16 5, i16 5, i16 5, i16 5, i16 5>
   %r = xor <8 x i16> %y, %sh0
   %sh1 = lshr <8 x i16> %r, <i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7, i16 7>
   ret <8 x i16> %sh1
 }
 
 
 define <16 x i8> @ashr_and(<16 x i8> %x, <16 x i8> %py, <16 x i8> %pz) {
 ; CHECK-LABEL: @ashr_and(
 ; CHECK-NEXT:    [[Y:%.*]] = srem <16 x i8> [[PY:%.*]], [[PZ:%.*]]
-; CHECK-NEXT:    [[SH0:%.*]] = ashr <16 x i8> [[X:%.*]], <i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3>
-; CHECK-NEXT:    [[R:%.*]] = and <16 x i8> [[Y]], [[SH0]]
-; CHECK-NEXT:    [[SH1:%.*]] = ashr <16 x i8> [[R]], <i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2>
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr <16 x i8> [[X:%.*]], <i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5>
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr <16 x i8> [[Y]], <i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2>
+; CHECK-NEXT:    [[SH1:%.*]] = and <16 x i8> [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret <16 x i8> [[SH1]]
 ;
   %y = srem <16 x i8> %py, %pz ; thwart complexity-based canonicalization
   %sh0 = ashr <16 x i8> %x, <i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3, i8 3>
   %r = and <16 x i8> %y, %sh0
   %sh1 = ashr <16 x i8> %r, <i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2, i8 2>
   ret <16 x i8> %sh1
 }
 
 define <2 x i64> @ashr_or(<2 x i64> %x, <2 x i64> %y) {
 ; CHECK-LABEL: @ashr_or(
-; CHECK-NEXT:    [[SH0:%.*]] = ashr <2 x i64> [[X:%.*]], <i64 5, i64 5>
-; CHECK-NEXT:    [[R:%.*]] = or <2 x i64> [[SH0]], [[Y:%.*]]
-; CHECK-NEXT:    [[SH1:%.*]] = ashr <2 x i64> [[R]], <i64 7, i64 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i64> [[X:%.*]], <i64 12, i64 12>
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr <2 x i64> [[Y:%.*]], <i64 7, i64 7>
+; CHECK-NEXT:    [[SH1:%.*]] = or <2 x i64> [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret <2 x i64> [[SH1]]
 ;
   %sh0 = ashr <2 x i64> %x, <i64 5, i64 5>
   %r = or <2 x i64> %sh0, %y
   %sh1 = ashr <2 x i64> %r, <i64 7, i64 7>
   ret <2 x i64> %sh1
 }
 
 define i32 @ashr_xor(i32 %x, i32 %py) {
 ; CHECK-LABEL: @ashr_xor(
 ; CHECK-NEXT:    [[Y:%.*]] = srem i32 [[PY:%.*]], 42
-; CHECK-NEXT:    [[SH0:%.*]] = ashr i32 [[X:%.*]], 5
-; CHECK-NEXT:    [[R:%.*]] = xor i32 [[Y]], [[SH0]]
-; CHECK-NEXT:    [[SH1:%.*]] = ashr i32 [[R]], 7
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[X:%.*]], 12
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr i32 [[Y]], 7
+; CHECK-NEXT:    [[SH1:%.*]] = xor i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[SH1]]
 ;
   %y = srem i32 %py, 42 ; thwart complexity-based canonicalization
   %sh0 = ashr i32 %x, 5
   %r = xor i32 %y, %sh0
   %sh1 = ashr i32 %r, 7
   ret i32 %sh1
 }