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

[spatel]

Summary

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 suggested 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).

[spatel]

Note: this patch is against trunk. It should patch cleanly independently of D48662 (but there would be offsets if that goes in 1st).

[lebedev.ri]

I don't have any notable feedback, so other than nits, looks good.
Maybe i'm missing something and you want to wait for a second review[er].

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1158

Would it be simpler to do

Value *OpV = Op0IsBinop ? Op0 : Op1
Value *OpC = Op0IsBinop ? Op1 : Op0

and use these later on?

1166

Note: ConstantExpr::getBinOpIdentity() does not handle div/rem.

1170–1174

Nit: i would like the comment to show an example here.

[spatel]

Patch updated:
No logic difference from the previous rev, just tried to improve the code structure/comments as suggested.
If there are no objections, I'll commit tomorrow.

lib/IR/Constants.cpp

return get(Instruction::LShr, C1, C2,
              isExact ? PossiblyExactOperator::IsExact : 0);
 }
 
 Constant *ConstantExpr::getAShr(Constant *C1, Constant *C2, bool isExact) {
   return get(Instruction::AShr, C1, C2,
              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:
     // Doesn't have an identity.
     return nullptr;
 
   case Instruction::Add:
   case Instruction::Or:
   case Instruction::Xor:
     return Constant::getNullValue(Ty);
 
   case Instruction::Mul:
     return ConstantInt::get(Ty, 1);
 
   case Instruction::And:
     return Constant::getAllOnesValue(Ty);
+
+  // FIXME: FAdd / FMul?
   }
 }
 
 Constant *ConstantExpr::getBinOpAbsorber(unsigned Opcode, Type *Ty) {
   switch (Opcode) {
   default:
     // Doesn't have an absorber.
     return nullptr;

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

BinopElts(BinaryOperator::BinaryOps Opc = (BinaryOperator::BinaryOps)0,
             Value *V0 = nullptr, Value *V1 = nullptr) :
       Opcode(Opc), Op0(V0), Op1(V1) {}
   operator bool() const { return Opcode != 0; }
 };
 
 /// Binops may be transformed into binops with different opcodes and operands.
 /// Reverse the usual canonicalization to enable folds with the non-canonical
 /// form of the binop. If a transform is possible, return the elements of the
 /// new binop. If not, return invalid elements.
 static BinopElts getAlternateBinop(BinaryOperator *BO, const DataLayout &DL) {
   Value *BO0 = BO->getOperand(0), *BO1 = BO->getOperand(1);
   Type *Ty = BO->getType();
   switch (BO->getOpcode()) {
     case Instruction::Shl: {
       // shl X, C --> mul X, (1 << C)
       Constant *C;
       if (match(BO1, m_Constant(C))) {
         Constant *ShlOne = ConstantExpr::getShl(ConstantInt::get(Ty, 1), C);
         return { Instruction::Mul, BO0, ShlOne };
       }
       break;
     }
     case Instruction::Or: {
       // or X, C --> add X, C (when X and C have no common bits set)
       const APInt *C;
       if (match(BO1, m_APInt(C)) && MaskedValueIsZero(BO0, *C, DL))
         return { Instruction::Add, BO0, BO1 };
       break;
     }
     default:
       break;
   }
   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)))
     return nullptr;
 
   Value *X, *Y;
   Constant *C0, *C1;
   bool ConstantsAreOp1;

test/Transforms/InstCombine/shuffle_select.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; Try to eliminate binops and shuffles when the shuffle is a select in disguise:
 ; PR37806 - https://bugs.llvm.org/show_bug.cgi?id=37806
 
 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>
   %s = shufflevector <4 x i32> %b, <4 x i32> %v, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   ret <4 x i32> %s
 }
 
 ; Constant operand 0 (LHS) could work for some non-commutative binops?

 }
 
 ; If any element of the shuffle mask operand is undef, that element of the result is undef.
 ; The shuffle is eliminated in this transform, but we can replace a constant element with undef.
 ; Preserve flags when possible.
 
 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>
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
   ret <4 x i32> %s
 }
 
 define <4 x i32> @shl(<4 x i32> %v) {

 ;
   %b = ashr <3 x i32> %v, <i32 11, i32 12, i32 13>
   %s = shufflevector <3 x i32> %b, <3 x i32> %v, <3 x i32> <i32 3, i32 1, i32 2>
   ret <3 x i32> %s
 }
 
 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>
   %s = shufflevector <3 x i42> %v, <3 x i42> %b, <3 x i32> <i32 0, i32 4, i32 undef>
   ret <3 x i42> %s
 }
 
 ; It doesn't matter if the intermediate op has extra uses.
 
 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]]
 ;
   %b = or <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
   %s = shufflevector <4 x i32> %b, <4 x i32> %v, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   call void @use_v4i32(<4 x i32> %b)
   ret <4 x i32> %s
 }
 
 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>
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 0, i32 5, i32 2, i32 3>
   ret <4 x i32> %s
 }
 
 define <4 x i32> @udiv(<4 x i32> %v) {