[CostModel] Recognise BROADCAST shuffle mask if the elements come from the second src

llvm-svn: 334620

1 files changed, 11 insertions, 4 deletions

diff --git a/llvm/lib/Analysis/TargetTransformInfo.cpp b/llvm/lib/Analysis/TargetTransformInfo.cpp
index a4bb9dd84c9..a470e9765e3 100644
--- a/llvm/lib/Analysis/TargetTransformInfo.cpp
+++ b/llvm/lib/Analysis/TargetTransformInfo.cpp
@@ -652,10 +652,17 @@ static bool isSingleSourceVectorMask(ArrayRef<int> Mask) {
 }
 
 static bool isZeroEltBroadcastVectorMask(ArrayRef<int> Mask) {
-  for (unsigned i = 0; i < Mask.size(); ++i)
-    if (Mask[i] > 0)
-      return false;
-  return true;
+  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) {