[Analysis] add optional index parameter to isSplatValue()

[spatel]

Summary

We want to allow splat value transforms to improve PR44588 and related bugs:
https://bugs.llvm.org/show_bug.cgi?id=44588
...but to do that, we need to know if values are splatted from the same, specific index (lane) rather than splatted from an arbitrary index.

We can improve the undef handling with 1-liner follow-ups because the Constant API optionally allow undefs now.

This change will cause a conflict with the planned update of shufflevector's mask in D72467, but that's a small fix-up depending on if it lands first or second.

[nikic]

I don't fully understand how this API is going to be used. The current API makes sense to me (are all elements in the vector the same?) but the new one is unclear, especially for the case where no shufflevector is involved. If I have a splat constant, which index does the value come from?

Do you already have a followup patch that uses this API that you could share?

[spatel]

In D73549#1845221, @nikic wrote:

I don't fully understand how this API is going to be used. The current API makes sense to me (are all elements in the vector the same?) but the new one is unclear, especially for the case where no shufflevector is involved. If I have a splat constant, which index does the value come from?

I agree it's not clear, so I'm certainly open to recommendations for improvement. For now, the index parameter is ignored on a splat constant. The TODO comment suggests a follow-up:

// FIXME: We can allow undefs, but if Index was specified, we may want to
//        check that the constant is defined at that index.

Do you already have a followup patch that uses this API that you could share?

Posted as D73703 - that has a test example where we are splatting from 2 different indexes, so we can't reassociate the binops.

[efriedma]

llvm/include/llvm/Analysis/VectorUtils.h

310

This description is really misleading.

The simple definition of a "splat", if we didn't have poison or undef, would simply be that we have to prove that each element is equal. (The index wouldn't be relevant.)

If you add poison, it becomes something like "Each element is either poisoned or equal to every other non-poisoned element. In addition, if an index is specified, either every element of the vector is poisoned, or the element at that index is not poisoned".

I haven't tried to concisely extend that to include undef, but the description should still look roughly like that.

[spatel]

Patch updated:
Changed documentation comment for isSplatValue().

[efriedma]

llvm/lib/Analysis/VectorUtils.cpp

362

Please use getMaskValue here.

[lebedev.ri]

This seems good to me.

[spatel]

Patch updated:
Use ShuffleVectorInst API to make the code more flexible (should not conflict with D72467); this also allows removing the assert from the previous rev (similar assert is included within getMaskValue()).

[efriedma]

LGTM

llvm/include/llvm/Analysis/VectorUtils.h

 /// from the vector.
 Value *findScalarElement(Value *V, unsigned EltNo);
 
 /// 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 the input value is known to be a vector with all identical
-/// elements (potentially including undefined elements).
+/// 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
+/// poisoned and equal to every other non-poisoned element.
 /// This may be more powerful than the related getSplatValue() because it is
 /// not limited by finding a scalar source value to a splatted vector.
-bool isSplatValue(const Value *V, unsigned Depth = 0);
+bool isSplatValue(const Value *V, int Index = -1, unsigned Depth = 0);
 
 /// Compute a map of integer instructions to their minimum legal type
 /// size.
 ///
 /// C semantics force sub-int-sized values (e.g. i8, i16) to be promoted to int
 /// type (e.g. i32) whenever arithmetic is performed on them.
 ///
 /// For targets with native i8 or i16 operations, usually InstCombine can shrink

llvm/lib/Analysis/VectorUtils.cpp

const llvm::Value *llvm::getSplatValue(const Value *V) {
 
   return nullptr;
 }
 
 // This setting is based on its counterpart in value tracking, but it could be
 // adjusted if needed.
 const unsigned MaxDepth = 6;
 
