[IR][ShuffleVector] Introduce `isReplicationMask()` matcher

[lebedev.ri]

Summary

Avid readers of this saga may recall from previous installments,
that replication mask replicates (lol) each of the VF elements
in a vector ReplicationFactor times. For example, the mask for
ReplicationFactor=3 and VF=4 is: <0,0,0,1,1,1,2,2,2,3,3,3>.
More importantly, replication mask is used by LoopVectorizer
when using masked interleaved memory operations.

As discussed in previous installments, while it is used by LV,
and we seem to support masked interleaved memory operations on X86,
it's support in cost model leaves a lot to be desired:
until basically yesterday even for AVX512 we had no cost model for it.

As it has been witnessed in the recent AVX2 X86TTIImpl::getInterleavedMemoryOpCost()
costmodel patches, while it is hard-enough to query the cost
of a particular assembly sequence [from llvm-mca],
afterwards the check lines LV costmodel tests must be updated manually.
This is, at the very least, boring.

Okay, now we have decent costmodel coverage for interleaving shuffles,
but now basically the same mind-killing sequence has to be performed
for replication mask. I think we can improve at least the second half
of the problem, by teaching the TargetTransformInfoImplCRTPBase::getUserCost()
to recognize Instruction::ShuffleVector that are repetition masks,
adding exhaustive test coverage using -cost-model -analyze + utils/update_analyze_test_checks.py

This way we can have good exhaustive coverage for cost model,
and only basic coverage for the LV costmodel.

This patch adds precise undef-aware isReplicationMask(), with exhaustive test coverage.
InstructionsTest.ShuffleMaskIsReplicationMask shows that it correctly detects all the known masks.
InstructionsTest.ShuffleMaskIsReplicationMask_Exhaustive_Correctness shows that if
we detected the replication mask with given params, then if we actually generate
a true replication mask with said params, it matches element-wise ignoring undef mask elements.

[lebedev.ri]

NFC, improve performance of isReplicationMask() matcher,
~halves runtime of InstructionsTest.ShuffleMaskIsReplicationMask_Exhaustive_Correctness test.

[lebedev.ri]

NFC, add an explicit test that changing mask elements to undef still results in a match (but not necessarily same tuple guess!).

[spatel]

LGTM - I don't have any experience with the test generator (so might want to wait on a 2nd opinion), but the code looks right.

llvm/include/llvm/IR/Instructions.h

2375

an -> a

[lebedev.ri]

In D113214#3111481, @spatel wrote:

LGTM

Thank you for the review!

I don't have any experience with the test generator (so might want to wait on a 2nd opinion), but the code looks right.

Which test generator, CombinationGenerator?

[spatel]

In D113214#3111487, @lebedev.ri wrote:

In D113214#3111481, @spatel wrote:

LGTM

Thank you for the review!

I don't have any experience with the test generator (so might want to wait on a 2nd opinion), but the code looks right.

Which test generator, CombinationGenerator?

Yes - I see that it's coming from exegesis, but I'm not familiar with that code.

[lebedev.ri]

In D113214#3111532, @spatel wrote:

In D113214#3111487, @lebedev.ri wrote:

In D113214#3111481, @spatel wrote:

LGTM

Thank you for the review!

I don't have any experience with the test generator (so might want to wait on a 2nd opinion), but the code looks right.

Which test generator, CombinationGenerator?

Yes - I see that it's coming from exegesis, but I'm not familiar with that code.

Ah, that's fine, i convinced the original author to write these new tests.

[rupprecht]

llvm/unittests/IR/InstructionsTest.cpp

1133

These two exhaustive test cases take an extremely long time to test (3+ minutes for each one). Would it make sense to use #ifdef EXPENSIVE_CHECKS for these two?

[lebedev.ri]

llvm/unittests/IR/InstructionsTest.cpp

1133

They take 1 and 2 seconds for me.
How long do tests in ConstantRangeTest.cpp take for you?

[rupprecht]

llvm/unittests/IR/InstructionsTest.cpp

1133

Weird. ConstantRangeTest takes 23 seconds in total, most test cases being <.5s but a few being ~2 seconds.

