[Analysis] add query to get splat value from array of ints

[spatel]

Summary

I was debug stepping through an x86 shuffle lowering and noticed we were doing an N^2 search for splat index. I didn't find the equivalent functionality anywhere else in LLVM, so here's a helper that takes an array of int and returns a splatted index while ignoring undefs (any negative value).

This might also be used inside existing ShuffleVectorInst/ShuffleVectorSDNode functions and/or help with D72467.

[RKSimon]

LGTM - The ShuffleVectorSDNode isSplatMask/getSplatIndex helpers could be cleaned up to use this as well.

llvm/include/llvm/Analysis/VectorUtils.h

 /// strides "a[i*stride]". Returns the symbolic stride, or null otherwise.
 Value *getStrideFromPointer(Value *Ptr, ScalarEvolution *SE, Loop *Lp);
 
 /// Given a vector and an element number, see if the scalar value is
 /// already around as a register, for example if it were inserted then extracted
 /// from the vector.
 Value *findScalarElement(Value *V, unsigned EltNo);
 
+/// If all non-negative \p Mask elements are the same value, return that value.
+/// If all elements are negative (undefined) or \p Mask contains different
+/// non-negative values, return -1.
+int getSplatIndex(ArrayRef<int> Mask);
+
 /// Get splat value if the input is a splat vector or return nullptr.
 /// The value may be extracted from a splat constants vector or from
 /// a sequence of instructions that broadcast a single value into a vector.
 const Value *getSplatValue(const Value *V);
 
 /// Return true if each element of the vector value \p V is poisoned or equal to
 /// every other non-poisoned element. If an index element is specified, either
 /// every element of the vector is poisoned or the element at that index is not

llvm/lib/Analysis/VectorUtils.cpp

if (match(V, m_Add(m_Value(Val), m_Constant(C))))
     if (Constant *Elt = C->getAggregateElement(EltNo))
       if (Elt->isNullValue())
         return findScalarElement(Val, EltNo);
 
   // Otherwise, we don't know.
   return nullptr;
 }
 
+int llvm::getSplatIndex(ArrayRef<int> Mask) {
+  int SplatIndex = -1;
+  for (int M : Mask) {
+    // Ignore invalid (undefined) mask elements.
+    if (M < 0)
+      continue;
+
+    // There can be only 1 non-negative mask element value if this is a splat.
+    if (SplatIndex != -1 && SplatIndex != M)
+      return -1;
+
+    // Initialize the splat index to the 1st non-negative mask element.
+    SplatIndex = M;
+  }
+  assert((SplatIndex == -1 || SplatIndex >= 0) && "Negative index?");
+  return SplatIndex;
+}
+
 /// Get splat value if the input is a splat vector or return nullptr.
 /// This function is not fully general. It checks only 2 cases:
 /// the input value is (1) a splat constant vector or (2) a sequence
 /// of instructions that broadcasts a scalar at element 0.
 const llvm::Value *llvm::getSplatValue(const Value *V) {
   if (isa<VectorType>(V->getType()))
     if (auto *C = dyn_cast<Constant>(V))
       return C->getSplatValue();

llvm/unittests/Analysis/VectorUtilsTest.cpp

TEST_F(BasicTest, isSplat) {
   Value *SplatC = IRB.CreateVectorSplat(5, ScalarC);
   EXPECT_TRUE(isSplatValue(SplatC));
 
   // FIXME: Constant splat analysis does not allow undef elements.
   Constant *SplatWithUndefC = ConstantVector::get({ScalarC, UndefScalar});
   EXPECT_FALSE(isSplatValue(SplatWithUndefC));
 }
 
+TEST_F(BasicTest, getSplatIndex) {
+  EXPECT_EQ(getSplatIndex({0,0,0}), 0);
+  EXPECT_EQ(getSplatIndex({1,0,0}), -1);     // no splat
+  EXPECT_EQ(getSplatIndex({0,1,1}), -1);     // no splat
+  EXPECT_EQ(getSplatIndex({42,42,42}), 42);  // array size is independent of splat index
+  EXPECT_EQ(getSplatIndex({42,42,-1}), 42);  // ignore negative
+  EXPECT_EQ(getSplatIndex({-1,42,-1}), 42);  // ignore negatives
+  EXPECT_EQ(getSplatIndex({-4,42,-42}), 42); // ignore all negatives
+  EXPECT_EQ(getSplatIndex({-4,-1,-42}), -1); // all negative values map to -1
+}
+
 TEST_F(VectorUtilsTest, isSplatValue_00) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> zeroinitializer\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_TRUE(isSplatValue(A));
 }