[VectorUtils] add IR-level analysis for widening of shuffle mask

[spatel]

Summary

This is similar to the recent move/addition of "scaleShuffleMask" (D76508), but there are a couple of differences:

1. The existing x86 helper (canWidenShuffleElements) always tries to divide-by-2, so it gets called iteratively and wouldn't handle the general case of non-pow-2 length.
2. The existing x86 code handles "SM_SentinelZero", but we don't have that in IR.

The motivation is to enable shuffle folds in instcombine/vector-combine that are similar to D76844 and D76727, but in the reverse-bitcast direction. Those patterns are visible in the tests for D40633.

[lebedev.ri]

This looks good to me, but do we want to have some roundtrip and/or exhaustive tests for this?

llvm/include/llvm/Analysis/VectorUtils.h

340

This reads weird, can seems out of place here.

[spatel]

Patch updated:

1. Improved code comments/asserts.
2. Changed logic - allow matching any negative (sentinel) value as long as that value repeats across all elements of a subsection. In the earlier version, we could match partial undef pieces in a widened element, but that means the round-trip with scaleShuffleMask is not always 1-to-1 (partial undefs would get remapped to defined values). This could get us into trouble with poison propagation, so let's avoid that. This way also allows usage from codegen because we would seamlessly handle SentinelZero or any other special mask constants.
3. Added round-trip unit tests to confirm that this is the opposite of scaleShuffleMask. I haven't tried creating exhaustive unit tests before, so not sure yet how to be both exhaustive and not blow up test timing on something like this where we would want to test across multiple dimensions (mask sizes, values, scale factor).

[lebedev.ri]

llvm/lib/Analysis/VectorUtils.cpp

401

As a preliminary step, rename this to be narrowShuffleMask() ?

llvm/unittests/Analysis/VectorUtilsTest.cpp

154–155

Oh hmm, i almost missed this. We indeed can't define undef shuffle mask elts here:

----------------------------------------
define <2 x i16> @t(<4 x i8> %x) {
%0:
  %t0 = shufflevector <4 x i8> %x, <4 x i8> undef, 4294967295, 1, 2, 3
  %t1 = bitcast <4 x i8> %t0 to <2 x i16>
  ret <2 x i16> %t1
}
=>
define <2 x i16> @t(<4 x i8> %x) {
%0:
  %t0 = bitcast <4 x i8> %x to <2 x i16>
  %t1 = shufflevector <2 x i16> %t0, <2 x i16> undef, 0, 1
  ret <2 x i16> %t1
}
Transformation doesn't verify!
ERROR: Target is more poisonous than source

Example:
<4 x i8> %x = < poison, poison, poison, poison >

Source:
<4 x i8> %t0 = < undef, poison, poison, poison >
<2 x i16> %t1 = < #x0000 (0)    [based on undef value], poison >

Target:
<2 x i16> %t0 = < poison, poison >
<2 x i16> %t1 = < poison, poison >
Source value: < #x0000 (0), poison >
Target value: < poison, poison >

Summary:
  0 correct transformations
  1 incorrect transformations
  0 Alive2 errors

[lebedev.ri]

llvm/lib/Analysis/VectorUtils.cpp

432–466

Ok, in light of not accepting partial sentinel values, what are your thoughts on the following then:

// Step through the input mask by splitting into Scale-sized subsections.
ScaledMask.clear();
ScaledMask.reserve(NumElts / Scale);