$ ./unittests/IR/IRTests --gtest_filter='InstructionsTest.ShuffleMask*'
Note: Google Test filter = InstructionsTest.ShuffleMask*
[==========] Running 4 tests from 1 test suite.
[----------] Global test environment set-up.
[----------] 4 tests from InstructionsTest
[ RUN      ] InstructionsTest.ShuffleMaskQueries
[       OK ] InstructionsTest.ShuffleMaskQueries (2 ms)
[ RUN      ] InstructionsTest.ShuffleMaskIsReplicationMask
[       OK ] InstructionsTest.ShuffleMaskIsReplicationMask (0 ms)
[ RUN      ] InstructionsTest.ShuffleMaskIsReplicationMask_undef
[       OK ] InstructionsTest.ShuffleMaskIsReplicationMask_undef (227553 ms)
[ RUN      ] InstructionsTest.ShuffleMaskIsReplicationMask_Exhaustive_Correctness
[       OK ] InstructionsTest.ShuffleMaskIsReplicationMask_Exhaustive_Correctness (203930 ms)
[----------] 4 tests from InstructionsTest (431486 ms total)

[----------] Global test environment tear-down
[==========] 4 tests from 1 test suite ran. (431486 ms total)
[  PASSED  ] 4 tests.

I'm running a debug build. I would expect performance to be degraded, but not this much. I'll see if release build makes a difference.

[rupprecht]

llvm/unittests/IR/InstructionsTest.cpp

1133

Yep, a release build is much faster:

$ ./unittests/IR/IRTests --gtest_filter='InstructionsTest.ShuffleMask*'
Note: Google Test filter = InstructionsTest.ShuffleMask*
[==========] Running 4 tests from 1 test suite.
[----------] Global test environment set-up.
[----------] 4 tests from InstructionsTest
[ RUN      ] InstructionsTest.ShuffleMaskQueries
[       OK ] InstructionsTest.ShuffleMaskQueries (1 ms)
[ RUN      ] InstructionsTest.ShuffleMaskIsReplicationMask
[       OK ] InstructionsTest.ShuffleMaskIsReplicationMask (0 ms)
[ RUN      ] InstructionsTest.ShuffleMaskIsReplicationMask_undef
[       OK ] InstructionsTest.ShuffleMaskIsReplicationMask_undef (4884 ms)
[ RUN      ] InstructionsTest.ShuffleMaskIsReplicationMask_Exhaustive_Correctness
[       OK ] InstructionsTest.ShuffleMaskIsReplicationMask_Exhaustive_Correctness (2030 ms)
[----------] 4 tests from InstructionsTest (6915 ms total)

[----------] Global test environment tear-down
[==========] 4 tests from 1 test suite ran. (6915 ms total)
[  PASSED  ] 4 tests.

I wonder what is so horrendously bad about this test case that the optimizer is able to get rid of so much?

[lebedev.ri]

llvm/unittests/IR/InstructionsTest.cpp

1133

Weird. ConstantRangeTest takes 23 seconds in total, most test cases being <.5s but a few being ~2 seconds.

In debug build, correct?

[lebedev.ri]

llvm/unittests/IR/InstructionsTest.cpp

1133

I don't have much sympathy for debug optimization-less build. You get what you asked for.
I've fine-tuned this in rG0b36431810076790b1f3362495b2eaebc9be96c7,
hopefully that no longer leads to bad run time for you.

[rupprecht]

llvm/unittests/IR/InstructionsTest.cpp

1133

Yes, 23 seconds for ConstantRangeTest is in a debug build configuration.

The runtime of this test is now less than a second for each test case, so more reasonable for running tests in a plain debug build. Thanks!

llvm/include/llvm/IR/Instructions.h

bool isInsertSubvectorMask(int &NumSubElts, int &Index) const {
     if (isa<ScalableVectorType>(getType()))
       return false;
 
     int NumSrcElts =
         cast<FixedVectorType>(Op<0>()->getType())->getNumElements();
     return isInsertSubvectorMask(ShuffleMask, NumSrcElts, NumSubElts, Index);
   }
 
+  /// Return true if this shuffle mask replicates each of the \p VF elements
+  /// in a vector \p ReplicationFactor times.
+  /// For example, the mask for \p ReplicationFactor=3 and \p VF=4 is:
+  ///   <0,0,0,1,1,1,2,2,2,3,3,3>
+  static bool isReplicationMask(ArrayRef<int> Mask, int &ReplicationFactor,
+                                int &VF);
+  static bool isReplicationMask(const Constant *Mask, int &ReplicationFactor,
+                                int &VF) {
+    assert(Mask->getType()->isVectorTy() && "Shuffle needs vector constant.");
+    // Not possible to express a shuffle mask for a scalable vector for this
+    // case.
+    if (isa<ScalableVectorType>(Mask->getType()))
+      return false;
+    SmallVector<int, 16> MaskAsInts;
+    getShuffleMask(Mask, MaskAsInts);
+    return isReplicationMask(MaskAsInts, ReplicationFactor, VF);
+  }
+
+  /// Return true if this shuffle mask is an replication mask.
+  bool isReplicationMask(int &ReplicationFactor, int &VF) const {
+    // Not possible to express a shuffle mask for a scalable vector for this
+    // case.
+    if (isa<ScalableVectorType>(getType()))
+      return false;
+
+    return isReplicationMask(ShuffleMask, ReplicationFactor, VF);
+  }
+
   /// 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/lib/IR/Instructions.cpp

