1//===- SLPShuffleAnalysis.h - SLP shuffle analysis base ---------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// Internal header used by SLPVectorizer.cpp. It defines the base class for
10// shuffle cost estimation and shuffle instruction emission. It does not depend
11// on BoUpSLP or any other SLP-private type.
12//
13//===----------------------------------------------------------------------===//
14
15#ifndef LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPSHUFFLEANALYSIS_H
16#define LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPSHUFFLEANALYSIS_H
17
18#include "SLPUtils.h"
19
20#include "llvm/ADT/ArrayRef.h"
21#include "llvm/ADT/STLExtras.h"
22#include "llvm/ADT/Sequence.h"
23#include "llvm/ADT/SmallBitVector.h"
24#include "llvm/ADT/SmallVector.h"
25#include "llvm/IR/Constants.h"
26#include "llvm/IR/DerivedTypes.h"
27#include "llvm/IR/Instructions.h"
28#include "llvm/Support/Casting.h"
29
30#include <algorithm>
31#include <cassert>
32
33namespace llvm::slpvectorizer {
34
35/// The base class for shuffle instruction emission and shuffle cost estimation.
36class BaseShuffleAnalysis {
37protected:
38 Type *ScalarTy = nullptr;
39
40 BaseShuffleAnalysis(Type *ScalarTy) : ScalarTy(ScalarTy) {}
41
42 /// V is expected to be a vectorized value.
43 /// When REVEC is disabled, there is no difference between VF and
44 /// VNumElements.
45 /// When REVEC is enabled, VF is VNumElements / ScalarTyNumElements.
46 /// e.g., if ScalarTy is <4 x Ty> and V1 is <8 x Ty>, 2 is returned instead
47 /// of 8.
48 unsigned getVF(Value *V) const {
49 assert(V && "V cannot be nullptr");
50 assert(isa<FixedVectorType>(V->getType()) &&
51 "V does not have FixedVectorType");
52 assert(ScalarTy && "ScalarTy cannot be nullptr");
53 unsigned ScalarTyNumElements = getNumElements(Ty: ScalarTy);
54 unsigned VNumElements =
55 cast<FixedVectorType>(Val: V->getType())->getNumElements();
56 assert(VNumElements > ScalarTyNumElements &&
57 "the number of elements of V is not large enough");
58 assert(VNumElements % ScalarTyNumElements == 0 &&
59 "the number of elements of V is not a vectorized value");
60 return VNumElements / ScalarTyNumElements;
61 }
62
63 /// Checks if the mask is an identity mask.
64 /// \param IsStrict if is true the function returns false if mask size does
65 /// not match vector size.
66 static bool isIdentityMask(ArrayRef<int> Mask, const FixedVectorType *VecTy,
67 bool IsStrict) {
68 int Limit = Mask.size();
69 int VF = VecTy->getNumElements();
70 int Index = -1;
71 if (VF == Limit && ShuffleVectorInst::isIdentityMask(Mask, NumSrcElts: Limit))
72 return true;
73 if (!IsStrict) {
74 // Consider extract subvector starting from index 0.
75 if (ShuffleVectorInst::isExtractSubvectorMask(Mask, NumSrcElts: VF, Index) &&
76 Index == 0)
77 return true;
78 // All VF-size submasks are identity (e.g.
79 // <poison,poison,poison,poison,0,1,2,poison,poison,1,2,3> etc. for VF 4).
80 if (Limit % VF == 0 && all_of(Range: seq<int>(Begin: 0, End: Limit / VF), P: [=](int Idx) {
81 ArrayRef<int> Slice = Mask.slice(N: Idx * VF, M: VF);
82 return all_of(Range&: Slice, P: equal_to(Arg: PoisonMaskElem)) ||
83 ShuffleVectorInst::isIdentityMask(Mask: Slice, NumSrcElts: VF);
84 }))
85 return true;
86 }
87 return false;
88 }
89
90 /// Tries to combine 2 different masks into single one.
91 /// \param LocalVF Vector length of the permuted input vector. \p Mask may
92 /// change the size of the vector, \p LocalVF is the original size of the
93 /// shuffled vector.
94 static void combineMasks(unsigned LocalVF, SmallVectorImpl<int> &Mask,
95 ArrayRef<int> ExtMask) {
96 unsigned VF = Mask.size();
97 SmallVector<int> NewMask(ExtMask.size(), PoisonMaskElem);
98 for (int I = 0, Sz = ExtMask.size(); I < Sz; ++I) {
99 if (ExtMask[I] == PoisonMaskElem)
100 continue;
101 int MaskedIdx = Mask[ExtMask[I] % VF];
102 NewMask[I] =
103 MaskedIdx == PoisonMaskElem ? PoisonMaskElem : MaskedIdx % LocalVF;
104 }
105 Mask.swap(RHS&: NewMask);
106 }
107
108 /// Looks through shuffles trying to reduce final number of shuffles in the
109 /// code. The function looks through the previously emitted shuffle
110 /// instructions and properly mark indices in mask as undef.
111 /// For example, given the code
112 /// \code
113 /// %s1 = shufflevector <2 x ty> %0, poison, <1, 0>
114 /// %s2 = shufflevector <2 x ty> %1, poison, <1, 0>
115 /// \endcode
116 /// and if need to emit shuffle of %s1 and %s2 with mask <1, 0, 3, 2>, it will
117 /// look through %s1 and %s2 and select vectors %0 and %1 with mask
118 /// <0, 1, 2, 3> for the shuffle.
119 /// If 2 operands are of different size, the smallest one will be resized and
120 /// the mask recalculated properly.
121 /// For example, given the code
122 /// \code
123 /// %s1 = shufflevector <2 x ty> %0, poison, <1, 0, 1, 0>
124 /// %s2 = shufflevector <2 x ty> %1, poison, <1, 0, 1, 0>
125 /// \endcode
126 /// and if need to emit shuffle of %s1 and %s2 with mask <1, 0, 5, 4>, it will
127 /// look through %s1 and %s2 and select vectors %0 and %1 with mask
128 /// <0, 1, 2, 3> for the shuffle.
129 /// So, it tries to transform permutations to simple vector merge, if
130 /// possible.
131 /// \param V The input vector which must be shuffled using the given \p Mask.
132 /// If the better candidate is found, \p V is set to this best candidate
133 /// vector.
134 /// \param Mask The input mask for the shuffle. If the best candidate is found
135 /// during looking-through-shuffles attempt, it is updated accordingly.
136 /// \param SinglePermute true if the shuffle operation is originally a
137 /// single-value-permutation. In this case the look-through-shuffles procedure
138 /// may look for resizing shuffles as the best candidates.
139 /// \return true if the shuffle results in the non-resizing identity shuffle
140 /// (and thus can be ignored), false - otherwise.
141 static bool peekThroughShuffles(Value *&V, SmallVectorImpl<int> &Mask,
142 bool SinglePermute) {
143 Value *Op = V;
144 ShuffleVectorInst *IdentityOp = nullptr;
145 SmallVector<int> IdentityMask;
146 while (auto *SV = dyn_cast<ShuffleVectorInst>(Val: Op)) {
147 // Exit if not a fixed vector type or changing size shuffle.
148 auto *SVTy = dyn_cast<FixedVectorType>(Val: SV->getType());
149 if (!SVTy)
150 break;
151 // Remember the identity or broadcast mask, if it is not a resizing
152 // shuffle. If no better candidates are found, this Op and Mask will be
153 // used in the final shuffle.
154 if (isIdentityMask(Mask, VecTy: SVTy, /*IsStrict=*/IsStrict: false)) {
155 if (!IdentityOp || !SinglePermute ||
156 (isIdentityMask(Mask, VecTy: SVTy, /*IsStrict=*/IsStrict: true) &&
157 !ShuffleVectorInst::isZeroEltSplatMask(Mask: IdentityMask,
158 NumSrcElts: IdentityMask.size()))) {
159 IdentityOp = SV;
160 // Store current mask in the IdentityMask so later we did not lost
161 // this info if IdentityOp is selected as the best candidate for the
162 // permutation.
163 IdentityMask.assign(RHS: Mask);
164 }
165 }
166 // Remember the broadcast mask. If no better candidates are found, this Op
167 // and Mask will be used in the final shuffle.
168 // Zero splat can be used as identity too, since it might be used with
169 // mask <0, 1, 2, ...>, i.e. identity mask without extra reshuffling.
170 // E.g. if need to shuffle the vector with the mask <3, 1, 2, 0>, which is
171 // expensive, the analysis founds out, that the source vector is just a
172 // broadcast, this original mask can be transformed to identity mask <0,
173 // 1, 2, 3>.
174 // \code
175 // %0 = shuffle %v, poison, zeroinitalizer
176 // %res = shuffle %0, poison, <3, 1, 2, 0>
177 // \endcode
178 // may be transformed to
179 // \code
180 // %0 = shuffle %v, poison, zeroinitalizer
181 // %res = shuffle %0, poison, <0, 1, 2, 3>
182 // \endcode
183 if (SV->isZeroEltSplat()) {
184 IdentityOp = SV;
185 IdentityMask.assign(RHS: Mask);
186 }
187 int LocalVF = Mask.size();
188 if (auto *SVOpTy =
189 dyn_cast<FixedVectorType>(Val: SV->getOperand(i_nocapture: 0)->getType()))
190 LocalVF = SVOpTy->getNumElements();
191 SmallVector<int> ExtMask(Mask.size(), PoisonMaskElem);
192 for (auto [Idx, I] : enumerate(First&: Mask)) {
193 if (I == PoisonMaskElem ||
194 static_cast<unsigned>(I) >= SV->getShuffleMask().size())
195 continue;
196 ExtMask[Idx] = SV->getMaskValue(Elt: I);
197 }
198 bool IsOp1Undef = isUndefVector</*isPoisonOnly=*/true>(
199 V: SV->getOperand(i_nocapture: 0),
200 UseMask: buildUseMask(VF: LocalVF, Mask: ExtMask, MaskArg: UseMask::FirstArg))
201 .all();
202 bool IsOp2Undef = isUndefVector</*isPoisonOnly=*/true>(
203 V: SV->getOperand(i_nocapture: 1),
204 UseMask: buildUseMask(VF: LocalVF, Mask: ExtMask, MaskArg: UseMask::SecondArg))
205 .all();
206 if (!IsOp1Undef && !IsOp2Undef) {
207 // Update mask and mark undef elems.
208 for (int &I : Mask) {
209 if (I == PoisonMaskElem)
210 continue;
211 if (SV->getMaskValue(Elt: I % SV->getShuffleMask().size()) ==
212 PoisonMaskElem)
213 I = PoisonMaskElem;
214 }
215 break;
216 }
217 SmallVector<int> ShuffleMask(SV->getShuffleMask());
218 combineMasks(LocalVF, Mask&: ShuffleMask, ExtMask: Mask);
219 Mask.swap(RHS&: ShuffleMask);
220 if (IsOp2Undef)
221 Op = SV->getOperand(i_nocapture: 0);
222 else
223 Op = SV->getOperand(i_nocapture: 1);
224 }
225 if (auto *OpTy = dyn_cast<FixedVectorType>(Val: Op->getType());
226 !OpTy || !isIdentityMask(Mask, VecTy: OpTy, IsStrict: SinglePermute) ||
227 ShuffleVectorInst::isZeroEltSplatMask(Mask, NumSrcElts: Mask.size())) {
228 if (IdentityOp) {
229 V = IdentityOp;
230 assert(Mask.size() == IdentityMask.size() &&
231 "Expected masks of same sizes.");
232 // Clear known poison elements.
233 for (auto [I, Idx] : enumerate(First&: Mask))
234 if (Idx == PoisonMaskElem)
235 IdentityMask[I] = PoisonMaskElem;
236 Mask.swap(RHS&: IdentityMask);
237 auto *Shuffle = dyn_cast<ShuffleVectorInst>(Val: V);
238 return SinglePermute &&
239 (isIdentityMask(Mask, VecTy: cast<FixedVectorType>(Val: V->getType()),
240 /*IsStrict=*/IsStrict: true) ||
241 (Shuffle && Mask.size() == Shuffle->getShuffleMask().size() &&
242 Shuffle->isZeroEltSplat() &&
243 ShuffleVectorInst::isZeroEltSplatMask(Mask, NumSrcElts: Mask.size()) &&
244 all_of(Range: enumerate(First&: Mask), P: [&](const auto &P) {
245 return P.value() == PoisonMaskElem ||
246 Shuffle->getShuffleMask()[P.index()] == 0;
247 })));
248 }
249 V = Op;
250 return false;
251 }
252 V = Op;
253 return true;
254 }
255
256 /// Smart shuffle instruction emission, walks through shuffles trees and
257 /// tries to find the best matching vector for the actual shuffle
258 /// instruction.
259 template <typename T, typename ShuffleBuilderTy, typename... Args>
260 static T createShuffle(Value *V1, Value *V2, ArrayRef<int> Mask,
261 ShuffleBuilderTy &Builder, Type *ScalarTy,
262 [[maybe_unused]] bool ReVec, Args... Arguments) {
263 assert(V1 && "Expected at least one vector value.");
264 unsigned ScalarTyNumElements = getNumElements(Ty: ScalarTy);
265 SmallVector<int> NewMask(Mask);
266 if (ScalarTyNumElements != 1) {
267 assert(ReVec && "FixedVectorType is not expected.");
268 transformScalarShuffleIndiciesToVector(VecTyNumElements: ScalarTyNumElements, Mask&: NewMask);
269 Mask = NewMask;
270 }
271 if (V2)
272 Builder.resizeToMatch(V1, V2);
273 int VF = Mask.size();
274 if (auto *FTy = dyn_cast<FixedVectorType>(Val: V1->getType()))
275 VF = FTy->getNumElements();
276 if (V2 && !isUndefVector</*IsPoisonOnly=*/true>(
277 V: V2, UseMask: buildUseMask(VF, Mask, MaskArg: UseMask::SecondArg))
278 .all()) {
279 // Peek through shuffles.
280 Value *Op1 = V1;
281 Value *Op2 = V2;
282 int VF =
283 cast<VectorType>(Val: V1->getType())->getElementCount().getKnownMinValue();
284 SmallVector<int> CombinedMask1(Mask.size(), PoisonMaskElem);
285 SmallVector<int> CombinedMask2(Mask.size(), PoisonMaskElem);
286 for (int I = 0, E = Mask.size(); I < E; ++I) {
287 if (Mask[I] < VF)
288 CombinedMask1[I] = Mask[I];
289 else
290 CombinedMask2[I] = Mask[I] - VF;
291 }
292 Value *PrevOp1;
293 Value *PrevOp2;
294 do {
295 PrevOp1 = Op1;
296 PrevOp2 = Op2;
297 (void)peekThroughShuffles(V&: Op1, Mask&: CombinedMask1, /*SinglePermute=*/SinglePermute: false);
298 (void)peekThroughShuffles(V&: Op2, Mask&: CombinedMask2, /*SinglePermute=*/SinglePermute: false);
299 // Check if we have 2 resizing shuffles - need to peek through operands
300 // again.
301 if (auto *SV1 = dyn_cast<ShuffleVectorInst>(Val: Op1))
302 if (auto *SV2 = dyn_cast<ShuffleVectorInst>(Val: Op2)) {
303 SmallVector<int> ExtMask1(Mask.size(), PoisonMaskElem);
304 for (auto [Idx, I] : enumerate(First&: CombinedMask1)) {
305 if (I == PoisonMaskElem)
306 continue;
307 ExtMask1[Idx] = SV1->getMaskValue(Elt: I);
308 }
309 SmallBitVector UseMask1 = buildUseMask(
310 VF: cast<FixedVectorType>(Val: SV1->getOperand(i_nocapture: 1)->getType())
311 ->getNumElements(),
312 Mask: ExtMask1, MaskArg: UseMask::SecondArg);
313 SmallVector<int> ExtMask2(CombinedMask2.size(), PoisonMaskElem);
314 for (auto [Idx, I] : enumerate(First&: CombinedMask2)) {
315 if (I == PoisonMaskElem)
316 continue;
317 ExtMask2[Idx] = SV2->getMaskValue(Elt: I);
318 }
319 SmallBitVector UseMask2 = buildUseMask(
320 VF: cast<FixedVectorType>(Val: SV2->getOperand(i_nocapture: 1)->getType())
321 ->getNumElements(),
322 Mask: ExtMask2, MaskArg: UseMask::SecondArg);
323 if (SV1->getOperand(i_nocapture: 0)->getType() ==
324 SV2->getOperand(i_nocapture: 0)->getType() &&
325 SV1->getOperand(i_nocapture: 0)->getType() != SV1->getType() &&
326 isUndefVector(V: SV1->getOperand(i_nocapture: 1), UseMask: UseMask1).all() &&
327 isUndefVector(V: SV2->getOperand(i_nocapture: 1), UseMask: UseMask2).all()) {
328 Op1 = SV1->getOperand(i_nocapture: 0);
329 Op2 = SV2->getOperand(i_nocapture: 0);
330 SmallVector<int> ShuffleMask1(SV1->getShuffleMask());
331 int LocalVF = ShuffleMask1.size();
332 if (auto *FTy = dyn_cast<FixedVectorType>(Val: Op1->getType()))
333 LocalVF = FTy->getNumElements();
334 combineMasks(LocalVF, Mask&: ShuffleMask1, ExtMask: CombinedMask1);
335 CombinedMask1.swap(RHS&: ShuffleMask1);
336 SmallVector<int> ShuffleMask2(SV2->getShuffleMask());
337 LocalVF = ShuffleMask2.size();
338 if (auto *FTy = dyn_cast<FixedVectorType>(Val: Op2->getType()))
339 LocalVF = FTy->getNumElements();
340 combineMasks(LocalVF, Mask&: ShuffleMask2, ExtMask: CombinedMask2);
341 CombinedMask2.swap(RHS&: ShuffleMask2);
342 }
343 }
344 } while (PrevOp1 != Op1 || PrevOp2 != Op2);
345 Builder.resizeToMatch(Op1, Op2);
346 VF = std::max(a: cast<VectorType>(Val: Op1->getType())
347 ->getElementCount()
348 .getKnownMinValue(),
349 b: cast<VectorType>(Val: Op2->getType())
350 ->getElementCount()
351 .getKnownMinValue());
352 for (int I = 0, E = Mask.size(); I < E; ++I) {
353 if (CombinedMask2[I] != PoisonMaskElem) {
354 assert(CombinedMask1[I] == PoisonMaskElem &&
355 "Expected undefined mask element");
356 CombinedMask1[I] = CombinedMask2[I] + (Op1 == Op2 ? 0 : VF);
357 }
358 }
359 if (Op1 == Op2 &&
360 (ShuffleVectorInst::isIdentityMask(Mask: CombinedMask1, NumSrcElts: VF) ||
361 (ShuffleVectorInst::isZeroEltSplatMask(Mask: CombinedMask1, NumSrcElts: VF) &&
362 isa<ShuffleVectorInst>(Val: Op1) &&
363 cast<ShuffleVectorInst>(Val: Op1)->getShuffleMask() ==
364 ArrayRef(CombinedMask1))))
365 return Builder.createIdentity(Op1);
366 return Builder.createShuffleVector(
367 Op1, Op1 == Op2 ? PoisonValue::get(T: Op1->getType()) : Op2,
368 CombinedMask1);
369 }
370 if (isa<PoisonValue>(Val: V1))
371 return Builder.createPoison(
372 cast<VectorType>(Val: V1->getType())->getElementType(), Mask.size());
373 bool IsIdentity = peekThroughShuffles(V&: V1, Mask&: NewMask, /*SinglePermute=*/SinglePermute: true);
374 assert(V1 && "Expected non-null value after looking through shuffles.");
375
376 if (!IsIdentity)
377 return Builder.createShuffleVector(V1, NewMask, Arguments...);
378 return Builder.createIdentity(V1);
379 }
380
381 /// Transforms mask \p CommonMask per given \p Mask to make proper set after
382 /// shuffle emission.
383 static void transformMaskAfterShuffle(MutableArrayRef<int> CommonMask,
384 ArrayRef<int> Mask) {
385 for (unsigned I : seq<unsigned>(Size: CommonMask.size()))
386 if (Mask[I] != PoisonMaskElem)
387 CommonMask[I] = I;
388 }
389};
390
391} // namespace llvm::slpvectorizer
392
393#endif // LLVM_LIB_TRANSFORMS_VECTORIZE_SLPVECTORIZER_SLPSHUFFLEANALYSIS_H
394