diff --git a/llvm/include/llvm/Analysis/VectorUtils.h b/llvm/include/llvm/Analysis/VectorUtils.h --- a/llvm/include/llvm/Analysis/VectorUtils.h +++ b/llvm/include/llvm/Analysis/VectorUtils.h @@ -342,6 +342,24 @@ void narrowShuffleMaskElts(int Scale, ArrayRef Mask, SmallVectorImpl &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 Mask, + SmallVectorImpl &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 --- a/llvm/lib/Analysis/VectorUtils.cpp +++ b/llvm/lib/Analysis/VectorUtils.cpp @@ -420,6 +420,57 @@ } } +bool llvm::widenShuffleMaskElts(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; + + ScaledMask.clear(); + ScaledMask.reserve(NumElts / Scale); + + // Step through the input mask by splitting into Scale-sized slices. + do { + ArrayRef 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 llvm::computeMinimumValueSizes(ArrayRef Blocks, DemandedBits &DB, const TargetTransformInfo *TTI) { diff --git a/llvm/unittests/Analysis/VectorUtilsTest.cpp b/llvm/unittests/Analysis/VectorUtilsTest.cpp --- a/llvm/unittests/Analysis/VectorUtilsTest.cpp +++ b/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, widenShuffleMaskElts) { + SmallVector WideMask; + SmallVector 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