[InstCombine] Fold binop of shifts with related amounts

[nikic]

Summary

Fold

binop(shift(ShiftedC1, ShAmt), shift(ShiftedC2, add(ShAmt, AddC)))
  ->
shift(binop(ShiftedC1, shift(ShiftedC2, AddC)), ShAmt)

where both shifts are the same and AddC is a valid shift amount.

Proofs: https://alive2.llvm.org/ce/z/PhVVeg

[nikic]

Conservatively bail on constant expressions.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2125

https://alive2.llvm.org/ce/z/Y46Cnq add with shl works as well. Fine with only handling BW for now but maybe InstructionCombining.cpp is a better place to put this so its easier to extend in the future.

2136

nit, since early out is an option, make this 2 seperate if statements to keep the expr a little less busy?

2146

static_cast instead of C-style

2152

nit: BinaryOperator::create(...) then drop the replaceInstUsesWith below?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2151

nit: Builder.CreateBinOp(ShiftOp, ShiftedC2, AddC) -> constantexpr.

llvm/test/Transforms/InstCombine/binop-of-displaced-shifts.ll

184

maybe a negative test addc is non-splat and one element > bitwidth.

[nikic]

Support adds, address comments.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

2151

IRBuilder is preferred unless the Constant * type is required. It will work transparently when shift constant expressions are removed.

[goldstein.w.n]

LGTM.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

static Instruction *canonicalizeLogicFirst(BinaryOperator &I,
   }
 
   auto *Add = cast<BinaryOperator>(Op0);
   Value *NewBinOp = Builder.CreateBinOp(OpC, X, ConstantInt::get(Ty, *C));
   return BinaryOperator::CreateWithCopiedFlags(Instruction::Add, NewBinOp,
                                                ConstantInt::get(Ty, *C2), Add);
 }
 
