[InstCombine] merge shuffled vector negate and multiply

[spatel]

Summary

Add the "(0 - X) --> X * -1" reverse identity to the list of alternate form binops.

We need a little hack to make the existing logic work because it does not expect to move constants from op0 to op1, but the code comment hopefully makes that clear. I don't think there are any other identities like that.

Fixes #54364 ( https://github.com/llvm/llvm-project/issues/54364 ).

[lebedev.ri]

Seems fine to me.

[RKSimon]

Thanks @spatel!

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

static BinopElts getAlternateBinop(BinaryOperator *BO, const DataLayout &DL) {
   }
   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;
   }
+  case Instruction::Sub:
+    // sub 0, X --> mul X, -1
+    if (match(BO0, m_ZeroInt()))
+      return {Instruction::Mul, BO1, ConstantInt::getAllOnesValue(Ty)};
+    break;
   default:
     break;
   }
   return {};
 }
 
 static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
   assert(Shuf.isSelect() && "Must have select-equivalent shuffle");
Instruction *InstCombinerImpl::foldSelectShuffle(ShuffleVectorInst &Shuf) {
   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;
 
+  // If one operand is "0 - X", allow that to be viewed as "X * -1"
+  // (ConstantsAreOp1) by getAlternateBinop below. If the neg is not paired
+  // with a multiply, we will exit because C0/C1 will not be set.
   Value *X, *Y;
-  Constant *C0, *C1;
+  Constant *C0 = nullptr, *C1 = nullptr;
   bool ConstantsAreOp1;
   if (match(B0, m_BinOp(m_Constant(C0), m_Value(X))) &&
       match(B1, m_BinOp(m_Constant(C1), m_Value(Y))))
     ConstantsAreOp1 = false;
-  else if (match(B0, m_BinOp(m_Value(X), m_Constant(C0))) &&
-           match(B1, m_BinOp(m_Value(Y), m_Constant(C1))))
+  else if (match(B0, m_CombineOr(m_BinOp(m_Value(X), m_Constant(C0)),
+                                 m_Neg(m_Value(X)))) &&
+           match(B1, m_CombineOr(m_BinOp(m_Value(Y), m_Constant(C1)),
+                                 m_Neg(m_Value(Y)))))
     ConstantsAreOp1 = true;
   else
     return nullptr;
 
   // We need matching binops to fold the lanes together.
   BinaryOperator::BinaryOps Opc0 = B0->getOpcode();
   BinaryOperator::BinaryOps Opc1 = B1->getOpcode();
   bool DropNSW = false;
if (BinopElts AltB0 = getAlternateBinop(B0, DL)) {
       C0 = cast<Constant>(AltB0.Op1);
     } else if (BinopElts AltB1 = getAlternateBinop(B1, DL)) {
       assert(isa<Constant>(AltB1.Op1) && "Expecting constant with alt binop");
       Opc1 = AltB1.Opcode;
       C1 = cast<Constant>(AltB1.Op1);
     }
   }
 
-  if (Opc0 != Opc1)
+  if (Opc0 != Opc1 || !C0 || !C1)
     return nullptr;
 
   // The opcodes must be the same. Use a new name to make that clear.
   BinaryOperator::BinaryOps BOpc = Opc0;
 
   // Select the constant elements needed for the single binop.
   ArrayRef<int> Mask = Shuf.getShuffleMask();
   Constant *NewC = ConstantExpr::getShuffleVector(C0, C1, Mask);

llvm/test/Transforms/InstCombine/shuffle_select.ll

