[InstCombine] enhance shuffle-of-binops to allow different variable ops (PR37806)
This was discussed in D48401 as another improvement for:
https://bugs.llvm.org/show_bug.cgi?id=37806
If we have 2 different variable values, then we shuffle (select) those lanes,
shuffle (select) the constants, and then perform the binop. This eliminates a binop.
The new shuffle uses the same shuffle mask as the existing shuffle, so there's no
danger of creating a difficult shuffle.
All of the earlier constraints still apply, but we also check for extra uses to
avoid creating more instructions than we'll remove.
Additionally, we're disallowing the fold for div/rem because that could expose a
UB hole.
Differential Revision: https://reviews.llvm.org/D48678
llvm-svn: 335974
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index d4fb086..dd579a6 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -1140,7 +1140,8 @@
return true;
}
-static Instruction *foldSelectShuffles(ShuffleVectorInst &Shuf) {
+static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
+ InstCombiner::BuilderTy &Builder) {
// Folds under here require the equivalent of a vector select.
if (!Shuf.isSelect())
return nullptr;
@@ -1150,16 +1151,14 @@
!match(Shuf.getOperand(1), m_BinOp(B1)))
return nullptr;
- // TODO: Fold the case with different variable operands (requires creating a
- // new shuffle and checking number of uses).
- Value *X;
+ Value *X, *Y;
Constant *C0, *C1;
bool ConstantsAreOp1;
if (match(B0, m_BinOp(m_Value(X), m_Constant(C0))) &&
- match(B1, m_BinOp(m_Specific(X), m_Constant(C1))))
+ match(B1, m_BinOp(m_Value(Y), m_Constant(C1))))
ConstantsAreOp1 = true;
else if (match(B0, m_BinOp(m_Constant(C0), m_Value(X))) &&
- match(B1, m_BinOp(m_Constant(C1), m_Specific(X))))
+ match(B1, m_BinOp(m_Constant(C1), m_Value(Y))))
ConstantsAreOp1 = false;
else
return nullptr;
@@ -1191,9 +1190,36 @@
// The opcodes must be the same. Use a new name to make that clear.
BinaryOperator::BinaryOps BOpc = Opc0;
- // Remove a binop and the shuffle by rearranging the constant:
- // shuffle (op X, C0), (op X, C1), M --> op X, C'
- // shuffle (op C0, X), (op C1, X), M --> op C', X
+ Value *V;
+ if (X == Y) {
+ // Remove a binop and the shuffle by rearranging the constant:
+ // shuffle (op V, C0), (op V, C1), M --> op V, C'
+ // shuffle (op C0, V), (op C1, V), M --> op C', V
+ V = X;
+ } else if (!Instruction::isIntDivRem(BOpc) &&
+ (B0->hasOneUse() || B1->hasOneUse())) {
+ // If there are 2 different variable operands, we must create a new shuffle
+ // (select) first, so check uses to ensure that we don't end up with more
+ // instructions than we started with.
+ //
+ // Note: In general, we do not create new shuffles in InstCombine because we
+ // do not know if a target can lower an arbitrary shuffle optimally. In this
+ // case, the shuffle uses the existing mask, so there is no additional risk.
+ //
+ // TODO: We are disallowing div/rem because a shuffle with an undef mask
+ // element would propagate an undef value to the div/rem. That's not
+ // safe in general because div/rem allow for undefined behavior. We can
+ // loosen this restriction (eg, check if the mask has no undefs or replace
+ // undef elements).
+
+ // Select the variable vectors first, then perform the binop:
+ // shuffle (op X, C0), (op Y, C1), M --> op (shuffle X, Y, M), C'
+ // shuffle (op C0, X), (op C1, Y), M --> op C', (shuffle X, Y, M)
+ V = Builder.CreateShuffleVector(X, Y, Shuf.getMask());
+ } else {
+ return nullptr;
+ }
+
Constant *NewC = ConstantExpr::getShuffleVector(C0, C1, Shuf.getMask());
// If the shuffle mask contains undef elements, then the new constant
@@ -1202,8 +1228,8 @@
if (Instruction::isIntDivRem(BOpc))
NewC = getSafeVectorConstantForIntDivRem(NewC);
- Instruction *NewBO = ConstantsAreOp1 ? BinaryOperator::Create(BOpc, X, NewC) :
- BinaryOperator::Create(BOpc, NewC, X);
+ Instruction *NewBO = ConstantsAreOp1 ? BinaryOperator::Create(BOpc, V, NewC) :
+ BinaryOperator::Create(BOpc, NewC, V);
// Flags are intersected from the 2 source binops.
NewBO->copyIRFlags(B0);
@@ -1223,7 +1249,7 @@
LHS, RHS, SVI.getMask(), SVI.getType(), SQ.getWithInstruction(&SVI)))
return replaceInstUsesWith(SVI, V);
- if (Instruction *I = foldSelectShuffles(SVI))
+ if (Instruction *I = foldSelectShuffle(SVI, Builder))
return I;
bool MadeChange = false;