Skip to content

Commit 06fbd77

Browse files
RKSimonRKSimon
committed
[WIP][VectorCombine] Fold "shuffle (binop (shuffle, shuffle)), undef" --> "binop (shuffle), (shuffle)"
Add foldPermuteOfBinops - to fold a permute (single source shuffle) through a binary op that is being fed by other shuffles. WIP - still need to add additional test coverage. Fixes #94546
1 parent bc999ee commit 06fbd77

File tree

5 files changed

+111
-20
lines changed

5 files changed

+111
-20
lines changed

‎llvm/lib/Transforms/Vectorize/VectorCombine.cpp‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ class VectorCombine {
112112
bool foldExtractedCmps(Instruction &I);
113113
bool foldSingleElementStore(Instruction &I);
114114
bool scalarizeLoadExtract(Instruction &I);
115+
bool foldPermuteOfBinops(Instruction &I);
115116
bool foldShuffleOfBinops(Instruction &I);
116117
bool foldShuffleOfCastops(Instruction &I);
117118
bool foldShuffleOfShuffles(Instruction &I);
@@ -1400,6 +1401,93 @@ bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
14001401
return true;
14011402
}
14021403

1404+
/// Try to convert "shuffle (binop (shuffle, shuffle)), undef"
1405+
/// --> "binop (shuffle), (shuffle)".
1406+
bool VectorCombine::foldPermuteOfBinops(Instruction &I) {
1407+
BinaryOperator *BinOp;
1408+
ArrayRef<int> OuterMask;
1409+
if (!match(&I,
1410+
m_Shuffle(m_OneUse(m_BinOp(BinOp)), m_Undef(), m_Mask(OuterMask))))
1411+
return false;
1412+
1413+
// Don't introduce poison into div/rem.
1414+
if (llvm::is_contained(OuterMask, PoisonMaskElem) && BinOp->isIntDivRem())
1415+
return false;
1416+
1417+
Value *Op00, *Op01;
1418+
ArrayRef<int> Mask0;
1419+
if (!match(BinOp->getOperand(0),
1420+
m_OneUse(m_Shuffle(m_Value(Op00), m_Value(Op01), m_Mask(Mask0)))))
1421+
return false;
1422+
1423+
Value *Op10, *Op11;
1424+
ArrayRef<int> Mask1;
1425+
if (!match(BinOp->getOperand(1),
1426+
m_OneUse(m_Shuffle(m_Value(Op10), m_Value(Op11), m_Mask(Mask1)))))
1427+
return false;
1428+
1429+
Instruction::BinaryOps Opcode = BinOp->getOpcode();
1430+
auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
1431+
auto *BinOpTy = dyn_cast<FixedVectorType>(BinOp->getType());
1432+
auto *Op0Ty = dyn_cast<FixedVectorType>(Op00->getType());
1433+
auto *Op1Ty = dyn_cast<FixedVectorType>(Op10->getType());
1434+
if (!ShuffleDstTy || !BinOpTy || !Op0Ty || !Op1Ty)
1435+
return false;
1436+
1437+
unsigned NumSrcElts = BinOpTy->getNumElements();
1438+
1439+
// Don't accept shuffles that reference the second (undef/poison) operand.
1440+
if (any_of(OuterMask, [NumSrcElts](int M) { return M >= (int)NumSrcElts; }))
1441+
return false;
1442+
1443+
// Merge outer / inner shuffles.
1444+
SmallVector<int> NewMask0, NewMask1;
1445+
for (int M : OuterMask) {
1446+
NewMask0.push_back(M >= 0 ? Mask0[M] : -1);
1447+
NewMask1.push_back(M >= 0 ? Mask1[M] : -1);
1448+
}
1449+
1450+
// Try to merge shuffles across the binop if the new shuffles are not costly.
1451+
TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput;
1452+
1453+
InstructionCost OldCost =
1454+
TTI.getArithmeticInstrCost(Opcode, BinOpTy, CostKind) +
1455+
TTI.getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc, BinOpTy,
1456+
OuterMask, CostKind, 0, nullptr, {BinOp}, &I) +
1457+
TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, Op0Ty, Mask0,
1458+
CostKind, 0, nullptr, {Op00, Op01},
1459+
cast<Instruction>(BinOp->getOperand(0))) +
1460+
TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, Op1Ty, Mask1,
1461+
CostKind, 0, nullptr, {Op10, Op11},
1462+
cast<Instruction>(BinOp->getOperand(1)));
1463+
1464+
InstructionCost NewCost =
1465+
TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, Op0Ty, NewMask0,
1466+
CostKind, 0, nullptr, {Op00, Op01}) +
1467+
TTI.getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc, Op1Ty, NewMask1,
1468+
CostKind, 0, nullptr, {Op10, Op11}) +
1469+
TTI.getArithmeticInstrCost(Opcode, ShuffleDstTy, CostKind);
1470+
1471+
LLVM_DEBUG(dbgs() << "Found a shuffle feeding a shuffled binop: " << I
1472+
<< "\n OldCost: " << OldCost << " vs NewCost: " << NewCost
1473+
<< "\n");
1474+
if (NewCost >= OldCost)
1475+
return false;
1476+
1477+
Value *Shuf0 = Builder.CreateShuffleVector(Op00, Op01, NewMask0);
1478+
Value *Shuf1 = Builder.CreateShuffleVector(Op10, Op11, NewMask1);
1479+
Value *NewBO = Builder.CreateBinOp(Opcode, Shuf0, Shuf1);
1480+
1481+
// Intersect flags from the old binops.
1482+
if (auto *NewInst = dyn_cast<Instruction>(NewBO))
1483+
NewInst->copyIRFlags(BinOp);
1484+
1485+
Worklist.pushValue(Shuf0);
1486+
Worklist.pushValue(Shuf1);
1487+
replaceValue(I, *NewBO);
1488+
return true;
1489+
}
1490+
14031491
/// Try to convert "shuffle (binop), (binop)" into "binop (shuffle), (shuffle)".
14041492
bool VectorCombine::foldShuffleOfBinops(Instruction &I) {
14051493
BinaryOperator *B0, *B1;
@@ -2736,6 +2824,7 @@ bool VectorCombine::run() {
27362824
MadeChange |= foldInsExtFNeg(I);
27372825
break;
27382826
case Instruction::ShuffleVector:
2827+
MadeChange |= foldPermuteOfBinops(I);
27392828
MadeChange |= foldShuffleOfBinops(I);
27402829
MadeChange |= foldShuffleOfCastops(I);
27412830
MadeChange |= foldShuffleOfShuffles(I);

‎llvm/test/Transforms/PhaseOrdering/X86/horiz-math-inseltpoison.ll‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -108,11 +108,10 @@ define <8 x float> @hadd_reverse_v8f32(<8 x float> %a, <8 x float> %b) #0 {
108108

109109
define <8 x float> @reverse_hadd_v8f32(<8 x float> %a, <8 x float> %b) #0 {
110110
; CHECK-LABEL: @reverse_hadd_v8f32(
111-
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <8 x float> [[A:%.*]], <8 x float> [[B:%.*]], <8 x i32> <i32 0, i32 2, i32 8, i32 10, i32 4, i32 6, i32 12, i32 14>
112-
; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <8 x float> [[A]], <8 x float> [[B]], <8 x i32> <i32 1, i32 3, i32 9, i32 11, i32 5, i32 7, i32 13, i32 15>
111+
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <8 x float> [[A:%.*]], <8 x float> [[B:%.*]], <8 x i32> <i32 14, i32 12, i32 6, i32 4, i32 10, i32 8, i32 2, i32 0>
112+
; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <8 x float> [[A]], <8 x float> [[B]], <8 x i32> <i32 15, i32 13, i32 7, i32 5, i32 11, i32 9, i32 3, i32 1>
113113
; CHECK-NEXT: [[TMP3:%.*]] = fadd <8 x float> [[TMP1]], [[TMP2]]
114-
; CHECK-NEXT: [[SHUFFLE:%.*]] = shufflevector <8 x float> [[TMP3]], <8 x float> poison, <8 x i32> <i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
115-
; CHECK-NEXT: ret <8 x float> [[SHUFFLE]]
114+
; CHECK-NEXT: ret <8 x float> [[TMP3]]
116115
;
117116
%vecext = extractelement <8 x float> %a, i32 0
118117
%vecext1 = extractelement <8 x float> %a, i32 1

‎llvm/test/Transforms/PhaseOrdering/X86/horiz-math.ll‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -108,11 +108,10 @@ define <8 x float> @hadd_reverse_v8f32(<8 x float> %a, <8 x float> %b) #0 {
108108

109109
define <8 x float> @reverse_hadd_v8f32(<8 x float> %a, <8 x float> %b) #0 {
110110
; CHECK-LABEL: @reverse_hadd_v8f32(
111-
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <8 x float> [[A:%.*]], <8 x float> [[B:%.*]], <8 x i32> <i32 0, i32 2, i32 8, i32 10, i32 4, i32 6, i32 12, i32 14>
112-
; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <8 x float> [[A]], <8 x float> [[B]], <8 x i32> <i32 1, i32 3, i32 9, i32 11, i32 5, i32 7, i32 13, i32 15>
111+
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <8 x float> [[A:%.*]], <8 x float> [[B:%.*]], <8 x i32> <i32 14, i32 12, i32 6, i32 4, i32 10, i32 8, i32 2, i32 0>
112+
; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <8 x float> [[A]], <8 x float> [[B]], <8 x i32> <i32 15, i32 13, i32 7, i32 5, i32 11, i32 9, i32 3, i32 1>
113113
; CHECK-NEXT: [[TMP3:%.*]] = fadd <8 x float> [[TMP1]], [[TMP2]]
114-
; CHECK-NEXT: [[SHUFFLE:%.*]] = shufflevector <8 x float> [[TMP3]], <8 x float> poison, <8 x i32> <i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
115-
; CHECK-NEXT: ret <8 x float> [[SHUFFLE]]
114+
; CHECK-NEXT: ret <8 x float> [[TMP3]]
116115
;
117116
%vecext = extractelement <8 x float> %a, i32 0
118117
%vecext1 = extractelement <8 x float> %a, i32 1

‎llvm/test/Transforms/PhaseOrdering/X86/pr50392.ll‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,9 @@ define <4 x double> @PR50392(<4 x double> %a, <4 x double> %b) {
3232
; AVX1-NEXT: ret <4 x double> [[SHUFFLE]]
3333
;
3434
; AVX2-LABEL: @PR50392(
35-
; AVX2-NEXT: [[TMP1:%.*]] = shufflevector <4 x double> [[A:%.*]], <4 x double> [[B:%.*]], <2 x i32> <i32 0, i32 4>
36-
; AVX2-NEXT: [[TMP2:%.*]] = shufflevector <4 x double> [[A]], <4 x double> [[B]], <2 x i32> <i32 1, i32 5>
37-
; AVX2-NEXT: [[TMP3:%.*]] = fadd <2 x double> [[TMP1]], [[TMP2]]
38-
; AVX2-NEXT: [[TMP4:%.*]] = shufflevector <2 x double> [[TMP3]], <2 x double> poison, <4 x i32> <i32 0, i32 poison, i32 1, i32 poison>
35+
; AVX2-NEXT: [[TMP1:%.*]] = shufflevector <4 x double> [[A:%.*]], <4 x double> [[B:%.*]], <4 x i32> <i32 0, i32 poison, i32 4, i32 poison>
36+
; AVX2-NEXT: [[TMP2:%.*]] = shufflevector <4 x double> [[A]], <4 x double> [[B]], <4 x i32> <i32 1, i32 poison, i32 5, i32 poison>
37+
; AVX2-NEXT: [[TMP4:%.*]] = fadd <4 x double> [[TMP1]], [[TMP2]]
3938
; AVX2-NEXT: [[SHIFT:%.*]] = shufflevector <4 x double> [[B]], <4 x double> poison, <4 x i32> <i32 poison, i32 poison, i32 3, i32 poison>
4039
; AVX2-NEXT: [[TMP5:%.*]] = fadd <4 x double> [[B]], [[SHIFT]]
4140
; AVX2-NEXT: [[SHUFFLE:%.*]] = shufflevector <4 x double> [[TMP4]], <4 x double> [[TMP5]], <4 x i32> <i32 0, i32 poison, i32 2, i32 6>

‎llvm/test/Transforms/PhaseOrdering/X86/pr94546.ll‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,18 @@ define <4 x double> @PR94546(<4 x double> %a, <4 x double> %b) {
1616
; SSE-NEXT: [[TMP4:%.*]] = shufflevector <2 x double> [[TMP3]], <2 x double> poison, <4 x i32> <i32 0, i32 poison, i32 poison, i32 1>
1717
; SSE-NEXT: ret <4 x double> [[TMP4]]
1818
;
19-
; AVX-LABEL: @PR94546(
20-
; AVX-NEXT: [[TMP1:%.*]] = shufflevector <4 x double> [[A:%.*]], <4 x double> [[B:%.*]], <2 x i32> <i32 0, i32 6>
21-
; AVX-NEXT: [[TMP2:%.*]] = shufflevector <4 x double> [[A]], <4 x double> [[B]], <2 x i32> <i32 1, i32 7>
22-
; AVX-NEXT: [[TMP3:%.*]] = fadd <2 x double> [[TMP1]], [[TMP2]]
23-
; AVX-NEXT: [[TMP4:%.*]] = shufflevector <2 x double> [[TMP3]], <2 x double> poison, <4 x i32> <i32 0, i32 poison, i32 poison, i32 1>
24-
; AVX-NEXT: ret <4 x double> [[TMP4]]
19+
; AVX1-LABEL: @PR94546(
20+
; AVX1-NEXT: [[TMP1:%.*]] = shufflevector <4 x double> [[A:%.*]], <4 x double> [[B:%.*]], <2 x i32> <i32 0, i32 6>
21+
; AVX1-NEXT: [[TMP2:%.*]] = shufflevector <4 x double> [[A]], <4 x double> [[B]], <2 x i32> <i32 1, i32 7>
22+
; AVX1-NEXT: [[TMP3:%.*]] = fadd <2 x double> [[TMP1]], [[TMP2]]
23+
; AVX1-NEXT: [[TMP4:%.*]] = shufflevector <2 x double> [[TMP3]], <2 x double> poison, <4 x i32> <i32 0, i32 poison, i32 poison, i32 1>
24+
; AVX1-NEXT: ret <4 x double> [[TMP4]]
25+
;
26+
; AVX2-LABEL: @PR94546(
27+
; AVX2-NEXT: [[TMP1:%.*]] = shufflevector <4 x double> [[A:%.*]], <4 x double> [[B:%.*]], <4 x i32> <i32 0, i32 poison, i32 poison, i32 6>
28+
; AVX2-NEXT: [[TMP2:%.*]] = shufflevector <4 x double> [[A]], <4 x double> [[B]], <4 x i32> <i32 1, i32 poison, i32 poison, i32 7>
29+
; AVX2-NEXT: [[TMP3:%.*]] = fadd <4 x double> [[TMP1]], [[TMP2]]
30+
; AVX2-NEXT: ret <4 x double> [[TMP3]]
2531
;
2632
%vecext = extractelement <4 x double> %a, i32 0
2733
%vecext1 = extractelement <4 x double> %a, i32 1
@@ -43,5 +49,4 @@ define <4 x double> @PR94546(<4 x double> %a, <4 x double> %b) {
4349
ret <4 x double> %shuffle
4450
}
4551
;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
46-
; AVX1: {{.*}}
47-
; AVX2: {{.*}}
52+
; AVX: {{.*}}

0 commit comments

Comments
 (0)