+// binop(shift(ShiftedC1, ShAmt), shift(ShiftedC2, add(ShAmt, AddC))) ->
+// shift(binop(ShiftedC1, shift(ShiftedC2, AddC)), ShAmt)
+// where both shifts are the same and AddC is a valid shift amount.
+static Value *foldBinOpOfDisplacedShifts(BinaryOperator &I,
+                                         IRBuilderBase &Builder) {
+  assert(I.isBitwiseLogicOp() && "Unexpected opcode");
+
+  Value *ShAmt;
+  Constant *ShiftedC1, *ShiftedC2, *AddC;
+  Type *Ty = I.getType();
+  unsigned BitWidth = Ty->getScalarSizeInBits();
+  if (!match(&I, m_c_BinOp(
+                     m_Shift(m_ImmConstant(ShiftedC1), m_Value(ShAmt)),
+                     m_Shift(m_ImmConstant(ShiftedC2),
+                             m_Add(m_Deferred(ShAmt), m_ImmConstant(AddC))))) ||
+      !match(AddC,
+             m_SpecificInt_ICMP(ICmpInst::ICMP_ULT, APInt(BitWidth, BitWidth))))
+    return nullptr;
+
+  // Avoid constant expressions.
+  auto *Op0Inst = dyn_cast<Instruction>(I.getOperand(0));
+  auto *Op1Inst = dyn_cast<Instruction>(I.getOperand(1));
+  if (!Op0Inst || !Op1Inst)
+    return nullptr;
+
+  // Both shifts must be the same.
+  Instruction::BinaryOps ShiftOp = (Instruction::BinaryOps)Op0Inst->getOpcode();
+  if (ShiftOp != Op1Inst->getOpcode())
+    return nullptr;
+
+  Value *NewC = Builder.CreateBinOp(
+      I.getOpcode(), ShiftedC1, Builder.CreateBinOp(ShiftOp, ShiftedC2, AddC));
+  return Builder.CreateBinOp(ShiftOp, NewC, ShAmt);
+}
+
 // FIXME: We use commutative matchers (m_c_*) for some, but not all, matches
 // here. We should standardize that construct where it is needed or choose some
 // other way to ensure that commutated variants of patterns are not missed.
 Instruction *InstCombinerImpl::visitAnd(BinaryOperator &I) {
   Type *Ty = I.getType();
 
   if (Value *V = simplifyAndInst(I.getOperand(0), I.getOperand(1),
                                  SQ.getWithInstruction(&I)))
if (Instruction *R = reassociateForUses(I, Builder))
     return R;
 
   if (Instruction *Canonicalized = canonicalizeLogicFirst(I, Builder))
     return Canonicalized;
 
   if (Instruction *Folded = foldLogicOfIsFPClass(I, Op0, Op1))
     return Folded;
 
+  if (Value *Res = foldBinOpOfDisplacedShifts(I, Builder))
+    return replaceInstUsesWith(I, Res);
+
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::matchBSwapOrBitReverse(Instruction &I,
                                                       bool MatchBSwaps,
                                                       bool MatchBitReversals) {
   SmallVector<Instruction *, 4> Insts;
   if (!recognizeBSwapOrBitReverseIdiom(&I, MatchBSwaps, MatchBitReversals,
if (Instruction *R = reassociateForUses(I, Builder))
     return R;
 
   if (Instruction *Canonicalized = canonicalizeLogicFirst(I, Builder))
     return Canonicalized;
 
   if (Instruction *Folded = foldLogicOfIsFPClass(I, Op0, Op1))
     return Folded;
 
+  if (Value *Res = foldBinOpOfDisplacedShifts(I, Builder))
+    return replaceInstUsesWith(I, Res);
+
   return nullptr;
 }
 
 /// A ^ B can be specified using other logic ops in a variety of patterns. We
 /// can fold these early and efficiently by morphing an existing instruction.
 static Instruction *foldXorToXor(BinaryOperator &I,
                                  InstCombiner::BuilderTy &Builder) {
   assert(I.getOpcode() == Instruction::Xor);
if (Instruction *Canonicalized = canonicalizeLogicFirst(I, Builder))
     return Canonicalized;
 
   if (Instruction *Folded = foldLogicOfIsFPClass(I, Op0, Op1))
     return Folded;
 
   if (Instruction *Folded = canonicalizeConditionalNegationViaMathToSelect(I))
     return Folded;
 
+  if (Value *Res = foldBinOpOfDisplacedShifts(I, Builder))
+    return replaceInstUsesWith(I, Res);
+
   return nullptr;
 }

llvm/test/Transforms/InstCombine/binop-of-displaced-shifts.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -S -passes=instcombine < %s | FileCheck %s
 
 define i8 @shl_or(i8 %x) {
 ; CHECK-LABEL: define i8 @shl_or
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl i8 16, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl i8 3, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl i8 22, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = shl i8 16, %x
   %add = add i8 %x, 1
   %shift2 = shl i8 3, %add
   %binop = or i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @lshr_or(i8 %x) {
 ; CHECK-LABEL: define i8 @lshr_or
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = lshr i8 16, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = lshr i8 3, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = lshr i8 17, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = lshr i8 16, %x
   %add = add i8 %x, 1
   %shift2 = lshr i8 3, %add
   %binop = or i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @ashr_or(i8 %x) {
 ; CHECK-LABEL: define i8 @ashr_or
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = ashr i8 -64, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = ashr i8 -128, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = ashr i8 -64, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = ashr i8 -64, %x
   %add = add i8 %x, 1
   %shift2 = ashr i8 -128, %add
   %binop = or i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @shl_xor(i8 %x) {
 ; CHECK-LABEL: define i8 @shl_xor
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl i8 16, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl i8 3, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = xor i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl i8 22, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = shl i8 16, %x
   %add = add i8 %x, 1
   %shift2 = shl i8 3, %add
   %binop = xor i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @lshr_xor(i8 %x) {
 ; CHECK-LABEL: define i8 @lshr_xor
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = lshr i8 16, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = lshr i8 3, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = xor i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = lshr i8 17, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = lshr i8 16, %x
   %add = add i8 %x, 1
   %shift2 = lshr i8 3, %add
   %binop = xor i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @ashr_xor(i8 %x) {
 ; CHECK-LABEL: define i8 @ashr_xor
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = ashr i8 -128, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = ashr i8 -64, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = xor i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = lshr i8 96, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = ashr i8 -128, %x
   %add = add i8 %x, 1
   %shift2 = ashr i8 -64, %add
   %binop = xor i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @shl_and(i8 %x) {
 ; CHECK-LABEL: define i8 @shl_and
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl i8 48, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl i8 8, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = and i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl i8 16, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = shl i8 48, %x
   %add = add i8 %x, 1
   %shift2 = shl i8 8, %add
   %binop = and i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @lshr_and(i8 %x) {
 ; CHECK-LABEL: define i8 @lshr_and
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = lshr i8 48, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = lshr i8 64, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = and i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = lshr i8 32, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = lshr i8 48, %x
   %add = add i8 %x, 1
   %shift2 = lshr i8 64, %add
   %binop = and i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @ashr_and(i8 %x) {
 ; CHECK-LABEL: define i8 @ashr_and
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = ashr i8 -64, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = ashr i8 -128, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = and i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = ashr i8 -64, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = ashr i8 -64, %x
   %add = add i8 %x, 1
   %shift2 = ashr i8 -128, %add
   %binop = and i8 %shift, %shift2
   ret i8 %binop
 }
 
 define i8 @shl_or_commuted(i8 %x) {
 ; CHECK-LABEL: define i8 @shl_or_commuted
 ; CHECK-SAME: (i8 [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl i8 16, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl i8 3, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or i8 [[SHIFT2]], [[SHIFT]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl i8 22, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = shl i8 16, %x
   %add = add i8 %x, 1
   %shift2 = shl i8 3, %add
   %binop = or i8 %shift2, %shift
   ret i8 %binop
 }
 
 define <2 x i8> @shl_or_splat(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @shl_or_splat
 ; CHECK-SAME: (<2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl <2 x i8> <i8 16, i8 16>, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 1>
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl <2 x i8> <i8 3, i8 3>, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or <2 x i8> [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl <2 x i8> <i8 22, i8 22>, [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[BINOP]]
 ;
   %shift = shl <2 x i8> <i8 16, i8 16>, %x
   %add = add <2 x i8> %x, <i8 1, i8 1>
   %shift2 = shl <2 x i8> <i8 3, i8 3>, %add
   %binop = or <2 x i8> %shift, %shift2
   ret <2 x i8> %binop
 }
 
 define <2 x i8> @shl_or_non_splat(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @shl_or_non_splat
 ; CHECK-SAME: (<2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl <2 x i8> <i8 16, i8 32>, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 2>
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl <2 x i8> <i8 3, i8 7>, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or <2 x i8> [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl <2 x i8> <i8 22, i8 60>, [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[BINOP]]
 ;
   %shift = shl <2 x i8> <i8 16, i8 32>, %x
   %add = add <2 x i8> %x, <i8 1, i8 2>
   %shift2 = shl <2 x i8> <i8 3, i8 7>, %add
   %binop = or <2 x i8> %shift, %shift2
   ret <2 x i8> %binop
 }
 
 define <2 x i8> @shl_or_undef_in_add(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @shl_or_undef_in_add
 ; CHECK-SAME: (<2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl <2 x i8> <i8 16, i8 16>, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 undef>
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl <2 x i8> <i8 3, i8 3>, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or <2 x i8> [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl <2 x i8> <i8 22, i8 poison>, [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[BINOP]]
 ;
   %shift = shl <2 x i8> <i8 16, i8 16>, %x
   %add = add <2 x i8> %x, <i8 1, i8 undef>
   %shift2 = shl <2 x i8> <i8 3, i8 3>, %add
   %binop = or <2 x i8> %shift, %shift2
   ret <2 x i8> %binop
 }
 
 define <2 x i8> @shl_or_undef_in_shift1(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @shl_or_undef_in_shift1
 ; CHECK-SAME: (<2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl <2 x i8> <i8 16, i8 undef>, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 1>
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl <2 x i8> <i8 3, i8 3>, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or <2 x i8> [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl <2 x i8> <i8 22, i8 -1>, [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[BINOP]]
 ;
   %shift = shl <2 x i8> <i8 16, i8 undef>, %x
   %add = add <2 x i8> %x, <i8 1, i8 1>
   %shift2 = shl <2 x i8> <i8 3, i8 3>, %add
   %binop = or <2 x i8> %shift, %shift2
   ret <2 x i8> %binop
 }
 
 define <2 x i8> @shl_or_undef_in_shift2(<2 x i8> %x) {
 ; CHECK-LABEL: define <2 x i8> @shl_or_undef_in_shift2
 ; CHECK-SAME: (<2 x i8> [[X:%.*]]) {
-; CHECK-NEXT:    [[SHIFT:%.*]] = shl <2 x i8> <i8 16, i8 16>, [[X]]
-; CHECK-NEXT:    [[ADD:%.*]] = add <2 x i8> [[X]], <i8 1, i8 1>
-; CHECK-NEXT:    [[SHIFT2:%.*]] = shl <2 x i8> <i8 3, i8 undef>, [[ADD]]
-; CHECK-NEXT:    [[BINOP:%.*]] = or <2 x i8> [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl <2 x i8> <i8 22, i8 16>, [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[BINOP]]
 ;
   %shift = shl <2 x i8> <i8 16, i8 16>, %x
   %add = add <2 x i8> %x, <i8 1, i8 1>
   %shift2 = shl <2 x i8> <i8 3, i8 undef>, %add
   %binop = or <2 x i8> %shift, %shift2
   ret <2 x i8> %binop
 }
 
 declare void @use(i8)
 
 define i8 @shl_or_multiuse(i8 %x) {
 ; CHECK-LABEL: define i8 @shl_or_multiuse
 ; CHECK-SAME: (i8 [[X:%.*]]) {
 ; CHECK-NEXT:    [[SHIFT:%.*]] = shl i8 16, [[X]]
 ; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[X]], 1
 ; CHECK-NEXT:    [[SHIFT2:%.*]] = shl i8 3, [[ADD]]
 ; CHECK-NEXT:    call void @use(i8 [[SHIFT]])
 ; CHECK-NEXT:    call void @use(i8 [[ADD]])
 ; CHECK-NEXT:    call void @use(i8 [[SHIFT2]])
-; CHECK-NEXT:    [[BINOP:%.*]] = or i8 [[SHIFT]], [[SHIFT2]]
+; CHECK-NEXT:    [[BINOP:%.*]] = shl i8 22, [[X]]
 ; CHECK-NEXT:    ret i8 [[BINOP]]
 ;
   %shift = shl i8 16, %x
   %add = add i8 %x, 1
   %shift2 = shl i8 3, %add
   call void @use(i8 %shift)
   call void @use(i8 %add)
   call void @use(i8 %shift2)