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

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" - we don't have
   that in IR, but this code should be safe to use with that or other
   special (negative) values.

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.

Differential Revision: https://reviews.llvm.org/D77881

diff --git a/llvm/include/llvm/Analysis/VectorUtils.h b/llvm/include/llvm/Analysis/VectorUtils.h
index d64ad12..6571c87 100644 (file)
--- a/llvm/include/llvm/Analysis/VectorUtils.h
+++ b/llvm/include/llvm/Analysis/VectorUtils.h
@@ -342,6 +342,24 @@ bool isSplatValue(const Value *V, int Index = -1, unsigned Depth = 0);
 void narrowShuffleMaskElts(int Scale, ArrayRef<int> Mask,
                            SmallVectorImpl<int> &ScaledMask);
 
+/// Try to transform a shuffle mask by replacing elements with the scaled index
+/// for an equivalent mask of widened elements. If all mask elements that would
+/// 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.
 ///

diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp
index a5fa3ec..31b600b 100644 (file)
--- a/llvm/lib/Analysis/VectorUtils.cpp
+++ b/llvm/lib/Analysis/VectorUtils.cpp
@@ -420,6 +420,57 @@ void llvm::narrowShuffleMaskElts(int Scale, ArrayRef<int> Mask,
   }
 }
 
+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) {

diff --git a/llvm/unittests/Analysis/VectorUtilsTest.cpp b/llvm/unittests/Analysis/VectorUtilsTest.cpp
index 24b1b0f..1f004a3 100644 (file)
--- a/llvm/unittests/Analysis/VectorUtilsTest.cpp
+++ b/llvm/unittests/Analysis/VectorUtilsTest.cpp
@@ -106,6 +106,62 @@ TEST_F(BasicTest, narrowShuffleMaskElts) {
   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