diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp --- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp +++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp @@ -963,6 +963,7 @@ class BoUpSLP { struct TreeEntry; struct ScheduleData; + class ShuffleInstructionBuilder; public: using ValueList = SmallVector; @@ -8114,67 +8115,205 @@ return Vec; } -namespace { -/// Merges shuffle masks and emits final shuffle instruction, if required. -class ShuffleInstructionBuilder { - IRBuilderBase &Builder; - const unsigned VF = 0; +/// Merges shuffle masks and emits final shuffle instruction, if required. It +/// supports shuffling of 2 input vectors. It implements lazy shuffles emission, +/// when the actual shuffle instruction is generated only if this is actually +/// required. Otherwise, the shuffle instruction emission is delayed till the +/// end of the process, to reduce the number of emitted instructions and further +/// analysis/transformations. +class BoUpSLP::ShuffleInstructionBuilder { bool IsFinalized = false; - SmallVector Mask; - /// Holds all of the instructions that we gathered. - SetVector &GatherShuffleSeq; - /// A list of blocks that we are going to CSE. - SetVector &CSEBlocks; + SmallVector CommonMask; + SmallVector InVectors; + IRBuilderBase &Builder; + BoUpSLP &R; + + class ShuffleIRBuilder { + IRBuilderBase &Builder; + /// Holds all of the instructions that we gathered. + SetVector &GatherShuffleExtractSeq; + /// A list of blocks that we are going to CSE. + SetVector &CSEBlocks; + + public: + ShuffleIRBuilder(IRBuilderBase &Builder, + SetVector &GatherShuffleExtractSeq, + SetVector &CSEBlocks) + : Builder(Builder), GatherShuffleExtractSeq(GatherShuffleExtractSeq), + CSEBlocks(CSEBlocks) {} + ~ShuffleIRBuilder() = default; + Value *createShuffleVector(Value *V1, Value *V2, ArrayRef Mask) { + Value *Vec = Builder.CreateShuffleVector(V1, V2, Mask); + if (auto *I = dyn_cast(Vec)) { + GatherShuffleExtractSeq.insert(I); + CSEBlocks.insert(I->getParent()); + } + return Vec; + } + Value *createShuffleVector(Value *V1, ArrayRef Mask) { + if (Mask.empty()) + return V1; + unsigned VF = Mask.size(); + unsigned LocalVF = cast(V1->getType())->getNumElements(); + if (VF == LocalVF && ShuffleVectorInst::isIdentityMask(Mask)) + return V1; + Value *Vec = Builder.CreateShuffleVector(V1, Mask); + if (auto *I = dyn_cast(Vec)) { + GatherShuffleExtractSeq.insert(I); + CSEBlocks.insert(I->getParent()); + } + return Vec; + } + }; + + /// Smart shuffle instruction emission, walks through shuffles trees and + /// tries to find the best matching vector for the actual shuffle + /// instruction. + Value *createShuffle(Value *V1, Value *V2, ArrayRef Mask) { + assert(V1 && "Expected at least one vector value."); + ShuffleIRBuilder ShuffleBuilder(Builder, R.GatherShuffleExtractSeq, + R.CSEBlocks); + if (V2) + return ShuffleBuilder.createShuffleVector(V1, V2, Mask); + return ShuffleBuilder.createShuffleVector(V1, Mask); + } public: - ShuffleInstructionBuilder(IRBuilderBase &Builder, unsigned VF, - SetVector &GatherShuffleSeq, - SetVector &CSEBlocks) - : Builder(Builder), VF(VF), GatherShuffleSeq(GatherShuffleSeq), - CSEBlocks(CSEBlocks) {} - - /// Adds a mask, inverting it before applying. - void addInversedMask(ArrayRef SubMask) { - if (SubMask.empty()) + ShuffleInstructionBuilder(IRBuilderBase &Builder, BoUpSLP &R) + : Builder(Builder), R(R) {} + + /// Adds 2 input vectors and the mask for their shuffling. + void add(Value *V1, Value *V2, ArrayRef Mask) { + assert(V1 && V2 && !Mask.empty() && "Expected non-empty input vectors."); + if (InVectors.empty()) { + InVectors.push_back(V1); + InVectors.push_back(V2); + CommonMask.assign(Mask.begin(), Mask.end()); return; - SmallVector NewMask; - inversePermutation(SubMask, NewMask); - addMask(NewMask); + } + Value *Vec = InVectors.front(); + if (InVectors.size() == 2) { + Vec = createShuffle(Vec, InVectors.back(), CommonMask); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (Mask[Idx] != UndefMaskElem) + CommonMask[Idx] = Idx; + } else if (cast(Vec->getType())->getNumElements() != + Mask.size()) { + Vec = createShuffle(Vec, nullptr, CommonMask); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (Mask[Idx] != UndefMaskElem) + CommonMask[Idx] = Idx; + } + V1 = createShuffle(V1, V2, Mask); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (Mask[Idx] != UndefMaskElem) + CommonMask[Idx] = Idx + Sz; + InVectors.front() = Vec; + if (InVectors.size() == 2) + InVectors.back() = V1; + else + InVectors.push_back(V1); + } + /// Adds another one input vector and the mask for the shuffling. + void add(Value *V1, ArrayRef Mask) { + if (InVectors.empty()) { + if (!isa(V1->getType())) { + V1 = createShuffle(V1, nullptr, CommonMask); + CommonMask.assign(Mask.size(), UndefMaskElem); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (Mask[Idx] != UndefMaskElem) + CommonMask[Idx] = Idx; + } + InVectors.push_back(V1); + CommonMask.assign(Mask.begin(), Mask.end()); + return; + } + const auto *It = find(InVectors, V1); + if (It == InVectors.end()) { + if (InVectors.size() == 2 || + InVectors.front()->getType() != V1->getType() || + !isa(V1->getType())) { + Value *V = InVectors.front(); + if (InVectors.size() == 2) { + V = createShuffle(InVectors.front(), InVectors.back(), CommonMask); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (CommonMask[Idx] != UndefMaskElem) + CommonMask[Idx] = Idx; + } else if (cast(V->getType())->getNumElements() != + CommonMask.size()) { + V = createShuffle(InVectors.front(), nullptr, CommonMask); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (CommonMask[Idx] != UndefMaskElem) + CommonMask[Idx] = Idx; + } + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (CommonMask[Idx] == UndefMaskElem && Mask[Idx] != UndefMaskElem) + CommonMask[Idx] = + V->getType() != V1->getType() + ? Idx + Sz + : Mask[Idx] + cast(V1->getType()) + ->getNumElements(); + if (V->getType() != V1->getType()) + V1 = createShuffle(V1, nullptr, Mask); + InVectors.front() = V; + if (InVectors.size() == 2) + InVectors.back() = V1; + else + InVectors.push_back(V1); + return; + } + // Check if second vector is required if the used elements are already + // used from the first one. + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (Mask[Idx] != UndefMaskElem && CommonMask[Idx] == UndefMaskElem) { + InVectors.push_back(V1); + break; + } + } + int VF = CommonMask.size(); + if (auto *FTy = dyn_cast(V1->getType())) + VF = FTy->getNumElements(); + for (unsigned Idx = 0, Sz = CommonMask.size(); Idx < Sz; ++Idx) + if (Mask[Idx] != UndefMaskElem && CommonMask[Idx] == UndefMaskElem) + CommonMask[Idx] = Mask[Idx] + (It == InVectors.begin() ? 0 : VF); } - - /// Functions adds masks, merging them into single one. - void addMask(ArrayRef SubMask) { - SmallVector NewMask(SubMask); - addMask(NewMask); + /// Adds another one input vector and the mask for the shuffling. + void add(Value *V1, ArrayRef Order) { + SmallVector NewMask; + inversePermutation(Order, NewMask); + add(V1, NewMask); } - - void addMask(ArrayRef SubMask) { ::addMask(Mask, SubMask); } - - Value *finalize(Value *V) { + /// Finalize emission of the shuffles. + Value * + finalize(ArrayRef ExtMask = std::nullopt) { IsFinalized = true; - unsigned ValueVF = cast(V->getType())->getNumElements(); - if (VF == ValueVF && Mask.empty()) - return V; - SmallVector NormalizedMask(VF, UndefMaskElem); - std::iota(NormalizedMask.begin(), NormalizedMask.end(), 0); - addMask(NormalizedMask); - - if (VF == ValueVF && ShuffleVectorInst::isIdentityMask(Mask)) - return V; - Value *Vec = Builder.CreateShuffleVector(V, Mask, "shuffle"); - if (auto *I = dyn_cast(Vec)) { - GatherShuffleSeq.insert(I); - CSEBlocks.insert(I->getParent()); + if (!ExtMask.empty()) { + if (CommonMask.empty()) { + CommonMask.assign(ExtMask.begin(), ExtMask.end()); + } else { + SmallVector NewMask(ExtMask.size(), UndefMaskElem); + for (int I = 0, Sz = ExtMask.size(); I < Sz; ++I) { + if (ExtMask[I] == UndefMaskElem) + continue; + NewMask[I] = CommonMask[ExtMask[I]]; + } + CommonMask.swap(NewMask); + } } - return Vec; + if (CommonMask.empty()) { + assert(InVectors.size() == 1 && "Expected only one vector with no mask"); + return InVectors.front(); + } + if (InVectors.size() == 2) + return createShuffle(InVectors.front(), InVectors.back(), CommonMask); + return createShuffle(InVectors.front(), nullptr, CommonMask); } ~ShuffleInstructionBuilder() { - assert((IsFinalized || Mask.empty()) && + assert((IsFinalized || CommonMask.empty()) && "Shuffle construction must be finalized."); } }; -} // namespace Value *BoUpSLP::vectorizeOperand(TreeEntry *E, unsigned NodeIdx) { ArrayRef VL = E->getOperand(NodeIdx); @@ -8275,8 +8414,7 @@ assert(E->State == TreeEntry::NeedToGather && "Expected gather node."); unsigned VF = E->getVectorFactor(); - ShuffleInstructionBuilder ShuffleBuilder(Builder, VF, GatherShuffleExtractSeq, - CSEBlocks); + ShuffleInstructionBuilder ShuffleBuilder(Builder, *this); SmallVector Gathered( VF, PoisonValue::get(E->Scalars.front()->getType())); bool NeedFreeze = false; @@ -8286,11 +8424,11 @@ inversePermutation(E->ReorderIndices, ReorderMask); if (!ReorderMask.empty()) reorderScalars(VL, ReorderMask); + SmallVector ReuseMask(VF, UndefMaskElem); if (!allConstant(VL)) { // For splats with can emit broadcasts instead of gathers, so try to find // such sequences. bool IsSplat = isSplat(VL) && (VL.size() > 2 || VL.front() == VL.back()); - SmallVector ReuseMask(VF, UndefMaskElem); SmallVector UndefPos; DenseMap UniquePositions; // Gather unique non-const values and all constant values. @@ -8356,14 +8494,14 @@ NeedFreeze = true; } } - ShuffleBuilder.addMask(ReuseMask); } else { + ReuseMask.clear(); copy(VL, Gathered.begin()); } // Gather unique scalars and all constants. Value *Vec = gather(Gathered); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - Vec = ShuffleBuilder.finalize(Vec); + ShuffleBuilder.add(Vec, ReuseMask); + Vec = ShuffleBuilder.finalize(E->ReuseShuffleIndices); if (NeedFreeze) Vec = Builder.CreateFreeze(Vec); return Vec; @@ -8377,10 +8515,20 @@ return E->VectorizedValue; } - bool NeedToShuffleReuses = !E->ReuseShuffleIndices.empty(); - unsigned VF = E->getVectorFactor(); - ShuffleInstructionBuilder ShuffleBuilder(Builder, VF, GatherShuffleExtractSeq, - CSEBlocks); + auto FinalShuffle = [&](Value *V) { + ShuffleInstructionBuilder ShuffleBuilder(Builder, *this); + if (E->State != TreeEntry::NeedToGather && + E->getOpcode() == Instruction::Store) { + ArrayRef Mask = + makeArrayRef(reinterpret_cast(E->ReorderIndices.begin()), + E->ReorderIndices.size()); + ShuffleBuilder.add(V, Mask); + } else { + ShuffleBuilder.add(V, E->ReorderIndices); + } + return ShuffleBuilder.finalize(E->ReuseShuffleIndices); + }; + if (E->State == TreeEntry::NeedToGather) { if (E->Idx > 0) { // We are in the middle of a vectorizable chain. We need to gather the @@ -8408,10 +8556,7 @@ } else { Vec = gather(E->Scalars); } - if (NeedToShuffleReuses) { - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - Vec = ShuffleBuilder.finalize(Vec); - } + Vec = FinalShuffle(Vec); E->VectorizedValue = Vec; return Vec; } @@ -8444,9 +8589,7 @@ Builder.SetInsertPoint(&*PH->getParent()->getFirstInsertionPt()); Builder.SetCurrentDebugLocation(PH->getDebugLoc()); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; @@ -8483,9 +8626,7 @@ case Instruction::ExtractElement: { Value *V = E->getSingleOperand(0); setInsertPointAfterBundle(E); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; return V; } @@ -8496,9 +8637,7 @@ Value *Ptr = Builder.CreateBitCast(LI->getOperand(0), PtrTy); LoadInst *V = Builder.CreateAlignedLoad(VecTy, Ptr, LI->getAlign()); Value *NewV = propagateMetadata(V, E->Scalars); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - NewV = ShuffleBuilder.finalize(NewV); + NewV = FinalShuffle(NewV); E->VectorizedValue = NewV; return NewV; } @@ -8621,9 +8760,7 @@ auto *CI = cast(VL0); Value *V = Builder.CreateCast(CI->getOpcode(), InVec, VecTy); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -8647,9 +8784,7 @@ CmpInst::Predicate P0 = cast(VL0)->getPredicate(); Value *V = Builder.CreateCmp(P0, L, R); propagateIRFlags(V, E->Scalars, VL0); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -8675,9 +8810,7 @@ } Value *V = Builder.CreateSelect(Cond, True, False); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -8699,9 +8832,7 @@ if (auto *I = dyn_cast(V)) V = propagateMetadata(I, E->Scalars); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -8746,9 +8877,7 @@ if (auto *I = dyn_cast(V)) V = propagateMetadata(I, E->Scalars); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -8793,9 +8922,7 @@ } Value *V = propagateMetadata(NewLI, E->Scalars); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; return V; @@ -8807,8 +8934,7 @@ setInsertPointAfterBundle(E); Value *VecValue = vectorizeOperand(E, 0); - ShuffleBuilder.addMask(E->ReorderIndices); - VecValue = ShuffleBuilder.finalize(VecValue); + VecValue = FinalShuffle(VecValue); Value *ScalarPtr = SI->getPointerOperand(); Value *VecPtr = Builder.CreateBitCast( @@ -8863,9 +8989,7 @@ V = propagateMetadata(I, GEPs); } - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -8942,9 +9066,7 @@ } propagateIRFlags(V, E->Scalars, VL0); - ShuffleBuilder.addInversedMask(E->ReorderIndices); - ShuffleBuilder.addMask(E->ReuseShuffleIndices); - V = ShuffleBuilder.finalize(V); + V = FinalShuffle(V); E->VectorizedValue = V; ++NumVectorInstructions; @@ -9026,7 +9148,6 @@ GatherShuffleExtractSeq.insert(I); CSEBlocks.insert(I->getParent()); } - V = ShuffleBuilder.finalize(V); E->VectorizedValue = V; ++NumVectorInstructions;