[Analysis] add isSplatValue() for vectors in IR

[spatel]

Summary

We have the related getSplatValue() already in IR (see code just above the proposed addition). But sometimes we only need to know that the value is a splat rather than capture the splatted scalar value. Also, we have an isSplatValue() function already in SDAG.

Motivation - recent bugs that would potentially benefit from improved splat analysis in IR:
https://bugs.llvm.org/show_bug.cgi?id=37428
https://bugs.llvm.org/show_bug.cgi?id=42174

[lebedev.ri]

Seems ok to me

llvm/lib/Analysis/VectorUtils.cpp

350

also fneg
also some intrinsics - overflow, saturating, ...

[RKSimon]

llvm/include/llvm/Analysis/VectorUtils.h

82

Make it clear that this is more powerful than getSplatValue ?

[spatel]

Patch updated:

1. Add text to header comment to indicate that this may be more powerful than getSplatValue().
2. Add TODO comment about recursing through several other types of operands.

[shawnl]

Looks fine

llvm/lib/Analysis/VectorUtils.cpp

354

We do not currently normalize vector selects to select.

see https://bugs.llvm.org/show_bug.cgi?id=41777

Should we? (as long as it doesn't require defining undefs?).

[spatel]

We have 2 text comment acknowledgements ("Seems fine", "Looks ok") - thanks for the prompt reviews!...can someone make it official by toggling the Phab state to 'Approved'? :)

llvm/lib/Analysis/VectorUtils.cpp

354

Let me know if I'm not getting the question right, but we have to be careful here to distinguish 'vector select' from 'bitwise select' (xxsel or vsel in PPC).

In IR, vector select is based on vector elements rather than bits:
http://llvm.org/docs/LangRef.html#select-instruction
(so it's more like the x86 'blend' instructions).

So to normalize to the IR vector select, we have to create something like this from a bitwise logic sequence

or (and x, cond), (and y, (not cond)) -->
%bitwise_cond = bitcast <4 x i32> %cond to <128 x i1>
%trueval = bitcast <4 x i32> %x to <128 x i1>
%falseval = bitcast <4 x i32> %y to <128 x i1>
%select = select <128 x i1> %bitwise_cond, <128 x i1> %trueval, <128 x i1> %falseval
%result = bitcast <128 x i1> %select to <4 x i32>

[RKSimon]

LGTM

[nikic]

llvm/lib/Analysis/VectorUtils.cpp

333

This seems potentially problematic, as there's no guarantee that all (independent) undef values will be chosen to be the same value down the line.

I don't know if this is a problem in the context where this is intended to be used though. Is isSplatValue() going to be used heuristically (for profitability) or is it necessary for transformation correctness?

[lebedev.ri]

llvm/lib/Analysis/VectorUtils.cpp

333

This is correct and is in line with other cases like this.
undef here basically means "this can have any value, we don't care", the uses won't care about that lane.
We can define undef to have some value, so we could just fill the undef lines with the splat value, and this get back to the original situation of a splat vector.

[spatel]

llvm/lib/Analysis/VectorUtils.cpp

333

We have both use cases represented in the motivating bugs:

1. PR37428 is a profitability transform that could potentially happen in CGP.
2. PR42174 is a canonicalization transform that could potentially happen in instcombine.

And given that we already have users of getSplatValue(), and it would be difficult to determine the safety on those, I was thinking that we'd add an optional 'AllowUndef' parameter to enable/disable that possibility. This would be similar to the SDAG 'wrapper' variant of this call:

/// Test whether \p V has a splatted value.
bool isSplatValue(SDValue V, bool AllowUndefs = false);

I'm not sure yet if we need/want to go further like the SDAG implementation and actually track the undef/demanded elements:

/// Test whether \p V has a splatted value for all the demanded elements.
///
/// On success \p UndefElts will indicate the elements that have UNDEF
/// values instead of the splat value, this is only guaranteed to be correct
/// for \p DemandedElts.
///
/// NOTE: The function will return true for a demanded splat of UNDEF values.
bool isSplatValue(SDValue V, const APInt &DemandedElts, APInt &UndefElts);

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).
+/// 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);
+
 /// 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
 /// the arithmetic type down again. However InstCombine refuses to create

llvm/lib/Analysis/VectorUtils.cpp

const llvm::Value *llvm::getSplatValue(const Value *V) {
   if (match(V, m_ShuffleVector(m_InsertElement(m_Value(), m_Value(Splat),
                                                m_ZeroInt()),
                                m_Value(), m_ZeroInt())))
     return Splat;
 
   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) {
+  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.
+    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;
+
+  // 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);
+
+  // 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);
+
+  // TODO: Add support for unary ops (fneg), casts, intrinsics (overflow ops).
+
+  return false;
+}
+
 MapVector<Instruction *, uint64_t>
 llvm::computeMinimumValueSizes(ArrayRef<BasicBlock *> Blocks, DemandedBits &DB,
                                const TargetTransformInfo *TTI) {
 
   // DemandedBits will give us every value's live-out bits. But we want
   // to ensure no extra casts would need to be inserted, so every DAG
   // of connected values must have the same minimum bitwidth.
   EquivalenceClasses<Value *> ECs;

llvm/unittests/Analysis/VectorUtilsTest.cpp

 //===- VectorUtilsTest.cpp - VectorUtils tests ------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/NoFolder.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/KnownBits.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
if (A == nullptr)
       report_fatal_error("@test must have an instruction %A");
   }
 
   LLVMContext Context;
   std::unique_ptr<Module> M;
   Instruction *A;
 };
 