for (ArrayRef<int> MaskSlice = Mask.take_front(Scale),
                   RemainingMaskElts = Mask.take_back(Mask.size() - Scale);
     !MaskSlice.empty(); MaskSlice = RemainingMaskElts.take_front(Scale),
                   RemainingMaskElts = RemainingMaskElts.take_back(
                       RemainingMaskElts.size() - Scale)) {
  assert((int)MaskSlice.size() == Scale && "Expected Scale-sized slice.");

  // The slice must be homogeneous.
  int OutputElt;

  if (MaskSlice.front() < 0) {
    // Negative values (undef or other "sentinel" values) must be equal across
    // the entire subsection.
    if (!is_splat(MaskSlice))
      return false;
    OutputElt = MaskSlice.front();
  } else {
    // A positive mask element must be cleanly divisible.
    if (MaskSlice.front() % Scale != 0)
      return false;
    // The elements of the subsection must be consecutive.
    auto ExpectedSlice =
        llvm::seq(MaskSlice.front(), MaskSlice.front() + Scale);
    assert(llvm::size(ExpectedSlice) == Scale && "Got wrong sequence.");
    if (!std::equal(adl_begin(MaskSlice), adl_end(MaskSlice),
                    adl_begin(ExpectedSlice)))
      return false;
    OutputElt = MaskSlice.front() / Scale;
  }

  // All narrow elements in this subsection map to the same wider element.
  ScaledMask.push_back(OutputElt);
}

https://godbolt.org/z/yoLNmG

This results in ~Scale less divisions.

[spatel]

llvm/lib/Analysis/VectorUtils.cpp

401

Yes, that would make more sense if we have both of these utils. Another possibility would be to make a single function with a ScaleUpOrDown bool param, but that's probably less readable in calling code.

432–466

Very C++. :)
I'm good with that...especially since you already wrote it and checked the optimization!

[lebedev.ri]

LG otherwise.

[spatel]

Patch updated:

1. Rebased after rG1318ddbc14c2 (name change for scaleShuffleMask), so this patch doesn't alter comments in the other function.
2. Adopted most of the implementation suggestions to improve readability/iteration processing . I couldn't take the 5-line 'for' statement, so used a take_front/drop_front pattern instead. Also used a straight 'for' loop to avoid adding an #include for llvm::seq(). Hopefully, that's same or better readability.

[spatel]

Patch updated again:
Tried to make the documentation comment easier to understand.

llvm/include/llvm/Analysis/VectorUtils.h

 /// scaled sequential indices for an equivalent mask of narrowed elements.
 /// Mask elements that are less than 0 (sentinel values) are repeated in the
 /// output mask.
 ///
 /// Example with Scale = 4:
 ///   <4 x i32> <3, 2, 0, -1> -->
 ///   <16 x i8> <12, 13, 14, 15, 8, 9, 10, 11, 0, 1, 2, 3, -1, -1, -1, -1>
 ///
-/// This is the reverse process of "canWidenShuffleElements", but can always
-/// succeed.
+/// This is the reverse process of "widenShuffleMask", but can always succeed.
 void scaleShuffleMask(size_t Scale, ArrayRef<int> Mask,
                       SmallVectorImpl<int> &ScaledMask);
 
+/// Try to reduce the length of a shuffle mask by replacing each element with
+/// the scaled index for an equivalent mask of widened elements.
+/// If all mask elements that map to a wider element of the new mask are
+/// undefined, that element of the new mask is undefined.
+///
+/// Example with Scale = 4:
+///   <16 x i8> <12, 13, 14, 15, 8, -1, 10, 11, 0, 1, 2, 3, -1, -1, -1, -1> -->
+///   <4 x i32> <3, 2, 0, -1>
+///
+/// This is the reverse process of scaleShuffleMask if it succeeds.
+bool widenShuffleMask(int Scale, ArrayRef<int> Mask,
+                      SmallVectorImpl<int> &ScaledMask);
+
 /// 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

if (match(V, m_Select(m_Value(X), m_Value(Y), m_Value(Z))))
     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;
 }
 
 void llvm::scaleShuffleMask(size_t Scale, ArrayRef<int> Mask,
                             SmallVectorImpl<int> &ScaledMask) {
   assert(Scale > 0 && "Unexpected scaling factor");
 
   // Fast-path: if no scaling, then it is just a copy.
   if (Scale == 1) {
     ScaledMask.assign(Mask.begin(), Mask.end());
     return;
   }
 
   ScaledMask.clear();
   for (int MaskElt : Mask)
     for (int ScaleElt = 0; ScaleElt != (int)Scale; ++ScaleElt)
       ScaledMask.push_back(MaskElt < 0 ? MaskElt : Scale * MaskElt + ScaleElt);
 }
 