bool ShuffleVectorInst::isConcat() const {
 
   // Use the mask length rather than the operands' vector lengths here. We
   // already know that the shuffle returns a vector twice as long as the inputs,
   // and neither of the inputs are undef vectors. If the mask picks consecutive
   // elements from both inputs, then this is a concatenation of the inputs.
   return isIdentityMaskImpl(getShuffleMask(), NumMaskElts);
 }
 
+static bool isReplicationMaskWithParams(ArrayRef<int> Mask,
+                                        int ReplicationFactor, int VF) {
+  assert(Mask.size() == (unsigned)ReplicationFactor * VF &&
+         "Unexpected mask size.");
+
+  for (int CurrElt : seq(0, VF)) {
+    ArrayRef<int> CurrSubMask = Mask.take_front(ReplicationFactor);
+    assert(CurrSubMask.size() == (unsigned)ReplicationFactor &&
+           "Run out of mask?");
+    Mask = Mask.drop_front(ReplicationFactor);
+    if (!all_of(CurrSubMask, [CurrElt](int MaskElt) {
+          return MaskElt == UndefMaskElem || MaskElt == CurrElt;
+        }))
+      return false;
+  }
+  assert(Mask.empty() && "Did not consume the whole mask?");
+
+  return true;
+}
+
+bool ShuffleVectorInst::isReplicationMask(ArrayRef<int> Mask,
+                                          int &ReplicationFactor, int &VF) {
+  // undef-less case is trivial.
+  if (none_of(Mask, [](int MaskElt) { return MaskElt == UndefMaskElem; })) {
+    ReplicationFactor =
+        Mask.take_while([](int MaskElt) { return MaskElt == 0; }).size();
+    if (ReplicationFactor == 0 || Mask.size() % ReplicationFactor != 0)
+      return false;
+    VF = Mask.size() / ReplicationFactor;
+    return isReplicationMaskWithParams(Mask, ReplicationFactor, VF);
+  }
+
+  // However, if the mask contains undef's, we have to enumerate possible tuples
+  // and pick one. There are bounds on replication factor: [1, mask size]
+  // (where RF=1 is an identity shuffle, RF=mask size is a broadcast shuffle)
+  // Additionally, mask size is a replication factor multiplied by vector size,
+  // which further significantly reduces the search space.
+
+  // Before doing that, let's perform basic sanity check first.
+  int Largest = -1;
+  for (int MaskElt : Mask) {
+    if (MaskElt == UndefMaskElem)
+      continue;
+    // Elements must be in non-decreasing order.
+    if (MaskElt < Largest)
+      return false;
+    Largest = std::max(Largest, MaskElt);
+  }
+
+  // Prefer larger replication factor if all else equal.
+  for (int PossibleReplicationFactor :
+       reverse(seq_inclusive<unsigned>(1, Mask.size()))) {
+    if (Mask.size() % PossibleReplicationFactor != 0)
+      continue;
+    int PossibleVF = Mask.size() / PossibleReplicationFactor;
+    if (!isReplicationMaskWithParams(Mask, PossibleReplicationFactor,
+                                     PossibleVF))
+      continue;
+    ReplicationFactor = PossibleReplicationFactor;
+    VF = PossibleVF;
+    return true;
+  }
+
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 //                             InsertValueInst Class
 //===----------------------------------------------------------------------===//
 
 void InsertValueInst::init(Value *Agg, Value *Val, ArrayRef<unsigned> Idxs,
                            const Twine &Name) {
   assert(getNumOperands() == 2 && "NumOperands not initialized?");
 

llvm/unittests/IR/InstructionsTest.cpp

 //===- llvm/unittest/IR/InstructionsTest.cpp - Instructions unit 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/AsmParser/Parser.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/ADT/CombinationGenerator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/VectorUtils.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
ShuffleVectorInst *Id15 =
                             Constant::getNullValue(VScaleV2Int32Ty),
                             Constant::getNullValue(VScaleV2Int32Ty));
   EXPECT_FALSE(Id15->isIdentityWithPadding());
   EXPECT_FALSE(Id15->isIdentityWithExtract());
   EXPECT_FALSE(Id15->isConcat());
   delete Id15;
 }
 