+struct BasicTest : public testing::Test {
+  LLVMContext Ctx;
+  std::unique_ptr<Module> M;
+  Function *F;
+  BasicBlock *BB;
+  IRBuilder<NoFolder> IRB;
+
+  BasicTest()
+      : M(new Module("VectorUtils", Ctx)),
+        F(Function::Create(
+            FunctionType::get(Type::getVoidTy(Ctx), /* IsVarArg */ false),
+            Function::ExternalLinkage, "f", M.get())),
+        BB(BasicBlock::Create(Ctx, "entry", F)), IRB(BB) {}
+};
+
+
 } // namespace
 
+TEST_F(BasicTest, isSplat) {
+  Value *UndefVec = UndefValue::get(VectorType::get(IRB.getInt8Ty(), 4));
+  EXPECT_TRUE(isSplatValue(UndefVec));
+
+  Constant *UndefScalar = UndefValue::get(IRB.getInt8Ty());
+  EXPECT_FALSE(isSplatValue(UndefScalar));
+
+  Constant *ScalarC = IRB.getInt8(42);
+  EXPECT_FALSE(isSplatValue(ScalarC));
+
+  Constant *OtherScalarC = IRB.getInt8(-42);
+  Constant *NonSplatC = ConstantVector::get({ScalarC, OtherScalarC});
+  EXPECT_FALSE(isSplatValue(NonSplatC));
+
+  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(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_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_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_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));
+}
+
+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_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_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");
+  EXPECT_FALSE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Binop_Not_Op1) {
+  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 1>\n"
+      "  %v1 = shufflevector <2 x i8> %x, <2 x i8> undef, <2 x i32> <i32 0, i32 1>\n"
+      "  %A = shl <2 x i8> %v0, %v1\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Select) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i1> %x, <2 x i8> %y, <2 x i8> %z) {\n"
+      "  %v0 = shufflevector <2 x i1> %x, <2 x i1> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %v1 = shufflevector <2 x i8> %y, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
+      "  %v2 = shufflevector <2 x i8> %z, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = select <2 x i1> %v0, <2 x i8> %v1, <2 x i8> %v2\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_TRUE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Select_ConstantOp) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i1> %x, <2 x i8> %y, <2 x i8> %z) {\n"
+      "  %v0 = shufflevector <2 x i1> %x, <2 x i1> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %v2 = shufflevector <2 x i8> %z, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = select <2 x i1> %v0, <2 x i8> <i8 42, i8 42>, <2 x i8> %v2\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_TRUE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Select_NotCond) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i1> %x, <2 x i8> %y, <2 x i8> %z) {\n"
+      "  %v1 = shufflevector <2 x i8> %y, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
+      "  %v2 = shufflevector <2 x i8> %z, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = select <2 x i1> %x, <2 x i8> %v1, <2 x i8> %v2\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Select_NotOp1) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i1> %x, <2 x i8> %y, <2 x i8> %z) {\n"
+      "  %v0 = shufflevector <2 x i1> %x, <2 x i1> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %v2 = shufflevector <2 x i8> %z, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %A = select <2 x i1> %v0, <2 x i8> %y, <2 x i8> %v2\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_Select_NotOp2) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i1> %x, <2 x i8> %y, <2 x i8> %z) {\n"
+      "  %v0 = shufflevector <2 x i1> %x, <2 x i1> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %v1 = shufflevector <2 x i8> %y, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
+      "  %A = select <2 x i1> %v0, <2 x i8> %v1, <2 x i8> %z\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_FALSE(isSplatValue(A));
+}
+
+TEST_F(VectorUtilsTest, isSplatValue_SelectBinop) {
+  parseAssembly(
+      "define <2 x i8> @test(<2 x i1> %x, <2 x i8> %y, <2 x i8> %z) {\n"
+      "  %v0 = shufflevector <2 x i1> %x, <2 x i1> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %v1 = shufflevector <2 x i8> %y, <2 x i8> undef, <2 x i32> <i32 0, i32 0>\n"
+      "  %v2 = shufflevector <2 x i8> %z, <2 x i8> undef, <2 x i32> <i32 1, i32 1>\n"
+      "  %bo = xor <2 x i8> %v1, %v2\n"
+      "  %A = select <2 x i1> %v0, <2 x i8> %bo, <2 x i8> %v2\n"
+      "  ret <2 x i8> %A\n"
+      "}\n");
+  EXPECT_TRUE(isSplatValue(A));
+}
+
 TEST_F(VectorUtilsTest, getSplatValueElt0) {
   parseAssembly(
       "define <2 x i8> @test(i8 %x) {\n"
       "  %ins = insertelement <2 x i8> undef, i8 %x, i32 0\n"
       "  %A = shufflevector <2 x i8> %ins, <2 x i8> undef, <2 x i32> zeroinitializer\n"
       "  ret <2 x i8> %A\n"
       "}\n");
   EXPECT_EQ(getSplatValue(A)->getName(), "x");