Skip to content

Commit

Permalink
[IR]Add NumSrcElts param to is..Mask static function in ShuffleVector…
Browse files Browse the repository at this point in the history 
…Inst.

Need to add NumSrcElts param to is..Mask functions in
ShuffleVectorInstruction class for better mask analysis. Mask.size() not
always matches the sizes of the permuted vector(s). Allows to better
estimate the cost in SLP and fix uses of the functions in other cases.

Differential Revision: https://reviews.llvm.org/D158449
  • Loading branch information
@alexey-bataev
alexey-bataev committed Sep 28, 2023
1 parent 6ca47eb commit c88c281
Show file tree
Hide file tree
Showing 14 changed files with 325 additions and 218 deletions.
19 changes: 8 additions & 11 deletions llvm/include/llvm/CodeGen/BasicTTIImpl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -935,31 +935,28 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
ArrayRef<int> Mask,
VectorType *Ty, int &Index,
VectorType *&SubTy) const {
int Limit = Mask.size() * 2;
if (Mask.empty() ||
// Extra check required by isSingleSourceMaskImpl function (called by
// ShuffleVectorInst::isSingleSourceMask).
any_of(Mask, [Limit](int I) { return I >= Limit; }))
if (Mask.empty())
return Kind;
int NumSrcElts = Ty->getElementCount().getKnownMinValue();
switch (Kind) {
case TTI::SK_PermuteSingleSrc:
if (ShuffleVectorInst::isReverseMask(Mask))
if (ShuffleVectorInst::isReverseMask(Mask, NumSrcElts))
return TTI::SK_Reverse;
if (ShuffleVectorInst::isZeroEltSplatMask(Mask))
if (ShuffleVectorInst::isZeroEltSplatMask(Mask, NumSrcElts))
return TTI::SK_Broadcast;
break;
case TTI::SK_PermuteTwoSrc: {
int NumSubElts;
if (Mask.size() > 2 && ShuffleVectorInst::isInsertSubvectorMask(
Mask, Mask.size(), NumSubElts, Index)) {
Mask, NumSrcElts, NumSubElts, Index)) {
SubTy = FixedVectorType::get(Ty->getElementType(), NumSubElts);
return TTI::SK_InsertSubvector;
}
if (ShuffleVectorInst::isSelectMask(Mask))
if (ShuffleVectorInst::isSelectMask(Mask, NumSrcElts))
return TTI::SK_Select;
if (ShuffleVectorInst::isTransposeMask(Mask))
if (ShuffleVectorInst::isTransposeMask(Mask, NumSrcElts))
return TTI::SK_Transpose;
if (ShuffleVectorInst::isSpliceMask(Mask, Index))
if (ShuffleVectorInst::isSpliceMask(Mask, NumSrcElts, Index))
return TTI::SK_Splice;
break;
}
Expand Down
70 changes: 39 additions & 31 deletions llvm/include/llvm/IR/Instructions.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2068,30 +2068,32 @@ class ShuffleVectorInst : public Instruction {
/// Return true if this shuffle mask chooses elements from exactly one source
/// vector.
/// Example: <7,5,undef,7>
/// This assumes that vector operands are the same length as the mask.
static bool isSingleSourceMask(ArrayRef<int> Mask);
static bool isSingleSourceMask(const Constant *Mask) {
/// This assumes that vector operands (of length \p NumSrcElts) are the same
/// length as the mask.
static bool isSingleSourceMask(ArrayRef<int> Mask, int NumSrcElts);
static bool isSingleSourceMask(const Constant *Mask, int NumSrcElts) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isSingleSourceMask(MaskAsInts);
return isSingleSourceMask(MaskAsInts, NumSrcElts);
}

/// Return true if this shuffle chooses elements from exactly one source
/// vector without changing the length of that vector.
/// Example: shufflevector <4 x n> A, <4 x n> B, <3,0,undef,3>
/// TODO: Optionally allow length-changing shuffles.
bool isSingleSource() const {
return !changesLength() && isSingleSourceMask(ShuffleMask);
return !changesLength() &&
isSingleSourceMask(ShuffleMask, ShuffleMask.size());
}

/// Return true if this shuffle mask chooses elements from exactly one source
/// vector without lane crossings. A shuffle using this mask is not
/// necessarily a no-op because it may change the number of elements from its
/// input vectors or it may provide demanded bits knowledge via undef lanes.
/// Example: <undef,undef,2,3>
static bool isIdentityMask(ArrayRef<int> Mask);
static bool isIdentityMask(const Constant *Mask) {
static bool isIdentityMask(ArrayRef<int> Mask, int NumSrcElts);
static bool isIdentityMask(const Constant *Mask, int NumSrcElts) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");

// Not possible to express a shuffle mask for a scalable vector for this
Expand All @@ -2101,7 +2103,7 @@ class ShuffleVectorInst : public Instruction {

SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isIdentityMask(MaskAsInts);
return isIdentityMask(MaskAsInts, NumSrcElts);
}

/// Return true if this shuffle chooses elements from exactly one source
Expand All @@ -2114,7 +2116,7 @@ class ShuffleVectorInst : public Instruction {
if (isa<ScalableVectorType>(getType()))
return false;

return !changesLength() && isIdentityMask(ShuffleMask);
return !changesLength() && isIdentityMask(ShuffleMask, ShuffleMask.size());
}

/// Return true if this shuffle lengthens exactly one source vector with
Expand All @@ -2138,12 +2140,12 @@ class ShuffleVectorInst : public Instruction {
/// In that case, the shuffle is better classified as an identity shuffle.
/// This assumes that vector operands are the same length as the mask
/// (a length-changing shuffle can never be equivalent to a vector select).
static bool isSelectMask(ArrayRef<int> Mask);
static bool isSelectMask(const Constant *Mask) {
static bool isSelectMask(ArrayRef<int> Mask, int NumSrcElts);
static bool isSelectMask(const Constant *Mask, int NumSrcElts) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isSelectMask(MaskAsInts);
return isSelectMask(MaskAsInts, NumSrcElts);
}

/// Return true if this shuffle chooses elements from its source vectors
Expand All @@ -2155,39 +2157,41 @@ class ShuffleVectorInst : public Instruction {
/// In that case, the shuffle is better classified as an identity shuffle.
/// TODO: Optionally allow length-changing shuffles.
bool isSelect() const {
return !changesLength() && isSelectMask(ShuffleMask);
return !changesLength() && isSelectMask(ShuffleMask, ShuffleMask.size());
}

/// Return true if this shuffle mask swaps the order of elements from exactly
/// one source vector.
/// Example: <7,6,undef,4>
/// This assumes that vector operands are the same length as the mask.
static bool isReverseMask(ArrayRef<int> Mask);
static bool isReverseMask(const Constant *Mask) {
/// This assumes that vector operands (of length \p NumSrcElts) are the same
/// length as the mask.
static bool isReverseMask(ArrayRef<int> Mask, int NumSrcElts);
static bool isReverseMask(const Constant *Mask, int NumSrcElts) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isReverseMask(MaskAsInts);
return isReverseMask(MaskAsInts, NumSrcElts);
}

/// Return true if this shuffle swaps the order of elements from exactly
/// one source vector.
/// Example: shufflevector <4 x n> A, <4 x n> B, <3,undef,1,undef>
/// TODO: Optionally allow length-changing shuffles.
bool isReverse() const {
return !changesLength() && isReverseMask(ShuffleMask);
return !changesLength() && isReverseMask(ShuffleMask, ShuffleMask.size());
}

/// Return true if this shuffle mask chooses all elements with the same value
/// as the first element of exactly one source vector.
/// Example: <4,undef,undef,4>
/// This assumes that vector operands are the same length as the mask.
static bool isZeroEltSplatMask(ArrayRef<int> Mask);
static bool isZeroEltSplatMask(const Constant *Mask) {
/// This assumes that vector operands (of length \p NumSrcElts) are the same
/// length as the mask.
static bool isZeroEltSplatMask(ArrayRef<int> Mask, int NumSrcElts);
static bool isZeroEltSplatMask(const Constant *Mask, int NumSrcElts) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isZeroEltSplatMask(MaskAsInts);
return isZeroEltSplatMask(MaskAsInts, NumSrcElts);
}

/// Return true if all elements of this shuffle are the same value as the
Expand All @@ -2197,7 +2201,8 @@ class ShuffleVectorInst : public Instruction {
/// TODO: Optionally allow length-changing shuffles.
/// TODO: Optionally allow splats from other elements.
bool isZeroEltSplat() const {
return !changesLength() && isZeroEltSplatMask(ShuffleMask);
return !changesLength() &&
isZeroEltSplatMask(ShuffleMask, ShuffleMask.size());
}

/// Return true if this shuffle mask is a transpose mask.
Expand Down Expand Up @@ -2232,12 +2237,12 @@ class ShuffleVectorInst : public Instruction {
/// ; Transposed matrix
/// t0 = < a, e, c, g > = shufflevector m0, m1 < 0, 4, 2, 6 >
/// t1 = < b, f, d, h > = shufflevector m0, m1 < 1, 5, 3, 7 >
static bool isTransposeMask(ArrayRef<int> Mask);
static bool isTransposeMask(const Constant *Mask) {
static bool isTransposeMask(ArrayRef<int> Mask, int NumSrcElts);
static bool isTransposeMask(const Constant *Mask, int NumSrcElts) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isTransposeMask(MaskAsInts);
return isTransposeMask(MaskAsInts, NumSrcElts);
}

/// Return true if this shuffle transposes the elements of its inputs without
Expand All @@ -2246,27 +2251,30 @@ class ShuffleVectorInst : public Instruction {
/// exact specification.
/// Example: shufflevector <4 x n> A, <4 x n> B, <0,4,2,6>
bool isTranspose() const {
return !changesLength() && isTransposeMask(ShuffleMask);
return !changesLength() && isTransposeMask(ShuffleMask, ShuffleMask.size());
}

/// Return true if this shuffle mask is a splice mask, concatenating the two
/// inputs together and then extracts an original width vector starting from
/// the splice index.
/// Example: shufflevector <4 x n> A, <4 x n> B, <1,2,3,4>
static bool isSpliceMask(ArrayRef<int> Mask, int &Index);
static bool isSpliceMask(const Constant *Mask, int &Index) {
/// This assumes that vector operands (of length \p NumSrcElts) are the same
/// length as the mask.
static bool isSpliceMask(ArrayRef<int> Mask, int NumSrcElts, int &Index);
static bool isSpliceMask(const Constant *Mask, int NumSrcElts, int &Index) {
assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
SmallVector<int, 16> MaskAsInts;
getShuffleMask(Mask, MaskAsInts);
return isSpliceMask(MaskAsInts, Index);
return isSpliceMask(MaskAsInts, NumSrcElts, Index);
}

/// Return true if this shuffle splices two inputs without changing the length
/// of the vectors. This operation concatenates the two inputs together and
/// then extracts an original width vector starting from the splice index.
/// Example: shufflevector <4 x n> A, <4 x n> B, <1,2,3,4>
bool isSplice(int &Index) const {
return !changesLength() && isSpliceMask(ShuffleMask, Index);
return !changesLength() &&
isSpliceMask(ShuffleMask, ShuffleMask.size(), Index);
}

/// Return true if this shuffle mask is an extract subvector mask.
Expand Down
2 changes: 1 addition & 1 deletion llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24119,7 +24119,7 @@ static SDValue foldExtractSubvectorFromShuffleVector(SDNode *N,

// Profitability check: only deal with extractions from the first subvector
// unless the mask becomes an identity mask.
if (!ShuffleVectorInst::isIdentityMask(NewMask) ||
if (!ShuffleVectorInst::isIdentityMask(NewMask, NewMask.size()) ||
any_of(NewMask, [](int M) { return M < 0; }))
for (auto &DemandedSubvector : DemandedSubvectors)
if (DemandedSubvector.second != 0)
Expand Down
75 changes: 44 additions & 31 deletions llvm/lib/IR/Instructions.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2136,10 +2136,10 @@ static bool isSingleSourceMaskImpl(ArrayRef<int> Mask, int NumOpElts) {
return UsesLHS || UsesRHS;
}

bool ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> Mask) {
bool ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> Mask, int NumSrcElts) {
// We don't have vector operand size information, so assume operands are the
// same size as the mask.
return isSingleSourceMaskImpl(Mask, Mask.size());
return isSingleSourceMaskImpl(Mask, NumSrcElts);
}

static bool isIdentityMaskImpl(ArrayRef<int> Mask, int NumOpElts) {
Expand All @@ -2154,65 +2154,75 @@ static bool isIdentityMaskImpl(ArrayRef<int> Mask, int NumOpElts) {
return true;
}

bool ShuffleVectorInst::isIdentityMask(ArrayRef<int> Mask) {
bool ShuffleVectorInst::isIdentityMask(ArrayRef<int> Mask, int NumSrcElts) {
if (Mask.size() != static_cast<unsigned>(NumSrcElts))
return false;
// We don't have vector operand size information, so assume operands are the
// same size as the mask.
return isIdentityMaskImpl(Mask, Mask.size());
return isIdentityMaskImpl(Mask, NumSrcElts);
}

bool ShuffleVectorInst::isReverseMask(ArrayRef<int> Mask) {
if (!isSingleSourceMask(Mask))
bool ShuffleVectorInst::isReverseMask(ArrayRef<int> Mask, int NumSrcElts) {
if (Mask.size() != static_cast<unsigned>(NumSrcElts))
return false;
if (!isSingleSourceMask(Mask, NumSrcElts))
return false;

// The number of elements in the mask must be at least 2.
int NumElts = Mask.size();
if (NumElts < 2)
if (NumSrcElts < 2)
return false;

for (int i = 0; i < NumElts; ++i) {
if (Mask[i] == -1)
for (int I = 0, E = Mask.size(); I < E; ++I) {
if (Mask[I] == -1)
continue;
if (Mask[i] != (NumElts - 1 - i) && Mask[i] != (NumElts + NumElts - 1 - i))
if (Mask[I] != (NumSrcElts - 1 - I) &&
Mask[I] != (NumSrcElts + NumSrcElts - 1 - I))
return false;
}
return true;
}

bool ShuffleVectorInst::isZeroEltSplatMask(ArrayRef<int> Mask) {
if (!isSingleSourceMask(Mask))
bool ShuffleVectorInst::isZeroEltSplatMask(ArrayRef<int> Mask, int NumSrcElts) {
if (Mask.size() != static_cast<unsigned>(NumSrcElts))
return false;
for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) {
if (Mask[i] == -1)
if (!isSingleSourceMask(Mask, NumSrcElts))
return false;
for (int I = 0, E = Mask.size(); I < E; ++I) {
if (Mask[I] == -1)
continue;
if (Mask[i] != 0 && Mask[i] != NumElts)
if (Mask[I] != 0 && Mask[I] != NumSrcElts)
return false;
}
return true;
}

bool ShuffleVectorInst::isSelectMask(ArrayRef<int> Mask) {
bool ShuffleVectorInst::isSelectMask(ArrayRef<int> Mask, int NumSrcElts) {
if (Mask.size() != static_cast<unsigned>(NumSrcElts))
return false;
// Select is differentiated from identity. It requires using both sources.
if (isSingleSourceMask(Mask))
if (isSingleSourceMask(Mask, NumSrcElts))
return false;
for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) {
if (Mask[i] == -1)
for (int I = 0, E = Mask.size(); I < E; ++I) {
if (Mask[I] == -1)
continue;
if (Mask[i] != i && Mask[i] != (NumElts + i))
if (Mask[I] != I && Mask[I] != (NumSrcElts + I))
return false;
}
return true;
}

bool ShuffleVectorInst::isTransposeMask(ArrayRef<int> Mask) {
bool ShuffleVectorInst::isTransposeMask(ArrayRef<int> Mask, int NumSrcElts) {
// Example masks that will return true:
// v1 = <a, b, c, d>
// v2 = <e, f, g, h>
// trn1 = shufflevector v1, v2 <0, 4, 2, 6> = <a, e, c, g>
// trn2 = shufflevector v1, v2 <1, 5, 3, 7> = <b, f, d, h>

if (Mask.size() != static_cast<unsigned>(NumSrcElts))
return false;
// 1. The number of elements in the mask must be a power-of-2 and at least 2.
int NumElts = Mask.size();
if (NumElts < 2 || !isPowerOf2_32(NumElts))
int Sz = Mask.size();
if (Sz < 2 || !isPowerOf2_32(Sz))
return false;

// 2. The first element of the mask must be either a 0 or a 1.
Expand All @@ -2221,23 +2231,26 @@ bool ShuffleVectorInst::isTransposeMask(ArrayRef<int> Mask) {

// 3. The difference between the first 2 elements must be equal to the
// number of elements in the mask.
if ((Mask[1] - Mask[0]) != NumElts)
if ((Mask[1] - Mask[0]) != NumSrcElts)
return false;

// 4. The difference between consecutive even-numbered and odd-numbered
// elements must be equal to 2.
for (int i = 2; i < NumElts; ++i) {
int MaskEltVal = Mask[i];
for (int I = 2; I < Sz; ++I) {
int MaskEltVal = Mask[I];
if (MaskEltVal == -1)
return false;
int MaskEltPrevVal = Mask[i - 2];
int MaskEltPrevVal = Mask[I - 2];
if (MaskEltVal - MaskEltPrevVal != 2)
return false;
}
return true;
}

bool ShuffleVectorInst::isSpliceMask(ArrayRef<int> Mask, int &Index) {
bool ShuffleVectorInst::isSpliceMask(ArrayRef<int> Mask, int NumSrcElts,
int &Index) {
if (Mask.size() != static_cast<unsigned>(NumSrcElts))
return false;
// Example: shufflevector <4 x n> A, <4 x n> B, <1,2,3,4>
int StartIndex = -1;
for (int I = 0, E = Mask.size(); I != E; ++I) {
Expand All @@ -2248,7 +2261,7 @@ bool ShuffleVectorInst::isSpliceMask(ArrayRef<int> Mask, int &Index) {
if (StartIndex == -1) {
// Don't support a StartIndex that begins in the second input, or if the
// first non-undef index would access below the StartIndex.
if (MaskEltVal < I || E <= (MaskEltVal - I))
if (MaskEltVal < I || NumSrcElts <= (MaskEltVal - I))
return false;

StartIndex = MaskEltVal - I;
Expand Down Expand Up @@ -2543,7 +2556,7 @@ bool ShuffleVectorInst::isOneUseSingleSourceMask(int VF) const {
// case.
if (isa<ScalableVectorType>(getType()))
return false;
if (!isSingleSourceMask(ShuffleMask))
if (!isSingleSourceMask(ShuffleMask, VF))
return false;

return isOneUseSingleSourceMask(ShuffleMask, VF);
Expand Down

0 comments on commit c88c281

Please sign in to comment.