+TEST(InstructionsTest, ShuffleMaskIsReplicationMask) {
+  for (int ReplicationFactor : seq_inclusive(1, 8)) {
+    for (int VF : seq_inclusive(1, 8)) {
+      const auto ReplicatedMask = createReplicatedMask(ReplicationFactor, VF);
+      int GuessedReplicationFactor = -1, GuessedVF = -1;
+      EXPECT_TRUE(ShuffleVectorInst::isReplicationMask(
+          ReplicatedMask, GuessedReplicationFactor, GuessedVF));
+      EXPECT_EQ(GuessedReplicationFactor, ReplicationFactor);
+      EXPECT_EQ(GuessedVF, VF);
+    }
+  }
+}
+
+TEST(InstructionsTest, ShuffleMaskIsReplicationMask_undef) {
+  for (int ReplicationFactor : seq_inclusive(1, 6)) {
+    for (int VF : seq_inclusive(1, 4)) {
+      const auto ReplicatedMask = createReplicatedMask(ReplicationFactor, VF);
+      int GuessedReplicationFactor = -1, GuessedVF = -1;
+
+      // If we change some mask elements to undef, we should still match.
+
+      SmallVector<SmallVector<bool>> ElementChoices(ReplicatedMask.size(),
+                                                    {false, true});
+
+      CombinationGenerator<bool, decltype(ElementChoices)::value_type,
+                           /*variable_smallsize=*/4>
+          G(ElementChoices);
+
+      G.generate([&](ArrayRef<bool> UndefOverrides) -> bool {
+        SmallVector<int> AdjustedMask;
+        AdjustedMask.reserve(ReplicatedMask.size());
+        for (auto I : zip(ReplicatedMask, UndefOverrides))
+          AdjustedMask.emplace_back(std::get<1>(I) ? -1 : std::get<0>(I));
+        assert(AdjustedMask.size() == ReplicatedMask.size() &&
+               "Size misprediction");
+
+        EXPECT_TRUE(ShuffleVectorInst::isReplicationMask(
+            AdjustedMask, GuessedReplicationFactor, GuessedVF));
+        // Do not check GuessedReplicationFactor and GuessedVF,
+        // with enough undef's we may deduce a different tuple.
+
+        return /*Abort=*/false;
+      });
+    }
+  }
+}
+
+TEST(InstructionsTest, ShuffleMaskIsReplicationMask_Exhaustive_Correctness) {
+  for (int ShufMaskNumElts : seq_inclusive(1, 8)) {
+    SmallVector<int> PossibleShufMaskElts;
+    PossibleShufMaskElts.reserve(ShufMaskNumElts + 2);
+    for (int PossibleShufMaskElt : seq_inclusive(-1, ShufMaskNumElts))
+      PossibleShufMaskElts.emplace_back(PossibleShufMaskElt);
+    assert(PossibleShufMaskElts.size() == ShufMaskNumElts + 2U &&
+           "Size misprediction");
+
+    SmallVector<SmallVector<int>> ElementChoices(ShufMaskNumElts,
+                                                 PossibleShufMaskElts);
+
+    CombinationGenerator<int, decltype(ElementChoices)::value_type,
+                         /*variable_smallsize=*/4>
+        G(ElementChoices);
+
+    G.generate([&](ArrayRef<int> Mask) -> bool {
+      int GuessedReplicationFactor = -1, GuessedVF = -1;
+      bool Match = ShuffleVectorInst::isReplicationMask(
+          Mask, GuessedReplicationFactor, GuessedVF);
+      if (!Match)
+        return /*Abort=*/false;
+
+      const auto ActualMask =
+          createReplicatedMask(GuessedReplicationFactor, GuessedVF);
+      EXPECT_EQ(Mask.size(), ActualMask.size());
+      for (auto I : zip(Mask, ActualMask)) {
+        int Elt = std::get<0>(I);
+        int ActualElt = std::get<0>(I);
+
+        if (Elt != -1)
+          EXPECT_EQ(Elt, ActualElt);
+      }
+
+      return /*Abort=*/false;
+    });
+  }
+}
+
 TEST(InstructionsTest, GetSplat) {
   // 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);