Index: llvm/include/llvm/Analysis/VectorUtils.h =================================================================== --- llvm/include/llvm/Analysis/VectorUtils.h +++ llvm/include/llvm/Analysis/VectorUtils.h @@ -337,11 +337,27 @@ /// <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 it always succeeds +/// because the indexes can always be multiplied (scaled up) to map to narrower +/// vector elements. void scaleShuffleMask(size_t Scale, ArrayRef Mask, SmallVectorImpl &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. This +/// transform is not always possible because indexes may not divide evenly +/// (scale down) to map to wider vector elements. +bool widenShuffleMask(int Scale, ArrayRef Mask, + SmallVectorImpl &ScaledMask); + /// Compute a map of integer instructions to their minimum legal type /// size. /// Index: llvm/lib/Analysis/VectorUtils.cpp =================================================================== --- llvm/lib/Analysis/VectorUtils.cpp +++ llvm/lib/Analysis/VectorUtils.cpp @@ -414,6 +414,64 @@ ScaledMask.push_back(MaskElt < 0 ? MaskElt : Scale * MaskElt + ScaleElt); } +bool llvm::widenShuffleMask(int Scale, ArrayRef Mask, + SmallVectorImpl &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; + for (int j = 0; j != Scale; ++j) { + int MaskElt = Mask[i + j]; + + // Negative values (undef or other "sentinel" values) must be equal across + // the entire subsection. + if (MaskElt < 0) { + // Initialize value for this subsection. + if (j == 0) + OutputElt = MaskElt; + else if (MaskElt != OutputElt) + return false; + continue; + } + + // A positive mask element must be offset and divisible. + if ((MaskElt - j) % Scale != 0) + return false; + + // Initialize the scaled output value based on the first element in this + // subsection. Otherwise, the input mask element must map to the same + // scaled output value for this subsection. + int ScaledElt = (MaskElt - j) / Scale; + if (j == 0) + OutputElt = ScaledElt; + else if (ScaledElt != OutputElt) + return false; + } + // All narrow elements in this subsection map to the same wider element. + ScaledMask.push_back(OutputElt); + } + + assert(ScaledMask.size() * Scale == Mask.size() && "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 llvm::computeMinimumValueSizes(ArrayRef Blocks, DemandedBits &DB, const TargetTransformInfo *TTI) { Index: llvm/unittests/Analysis/VectorUtilsTest.cpp =================================================================== --- llvm/unittests/Analysis/VectorUtilsTest.cpp +++ llvm/unittests/Analysis/VectorUtilsTest.cpp @@ -106,6 +106,62 @@ 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 WideMask; + SmallVector NarrowMask; + + // scale == 1 is a copy + EXPECT_TRUE(widenShuffleMask(1, {3,2,0,-1}, WideMask)); + EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({3,2,0,-1})); + + // back to original mask + scaleShuffleMask(1, makeArrayRef(WideMask), NarrowMask); + EXPECT_EQ(makeArrayRef(NarrowMask), 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})); + + // back to original mask + scaleShuffleMask(3, makeArrayRef(WideMask), NarrowMask); + EXPECT_EQ(makeArrayRef(NarrowMask), makeArrayRef({0,1,2})); + + // 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})); + + // back to original mask + scaleShuffleMask(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(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 + EXPECT_TRUE(widenShuffleMask(3, {6,7,8,0,1,2,-1,-1,-1}, WideMask)); + EXPECT_EQ(makeArrayRef(WideMask), makeArrayRef({2,0,-1})); + + // back to original mask + scaleShuffleMask(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(widenShuffleMask(3, {-1,7,8,0,-1,2,-1,-1,-1}, WideMask)); + + // negative indexes must match across a wide element + EXPECT_FALSE(widenShuffleMask(2, {-1,-2,-1,-1}, WideMask)); + + // negative indexes must match across a wide element + EXPECT_TRUE(widenShuffleMask(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