diff --git a/llvm/include/llvm/Analysis/VectorUtils.h b/llvm/include/llvm/Analysis/VectorUtils.h --- a/llvm/include/llvm/Analysis/VectorUtils.h +++ b/llvm/include/llvm/Analysis/VectorUtils.h @@ -533,6 +533,12 @@ llvm::SmallVector createSequentialMask(unsigned Start, unsigned NumInts, unsigned NumUndefs); +/// Given a shuffle mask for a binary shuffle, create the equivalent shuffle +/// mask assuming both operands are identical. This assumes that the unary +/// shuffle will use elements from operand 0 (operand 1 will be unused). +llvm::SmallVector createUnaryMask(ArrayRef Mask, + unsigned NumElts); + /// Concatenate a list of vectors. /// /// This function generates code that concatenate the vectors in \p Vecs into a diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp --- a/llvm/lib/Analysis/VectorUtils.cpp +++ b/llvm/lib/Analysis/VectorUtils.cpp @@ -830,6 +830,23 @@ return Mask; } +llvm::SmallVector llvm::createUnaryMask(ArrayRef Mask, + unsigned NumElts) { + // Avoid casts in the loop and make sure we have a reasonable number. + int NumEltsSigned = NumElts; + assert(NumEltsSigned > 0 && "Expected smaller or non-zero element count"); + + // If the mask chooses an element from operand 1, reduce it to choose from the + // corresponding element of operand 0. Undef mask elements are unchanged. + SmallVector UnaryMask; + for (int MaskElt : Mask) { + assert((MaskElt < NumEltsSigned * 2) && "Expected valid shuffle mask"); + int UnaryElt = MaskElt >= NumEltsSigned ? MaskElt - NumEltsSigned : MaskElt; + UnaryMask.push_back(UnaryElt); + } + return UnaryMask; +} + /// A helper function for concatenating vectors. This function concatenates two /// vectors having the same element type. If the second vector has fewer /// elements than the first, it is padded with undefs. diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp --- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp @@ -21254,15 +21254,9 @@ ShuffleVectorSDNode *SVN = cast(N); // Canonicalize shuffle v, v -> v, undef - if (N0 == N1) { - SmallVector NewMask; - for (unsigned i = 0; i != NumElts; ++i) { - int Idx = SVN->getMaskElt(i); - if (Idx >= (int)NumElts) Idx -= NumElts; - NewMask.push_back(Idx); - } - return DAG.getVectorShuffle(VT, SDLoc(N), N0, DAG.getUNDEF(VT), NewMask); - } + if (N0 == N1) + return DAG.getVectorShuffle(VT, SDLoc(N), N0, DAG.getUNDEF(VT), + createUnaryMask(SVN->getMask(), NumElts)); // Canonicalize shuffle undef, v -> v, undef. Commute the shuffle mask. if (N0.isUndef()) diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp --- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp +++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp @@ -2484,16 +2484,7 @@ if (LHS == RHS) { assert(!match(RHS, m_Undef()) && "Shuffle with 2 undef ops not simplified?"); - // Remap any references to RHS to use LHS. - SmallVector Elts; - for (unsigned i = 0; i != VWidth; ++i) { - // Propagate undef elements or force mask to LHS. - if (Mask[i] < 0) - Elts.push_back(UndefMaskElem); - else - Elts.push_back(Mask[i] % LHSWidth); - } - return new ShuffleVectorInst(LHS, Elts); + return new ShuffleVectorInst(LHS, createUnaryMask(Mask, LHSWidth)); } // shuffle undef, x, mask --> shuffle x, undef, mask'