;
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 4, i32 undef, i32 2, i32 3>
   ret <4 x i32> %t3
 }
 
 ; Negate can be converted to mul to enable the fold.
 
 define <4 x i32> @mul_neg(<4 x i32> %x) {
 ; CHECK-LABEL: @mul_neg(
-; CHECK-NEXT:    [[M:%.*]] = mul <4 x i32> [[X:%.*]], <i32 257, i32 -3, i32 poison, i32 -9>
-; CHECK-NEXT:    [[N:%.*]] = sub <4 x i32> <i32 poison, i32 poison, i32 0, i32 poison>, [[X]]
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i32> [[M]], <4 x i32> [[N]], <4 x i32> <i32 0, i32 1, i32 6, i32 3>
-; CHECK-NEXT:    ret <4 x i32> [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = mul <4 x i32> [[X:%.*]], <i32 257, i32 -3, i32 -1, i32 -9>
+; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %m = mul <4 x i32> %x, <i32 257, i32 -3, i32 poison, i32 -9>
   %n = sub <4 x i32> <i32 poison, i32 poison, i32 0, i32 poison>, %x
   %r = shufflevector <4 x i32> %m, <4 x i32> %n, <4 x i32> <i32 0, i32 1, i32 6, i32 3>
   ret <4 x i32> %r
 }
 
 define <3 x i79> @neg_mul(<3 x i79> %x) {
 ; CHECK-LABEL: @neg_mul(
-; CHECK-NEXT:    [[N:%.*]] = sub nsw <3 x i79> <i79 0, i79 poison, i79 0>, [[X:%.*]]
-; CHECK-NEXT:    [[M:%.*]] = mul nsw <3 x i79> [[X]], <i79 poison, i79 -3, i79 poison>
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <3 x i79> [[N]], <3 x i79> [[M]], <3 x i32> <i32 0, i32 4, i32 2>
-; CHECK-NEXT:    ret <3 x i79> [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = mul nsw <3 x i79> [[X:%.*]], <i79 -1, i79 -3, i79 -1>
+; CHECK-NEXT:    ret <3 x i79> [[TMP1]]
 ;
   %n = sub nsw <3 x i79> <i79 0, i79 poison, i79 0>, %x
   %m = mul nsw <3 x i79> %x, <i79 poison, i79 -3, i79 poison>
   %r = shufflevector <3 x i79> %n, <3 x i79> %m, <3 x i32> <i32 0, i32 4, i32 2>
   ret <3 x i79> %r
 }
 
 define <4 x i32> @mul_neg_2_vars(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @mul_neg_2_vars(
-; CHECK-NEXT:    [[M:%.*]] = mul nuw <4 x i32> [[X:%.*]], <i32 42, i32 poison, i32 poison, i32 6>
-; CHECK-NEXT:    [[N:%.*]] = sub nsw <4 x i32> <i32 poison, i32 0, i32 0, i32 poison>, [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i32> [[M]], <4 x i32> [[N]], <4 x i32> <i32 0, i32 5, i32 6, i32 3>
-; CHECK-NEXT:    ret <4 x i32> [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[X:%.*]], <4 x i32> [[Y:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 3>
+; CHECK-NEXT:    [[TMP2:%.*]] = mul <4 x i32> [[TMP1]], <i32 42, i32 -1, i32 -1, i32 6>
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
 ;
   %m = mul nuw <4 x i32> %x, <i32 42, i32 poison, i32 poison, i32 6>
   %n = sub nsw <4 x i32> <i32 poison, i32 0, i32 0, i32 poison>, %y
   %r = shufflevector <4 x i32> %m, <4 x i32> %n, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
   ret <4 x i32> %r
 }
 
 define <4 x i32> @neg_mul_2_vars(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @neg_mul_2_vars(
-; CHECK-NEXT:    [[N:%.*]] = sub nsw <4 x i32> <i32 0, i32 poison, i32 0, i32 poison>, [[Y:%.*]]
-; CHECK-NEXT:    [[M:%.*]] = mul nuw nsw <4 x i32> [[X:%.*]], <i32 poison, i32 42, i32 poison, i32 6>
-; CHECK-NEXT:    [[R:%.*]] = shufflevector <4 x i32> [[N]], <4 x i32> [[M]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
-; CHECK-NEXT:    ret <4 x i32> [[R]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[Y:%.*]], <4 x i32> [[X:%.*]], <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[TMP2:%.*]] = mul nsw <4 x i32> [[TMP1]], <i32 -1, i32 42, i32 -1, i32 6>
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
 ;
   %n = sub nsw <4 x i32> <i32 0, i32 poison, i32 0, i32 poison>, %y
   %m = mul nuw nsw <4 x i32> %x, <i32 poison, i32 42, i32 poison, i32 6>
   %r = shufflevector <4 x i32> %n, <4 x i32> %m, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
   ret <4 x i32> %r
 }
 
 ; Or with constant can be converted to add to enable the fold.