[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

338–339

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

400–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

438–472

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

400–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.

438–472

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

 
 /// Replace each shuffle mask index with the 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 widening shuffle mask elements, but it always
 /// succeeds because the indexes can always be multiplied (scaled up) to map to
 /// narrower vector elements.
 void narrowShuffleMaskElts(int Scale, ArrayRef<int> Mask,
                            SmallVectorImpl<int> &ScaledMask);
 
+/// Try to replace each shuffle mask index 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
+/// the same negative number (sentinel value), that element of the new mask is
+/// the same value. If any element in a given slice is negative and some other
+/// element in that slice is not the same value, return false (partial matches
+/// with sentinel values are not allowed).
+///
+/// Example with Scale = 4:
+///   <16 x i8> <12, 13, 14, 15, 8, 9, 10, 11, 0, 1, 2, 3, -1, -1, -1, -1> -->
+///   <4 x i32> <3, 2, 0, -1>
+///
+/// This is the reverse process of narrowing shuffle mask elements if it
+/// succeeds. This transform is not always possible because indexes may not
+/// divide evenly (scale down) to map to wider vector elements.
+bool widenShuffleMaskElts(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

bool llvm::isSplatValue(const Value *V, int Index, unsigned Depth) {
   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::narrowShuffleMaskElts(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;
   }
 
   ScaledMask.clear();
   for (int MaskElt : Mask) {
     if (MaskElt >= 0) {
       assert(((uint64_t)Scale * MaskElt + (Scale - 1)) <=
                  std::numeric_limits<int32_t>::max() &&
              "Overflowed 32-bits");
     }
     for (int SliceElt = 0; SliceElt != Scale; ++SliceElt)
       ScaledMask.push_back(MaskElt < 0 ? MaskElt : Scale * MaskElt + SliceElt);
   }
 }
 
+bool llvm::widenShuffleMaskElts(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;
+
+  ScaledMask.clear();
+  ScaledMask.reserve(NumElts / Scale);
+
+  // Step through the input mask by splitting into Scale-sized slices.
+  do {
+    ArrayRef<int> MaskSlice = Mask.take_front(Scale);
+    assert((int)MaskSlice.size() == Scale && "Expected Scale-sized slice.");
+
+    // The first element of the slice determines how we evaluate this slice.
+    int SliceFront = MaskSlice.front();
+    if (SliceFront < 0) {
+      // Negative values (undef or other "sentinel" values) must be equal across
+      // the entire slice.
+      if (!is_splat(MaskSlice))
+        return false;
+      ScaledMask.push_back(SliceFront);
+    } else {
+      // A positive mask element must be cleanly divisible.
+      if (SliceFront % Scale != 0)
+        return false;
+      // Elements of the slice must be consecutive.
+      for (int i = 1; i < Scale; ++i)
+        if (MaskSlice[i] != SliceFront + i)
+          return false;
+      ScaledMask.push_back(SliceFront / Scale);
+    }
+    Mask = Mask.drop_front(Scale);
+  } while (!Mask.empty());
+
+  assert((int)ScaledMask.size() * Scale == NumElts && "Unexpected scaled mask");
+
+  // 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, narrowShuffleMaskElts) {
   SmallVector<int, 16> ScaledMask;
   narrowShuffleMaskElts(1, {3,2,0,-2}, ScaledMask);
   EXPECT_EQ(makeArrayRef(ScaledMask), makeArrayRef({3,2,0,-2}));
   narrowShuffleMaskElts(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, widenShuffleMaskElts) {
+  SmallVector<int, 16> WideMask;
+  SmallVector<int, 16> NarrowMask;
+
+  // scale == 1 is a copy
+  EXPECT_TRUE(widenShuffleMaskElts(1, {3,2,0,-1}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({3,2,0,-1}));
+
+  // back to original mask
+  narrowShuffleMaskElts(1, makeArrayRef(WideMask), NarrowMask);
+  EXPECT_EQ(makeArrayRef(NarrowMask), makeArrayRef({3,2,0,-1}));
+
+  // can't widen non-consecutive 3/2
+  EXPECT_FALSE(widenShuffleMaskElts(2, {3,2,0,-1}, WideMask));
+
+  // can't widen if not evenly divisible
+  EXPECT_FALSE(widenShuffleMaskElts(2, {0,1,2}, WideMask));
+
+  // can always widen identity to single element
+  EXPECT_TRUE(widenShuffleMaskElts(3, {0,1,2}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({0}));
+
+  // back to original mask
+  narrowShuffleMaskElts(3, makeArrayRef(WideMask), NarrowMask);
+  EXPECT_EQ(makeArrayRef(NarrowMask), makeArrayRef({0,1,2}));
+
+  // groups of 4 must be consecutive/undef
+  EXPECT_TRUE(widenShuffleMaskElts(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}));
+
+  // back to original mask
+  narrowShuffleMaskElts(4, makeArrayRef(WideMask), NarrowMask);
+  EXPECT_EQ(makeArrayRef(NarrowMask), makeArrayRef({12,13,14,15,8,9,10,11,0,1,2,3,-1,-1,-1,-1}));
+
+  // groups of 2 must be consecutive/undef
+  EXPECT_FALSE(widenShuffleMaskElts(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
+  EXPECT_TRUE(widenShuffleMaskElts(3, {6,7,8,0,1,2,-1,-1,-1}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({2,0,-1}));
+
+  // back to original mask
+  narrowShuffleMaskElts(3, makeArrayRef(WideMask), NarrowMask);
+  EXPECT_EQ(makeArrayRef(NarrowMask), makeArrayRef({6,7,8,0,1,2,-1,-1,-1}));
+
+  // groups of 3 must be consecutive/undef (partial undefs are not ok)
+  EXPECT_FALSE(widenShuffleMaskElts(3, {-1,7,8,0,-1,2,-1,-1,-1}, WideMask));
+
+  // negative indexes must match across a wide element
+  EXPECT_FALSE(widenShuffleMaskElts(2, {-1,-2,-1,-1}, WideMask));
+
+  // negative indexes must match across a wide element
+  EXPECT_TRUE(widenShuffleMaskElts(2, {-2,-2,-3,-3}, WideMask));
+  EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({-2,-3}));
+}
+
 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