Index: llvm/include/llvm/Analysis/VectorUtils.h =================================================================== --- llvm/include/llvm/Analysis/VectorUtils.h +++ llvm/include/llvm/Analysis/VectorUtils.h @@ -337,11 +337,23 @@ /// <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 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. +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,58 @@ 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 = 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 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,37 @@ 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; + // 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