Skip to content

Commit a4d1026

Browse files
davemgreendavemgreen
authored
[VectorCombine] Add foldShuffleToIdentity (#88693)
This patch adds a basic version of a combine that attempts to remove shuffles that when combined simplify away to an identity shuffle. For example: %ab = shufflevector <8 x half> %a, <8 x half> poison, <4 x i32> <i32 3, i32 2, i32 1, i32 0> %at = shufflevector <8 x half> %a, <8 x half> poison, <4 x i32> <i32 7, i32 6, i32 5, i32 4> %abt = fneg <4 x half> %at %abb = fneg <4 x half> %ab %r = shufflevector <4 x half> %abt, <4 x half> %abb, <8 x i32> <i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0> By looking through the shuffles and fneg, it can be simplified to: %r = fneg <8 x half> %a The code tracks each lane starting from the original shuffle, keeping a track of a vector of {src, idx}. As we propagate up through the instructions we will either look through intermediate instructions (binops and unops) or see a collections of lanes that all have the same src and incrementing idx (an identity). We can also see a single value with identical lanes, which we can treat like a splat. Only the basic version is added here, handling identities, splats, binops and unops. In follow-up patches other instructions can be added such as constants, intrinsics, cmp/sel and zext/sext/trunc.
1 parent 46c2d93 commit a4d1026

File tree

4 files changed

+212
-146
lines changed

4 files changed

+212
-146
lines changed

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

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ class VectorCombine {
114114
bool foldShuffleOfBinops(Instruction &I);
115115
bool foldShuffleOfCastops(Instruction &I);
116116
bool foldShuffleOfShuffles(Instruction &I);
117+
bool foldShuffleToIdentity(Instruction &I);
117118
bool foldShuffleFromReductions(Instruction &I);
118119
bool foldTruncFromReductions(Instruction &I);
119120
bool foldSelectShuffle(Instruction &I, bool FromReduction = false);
@@ -1667,6 +1668,151 @@ bool VectorCombine::foldShuffleOfShuffles(Instruction &I) {
16671668
return true;
16681669
}
16691670

1671+
// Starting from a shuffle, look up through operands tracking the shuffled index
1672+
// of each lane. If we can simplify away the shuffles to identities then
1673+
// do so.
1674+
bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
1675+
auto *Ty = dyn_cast<FixedVectorType>(I.getType());
1676+
if (!Ty || !isa<Instruction>(I.getOperand(0)) ||
1677+
!isa<Instruction>(I.getOperand(1)))
1678+
return false;
1679+
1680+
using InstLane = std::pair<Value *, int>;
1681+
1682+
auto LookThroughShuffles = [](Value *V, int Lane) -> InstLane {
1683+
while (auto *SV = dyn_cast<ShuffleVectorInst>(V)) {
1684+
unsigned NumElts =
1685+
cast<FixedVectorType>(SV->getOperand(0)->getType())->getNumElements();
1686+
int M = SV->getMaskValue(Lane);
1687+
if (M < 0)
1688+
return {nullptr, PoisonMaskElem};
1689+
else if (M < (int)NumElts) {
1690+
V = SV->getOperand(0);
1691+
Lane = M;
1692+
} else {
1693+
V = SV->getOperand(1);
1694+
Lane = M - NumElts;
1695+
}
1696+
}
1697+
return InstLane{V, Lane};
1698+
};
1699+
1700+
auto GenerateInstLaneVectorFromOperand =
1701+
[&LookThroughShuffles](ArrayRef<InstLane> Item, int Op) {
1702+
SmallVector<InstLane> NItem;
1703+
for (InstLane V : Item) {
1704+
NItem.emplace_back(
1705+
!V.first
1706+
? InstLane{nullptr, PoisonMaskElem}
1707+
: LookThroughShuffles(
1708+
cast<Instruction>(V.first)->getOperand(Op), V.second));
1709+
}
1710+
return NItem;
1711+
};
1712+
1713+
SmallVector<InstLane> Start(Ty->getNumElements());
1714+
for (unsigned M = 0, E = Ty->getNumElements(); M < E; ++M)
1715+
Start[M] = LookThroughShuffles(&I, M);
1716+
1717+
SmallVector<SmallVector<InstLane>> Worklist;
1718+
Worklist.push_back(Start);
1719+
SmallPtrSet<Value *, 4> IdentityLeafs, SplatLeafs;
1720+
unsigned NumVisited = 0;
1721+
1722+
while (!Worklist.empty()) {
1723+
SmallVector<InstLane> Item = Worklist.pop_back_val();
1724+
if (++NumVisited > MaxInstrsToScan)
1725+
return false;
1726+
1727+
// If we found an undef first lane then bail out to keep things simple.
1728+
if (!Item[0].first)
1729+
return false;
1730+
1731+
// Look for an identity value.
1732+
if (Item[0].second == 0 && Item[0].first->getType() == Ty &&
1733+
all_of(drop_begin(enumerate(Item)), [&](const auto &E) {
1734+
return !E.value().first || (E.value().first == Item[0].first &&
1735+
E.value().second == (int)E.index());
1736+
})) {
1737+
IdentityLeafs.insert(Item[0].first);
1738+
continue;
1739+
}
1740+
// Look for a splat value.
1741+
if (all_of(drop_begin(Item), [&](InstLane &IL) {
1742+
return !IL.first ||
1743+
(IL.first == Item[0].first && IL.second == Item[0].second);
1744+
})) {
1745+
SplatLeafs.insert(Item[0].first);
1746+
continue;
1747+
}
1748+
1749+
// We need each element to be the same type of value, and check that each
1750+
// element has a single use.
1751+
if (!all_of(drop_begin(Item), [&](InstLane IL) {
1752+
if (!IL.first)
1753+
return true;
1754+
if (auto *I = dyn_cast<Instruction>(IL.first); I && !I->hasOneUse())
1755+
return false;
1756+
if (IL.first->getValueID() != Item[0].first->getValueID())
1757+
return false;
1758+
auto *II = dyn_cast<IntrinsicInst>(IL.first);
1759+
return !II ||
1760+
II->getIntrinsicID() ==
1761+
cast<IntrinsicInst>(Item[0].first)->getIntrinsicID();
1762+
}))
1763+
return false;
1764+
1765+
// Check the operator is one that we support. We exclude div/rem in case
1766+
// they hit UB from poison lanes.
1767+
if (isa<BinaryOperator>(Item[0].first) &&
1768+
!cast<BinaryOperator>(Item[0].first)->isIntDivRem()) {
1769+
Worklist.push_back(GenerateInstLaneVectorFromOperand(Item, 0));
1770+
Worklist.push_back(GenerateInstLaneVectorFromOperand(Item, 1));
1771+
} else if (isa<UnaryOperator>(Item[0].first)) {
1772+
Worklist.push_back(GenerateInstLaneVectorFromOperand(Item, 0));
1773+
} else {
1774+
return false;
1775+
}
1776+
}
1777+
1778+
// If we got this far, we know the shuffles are superfluous and can be
1779+
// removed. Scan through again and generate the new tree of instructions.
1780+
std::function<Value *(ArrayRef<InstLane>)> Generate =
1781+
[&](ArrayRef<InstLane> Item) -> Value * {
1782+
if (IdentityLeafs.contains(Item[0].first) &&
1783+
all_of(drop_begin(enumerate(Item)), [&](const auto &E) {
1784+
return !E.value().first || (E.value().first == Item[0].first &&
1785+
E.value().second == (int)E.index());
1786+
})) {
1787+
return Item[0].first;
1788+
}
1789+
if (SplatLeafs.contains(Item[0].first)) {
1790+
if (auto ILI = dyn_cast<Instruction>(Item[0].first))
1791+
Builder.SetInsertPoint(*ILI->getInsertionPointAfterDef());
1792+
else if (isa<Argument>(Item[0].first))
1793+
Builder.SetInsertPointPastAllocas(I.getParent()->getParent());
1794+
SmallVector<int, 16> Mask(Ty->getNumElements(), Item[0].second);
1795+
return Builder.CreateShuffleVector(Item[0].first, Mask);
1796+
}
1797+
1798+
auto *I = cast<Instruction>(Item[0].first);
1799+
SmallVector<Value *> Ops(I->getNumOperands());
1800+
for (unsigned Idx = 0, E = I->getNumOperands(); Idx < E; Idx++)
1801+
Ops[Idx] = Generate(GenerateInstLaneVectorFromOperand(Item, Idx));
1802+
Builder.SetInsertPoint(I);
1803+
if (auto BI = dyn_cast<BinaryOperator>(I))
1804+
return Builder.CreateBinOp((Instruction::BinaryOps)BI->getOpcode(),
1805+
Ops[0], Ops[1]);
1806+
assert(isa<UnaryInstruction>(I) &&
1807+
"Unexpected instruction type in Generate");
1808+
return Builder.CreateUnOp((Instruction::UnaryOps)I->getOpcode(), Ops[0]);
1809+
};
1810+
1811+
Value *V = Generate(Start);
1812+
replaceValue(I, *V);
1813+
return true;
1814+
}
1815+
16701816
/// Given a commutative reduction, the order of the input lanes does not alter
16711817
/// the results. We can use this to remove certain shuffles feeding the
16721818
/// reduction, removing the need to shuffle at all.
@@ -2224,6 +2370,7 @@ bool VectorCombine::run() {
22242370
MadeChange |= foldShuffleOfCastops(I);
22252371
MadeChange |= foldShuffleOfShuffles(I);
22262372
MadeChange |= foldSelectShuffle(I);
2373+
MadeChange |= foldShuffleToIdentity(I);
22272374
break;
22282375
case Instruction::BitCast:
22292376
MadeChange |= foldBitcastShuffle(I);

‎llvm/test/Transforms/PhaseOrdering/AArch64/interleavevectorization.ll

Lines changed: 10 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -22,19 +22,13 @@ define void @add4(ptr noalias noundef %x, ptr noalias noundef %y, i32 noundef %n
2222
; CHECK-NEXT: [[WIDE_VEC:%.*]] = load <32 x i16>, ptr [[TMP0]], align 2
2323
; CHECK-NEXT: [[TMP1:%.*]] = getelementptr inbounds i16, ptr [[X]], i64 [[OFFSET_IDX]]
2424
; CHECK-NEXT: [[WIDE_VEC24:%.*]] = load <32 x i16>, ptr [[TMP1]], align 2
25-
; CHECK-NEXT: [[TMP2:%.*]] = add <32 x i16> [[WIDE_VEC24]], [[WIDE_VEC]]
26-
; CHECK-NEXT: [[TMP3:%.*]] = add <32 x i16> [[WIDE_VEC24]], [[WIDE_VEC]]
27-
; CHECK-NEXT: [[TMP4:%.*]] = add <32 x i16> [[WIDE_VEC24]], [[WIDE_VEC]]
28-
; CHECK-NEXT: [[TMP5:%.*]] = or disjoint i64 [[OFFSET_IDX]], 3
29-
; CHECK-NEXT: [[TMP6:%.*]] = add <32 x i16> [[WIDE_VEC24]], [[WIDE_VEC]]
30-
; CHECK-NEXT: [[GEP:%.*]] = getelementptr i16, ptr [[INVARIANT_GEP]], i64 [[TMP5]]
31-
; CHECK-NEXT: [[TMP7:%.*]] = shufflevector <32 x i16> [[TMP2]], <32 x i16> [[TMP3]], <16 x i32> <i32 0, i32 4, i32 8, i32 12, i32 16, i32 20, i32 24, i32 28, i32 33, i32 37, i32 41, i32 45, i32 49, i32 53, i32 57, i32 61>
32-
; CHECK-NEXT: [[TMP8:%.*]] = shufflevector <32 x i16> [[TMP4]], <32 x i16> [[TMP6]], <16 x i32> <i32 2, i32 6, i32 10, i32 14, i32 18, i32 22, i32 26, i32 30, i32 35, i32 39, i32 43, i32 47, i32 51, i32 55, i32 59, i32 63>
33-
; CHECK-NEXT: [[INTERLEAVED_VEC:%.*]] = shufflevector <16 x i16> [[TMP7]], <16 x i16> [[TMP8]], <32 x i32> <i32 0, i32 8, i32 16, i32 24, i32 1, i32 9, i32 17, i32 25, i32 2, i32 10, i32 18, i32 26, i32 3, i32 11, i32 19, i32 27, i32 4, i32 12, i32 20, i32 28, i32 5, i32 13, i32 21, i32 29, i32 6, i32 14, i32 22, i32 30, i32 7, i32 15, i32 23, i32 31>
25+
; CHECK-NEXT: [[INTERLEAVED_VEC:%.*]] = add <32 x i16> [[WIDE_VEC24]], [[WIDE_VEC]]
26+
; CHECK-NEXT: [[TMP2:%.*]] = or disjoint i64 [[OFFSET_IDX]], 3
27+
; CHECK-NEXT: [[GEP:%.*]] = getelementptr i16, ptr [[INVARIANT_GEP]], i64 [[TMP2]]
3428
; CHECK-NEXT: store <32 x i16> [[INTERLEAVED_VEC]], ptr [[GEP]], align 2
3529
; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 8
36-
; CHECK-NEXT: [[TMP9:%.*]] = icmp eq i64 [[INDEX_NEXT]], 256
37-
; CHECK-NEXT: br i1 [[TMP9]], label [[FOR_END:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
30+
; CHECK-NEXT: [[TMP3:%.*]] = icmp eq i64 [[INDEX_NEXT]], 256
31+
; CHECK-NEXT: br i1 [[TMP3]], label [[FOR_END:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
3832
; CHECK: for.end:
3933
; CHECK-NEXT: ret void
4034
;
@@ -412,22 +406,13 @@ define void @addmul(ptr noalias noundef %x, ptr noundef %y, ptr noundef %z, i32
412406
; CHECK-NEXT: [[TMP2:%.*]] = mul <32 x i16> [[WIDE_VEC31]], [[WIDE_VEC]]
413407
; CHECK-NEXT: [[TMP3:%.*]] = getelementptr inbounds i16, ptr [[X]], i64 [[OFFSET_IDX]]
414408
; CHECK-NEXT: [[WIDE_VEC36:%.*]] = load <32 x i16>, ptr [[TMP3]], align 2
415-
; CHECK-NEXT: [[TMP4:%.*]] = add <32 x i16> [[TMP2]], [[WIDE_VEC36]]
416-
; CHECK-NEXT: [[TMP5:%.*]] = mul <32 x i16> [[WIDE_VEC31]], [[WIDE_VEC]]
417-
; CHECK-NEXT: [[TMP6:%.*]] = add <32 x i16> [[TMP5]], [[WIDE_VEC36]]
418-
; CHECK-NEXT: [[TMP7:%.*]] = mul <32 x i16> [[WIDE_VEC31]], [[WIDE_VEC]]
419-
; CHECK-NEXT: [[TMP8:%.*]] = add <32 x i16> [[TMP7]], [[WIDE_VEC36]]
420-
; CHECK-NEXT: [[TMP9:%.*]] = or disjoint i64 [[OFFSET_IDX]], 3
421-
; CHECK-NEXT: [[TMP10:%.*]] = mul <32 x i16> [[WIDE_VEC31]], [[WIDE_VEC]]
422-
; CHECK-NEXT: [[TMP11:%.*]] = add <32 x i16> [[TMP10]], [[WIDE_VEC36]]
423-
; CHECK-NEXT: [[GEP:%.*]] = getelementptr i16, ptr [[INVARIANT_GEP]], i64 [[TMP9]]
424-
; CHECK-NEXT: [[TMP12:%.*]] = shufflevector <32 x i16> [[TMP4]], <32 x i16> [[TMP6]], <16 x i32> <i32 0, i32 4, i32 8, i32 12, i32 16, i32 20, i32 24, i32 28, i32 33, i32 37, i32 41, i32 45, i32 49, i32 53, i32 57, i32 61>
425-
; CHECK-NEXT: [[TMP13:%.*]] = shufflevector <32 x i16> [[TMP8]], <32 x i16> [[TMP11]], <16 x i32> <i32 2, i32 6, i32 10, i32 14, i32 18, i32 22, i32 26, i32 30, i32 35, i32 39, i32 43, i32 47, i32 51, i32 55, i32 59, i32 63>
426-
; CHECK-NEXT: [[INTERLEAVED_VEC:%.*]] = shufflevector <16 x i16> [[TMP12]], <16 x i16> [[TMP13]], <32 x i32> <i32 0, i32 8, i32 16, i32 24, i32 1, i32 9, i32 17, i32 25, i32 2, i32 10, i32 18, i32 26, i32 3, i32 11, i32 19, i32 27, i32 4, i32 12, i32 20, i32 28, i32 5, i32 13, i32 21, i32 29, i32 6, i32 14, i32 22, i32 30, i32 7, i32 15, i32 23, i32 31>
409+
; CHECK-NEXT: [[INTERLEAVED_VEC:%.*]] = add <32 x i16> [[TMP2]], [[WIDE_VEC36]]
410+
; CHECK-NEXT: [[TMP4:%.*]] = or disjoint i64 [[OFFSET_IDX]], 3
411+
; CHECK-NEXT: [[GEP:%.*]] = getelementptr i16, ptr [[INVARIANT_GEP]], i64 [[TMP4]]
427412
; CHECK-NEXT: store <32 x i16> [[INTERLEAVED_VEC]], ptr [[GEP]], align 2
428413
; CHECK-NEXT: [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 8
429-
; CHECK-NEXT: [[TMP14:%.*]] = icmp eq i64 [[INDEX_NEXT]], 256
430-
; CHECK-NEXT: br i1 [[TMP14]], label [[FOR_END:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP5:![0-9]+]]
414+
; CHECK-NEXT: [[TMP5:%.*]] = icmp eq i64 [[INDEX_NEXT]], 256
415+
; CHECK-NEXT: br i1 [[TMP5]], label [[FOR_END:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP5:![0-9]+]]
431416
; CHECK: for.end:
432417
; CHECK-NEXT: ret void
433418
;

0 commit comments

Comments
 (0)