-bool llvm::isSplatValue(const Value *V, unsigned Depth) {
+bool llvm::isSplatValue(const Value *V, int Index, unsigned Depth) {
   assert(Depth <= MaxDepth && "Limit Search Depth");
 
   if (isa<VectorType>(V->getType())) {
     if (isa<UndefValue>(V))
       return true;
-    // FIXME: Constant splat analysis does not allow undef elements.
+    // FIXME: We can allow undefs, but if Index was specified, we may want to
+    //        check that the constant is defined at that index.
     if (auto *C = dyn_cast<Constant>(V))
       return C->getSplatValue() != nullptr;
   }
 
-  // FIXME: Constant splat analysis does not allow undef elements.
-  Constant *Mask;
-  if (match(V, m_ShuffleVector(m_Value(), m_Value(), m_Constant(Mask))))
-    return Mask->getSplatValue() != nullptr;
+  if (auto *Shuf = dyn_cast<ShuffleVectorInst>(V)) {
+    // FIXME: We can safely allow undefs here. If Index was specified, we will
+    //        check that the mask elt is defined at the required index.
+    if (!Shuf->getMask()->getSplatValue())
+      return false;
+
+    // Match any index.
+    if (Index == -1)
+      return true;
+
+    // Match a specific element. The mask should be defined at and match the
+    // specified index.
+    return Shuf->getMaskValue(Index) == Index;
+  }
 
   // The remaining tests are all recursive, so bail out if we hit the limit.
   if (Depth++ == MaxDepth)
     return false;
 
   // If both operands of a binop are splats, the result is a splat.
   Value *X, *Y, *Z;
   if (match(V, m_BinOp(m_Value(X), m_Value(Y))))
-    return isSplatValue(X, Depth) && isSplatValue(Y, Depth);
+    return isSplatValue(X, Index, Depth) && isSplatValue(Y, Index, Depth);
 
   // If all operands of a select are splats, the result is a splat.
   if (match(V, m_Select(m_Value(X), m_Value(Y), m_Value(Z))))
-    return isSplatValue(X, Depth) && isSplatValue(Y, Depth) &&
-           isSplatValue(Z, Depth);
+    return isSplatValue(X, Index, Depth) && isSplatValue(Y, Index, Depth) &&
+           isSplatValue(Z, Index, Depth);
 
   // TODO: Add support for unary ops (fneg), casts, intrinsics (overflow ops).
 
   return false;
 }
 
 MapVector<Instruction *, uint64_t>
 llvm::computeMinimumValueSizes(ArrayRef<BasicBlock *> Blocks, DemandedBits &DB,

llvm/unittests/Analysis/VectorUtilsTest.cpp

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));
 }
 
+TEST_F(VectorUtilsTest, isSplatValue_00_index0) {
+  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, 0));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_00_index1) {
+  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_FALSE(isSplatValue(A, 1));
+}
+
 TEST_F(VectorUtilsTest, isSplatValue_11) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_TRUE(isSplatValue(A));
 }
 
+TEST_F(VectorUtilsTest, isSplatValue_11_index0) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 0));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_11_index1) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_TRUE(isSplatValue(A, 1));
+}
+
 TEST_F(VectorUtilsTest, isSplatValue_01) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 1>\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_FALSE(isSplatValue(A));
 }
 
-// FIXME: Constant (mask) splat analysis does not allow undef elements.
+TEST_F(VectorUtilsTest, isSplatValue_01_index0) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 1>\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 0));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_01_index1) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 1>\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 1));
+}
+
+// FIXME: Allow undef matching with Constant (mask) splat analysis.
 
 TEST_F(VectorUtilsTest, isSplatValue_0u) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 undef>\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_FALSE(isSplatValue(A));
 }
 
+// FIXME: Allow undef matching with Constant (mask) splat analysis.
+
+TEST_F(VectorUtilsTest, isSplatValue_0u_index0) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 undef>\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 0));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_0u_index1) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %A = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 undef>\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 1));
+}
+
 TEST_F(VectorUtilsTest, isSplatValue_Binop) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %v0 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
       "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
       "  %A = udiv <2 x i8> %v0, %v1\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_TRUE(isSplatValue(A));
 }
 
+TEST_F(VectorUtilsTest, isSplatValue_Binop_index0) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %v0 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
+      "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = udiv <2 x i8> %v0, %v1\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 0));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Binop_index1) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %v0 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
+      "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = udiv <2 x i8> %v0, %v1\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 1));
+}
+
 TEST_F(VectorUtilsTest, isSplatValue_Binop_ConstantOp0) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
       "  %A = ashr <2 x i8> <i8 42, i8 42>, %v1\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_TRUE(isSplatValue(A));
 }
 
+TEST_F(VectorUtilsTest, isSplatValue_Binop_ConstantOp0_index0) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = ashr <2 x i8> <i8 42, i8 42>, %v1\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A, 0));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Binop_ConstantOp0_index1) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i8> %x) {\n"
+      "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = ashr <2 x i8> <i8 42, i8 42>, %v1\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_TRUE(isSplatValue(A, 1));
+}
+
 TEST_F(VectorUtilsTest, isSplatValue_Binop_Not_Op0) {
   parseAssembly(
       "define <2 x i8> @test(<2 x i8> %x) {\n"
       "  %v0 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 0>\n"
       "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
       "  %A = add <2 x i8> %v0, %v1\n"
       "  ret <2 x i8> %A\n"
       "}\n");