From a1bdb016564d3c0867e83f3bb7494bfe8128f4b6 Mon Sep 17 00:00:00 2001 From: David Green Date: Mon, 17 Jun 2024 15:25:33 +0100 Subject: [PATCH] [VectorCombine] Change shuffleToIdentity to use Use. NFC When looking up through shuffles, a Value can be multiple different leaf types (for example an identity from one position, a splat from another). We currently detect this by recalculating which type of leaf it is when generating, but as more types of leafs are added (#94954) this doesn't scale very well. This patch switches it to use Use, not Value, to more accurately detect which type of leaf each Use should have. --- .../Transforms/Vectorize/VectorCombine.cpp | 108 +++++++++--------- .../AArch64/shuffletoidentity.ll | 13 +++ 2 files changed, 65 insertions(+), 56 deletions(-) diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp index e608c7fb604681..5b9fe1c9c9854d 100644 --- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp +++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp @@ -1669,33 +1669,34 @@ bool VectorCombine::foldShuffleOfShuffles(Instruction &I) { return true; } -using InstLane = std::pair; +using InstLane = std::pair; -static InstLane lookThroughShuffles(Value *V, int Lane) { - while (auto *SV = dyn_cast(V)) { +static InstLane lookThroughShuffles(Use *U, int Lane) { + while (auto *SV = dyn_cast(U->get())) { unsigned NumElts = cast(SV->getOperand(0)->getType())->getNumElements(); int M = SV->getMaskValue(Lane); if (M < 0) return {nullptr, PoisonMaskElem}; if (static_cast(M) < NumElts) { - V = SV->getOperand(0); + U = &SV->getOperandUse(0); Lane = M; } else { - V = SV->getOperand(1); + U = &SV->getOperandUse(1); Lane = M - NumElts; } } - return InstLane{V, Lane}; + return InstLane{U, Lane}; } static SmallVector generateInstLaneVectorFromOperand(ArrayRef Item, int Op) { SmallVector NItem; for (InstLane IL : Item) { - auto [V, Lane] = IL; + auto [U, Lane] = IL; InstLane OpLane = - V ? lookThroughShuffles(cast(V)->getOperand(Op), Lane) + U ? lookThroughShuffles(&cast(U->get())->getOperandUse(Op), + Lane) : InstLane{nullptr, PoisonMaskElem}; NItem.emplace_back(OpLane); } @@ -1703,29 +1704,24 @@ generateInstLaneVectorFromOperand(ArrayRef Item, int Op) { } static Value *generateNewInstTree(ArrayRef Item, FixedVectorType *Ty, - const SmallPtrSet &IdentityLeafs, - const SmallPtrSet &SplatLeafs, + const SmallPtrSet &IdentityLeafs, + const SmallPtrSet &SplatLeafs, IRBuilder<> &Builder) { - auto [FrontV, FrontLane] = Item.front(); - - if (IdentityLeafs.contains(FrontV) && - all_of(drop_begin(enumerate(Item)), [Item](const auto &E) { - Value *FrontV = Item.front().first; - auto [V, Lane] = E.value(); - return !V || (V == FrontV && Lane == (int)E.index()); - })) { - return FrontV; + auto [FrontU, FrontLane] = Item.front(); + + if (IdentityLeafs.contains(FrontU)) { + return FrontU->get(); } - if (SplatLeafs.contains(FrontV)) { - if (auto *ILI = dyn_cast(FrontV)) + if (SplatLeafs.contains(FrontU)) { + if (auto *ILI = dyn_cast(FrontU)) Builder.SetInsertPoint(*ILI->getInsertionPointAfterDef()); - else if (auto *Arg = dyn_cast(FrontV)) + else if (auto *Arg = dyn_cast(FrontU)) Builder.SetInsertPointPastAllocas(Arg->getParent()); SmallVector Mask(Ty->getNumElements(), FrontLane); - return Builder.CreateShuffleVector(FrontV, Mask); + return Builder.CreateShuffleVector(FrontU->get(), Mask); } - auto *I = cast(FrontV); + auto *I = cast(FrontU->get()); auto *II = dyn_cast(I); unsigned NumOps = I->getNumOperands() - (II ? 1 : 0); SmallVector Ops(NumOps); @@ -1741,7 +1737,7 @@ static Value *generateNewInstTree(ArrayRef Item, FixedVectorType *Ty, SmallVector ValueList; for (const auto &Lane : Item) if (Lane.first) - ValueList.push_back(Lane.first); + ValueList.push_back(Lane.first->get()); Builder.SetInsertPoint(I); Type *DstTy = @@ -1785,16 +1781,16 @@ static Value *generateNewInstTree(ArrayRef Item, FixedVectorType *Ty, // do so. bool VectorCombine::foldShuffleToIdentity(Instruction &I) { auto *Ty = dyn_cast(I.getType()); - if (!Ty) + if (!Ty || I.use_empty()) return false; SmallVector Start(Ty->getNumElements()); for (unsigned M = 0, E = Ty->getNumElements(); M < E; ++M) - Start[M] = lookThroughShuffles(&I, M); + Start[M] = lookThroughShuffles(&*I.use_begin(), M); SmallVector> Worklist; Worklist.push_back(Start); - SmallPtrSet IdentityLeafs, SplatLeafs; + SmallPtrSet IdentityLeafs, SplatLeafs; unsigned NumVisited = 0; while (!Worklist.empty()) { @@ -1802,52 +1798,52 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) { return false; SmallVector Item = Worklist.pop_back_val(); - auto [FrontV, FrontLane] = Item.front(); + auto [FrontU, FrontLane] = Item.front(); // If we found an undef first lane then bail out to keep things simple. - if (!FrontV) + if (!FrontU) return false; // Look for an identity value. - if (!FrontLane && - cast(FrontV->getType())->getNumElements() == + if (FrontLane == 0 && + cast(FrontU->get()->getType())->getNumElements() == Ty->getNumElements() && all_of(drop_begin(enumerate(Item)), [Item](const auto &E) { - Value *FrontV = Item.front().first; - return !E.value().first || (E.value().first == FrontV && + Value *FrontV = Item.front().first->get(); + return !E.value().first || (E.value().first->get() == FrontV && E.value().second == (int)E.index()); })) { - IdentityLeafs.insert(FrontV); + IdentityLeafs.insert(FrontU); continue; } // Look for constants, for the moment only supporting constant splats. - if (auto *C = dyn_cast(FrontV); + if (auto *C = dyn_cast(FrontU); C && C->getSplatValue() && all_of(drop_begin(Item), [Item](InstLane &IL) { - Value *FrontV = Item.front().first; - Value *V = IL.first; - return !V || V == FrontV; + Value *FrontV = Item.front().first->get(); + Use *U = IL.first; + return !U || U->get() == FrontV; })) { - SplatLeafs.insert(FrontV); + SplatLeafs.insert(FrontU); continue; } // Look for a splat value. if (all_of(drop_begin(Item), [Item](InstLane &IL) { - auto [FrontV, FrontLane] = Item.front(); - auto [V, Lane] = IL; - return !V || (V == FrontV && Lane == FrontLane); + auto [FrontU, FrontLane] = Item.front(); + auto [U, Lane] = IL; + return !U || (U->get() == FrontU->get() && Lane == FrontLane); })) { - SplatLeafs.insert(FrontV); + SplatLeafs.insert(FrontU); continue; } // We need each element to be the same type of value, and check that each // element has a single use. if (!all_of(drop_begin(Item), [Item](InstLane IL) { - Value *FrontV = Item.front().first; - Value *V = IL.first; - if (!V) + Value *FrontV = Item.front().first->get(); + if (!IL.first) return true; + Value *V = IL.first->get(); if (auto *I = dyn_cast(V); I && !I->hasOneUse()) return false; if (V->getValueID() != FrontV->getValueID()) @@ -1869,25 +1865,25 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) { // Check the operator is one that we support. We exclude div/rem in case // they hit UB from poison lanes. - if ((isa(FrontV) && - !cast(FrontV)->isIntDivRem()) || - isa(FrontV)) { + if ((isa(FrontU) && + !cast(FrontU)->isIntDivRem()) || + isa(FrontU)) { Worklist.push_back(generateInstLaneVectorFromOperand(Item, 0)); Worklist.push_back(generateInstLaneVectorFromOperand(Item, 1)); - } else if (isa(FrontV)) { + } else if (isa(FrontU)) { Worklist.push_back(generateInstLaneVectorFromOperand(Item, 0)); - } else if (isa(FrontV)) { + } else if (isa(FrontU)) { Worklist.push_back(generateInstLaneVectorFromOperand(Item, 0)); Worklist.push_back(generateInstLaneVectorFromOperand(Item, 1)); Worklist.push_back(generateInstLaneVectorFromOperand(Item, 2)); - } else if (auto *II = dyn_cast(FrontV); + } else if (auto *II = dyn_cast(FrontU); II && isTriviallyVectorizable(II->getIntrinsicID())) { for (unsigned Op = 0, E = II->getNumOperands() - 1; Op < E; Op++) { if (isVectorIntrinsicWithScalarOpAtArg(II->getIntrinsicID(), Op)) { if (!all_of(drop_begin(Item), [Item, Op](InstLane &IL) { - Value *FrontV = Item.front().first; - Value *V = IL.first; - return !V || (cast(V)->getOperand(Op) == + Value *FrontV = Item.front().first->get(); + Use *U = IL.first; + return !U || (cast(U->get())->getOperand(Op) == cast(FrontV)->getOperand(Op)); })) return false; diff --git a/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll b/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll index c020aeb6483f5c..9ad042c112b4c0 100644 --- a/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll +++ b/llvm/test/Transforms/VectorCombine/AArch64/shuffletoidentity.ll @@ -202,6 +202,19 @@ define <8 x i8> @abs_different(<8 x i8> %a) { ret <8 x i8> %r } +define <4 x i32> @poison_intrinsic(<2 x i16> %l256) { +; CHECK-LABEL: @poison_intrinsic( +; CHECK-NEXT: [[L266:%.*]] = call <2 x i16> @llvm.abs.v2i16(<2 x i16> [[L256:%.*]], i1 false) +; CHECK-NEXT: [[L267:%.*]] = zext <2 x i16> [[L266]] to <2 x i32> +; CHECK-NEXT: [[L271:%.*]] = shufflevector <2 x i32> [[L267]], <2 x i32> poison, <4 x i32> +; CHECK-NEXT: ret <4 x i32> [[L271]] +; + %l266 = call <2 x i16> @llvm.abs.v2i16(<2 x i16> %l256, i1 false) + %l267 = zext <2 x i16> %l266 to <2 x i32> + %l271 = shufflevector <2 x i32> %l267, <2 x i32> poison, <4 x i32> + ret <4 x i32> %l271 +} + define <8 x half> @splat0(<8 x half> %a, <8 x half> %b) { ; CHECK-LABEL: @splat0( ; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <8 x half> [[B:%.*]], <8 x half> poison, <8 x i32> zeroinitializer