From ac619a09ecb228834c200afbae4a4d7aeaa5576a Mon Sep 17 00:00:00 2001 From: Sanjay Patel Date: Thu, 30 Aug 2018 15:05:38 +0000 Subject: [PATCH] [IR] add shuffle queries for identity extend/extract This was one of the potential follow-ups suggested in D48236, and these will be used to make matching the patterns in PR38691 cleaner: https://bugs.llvm.org/show_bug.cgi?id=38691 About the vocabulary: in the DAG, these would be concat_vector with an undef operand or extract_subvector. Alternate names are discussed in the review, but I think these are familiar/good enough to proceed. Once we have uses of them in code, we might adjust if there are better options. https://reviews.llvm.org/D51392 llvm-svn: 341075 --- llvm/include/llvm/IR/Instructions.h | 13 +++-- llvm/lib/IR/Instructions.cpp | 57 ++++++++++++++++---- llvm/unittests/IR/InstructionsTest.cpp | 72 ++++++++++++++++++++++++++ 3 files changed, 130 insertions(+), 12 deletions(-) diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h index 140a6f204245..dea7fcd9b64c 100644 --- a/llvm/include/llvm/IR/Instructions.h +++ b/llvm/include/llvm/IR/Instructions.h @@ -2456,7 +2456,7 @@ public: } /// Return true if this shuffle returns a vector with a different number of - /// elements than its source elements. + /// elements than its source vectors. /// Example: shufflevector <4 x n> A, <4 x n> B, <1,2> bool changesLength() const { unsigned NumSourceElts = Op<0>()->getType()->getVectorNumElements(); @@ -2497,15 +2497,22 @@ public: return isIdentityMask(MaskAsInts); } - /// Return true if this shuffle mask chooses elements from exactly one source + /// Return true if this shuffle chooses elements from exactly one source /// vector without lane crossings and does not change the number of elements /// from its input vectors. /// Example: shufflevector <4 x n> A, <4 x n> B, <4,undef,6,undef> - /// TODO: Optionally allow length-changing shuffles. bool isIdentity() const { return !changesLength() && isIdentityMask(getShuffleMask()); } + /// Return true if this shuffle lengthens exactly one source vector with + /// undefs in the high elements. + bool isIdentityWithPadding() const; + + /// Return true if this shuffle extracts the first N elements of exactly one + /// source vector. + bool isIdentityWithExtract() const; + /// Return true if this shuffle mask chooses elements from its source vectors /// without lane crossings. A shuffle using this mask would be /// equivalent to a vector select with a constant condition operand. diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp index 1a2752deae15..2bf9f0b12e71 100644 --- a/llvm/lib/IR/Instructions.cpp +++ b/llvm/lib/IR/Instructions.cpp @@ -1660,17 +1660,17 @@ void ShuffleVectorInst::getShuffleMask(const Constant *Mask, } } -bool ShuffleVectorInst::isSingleSourceMask(ArrayRef Mask) { +static bool isSingleSourceMaskImpl(ArrayRef Mask, int NumOpElts) { assert(!Mask.empty() && "Shuffle mask must contain elements"); bool UsesLHS = false; bool UsesRHS = false; - for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) { + for (int i = 0, NumMaskElts = Mask.size(); i < NumMaskElts; ++i) { if (Mask[i] == -1) continue; - assert(Mask[i] >= 0 && Mask[i] < (NumElts * 2) && + assert(Mask[i] >= 0 && Mask[i] < (NumOpElts * 2) && "Out-of-bounds shuffle mask element"); - UsesLHS |= (Mask[i] < NumElts); - UsesRHS |= (Mask[i] >= NumElts); + UsesLHS |= (Mask[i] < NumOpElts); + UsesRHS |= (Mask[i] >= NumOpElts); if (UsesLHS && UsesRHS) return false; } @@ -1678,18 +1678,30 @@ bool ShuffleVectorInst::isSingleSourceMask(ArrayRef Mask) { return true; } -bool ShuffleVectorInst::isIdentityMask(ArrayRef Mask) { - if (!isSingleSourceMask(Mask)) +bool ShuffleVectorInst::isSingleSourceMask(ArrayRef Mask) { + // We don't have vector operand size information, so assume operands are the + // same size as the mask. + return isSingleSourceMaskImpl(Mask, Mask.size()); +} + +static bool isIdentityMaskImpl(ArrayRef Mask, int NumOpElts) { + if (!isSingleSourceMaskImpl(Mask, NumOpElts)) return false; - for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) { + for (int i = 0, NumMaskElts = Mask.size(); i < NumMaskElts; ++i) { if (Mask[i] == -1) continue; - if (Mask[i] != i && Mask[i] != (NumElts + i)) + if (Mask[i] != i && Mask[i] != (NumOpElts + i)) return false; } return true; } +bool ShuffleVectorInst::isIdentityMask(ArrayRef Mask) { + // We don't have vector operand size information, so assume operands are the + // same size as the mask. + return isIdentityMaskImpl(Mask, Mask.size()); +} + bool ShuffleVectorInst::isReverseMask(ArrayRef Mask) { if (!isSingleSourceMask(Mask)) return false; @@ -1761,6 +1773,33 @@ bool ShuffleVectorInst::isTransposeMask(ArrayRef Mask) { return true; } +bool ShuffleVectorInst::isIdentityWithPadding() const { + int NumOpElts = Op<0>()->getType()->getVectorNumElements(); + int NumMaskElts = getType()->getVectorNumElements(); + if (NumMaskElts <= NumOpElts) + return false; + + // The first part of the mask must choose elements from exactly 1 source op. + ArrayRef Mask = getShuffleMask(); + if (!isIdentityMaskImpl(Mask, NumOpElts)) + return false; + + // All extending must be with undef elements. + for (int i = NumOpElts; i < NumMaskElts; ++i) + if (Mask[i] != -1) + return false; + + return true; +} + +bool ShuffleVectorInst::isIdentityWithExtract() const { + int NumOpElts = Op<0>()->getType()->getVectorNumElements(); + int NumMaskElts = getType()->getVectorNumElements(); + if (NumMaskElts >= NumOpElts) + return false; + + return isIdentityMaskImpl(getShuffleMask(), NumOpElts); +} //===----------------------------------------------------------------------===// // InsertValueInst Class diff --git a/llvm/unittests/IR/InstructionsTest.cpp b/llvm/unittests/IR/InstructionsTest.cpp index 6bac59620d9e..d153f999d654 100644 --- a/llvm/unittests/IR/InstructionsTest.cpp +++ b/llvm/unittests/IR/InstructionsTest.cpp @@ -837,6 +837,78 @@ TEST(InstructionsTest, ShuffleMaskQueries) { EXPECT_TRUE(ShuffleVectorInst::isTransposeMask(ConstantVector::get({C1, C5, C3, C7}))); EXPECT_TRUE(ShuffleVectorInst::isTransposeMask(ConstantVector::get({C1, C3}))); + + // Identity with undef elts. + ShuffleVectorInst *Id1 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C3, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, CU, CU})); + EXPECT_TRUE(Id1->isIdentity()); + EXPECT_FALSE(Id1->isIdentityWithPadding()); + EXPECT_FALSE(Id1->isIdentityWithExtract()); + delete Id1; + + // Result has less elements than operands. + ShuffleVectorInst *Id2 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2})); + EXPECT_FALSE(Id2->isIdentity()); + EXPECT_FALSE(Id2->isIdentityWithPadding()); + EXPECT_TRUE(Id2->isIdentityWithExtract()); + delete Id2; + + // Result has less elements than operands; choose from Op1. + ShuffleVectorInst *Id3 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C4, CU, C6})); + EXPECT_FALSE(Id3->isIdentity()); + EXPECT_FALSE(Id3->isIdentityWithPadding()); + EXPECT_TRUE(Id3->isIdentityWithExtract()); + delete Id3; + + // Result has less elements than operands; choose from Op0 and Op1 is not identity. + ShuffleVectorInst *Id4 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C4, C1, C6})); + EXPECT_FALSE(Id4->isIdentity()); + EXPECT_FALSE(Id4->isIdentityWithPadding()); + EXPECT_FALSE(Id4->isIdentityWithExtract()); + delete Id4; + + // Result has more elements than operands, and extra elements are undef. + ShuffleVectorInst *Id5 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({CU, C1, C2, C3, CU, CU})); + EXPECT_FALSE(Id5->isIdentity()); + EXPECT_TRUE(Id5->isIdentityWithPadding()); + EXPECT_FALSE(Id5->isIdentityWithExtract()); + delete Id5; + + // Result has more elements than operands, and extra elements are undef; choose from Op1. + ShuffleVectorInst *Id6 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C4, C5, C6, CU, CU, CU})); + EXPECT_FALSE(Id6->isIdentity()); + EXPECT_TRUE(Id6->isIdentityWithPadding()); + EXPECT_FALSE(Id6->isIdentityWithExtract()); + delete Id6; + + // Result has more elements than operands, but extra elements are not undef. + ShuffleVectorInst *Id7 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3, CU, C1})); + EXPECT_FALSE(Id7->isIdentity()); + EXPECT_FALSE(Id7->isIdentityWithPadding()); + EXPECT_FALSE(Id7->isIdentityWithExtract()); + delete Id7; + + // Result has more elements than operands; choose from Op0 and Op1 is not identity. + ShuffleVectorInst *Id8 = new ShuffleVectorInst(ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C0, C1, C2, C3}), + ConstantVector::get({C4, CU, C2, C3, CU, CU})); + EXPECT_FALSE(Id8->isIdentity()); + EXPECT_FALSE(Id8->isIdentityWithPadding()); + EXPECT_FALSE(Id8->isIdentityWithExtract()); + delete Id8; } TEST(InstructionsTest, SkipDebug) {