[IR] add shuffle queries for identity extend/extract

[spatel]

Summary

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

I'm not sure if we have preferred/existing vocabulary for these operations. In the DAG, I think these would be concat_vector with an undef operand or extract_subvector. Alternate naming suggestions welcome.

[lebedev.ri]

I think this looks good. Naming suggestion included.

include/llvm/IR/Instructions.h

2508–2514

isIdentityWithExtend() sounds very much like [[ https://llvm.org/docs/LangRef.html#zext-to-instruction | zext ]]-but-with-undef.
So maybe isIdentityWithExtract() can follow suit and be named isIdentityWithTruncate()?

[spatel]

include/llvm/IR/Instructions.h

2508–2514

Hmm...that actually makes me less positive about 'extend' because we don't want to confuse these with the per-element zext/sext/trunc ops.

How about 'widen' and 'narrow'? That's actually how I've referred to these ops in a draft patch that would use them.

[lebedev.ri]

include/llvm/IR/Instructions.h

2508–2514

Yet another alternative name for isIdentityWithExtend(): isIdentityWithUndefPadding().

[spatel]

include/llvm/IR/Instructions.h

2508–2514

More exact description works for me. Let's go with that unless/until we think of something better.

[lebedev.ri]

include/llvm/IR/Instructions.h

2508–2514

In the same line of thought, isIdentityWithExtract() could be isIdentityPartial().
It won't be identity if the parital-ness is not at the end, so i'd think it is as descriptive.

[spatel]

I messed up linking the commit to the review - it was rL341075.

I went with isIdentityWithPadding() and kept isIdentityWithExtract() for now. I'll post a patch that uses these, and we can adjust the names if those don't feel right.

include/llvm/IR/Instructions.h

public:
 
   SmallVector<int, 16> getShuffleMask() const {
     SmallVector<int, 16> Mask;
     getShuffleMask(Mask);
     return Mask;
   }
 
   /// 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();
     unsigned NumMaskElts = getMask()->getType()->getVectorNumElements();
     return NumSourceElts != NumMaskElts;
   }
 
   /// Return true if this shuffle mask chooses elements from exactly one source
public:
   static bool isIdentityMask(ArrayRef<int> Mask);
   static bool isIdentityMask(const Constant *Mask) {
     assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
     SmallVector<int, 16> MaskAsInts;
     getShuffleMask(Mask, MaskAsInts);
     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 extends exactly one source vector with undefs
+  /// in the high elements.
+  bool isIdentityWithExtend() 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.
   /// Example: <4,1,6,undef>
   /// This returns false if the mask does not choose from both input vectors.
   /// 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).

lib/IR/Instructions.cpp

void ShuffleVectorInst::getShuffleMask(const Constant *Mask,
   }
   for (unsigned i = 0; i != NumElts; ++i) {
     Constant *C = Mask->getAggregateElement(i);
     Result.push_back(isa<UndefValue>(C) ? -1 :
                      cast<ConstantInt>(C)->getZExtValue());
   }
 }
 
-bool ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> Mask) {
+static bool isSingleSourceMaskImpl(ArrayRef<int> 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;
   }
   assert((UsesLHS ^ UsesRHS) && "Should have selected from exactly 1 source");
   return true;
 }
 
-bool ShuffleVectorInst::isIdentityMask(ArrayRef<int> Mask) {
-  if (!isSingleSourceMask(Mask))
+bool ShuffleVectorInst::isSingleSourceMask(ArrayRef<int> 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<int> 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<int> 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<int> Mask) {
   if (!isSingleSourceMask(Mask))
     return false;
   for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) {
     if (Mask[i] == -1)
       continue;
     if (Mask[i] != (NumElts - 1 - i) && Mask[i] != (NumElts + NumElts - 1 - i))
       return false;
if (MaskEltVal == -1)
       return false;
     int MaskEltPrevVal = Mask[i - 2];
     if (MaskEltVal - MaskEltPrevVal != 2)
       return false;
   }
   return true;
 }
 
+bool ShuffleVectorInst::isIdentityWithExtend() 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<int> 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
 //===----------------------------------------------------------------------===//
 
 void InsertValueInst::init(Value *Agg, Value *Val, ArrayRef<unsigned> Idxs,
                            const Twine &Name) {
   assert(getNumOperands() == 2 && "NumOperands not initialized?");

unittests/IR/InstructionsTest.cpp

TEST(InstructionsTest, ShuffleMaskQueries) {
   EXPECT_TRUE(ShuffleVectorInst::isSingleSourceMask(ConstantVector::get({C7, C5, CU, C7})));
   EXPECT_TRUE(ShuffleVectorInst::isSingleSourceMask(ConstantVector::get({C3, C0, CU, C3})));
 
   EXPECT_TRUE(ShuffleVectorInst::isZeroEltSplatMask(ConstantVector::get({C4, CU, CU, C4})));
   EXPECT_TRUE(ShuffleVectorInst::isZeroEltSplatMask(ConstantVector::get({CU, C0, CU, C0})));
 
   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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  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->isIdentityWithExtend());
+  EXPECT_FALSE(Id8->isIdentityWithExtract());
+  delete Id8;
 }
 
 TEST(InstructionsTest, SkipDebug) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C,
                                       R"(
       declare void @llvm.dbg.value(metadata, metadata, metadata)