[IR] move shuffle mask queries from TTI to ShuffleVectorInst

[spatel]

Summary

The optimizer is getting smarter (eg, D47986) about differentiating shuffles based on its mask values, so we should make queries on the mask constant operand generally available to avoid code duplication.

We'll probably use this soon in the vectorizers and instcombine (D48023 and https://bugs.llvm.org/show_bug.cgi?id=37806).

We might clean up TTI a bit more once all of its current 'SK_*' options are covered.

[lebedev.ri]

lib/IR/Instructions.cpp

1693–1694 ↗

(On Diff #151550)

Negative values represent undef?

1704–1705 ↗

(On Diff #151550)

Just as a sanity check, i would have expected to see an assert here:

assert((FoundLHS ^ FoundRHS) && "Should have selected from exactly one source");

1735–1738 ↗

(On Diff #151550)

This method was declared with:

/// 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.

You are sure this does what that comment says it does?

1739–1740 ↗

(On Diff #151550)

Just as a sanity check, i would have expected to see an assert here:

assert((ReverseLHS ^ ReverseRHS) && "Should have selected from exactly one source");

1757–1758 ↗

(On Diff #151550)

Just as a sanity check, i would have expected to see an assert here:

assert(!(ShuffleLHS && ShuffleRHS) && "Should not have selected from both sources");

1773–1774 ↗

(On Diff #151550)

Just as a sanity check, i would have expected to see an assert here:

assert((SplatLHS || SplatRHS) && "Should have splatted from at least one of sources");

1783–1787 ↗

(On Diff #151550)

unsigned NumSourceElts = Op<0>()->getType()->getVectorNumElements();

But this does not seem to take the NumSourceElts & 1 into account, like the comment says?

[spatel]

lib/IR/Instructions.cpp

1693–1694 ↗

(On Diff #151550)

Good question/catch! Officially, only -1 represents undef:

/// Return the shuffle mask value for the specified element of the mask.
/// Return -1 if the element is undef.
static int getMaskValue(Constant *Mask, unsigned Elt);

...and I wish we weren't using a magic number rather than a named constant, but that's a separate patch.

This code was copied/translated from the existing code, but I'll fix it.

1704–1705 ↗

(On Diff #151550)

Agreed, and we should assert other assumptions (eg, the input Constant is a vector).

1735–1738 ↗

(On Diff #151550)

This was copied from the existing code, so unless I introduced a typo, that's the correct logic.

Note that this patch is intended to be no-functional-change for any existing users of the TTI code, so there are existing regression tests in place (though if typical, they are far from thorough).

I'll definitely add more unit tests to exercise this further.

1783–1787 ↗

(On Diff #151550)

I don't understand this comment. Can you provide more detail/example of the possible logic bug?

[lebedev.ri]

lib/IR/Instructions.cpp

1704–1705 ↗

(On Diff #151550)

Yes, assertions[/contracts] are nice.

1735–1738 ↗

(On Diff #151550)

I understand that it is meant to be a NFC, but since all the relevant changes
are nicely layed out in a single diff, it is really easy to look them through,
and note possible discrepancies.

Otherwise, if this is fully NFC, you could just commit as is :)

1783–1787 ↗

(On Diff #151550)

I can try.
The comment states:

// 2. The first element of the mask must be either a zero (for the
// even-numbered vector elements) or a one (for the odd-numbered vector
// elements).

But the check is:

if (getMaskValue(Mask, 0) != 0 && getMaskValue(Mask, 0) != 1)
  return false;

I.e. if the first element is not 0 and not 1, we bailout.
But the comment suggests that the logic should be:

unsigned NumSourceElts = Op<0>()->getType()->getVectorNumElements();
if (!(((NumSourceElts % 2 == 0) && (getMaskValue(Mask, 0) == 0)) ||
      ((NumSourceElts % 2 == 1) && (getMaskValue(Mask, 0) == 1))))
    return false;

Is this any clearer?

Or, the comment is wrong.

[lebedev.ri]

lib/IR/Instructions.cpp

1783–1787 ↗

(On Diff #151550)

In other words, either the comment should not talk about even/odd -numbered vector elements,
or the check should actually consider even/odd -ness.

[spatel]

Patch updated:

1. Fixed undef checks to be "x == -1" rather than "x < 0".
2. Added assertions for inputs and logic.
3. Added more unit tests (the last group is based on the mask examples that I added to the header file code comments...would be extra embarrassing to get those wrong!).
4. Reduced confusing comment for transpose (and restated the canonical TRN1 / TRN2 examples for clarity).
5. Changed all 'unsigned' types to 'int' so we don't have to cast (this is always an annoyance with mask values).
6. Changed variable names and added locals to try to make the logic clearer and more uniform.

[lebedev.ri]

Ok, this looks about right.
A few high-level remarks:

• There is an obvious pattern in these is*Mask() functions. This makes them slightly self-duplicating. I wonder if it would be possible to refactor them somehow to reduce duplication.
• There is some correlation between those functions. I would personally add an exhaustive test (with googletest's Combine, if it is not disabled?) to check for those correlations.

include/llvm/IR/Instructions.h

2467 ↗

(On Diff #151644)

All these functions should be able to take const Constant *Mask.

lib/IR/Instructions.cpp

1740–1741 ↗

(On Diff #151644)

This is slightly inconsistent, why not use helper variables here?

bool IsLHS = MaskEltVal == (LastEltIndex - i);
bool IsRHS = MaskEltVal == (NumElts + LastEltIndex - i);
UsesLHS &= IsLHS;
UsesRHS &= IsRHS;

[RKSimon]

I haven't had a chance to go through this yet, but it'd be useful to be able to call these through ArrayRef<int> style shuffle masks (see the TODO I added for SLP to D48023)

Plus, we probably need to think about input/output size mismatches a little more - Identity and Broadcast for instance don't really care - adding insert/extract subvector pattern matches would be a start.

[spatel]

Patch updated:
Redefine all of the Constant* functions in terms of ArrayRef<int> parent implementations. This makes the code slightly cleaner and might open up the possibility of removing some redundant code from the backend.

I know we can do better to share code, test, and add functionality, but I'd like to get this committed so I'm not holding up functional changes and existing cleanup that is gated on this first step.

[spatel]

Uploaded the wrong draft - went overboard on the 'const &' on that last one. An ArrayRef is a const ref.

[lebedev.ri]

include/llvm/IR/Instructions.h

2467 ↗

(On Diff #151650)

i think it will be less confusing to replace -1 with undef.

2471 ↗

(On Diff #151650)

From where does this 16 come from?
Magical expected maximal number of elements in one vector? (AVX512 / 16 elements = 32 bit/element)
Might be better to static constexpr int it beforehand.

2487 ↗

(On Diff #151650)

i think it will be less confusing to replace -1 with undef.

2513 ↗

(On Diff #151650)

i think it will be less confusing to replace -1 with undef.

2532 ↗

(On Diff #151650)

i think it will be less confusing to replace -1 with undef.

2551 ↗

(On Diff #151650)

i think it will be less confusing to replace -1 with undef.

lib/IR/Instructions.cpp

1726–1743 ↗

(On Diff #151650)

Hm, why is this so different from isIdentityMask()?
I'm not sure why it couldn't be a copy of isIdentityMask() with just changed Is{LR}HS logic.

1763–1777 ↗

(On Diff #151650)

Same here. The foundation looks very similar, yet the approach seems different.

[RKSimon]

LGTM with a slight addition to the comments - SLP for example supports alternate/permute cases for different sized inputs

include/llvm/IR/Instructions.h

2514 ↗

(On Diff #151650)

Make this comment in each case a TODO - at some point we're going to have to (optionally) support length changing shuffle patterns

[spatel]

include/llvm/IR/Instructions.h

2471 ↗

(On Diff #151650)

Yes - copy/paste of magic constant. This is typical SmallVector usage throughout LLVM. OK, if we make that a follow-up cleanup and zap all of them (at least in this file)?

lib/IR/Instructions.cpp

1726–1743 ↗

(On Diff #151650)

The symmetry will be clear if we have everyone call isSingleSourceMask as a preliminary step. Then the basic logic will be identical across everything except for the isTranspose oddball.

[spatel]

Patch updated:

1. Improved header comment examples (use 'undef' rather than -1).
2. Reimplemented logic to maximize code reuse (make use of isSingleSourceMask() internally).
3. Added 'TODO' notes about allowing size-changing shuffles.

[lebedev.ri]

Nice!
I like this.

include/llvm/IR/Instructions.h

2471 ↗

(On Diff #151650)

Yep, absolutely, just pointing out.

lib/IR/Instructions.cpp

1691 ↗

(On Diff #151942)

Oh, you can define multiple variables in init of the for loop?
Somehow i never realized that..

[spatel]

lib/IR/Instructions.cpp

1691 ↗

(On Diff #151942)

Yep - http://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop

llvm/trunk/include/llvm/IR/Instructions.h

public:
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Value);
 
   Constant *getMask() const {
     return cast<Constant>(getOperand(2));
   }
 
   /// Return the shuffle mask value for the specified element of the mask.
   /// Return -1 if the element is undef.
-  static int getMaskValue(Constant *Mask, unsigned Elt);
+  static int getMaskValue(const Constant *Mask, unsigned Elt);
 
   /// Return the shuffle mask value of this instruction for the given element
   /// index. Return -1 if the element is undef.
   int getMaskValue(unsigned Elt) const {
     return getMaskValue(getMask(), Elt);
   }
 
   /// Convert the input shuffle mask operand to a vector of integers. Undefined
   /// elements of the mask are returned as -1.
-  static void getShuffleMask(Constant *Mask, SmallVectorImpl<int> &Result);
+  static void getShuffleMask(const Constant *Mask,
+                             SmallVectorImpl<int> &Result);
 
   /// Return the mask for this instruction as a vector of integers. Undefined
   /// elements of the mask are returned as -1.
   void getShuffleMask(SmallVectorImpl<int> &Result) const {
     return getShuffleMask(getMask(), Result);
   }
 
   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.
+  /// 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
+  /// 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) {
+    assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
+    SmallVector<int, 16> MaskAsInts;
+    getShuffleMask(Mask, MaskAsInts);
+    return isSingleSourceMask(MaskAsInts);
+  }
+
+  /// 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(getMask());
+  }
+
+  /// 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) {
+    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
+  /// 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 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).
+  static bool isSelectMask(ArrayRef<int> Mask);
+  static bool isSelectMask(const Constant *Mask) {
+    assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
+    SmallVector<int, 16> MaskAsInts;
+    getShuffleMask(Mask, MaskAsInts);
+    return isSelectMask(MaskAsInts);
+  }
+
+  /// Return true if this shuffle chooses elements from its source vectors
+  /// without lane crossings and all operands have the same number of elements.
+  /// In other words, this shuffle is equivalent to a vector select with a
+  /// constant condition operand.
+  /// Example: shufflevector <4 x n> A, <4 x n> B, <undef,1,6,3>
+  /// 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.
+  /// TODO: Optionally allow length-changing shuffles.
+  bool isSelect() const {
+    return !changesLength() && isSelectMask(getMask());
+  }
+
+  /// 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) {
+    assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
+    SmallVector<int, 16> MaskAsInts;
+    getShuffleMask(Mask, MaskAsInts);
+    return isReverseMask(MaskAsInts);
+  }
+
+  /// 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(getMask());
+  }
+
+  /// 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) {
+    assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
+    SmallVector<int, 16> MaskAsInts;
+    getShuffleMask(Mask, MaskAsInts);
+    return isZeroEltSplatMask(MaskAsInts);
+  }
+
+  /// Return true if all elements of this shuffle are the same value as the
+  /// first element of exactly one source vector without changing the length
+  /// of that vector.
+  /// Example: shufflevector <4 x n> A, <4 x n> B, <undef,0,undef,0>
+  /// TODO: Optionally allow length-changing shuffles.
+  /// TODO: Optionally allow splats from other elements.
+  bool isZeroEltSplat() const {
+    return !changesLength() && isZeroEltSplatMask(getMask());
+  }
+
+  /// Return true if this shuffle mask is a transpose mask.
+  /// Transpose vector masks transpose a 2xn matrix. They read corresponding
+  /// even- or odd-numbered vector elements from two n-dimensional source
+  /// vectors and write each result into consecutive elements of an
+  /// n-dimensional destination vector. Two shuffles are necessary to complete
+  /// the transpose, one for the even elements and another for the odd elements.
+  /// This description closely follows how the TRN1 and TRN2 AArch64
+  /// instructions operate.
+  ///
+  /// For example, a simple 2x2 matrix can be transposed with:
+  ///
+  ///   ; Original matrix
+  ///   m0 = <a, b>
+  ///   m1 = <c, d>
+  ///
+  ///   ; Transposed matrix
+  ///   t0 = <a, c> = shufflevector m0, m1, <0, 2>
+  ///   t1 = <b, d> = shufflevector m0, m1, <1, 3>
+  ///
+  /// For matrices having greater than n columns, the resulting nx2 transposed
+  /// matrix is stored in two result vectors such that one vector contains
+  /// interleaved elements from all the even-numbered rows and the other vector
+  /// contains interleaved elements from all the odd-numbered rows. For example,
+  /// a 2x4 matrix can be transposed with:
+  ///
+  ///   ; Original matrix
+  ///   m0 = <a, b, c, d>
+  ///   m1 = <e, f, g, h>
+  ///
+  ///   ; 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) {
+    assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
+    SmallVector<int, 16> MaskAsInts;
+    getShuffleMask(Mask, MaskAsInts);
+    return isTransposeMask(MaskAsInts);
+  }
+
+  /// Return true if this shuffle transposes the elements of its inputs without
+  /// changing the length of the vectors. This operation may also be known as a
+  /// merge or interleave. See the description for isTransposeMask() for the
+  /// exact specification.
+  /// Example: shufflevector <4 x n> A, <4 x n> B, <0,4,2,6>
+  bool isTranspose() const {
+    return !changesLength() && isTransposeMask(getMask());
+  }
+
   /// Change values in a shuffle permute mask assuming the two vector operands
   /// of length InVecNumElts have swapped position.
   static void commuteShuffleMask(MutableArrayRef<int> Mask,
                                  unsigned InVecNumElts) {
     for (int &Idx : Mask) {
       if (Idx == -1)
         continue;
       Idx = Idx < (int)InVecNumElts ? Idx + InVecNumElts : Idx - InVecNumElts;

llvm/trunk/lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::shouldExpandReduction(const IntrinsicInst *II) const {
   return TTIImpl->shouldExpandReduction(II);
 }
 
 int TargetTransformInfo::getInstructionLatency(const Instruction *I) const {
   return TTIImpl->getInstructionLatency(I);
 }
 
-static bool isReverseVectorMask(ArrayRef<int> Mask) {
-  bool ReverseLHS = true;
-  bool ReverseRHS = true;
-  unsigned MaskSize = Mask.size();
-
-  for (unsigned i = 0; i < MaskSize && (ReverseLHS || ReverseRHS); ++i) {
-    if (Mask[i] < 0)
-      continue;
-    ReverseLHS &= (Mask[i] == (int)(MaskSize - 1 - i));
-    ReverseRHS &= (Mask[i] == (int)(MaskSize + MaskSize - 1 - i));
-  }
-  return ReverseLHS || ReverseRHS;
-}
-
-static bool isSingleSourceVectorMask(ArrayRef<int> Mask) {
-  bool ShuffleLHS = false;
-  bool ShuffleRHS = false;
-  unsigned MaskSize = Mask.size();
-
-  for (unsigned i = 0; i < MaskSize && !(ShuffleLHS && ShuffleRHS); ++i) {
-    if (Mask[i] < 0)
-      continue;
-    if ((unsigned)Mask[i] >= MaskSize)
-      ShuffleRHS = true;
-    else
-      ShuffleLHS = true;
-  }
-  return !(ShuffleLHS && ShuffleRHS);
-}
-
-static bool isZeroEltBroadcastVectorMask(ArrayRef<int> Mask) {
-  bool BroadcastLHS = true;
-  bool BroadcastRHS = true;
-  unsigned MaskSize = Mask.size();
-
-  for (unsigned i = 0; i < MaskSize && (BroadcastLHS || BroadcastRHS); ++i) {
-    if (Mask[i] < 0)
-      continue;
-    BroadcastLHS &= (Mask[i] == 0);
-    BroadcastRHS &= (Mask[i] == (int)MaskSize);
-  }
-  return BroadcastLHS || BroadcastRHS;
-}
-
-static bool isIdentityVectorMask(ArrayRef<int> Mask) {
-  bool IdentityLHS = true;
-  bool IdentityRHS = true;
-  unsigned MaskSize = Mask.size();
-
-  // Example: shufflevector A, B, <0,1,u,3>
-  // Example: shufflevector A, B, <4,u,6,u>
-  for (unsigned i = 0; i < MaskSize && (IdentityLHS || IdentityRHS); ++i) {
-    if (Mask[i] < 0)
-      continue;
-    IdentityLHS &= (Mask[i] == (int)i);
-    IdentityRHS &= (Mask[i] == (int)(i + MaskSize));
-  }
-  return IdentityLHS || IdentityRHS;
-}
-
-static bool isSelectVectorMask(ArrayRef<int> Mask) {
-  bool IsSelect = true;
-  bool FoundLHS = false;
-  bool FoundRHS = false;
-  unsigned MaskSize = Mask.size();
-
-  // Example: shufflevector A, B, <0,1,6,3>
-  // Example: shufflevector A, B, <4,1,6,3>
-  for (unsigned i = 0; i < MaskSize && IsSelect; ++i) {
-    if (Mask[i] < 0)
-      continue;
-    bool IsLHS = (Mask[i] == (int)i);
-    bool IsRHS = (Mask[i] == (int)(i + MaskSize));
-    FoundLHS |= IsLHS;
-    FoundRHS |= IsRHS;
-    IsSelect = IsLHS || IsRHS;
-  }
-  // If we don't use both vectors this is really an Identity mask.
-  return IsSelect && FoundLHS && FoundRHS;
-}
-
-static bool isTransposeVectorMask(ArrayRef<int> Mask) {
-  // Transpose vector masks transpose a 2xn matrix. They read corresponding
-  // even- or odd-numbered vector elements from two n-dimensional source
-  // vectors and write each result into consecutive elements of an
-  // n-dimensional destination vector. Two shuffles are necessary to complete
-  // the transpose, one for the even elements and another for the odd elements.
-  // This description closely follows how the TRN1 and TRN2 AArch64
-  // instructions operate.
-  //
-  // For example, a simple 2x2 matrix can be transposed with:
-  //
-  //   ; Original matrix
-  //   m0 = <a, b>
-  //   m1 = <c, d>
-  //
-  //   ; Transposed matrix
-  //   t0 = <a, c> = shufflevector m0, m1, <0, 2>
-  //   t1 = <b, d> = shufflevector m0, m1, <1, 3>
-  //
-  // For matrices having greater than n columns, the resulting nx2 transposed
-  // matrix is stored in two result vectors such that one vector contains
-  // interleaved elements from all the even-numbered rows and the other vector
-  // contains interleaved elements from all the odd-numbered rows. For example,
-  // a 2x4 matrix can be transposed with:
-  //
-  //   ; Original matrix
-  //   m0 = <a, b, c, d>
-  //   m1 = <e, f, g, h>
-  //
-  //   ; 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>
-  //
-  // The above explanation places limitations on what valid transpose masks can
-  // look like. These limitations are defined by the checks below.
-  //
-  // 1. The number of elements in the mask must be a power of two.
-  if (!isPowerOf2_32(Mask.size()))
-    return false;
-
-  // 2. The first element of the mask must be either a zero (for the
-  // even-numbered vector elements) or a one (for the odd-numbered vector
-  // elements).
-  if (Mask[0] != 0 && Mask[0] != 1)
-    return false;
-
-  // 3. The difference between the first two elements must be equal to the
-  // number of elements in the mask.
-  if (Mask[1] - Mask[0] != (int)Mask.size())
-    return false;
-
-  // 4. The difference between consecutive even-numbered and odd-numbered
-  // elements must be equal to two.
-  for (int I = 2; I < (int)Mask.size(); ++I)
-    if (Mask[I] - Mask[I - 2] != 2)
-      return false;
-
-  return true;
-}
-
 static TargetTransformInfo::OperandValueKind
 getOperandInfo(Value *V, TargetTransformInfo::OperandValueProperties &OpProps) {
   TargetTransformInfo::OperandValueKind OpInfo =
       TargetTransformInfo::OK_AnyValue;
   OpProps = TargetTransformInfo::OP_None;
 
   const Value *Splat = getSplatValue(V);
 
case Instruction::InsertElement: {
     unsigned Idx = -1; 
     if (CI)
       Idx = CI->getZExtValue();
     return getVectorInstrCost(I->getOpcode(),
                                    IE->getType(), Idx);
   }
   case Instruction::ShuffleVector: {
     const ShuffleVectorInst *Shuffle = cast<ShuffleVectorInst>(I);
-    Type *VecTypOp0 = Shuffle->getOperand(0)->getType();
-    unsigned NumVecElems = VecTypOp0->getVectorNumElements();
-    SmallVector<int, 16> Mask = Shuffle->getShuffleMask();
+    // TODO: Identify and add costs for insert/extract subvector, etc.
+    if (Shuffle->changesLength())
+      return -1;
     
-    if (NumVecElems == Mask.size()) {
-      if (isIdentityVectorMask(Mask))
-        return 0;
+    if (Shuffle->isIdentity())
 
-      if (isReverseVectorMask(Mask))
-        return TTIImpl->getShuffleCost(TargetTransformInfo::SK_Reverse,
-                                       VecTypOp0, 0, nullptr);
-
-      if (isSelectVectorMask(Mask))
-        return TTIImpl->getShuffleCost(TargetTransformInfo::SK_Select,
-                                       VecTypOp0, 0, nullptr);
-
-      if (isTransposeVectorMask(Mask))
-        return TTIImpl->getShuffleCost(TargetTransformInfo::SK_Transpose,
-                                       VecTypOp0, 0, nullptr);
-
-      if (isZeroEltBroadcastVectorMask(Mask))
-        return TTIImpl->getShuffleCost(TargetTransformInfo::SK_Broadcast,
-                                       VecTypOp0, 0, nullptr);
-
-      if (isSingleSourceVectorMask(Mask))
-        return TTIImpl->getShuffleCost(TargetTransformInfo::SK_PermuteSingleSrc,
-                                       VecTypOp0, 0, nullptr);
+    Type *Ty = Shuffle->getType();
+    if (Shuffle->isReverse())
+      return TTIImpl->getShuffleCost(SK_Reverse, Ty, 0, nullptr);
 
-      return TTIImpl->getShuffleCost(TargetTransformInfo::SK_PermuteTwoSrc,
-                                     VecTypOp0, 0, nullptr);
-    }
+    if (Shuffle->isSelect())
+      return TTIImpl->getShuffleCost(SK_Select, Ty, 0, nullptr);
 
-    return -1;
+    if (Shuffle->isTranspose())
+      return TTIImpl->getShuffleCost(SK_Transpose, Ty, 0, nullptr);
+
+    if (Shuffle->isZeroEltSplat())
+      return TTIImpl->getShuffleCost(SK_Broadcast, Ty, 0, nullptr);
+
+    if (Shuffle->isSingleSource())
+      return TTIImpl->getShuffleCost(SK_PermuteSingleSrc, Ty, 0, nullptr);
+
+    return TTIImpl->getShuffleCost(SK_PermuteTwoSrc, Ty, 0, nullptr);
   }
   case Instruction::Call:
     if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
       SmallVector<Value *, 4> Args(II->arg_operands());
 
       FastMathFlags FMF;
       if (auto *FPMO = dyn_cast<FPMathOperator>(II))
         FMF = FPMO->getFastMathFlags();

llvm/trunk/lib/IR/Instructions.cpp

bool ShuffleVectorInst::isValidOperands(const Value *V1, const Value *V2,
   // ugliness to allow such a mask pass.
   if (const auto *CE = dyn_cast<ConstantExpr>(Mask))
     if (CE->getOpcode() == Instruction::UserOp1)
       return true;
 
   return false;
 }
 
-int ShuffleVectorInst::getMaskValue(Constant *Mask, unsigned i) {
+int ShuffleVectorInst::getMaskValue(const Constant *Mask, unsigned i) {
   assert(i < Mask->getType()->getVectorNumElements() && "Index out of range");
   if (auto *CDS = dyn_cast<ConstantDataSequential>(Mask))
     return CDS->getElementAsInteger(i);
   Constant *C = Mask->getAggregateElement(i);
   if (isa<UndefValue>(C))
     return -1;
   return cast<ConstantInt>(C)->getZExtValue();
 }
 
-void ShuffleVectorInst::getShuffleMask(Constant *Mask,
+void ShuffleVectorInst::getShuffleMask(const Constant *Mask,
                                        SmallVectorImpl<int> &Result) {
   unsigned NumElts = Mask->getType()->getVectorNumElements();
   
   if (auto *CDS = dyn_cast<ConstantDataSequential>(Mask)) {
     for (unsigned i = 0; i != NumElts; ++i)
       Result.push_back(CDS->getElementAsInteger(i));
     return;
   }    
   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) {
+  assert(!Mask.empty() && "Shuffle mask must contain elements");
+  bool UsesLHS = false;
+  bool UsesRHS = false;
+  for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) {
+    if (Mask[i] == -1)
+      continue;
+    assert(Mask[i] >= 0 && Mask[i] < (NumElts * 2) &&
+           "Out-of-bounds shuffle mask element");
+    UsesLHS |= (Mask[i] < NumElts);
+    UsesRHS |= (Mask[i] >= NumElts);
+    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))
+    return false;
+  for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) {
+    if (Mask[i] == -1)
+      continue;
+    if (Mask[i] != i && Mask[i] != (NumElts + i))
+      return false;
+  }
+  return true;
+}
+
+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;
+  }
+  return true;
+}
+
+bool ShuffleVectorInst::isZeroEltSplatMask(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] != 0 && Mask[i] != NumElts)
+      return false;
+  }
+  return true;
+}
+
+bool ShuffleVectorInst::isSelectMask(ArrayRef<int> Mask) {
+  // Select is differentiated from identity. It requires using both sources.
+  if (isSingleSourceMask(Mask))
+    return false;
+  for (int i = 0, NumElts = Mask.size(); i < NumElts; ++i) {
+    if (Mask[i] == -1)
+      continue;
+    if (Mask[i] != i && Mask[i] != (NumElts + i))
+      return false;
+  }
+  return true;
+}
+
+bool ShuffleVectorInst::isTransposeMask(ArrayRef<int> Mask) {
+  // 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>
+
+  // 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))
+    return false;
+
+  // 2. The first element of the mask must be either a 0 or a 1.
+  if (Mask[0] != 0 && Mask[0] != 1)
+    return false;
+
+  // 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)
+    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];
+    if (MaskEltVal == -1)
+      return false;
+    int MaskEltPrevVal = Mask[i - 2];
+    if (MaskEltVal - MaskEltPrevVal != 2)
+      return false;
+  }
+  return true;
+}
+
+
 //===----------------------------------------------------------------------===//
 //                             InsertValueInst Class
 //===----------------------------------------------------------------------===//
 
 void InsertValueInst::init(Value *Agg, Value *Val, ArrayRef<unsigned> Idxs, 
                            const Twine &Name) {
   assert(getNumOperands() == 2 && "NumOperands not initialized?");
 

llvm/trunk/unittests/IR/InstructionsTest.cpp

 }
 
 TEST(InstructionsTest, CommuteShuffleMask) {
   SmallVector<int, 16> Indices({-1, 0, 7});
   ShuffleVectorInst::commuteShuffleMask(Indices, 4);
   EXPECT_THAT(Indices, testing::ContainerEq(ArrayRef<int>({-1, 4, 3})));
 }
 
+TEST(InstructionsTest, ShuffleMaskQueries) {
+  // Create the elements for various constant vectors.
+  LLVMContext Ctx;
+  Type *Int32Ty = Type::getInt32Ty(Ctx);
+  Constant *CU = UndefValue::get(Int32Ty);
+  Constant *C0 = ConstantInt::get(Int32Ty, 0);
+  Constant *C1 = ConstantInt::get(Int32Ty, 1);
+  Constant *C2 = ConstantInt::get(Int32Ty, 2);
+  Constant *C3 = ConstantInt::get(Int32Ty, 3);
+  Constant *C4 = ConstantInt::get(Int32Ty, 4);
+  Constant *C5 = ConstantInt::get(Int32Ty, 5);
+  Constant *C6 = ConstantInt::get(Int32Ty, 6);
+  Constant *C7 = ConstantInt::get(Int32Ty, 7);
+
+  Constant *Identity = ConstantVector::get({C0, CU, C2, C3, C4});
+  EXPECT_TRUE(ShuffleVectorInst::isIdentityMask(Identity));
+  EXPECT_FALSE(ShuffleVectorInst::isSelectMask(Identity)); // identity is distinguished from select
+  EXPECT_FALSE(ShuffleVectorInst::isReverseMask(Identity));
+  EXPECT_TRUE(ShuffleVectorInst::isSingleSourceMask(Identity)); // identity is always single source
+  EXPECT_FALSE(ShuffleVectorInst::isZeroEltSplatMask(Identity));
+  EXPECT_FALSE(ShuffleVectorInst::isTransposeMask(Identity));
+
+  Constant *Select = ConstantVector::get({CU, C1, C5});
+  EXPECT_FALSE(ShuffleVectorInst::isIdentityMask(Select));
+  EXPECT_TRUE(ShuffleVectorInst::isSelectMask(Select));
+  EXPECT_FALSE(ShuffleVectorInst::isReverseMask(Select));
+  EXPECT_FALSE(ShuffleVectorInst::isSingleSourceMask(Select));
+  EXPECT_FALSE(ShuffleVectorInst::isZeroEltSplatMask(Select));
+  EXPECT_FALSE(ShuffleVectorInst::isTransposeMask(Select));
+  
+  Constant *Reverse = ConstantVector::get({C3, C2, C1, CU});
+  EXPECT_FALSE(ShuffleVectorInst::isIdentityMask(Reverse));
+  EXPECT_FALSE(ShuffleVectorInst::isSelectMask(Reverse));
+  EXPECT_TRUE(ShuffleVectorInst::isReverseMask(Reverse));
+  EXPECT_TRUE(ShuffleVectorInst::isSingleSourceMask(Reverse)); // reverse is always single source
+  EXPECT_FALSE(ShuffleVectorInst::isZeroEltSplatMask(Reverse));
+  EXPECT_FALSE(ShuffleVectorInst::isTransposeMask(Reverse));
+
+  Constant *SingleSource = ConstantVector::get({C2, C2, C0, CU});
+  EXPECT_FALSE(ShuffleVectorInst::isIdentityMask(SingleSource));
+  EXPECT_FALSE(ShuffleVectorInst::isSelectMask(SingleSource));
+  EXPECT_FALSE(ShuffleVectorInst::isReverseMask(SingleSource));
+  EXPECT_TRUE(ShuffleVectorInst::isSingleSourceMask(SingleSource));
+  EXPECT_FALSE(ShuffleVectorInst::isZeroEltSplatMask(SingleSource));
+  EXPECT_FALSE(ShuffleVectorInst::isTransposeMask(SingleSource));
+
+  Constant *ZeroEltSplat = ConstantVector::get({C0, C0, CU, C0});
+  EXPECT_FALSE(ShuffleVectorInst::isIdentityMask(ZeroEltSplat));
+  EXPECT_FALSE(ShuffleVectorInst::isSelectMask(ZeroEltSplat));
+  EXPECT_FALSE(ShuffleVectorInst::isReverseMask(ZeroEltSplat));
+  EXPECT_TRUE(ShuffleVectorInst::isSingleSourceMask(ZeroEltSplat)); // 0-splat is always single source
+  EXPECT_TRUE(ShuffleVectorInst::isZeroEltSplatMask(ZeroEltSplat));
+  EXPECT_FALSE(ShuffleVectorInst::isTransposeMask(ZeroEltSplat));
+
+  Constant *Transpose = ConstantVector::get({C0, C4, C2, C6});
+  EXPECT_FALSE(ShuffleVectorInst::isIdentityMask(Transpose));
+  EXPECT_FALSE(ShuffleVectorInst::isSelectMask(Transpose));
+  EXPECT_FALSE(ShuffleVectorInst::isReverseMask(Transpose));
+  EXPECT_FALSE(ShuffleVectorInst::isSingleSourceMask(Transpose));
+  EXPECT_FALSE(ShuffleVectorInst::isZeroEltSplatMask(Transpose));
+  EXPECT_TRUE(ShuffleVectorInst::isTransposeMask(Transpose));
+
+  // More tests to make sure the logic is/stays correct...
+  EXPECT_TRUE(ShuffleVectorInst::isIdentityMask(ConstantVector::get({CU, C1, CU, C3})));
+  EXPECT_TRUE(ShuffleVectorInst::isIdentityMask(ConstantVector::get({C4, CU, C6, CU})));
+
+  EXPECT_TRUE(ShuffleVectorInst::isSelectMask(ConstantVector::get({C4, C1, C6, CU})));
+  EXPECT_TRUE(ShuffleVectorInst::isSelectMask(ConstantVector::get({CU, C1, C6, C3})));
+
+  EXPECT_TRUE(ShuffleVectorInst::isReverseMask(ConstantVector::get({C7, C6, CU, C4})));
+  EXPECT_TRUE(ShuffleVectorInst::isReverseMask(ConstantVector::get({C3, CU, C1, CU})));
+
+  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})));
+}
+
 } // end anonymous namespace
 } // end namespace llvm