+bool llvm::widenShuffleMask(int Scale, ArrayRef<int> Mask,
+                            SmallVectorImpl<int> &ScaledMask) {
+  assert(Scale > 0 && "Unexpected scaling factor");
+
+  // Fast-path: if no scaling, then it is just a copy.
+  if (Scale == 1) {
+    ScaledMask.assign(Mask.begin(), Mask.end());
+    return true;
+  }
+
+  // We must map the original elements down evenly to a type with less elements.
+  int NumElts = Mask.size();
+  if (NumElts % Scale != 0)
+    return false;
+
+  // Step through the input mask by splitting into Scale-sized subsections.
+  ScaledMask.clear();
+  for (int i = 0, e = NumElts; i != e; i += Scale) {
+    // Check if the elements that map to a widened mask element are consecutive.
+    int OutputElt = UndefMaskElem;
+    for (int j = 0; j != Scale; ++j) {
+      int MaskElt = Mask[i + j];
+      if (MaskElt == UndefMaskElem)
+        continue;
+
+      // This could be relaxed to deal with codegen mask sentinel constants
+      // (for example - zero element).
+      if (MaskElt < 0)
+        return false;
+
+      // The input mask element is defined, so it must be offset and divisible.
+      if ((MaskElt - j) % Scale != 0)
+        return false;
+
+      // Initialize the scaled output value if it is not set. Otherwise, the
+      // input mask element must map to the same scaled output value for this
+      // subsection.
+      int ScaledElt = (MaskElt - j) / Scale;
+      if (OutputElt == UndefMaskElem)
+        OutputElt = ScaledElt;
+      else if (ScaledElt != OutputElt)
+        return false;
+    }
+    // All narrow elements in this subsection map to the same wider element.
+    ScaledMask.push_back(OutputElt);
+  }
+
+  // All elements of the original mask can be scaled down to map to the elements
+  // of a mask with wider elements.
+  return true;
+}
+
 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

 TEST_F(BasicTest, scaleShuffleMask) {
   SmallVector<int, 16> ScaledMask;
   scaleShuffleMask(1, {3,2,0,-2}, ScaledMask);
   EXPECT_EQ(makeArrayRef(ScaledMask), makeArrayRef({3,2,0,-2}));
   scaleShuffleMask(4, {3,2,0,-1}, ScaledMask);
   EXPECT_EQ(makeArrayRef(ScaledMask), makeArrayRef({12,13,14,15,8,9,10,11,0,1,2,3,-1,-1,-1,-1}));
 }
 
+TEST_F(BasicTest, widenShuffleMask) {
+  SmallVector<int, 16> WideMask;
+  // scale == 1 is a copy
+  EXPECT_TRUE(widenShuffleMask(1, {3,2,0,-1}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({3,2,0,-1}));
+
+  // can't widen non-consecutive 3/2
+  EXPECT_FALSE(widenShuffleMask(2, {3,2,0,-1}, WideMask));
+
+  // can't widen if not evenly divisible
+  EXPECT_FALSE(widenShuffleMask(2, {0,1,2}, WideMask));
+
+  // can always widen identity to single element
+  EXPECT_TRUE(widenShuffleMask(3, {0,1,2}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({0}));
+
+  // groups of 4 must be consecutive/undef
+  EXPECT_TRUE(widenShuffleMask(4, {12,13,14,15,8,9,10,11,0,1,2,3,-1,-1,-1,-1}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({3,2,0,-1}));
+
+  // groups of 2 must be consecutive/undef
+  EXPECT_FALSE(widenShuffleMask(2, {12,12,14,15,8,9,10,11,0,1,2,3,-1,-1,-1,-1}, WideMask));
+
+  // groups of 3 must be consecutive/undef (partial undefs are ok)
+  EXPECT_TRUE(widenShuffleMask(3, {-1,7,8,0,-1,2,-1,-1,-1}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({2,0,-1}));
+
+  // unknown/unhandled negative index
+  EXPECT_FALSE(widenShuffleMask(2, {-1,-2,-1,-1}, WideMask));
+}
+
 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));
 }
 
 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"