[InstCombine] fold shuffle-with-binop and common value

This is the last significant change suggested in PR37806:
https://bugs.llvm.org/show_bug.cgi?id=37806#c5
...though there are several follow-ups noted in the code comments 
in this patch to complete this transform.

It's possible that a binop feeding a select-shuffle has been eliminated 
by earlier transforms (or the code was just written like this in the 1st 
place), so we'll fail to match the patterns that have 2 binops from: 
D48401, 
D48678, 
D48662, 
D48485.

In that case, we can try to materialize identity constants for the remaining
binop to fill in the "ghost" lanes of the vector (where we just want to pass 
through the original values of the source operand).

I added comments to ConstantExpr::getBinOpIdentity() to show planned follow-ups. 
For now, we only handle the 5 commutative integer binops (add/mul/and/or/xor).

Differential Revision: https://reviews.llvm.org/D48830

llvm-svn: 336196

llvm/lib/IR/Constants.cpp

@@ -2261,6 +2261,9 @@ Constant *ConstantExpr::getAShr(Constant *C1, Constant *C2, bool isExact) {
              isExact ? PossiblyExactOperator::IsExact : 0);
 }
 
+// FIXME: Add a parameter to specify the operand number for non-commutative ops.
+// For example, the operand 1 identity constant for any shift is the null value
+// because shift-by-0 always returns operand 0.
 Constant *ConstantExpr::getBinOpIdentity(unsigned Opcode, Type *Ty) {
   switch (Opcode) {
   default:
@@ -2277,6 +2280,8 @@ Constant *ConstantExpr::getBinOpIdentity(unsigned Opcode, Type *Ty) {
 
   case Instruction::And:
     return Constant::getAllOnesValue(Ty);
+
+  // FIXME: FAdd / FMul?
   }
 }
 

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -1182,13 +1182,60 @@ static BinopElts getAlternateBinop(BinaryOperator *BO, const DataLayout &DL) {
   return {};
 }
 
+static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
+  assert(Shuf.isSelect() && "Must have select-equivalent shuffle");
+
+  // Are we shuffling together some value and that same value after it has been
+  // modified by a binop with a constant?
+  Value *Op0 = Shuf.getOperand(0), *Op1 = Shuf.getOperand(1);
+  Constant *C;
+  bool Op0IsBinop;
+  if (match(Op0, m_BinOp(m_Specific(Op1), m_Constant(C))))
+    Op0IsBinop = true;
+  else if (match(Op1, m_BinOp(m_Specific(Op0), m_Constant(C))))
+    Op0IsBinop = false;
+  else
+    return nullptr;
+
+  auto *BO = cast<BinaryOperator>(Op0IsBinop ? Op0 : Op1);
+  Value *X = Op0IsBinop ? Op1 : Op0;
+  // TODO: Allow div/rem by accounting for potential UB due to undef elements.
+  if (BO->isIntDivRem())
+    return nullptr;
+
+  // The identity constant for a binop leaves a variable operand unchanged. For
+  // a vector, this is a splat of something like 0, -1, or 1.
+  // If there's no identity constant for this binop, we're done.
+  BinaryOperator::BinaryOps BOpcode = BO->getOpcode();
+  Constant *IdC = ConstantExpr::getBinOpIdentity(BOpcode, Shuf.getType());
+  if (!IdC)
+    return nullptr;
+
+  // Shuffle identity constants into the lanes that return the original value.
+  // Example: shuf (mul X, {-1,-2,-3,-4}), X, {0,5,6,3} --> mul X, {-1,1,1,-4}
+  // Example: shuf X, (add X, {-1,-2,-3,-4}), {0,1,6,7} --> add X, {0,0,-3,-4}
+  // The existing binop constant vector remains in the same operand position.
+  Constant *Mask = Shuf.getMask();
+  Constant *NewC = Op0IsBinop ? ConstantExpr::getShuffleVector(C, IdC, Mask) :
+                                ConstantExpr::getShuffleVector(IdC, C, Mask);
+
+  // shuf (bop X, C), X, M --> bop X, C'
+  // shuf X, (bop X, C), M --> bop X, C'
+  Instruction *NewBO = BinaryOperator::Create(BOpcode, X, NewC);
+  NewBO->copyIRFlags(BO);
+  return NewBO;
+}
+
 /// Try to fold shuffles that are the equivalent of a vector select.
 static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
                                       InstCombiner::BuilderTy &Builder,
                                       const DataLayout &DL) {
   if (!Shuf.isSelect())
     return nullptr;
 
+  if (Instruction *I = foldSelectShuffleWith1Binop(Shuf))
+    return I;
+
   BinaryOperator *B0, *B1;
   if (!match(Shuf.getOperand(0), m_BinOp(B0)) ||
       !match(Shuf.getOperand(1), m_BinOp(B1)))

llvm/test/Transforms/InstCombine/shuffle_select.ll

@@ -6,8 +6,7 @@
 
 define <4 x i32> @add(<4 x i32> %v) {
 ; CHECK-LABEL: @add(
-; CHECK-NEXT:    [[B:%.*]] = add <4 x i32> [[V:%.*]], <i32 11, i32 undef, i32 13, i32 undef>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[B]], <4 x i32> [[V]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[S:%.*]] = add <4 x i32> [[V:%.*]], <i32 11, i32 0, i32 13, i32 0>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = add <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
@@ -34,8 +33,7 @@ define <4 x i32> @sub(<4 x i32> %v) {
 
 define <4 x i32> @mul(<4 x i32> %v) {
 ; CHECK-LABEL: @mul(
-; CHECK-NEXT:    [[B:%.*]] = mul nuw nsw <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 undef, i32 14>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[V]], <4 x i32> [[B]], <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[S:%.*]] = mul nuw nsw <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 1, i32 14>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = mul nsw nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
@@ -80,8 +78,7 @@ define <3 x i32> @ashr(<3 x i32> %v) {
 
 define <3 x i42> @and(<3 x i42> %v) {
 ; CHECK-LABEL: @and(
-; CHECK-NEXT:    [[B:%.*]] = and <3 x i42> [[V:%.*]], <i42 undef, i42 12, i42 undef>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <3 x i42> [[V]], <3 x i42> [[B]], <3 x i32> <i32 0, i32 4, i32 undef>
+; CHECK-NEXT:    [[S:%.*]] = and <3 x i42> [[V:%.*]], <i42 -1, i42 12, i42 undef>
 ; CHECK-NEXT:    ret <3 x i42> [[S]]
 ;
   %b = and <3 x i42> %v, <i42 11, i42 12, i42 13>
@@ -96,7 +93,7 @@ declare void @use_v4i32(<4 x i32>)
 define <4 x i32> @or(<4 x i32> %v) {
 ; CHECK-LABEL: @or(
 ; CHECK-NEXT:    [[B:%.*]] = or <4 x i32> [[V:%.*]], <i32 11, i32 12, i32 13, i32 14>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[B]], <4 x i32> [[V]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[S:%.*]] = or <4 x i32> [[V]], <i32 0, i32 0, i32 13, i32 14>
 ; CHECK-NEXT:    call void @use_v4i32(<4 x i32> [[B]])
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
@@ -108,8 +105,7 @@ define <4 x i32> @or(<4 x i32> %v) {
 
 define <4 x i32> @xor(<4 x i32> %v) {
 ; CHECK-LABEL: @xor(
-; CHECK-NEXT:    [[B:%.*]] = xor <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 undef, i32 undef>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[V]], <4 x i32> [[B]], <4 x i32> <i32 0, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[S:%.*]] = xor <4 x i32> [[V:%.*]], <i32 0, i32 12, i32 0, i32 0>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = xor <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>