1//===-- X86PartialReduction.cpp -------------------------------------------===//
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// This pass looks for add instructions used by a horizontal reduction to see
10// if we might be able to use pmaddwd or psadbw. Some cases of this require
11// cross basic block knowledge and can't be done in SelectionDAG.
12//
13//===----------------------------------------------------------------------===//
14
15#include "X86.h"
16#include "X86TargetMachine.h"
17#include "llvm/Analysis/ValueTracking.h"
18#include "llvm/CodeGen/TargetPassConfig.h"
19#include "llvm/IR/Analysis.h"
20#include "llvm/IR/Constants.h"
21#include "llvm/IR/DerivedTypes.h"
22#include "llvm/IR/IRBuilder.h"
23#include "llvm/IR/Instructions.h"
24#include "llvm/IR/IntrinsicsX86.h"
25#include "llvm/IR/PassManager.h"
26#include "llvm/IR/PatternMatch.h"
27#include "llvm/Pass.h"
28#include "llvm/Support/KnownBits.h"
29
30using namespace llvm;
31
32#define DEBUG_TYPE "x86-partial-reduction"
33
34namespace {
35
36class X86PartialReduction {
37 const X86TargetMachine *TM;
38 const DataLayout *DL = nullptr;
39 const X86Subtarget *ST = nullptr;
40
41public:
42 X86PartialReduction(const X86TargetMachine *TM) : TM(TM) {}
43 bool run(Function &F);
44
45private:
46 bool tryMAddReplacement(Instruction *Op, bool ReduceInOneBB);
47 bool trySADReplacement(Instruction *Op);
48 bool tryByteSumReplacement(Instruction *Op);
49};
50
51class X86PartialReductionLegacy : public FunctionPass {
52public:
53 static char ID; // Pass identification, replacement for typeid.
54
55 X86PartialReductionLegacy() : FunctionPass(ID) {}
56
57 bool runOnFunction(Function &F) override;
58
59 void getAnalysisUsage(AnalysisUsage &AU) const override {
60 AU.setPreservesCFG();
61 }
62
63 StringRef getPassName() const override { return "X86 Partial Reduction"; }
64};
65}
66
67FunctionPass *llvm::createX86PartialReductionLegacyPass() {
68 return new X86PartialReductionLegacy();
69}
70
71char X86PartialReductionLegacy::ID = 0;
72
73INITIALIZE_PASS(X86PartialReductionLegacy, DEBUG_TYPE, "X86 Partial Reduction",
74 false, false)
75
76// This function should be aligned with detectExtMul() in X86ISelLowering.cpp.
77static bool matchVPDPBUSDPattern(const X86Subtarget *ST, BinaryOperator *Mul,
78 const DataLayout *DL) {
79 if (!ST->hasVNNI() && !ST->hasAVXVNNI())
80 return false;
81
82 Value *LHS = Mul->getOperand(i_nocapture: 0);
83 Value *RHS = Mul->getOperand(i_nocapture: 1);
84
85 if (isa<SExtInst>(Val: LHS))
86 std::swap(a&: LHS, b&: RHS);
87
88 auto IsFreeTruncation = [&](Value *Op) {
89 if (auto *Cast = dyn_cast<CastInst>(Val: Op)) {
90 if (Cast->getParent() == Mul->getParent() &&
91 (Cast->getOpcode() == Instruction::SExt ||
92 Cast->getOpcode() == Instruction::ZExt) &&
93 Cast->getOperand(i_nocapture: 0)->getType()->getScalarSizeInBits() <= 8)
94 return true;
95 }
96
97 return isa<Constant>(Val: Op);
98 };
99
100 // (dpbusd (zext a), (sext, b)). Since the first operand should be unsigned
101 // value, we need to check LHS is zero extended value. RHS should be signed
102 // value, so we just check the signed bits.
103 if ((IsFreeTruncation(LHS) &&
104 computeKnownBits(V: LHS, DL: *DL).countMaxActiveBits() <= 8) &&
105 (IsFreeTruncation(RHS) && ComputeMaxSignificantBits(Op: RHS, DL: *DL) <= 8))
106 return true;
107
108 return false;
109}
110
111bool X86PartialReduction::tryMAddReplacement(Instruction *Op,
112 bool ReduceInOneBB) {
113 if (!ST->hasSSE2())
114 return false;
115
116 // Need at least 8 elements.
117 if (cast<FixedVectorType>(Val: Op->getType())->getNumElements() < 8)
118 return false;
119
120 // Element type should be i32.
121 if (!cast<VectorType>(Val: Op->getType())->getElementType()->isIntegerTy(BitWidth: 32))
122 return false;
123
124 auto *Mul = dyn_cast<BinaryOperator>(Val: Op);
125 if (!Mul || Mul->getOpcode() != Instruction::Mul)
126 return false;
127
128 Value *LHS = Mul->getOperand(i_nocapture: 0);
129 Value *RHS = Mul->getOperand(i_nocapture: 1);
130
131 // If the target support VNNI, leave it to ISel to combine reduce operation
132 // to VNNI instruction.
133 // TODO: we can support transforming reduce to VNNI intrinsic for across block
134 // in this pass.
135 if (ReduceInOneBB && matchVPDPBUSDPattern(ST, Mul, DL))
136 return false;
137
138 // LHS and RHS should be only used once or if they are the same then only
139 // used twice. Only check this when SSE4.1 is enabled and we have zext/sext
140 // instructions, otherwise we use punpck to emulate zero extend in stages. The
141 // trunc/ we need to do likely won't introduce new instructions in that case.
142 if (ST->hasSSE41()) {
143 if (LHS == RHS) {
144 if (!isa<Constant>(Val: LHS) && !LHS->hasNUses(N: 2))
145 return false;
146 } else {
147 if (!isa<Constant>(Val: LHS) && !LHS->hasOneUse())
148 return false;
149 if (!isa<Constant>(Val: RHS) && !RHS->hasOneUse())
150 return false;
151 }
152 }
153
154 auto CanShrinkOp = [&](Value *Op) {
155 auto IsFreeTruncation = [&](Value *Op) {
156 if (auto *Cast = dyn_cast<CastInst>(Val: Op)) {
157 if (Cast->getParent() == Mul->getParent() &&
158 (Cast->getOpcode() == Instruction::SExt ||
159 Cast->getOpcode() == Instruction::ZExt) &&
160 Cast->getOperand(i_nocapture: 0)->getType()->getScalarSizeInBits() <= 16)
161 return true;
162 }
163
164 return isa<Constant>(Val: Op);
165 };
166
167 // If the operation can be freely truncated and has enough sign bits we
168 // can shrink.
169 if (IsFreeTruncation(Op) && ComputeNumSignBits(Op, DL: *DL, AC: nullptr, CxtI: Mul) > 16)
170 return true;
171
172 // SelectionDAG has limited support for truncating through an add or sub if
173 // the inputs are freely truncatable.
174 if (auto *BO = dyn_cast<BinaryOperator>(Val: Op)) {
175 if (BO->getParent() == Mul->getParent() &&
176 IsFreeTruncation(BO->getOperand(i_nocapture: 0)) &&
177 IsFreeTruncation(BO->getOperand(i_nocapture: 1)) &&
178 ComputeNumSignBits(Op, DL: *DL, AC: nullptr, CxtI: Mul) > 16)
179 return true;
180 }
181
182 return false;
183 };
184
185 // Both Ops need to be shrinkable.
186 if (!CanShrinkOp(LHS) && !CanShrinkOp(RHS))
187 return false;
188
189 IRBuilder<> Builder(Mul);
190
191 auto *MulTy = cast<FixedVectorType>(Val: Op->getType());
192 unsigned NumElts = MulTy->getNumElements();
193
194 // Extract even elements and odd elements and add them together. This will
195 // be pattern matched by SelectionDAG to pmaddwd. This instruction will be
196 // half the original width.
197 SmallVector<int, 16> EvenMask(NumElts / 2);
198 SmallVector<int, 16> OddMask(NumElts / 2);
199 for (int i = 0, e = NumElts / 2; i != e; ++i) {
200 EvenMask[i] = i * 2;
201 OddMask[i] = i * 2 + 1;
202 }
203 // Creating a new mul so the replaceAllUsesWith below doesn't replace the
204 // uses in the shuffles we're creating.
205 Value *NewMul = Builder.CreateMul(LHS: Mul->getOperand(i_nocapture: 0), RHS: Mul->getOperand(i_nocapture: 1));
206 Value *EvenElts = Builder.CreateShuffleVector(V1: NewMul, V2: NewMul, Mask: EvenMask);
207 Value *OddElts = Builder.CreateShuffleVector(V1: NewMul, V2: NewMul, Mask: OddMask);
208 Value *MAdd = Builder.CreateAdd(LHS: EvenElts, RHS: OddElts);
209
210 // Concatenate zeroes to extend back to the original type.
211 SmallVector<int, 32> ConcatMask(NumElts);
212 std::iota(first: ConcatMask.begin(), last: ConcatMask.end(), value: 0);
213 Value *Zero = Constant::getNullValue(Ty: MAdd->getType());
214 Value *Concat = Builder.CreateShuffleVector(V1: MAdd, V2: Zero, Mask: ConcatMask);
215
216 Mul->replaceAllUsesWith(V: Concat);
217 Mul->eraseFromParent();
218
219 return true;
220}
221
222bool X86PartialReduction::trySADReplacement(Instruction *Op) {
223 if (!ST->hasSSE2())
224 return false;
225
226 // TODO: There's nothing special about i32, any integer type above i16 should
227 // work just as well.
228 if (!cast<VectorType>(Val: Op->getType())->getElementType()->isIntegerTy(BitWidth: 32))
229 return false;
230
231 Value *LHS;
232 if (match(V: Op, P: PatternMatch::m_Intrinsic<Intrinsic::abs>())) {
233 LHS = Op->getOperand(i: 0);
234 } else {
235 // Operand should be a select.
236 auto *SI = dyn_cast<SelectInst>(Val: Op);
237 if (!SI)
238 return false;
239
240 Value *RHS;
241 // Select needs to implement absolute value.
242 auto SPR = matchSelectPattern(V: SI, LHS, RHS);
243 if (SPR.Flavor != SPF_ABS)
244 return false;
245 }
246
247 // Need a subtract of two values.
248 auto *Sub = dyn_cast<BinaryOperator>(Val: LHS);
249 if (!Sub || Sub->getOpcode() != Instruction::Sub)
250 return false;
251
252 // Look for zero extend from i8.
253 auto getZeroExtendedVal = [](Value *Op) -> Value * {
254 if (auto *ZExt = dyn_cast<ZExtInst>(Val: Op))
255 if (cast<VectorType>(Val: ZExt->getOperand(i_nocapture: 0)->getType())
256 ->getElementType()
257 ->isIntegerTy(BitWidth: 8))
258 return ZExt->getOperand(i_nocapture: 0);
259
260 return nullptr;
261 };
262
263 // Both operands of the subtract should be extends from vXi8.
264 Value *Op0 = getZeroExtendedVal(Sub->getOperand(i_nocapture: 0));
265 Value *Op1 = getZeroExtendedVal(Sub->getOperand(i_nocapture: 1));
266 if (!Op0 || !Op1)
267 return false;
268
269 IRBuilder<> Builder(Op);
270
271 auto *OpTy = cast<FixedVectorType>(Val: Op->getType());
272 unsigned NumElts = OpTy->getNumElements();
273
274 unsigned IntrinsicNumElts;
275 Intrinsic::ID IID;
276 if (ST->useBWIRegs() && NumElts >= 64) {
277 IID = Intrinsic::x86_avx512_psad_bw_512;
278 IntrinsicNumElts = 64;
279 } else if (ST->hasAVX2() && NumElts >= 32) {
280 IID = Intrinsic::x86_avx2_psad_bw;
281 IntrinsicNumElts = 32;
282 } else {
283 IID = Intrinsic::x86_sse2_psad_bw;
284 IntrinsicNumElts = 16;
285 }
286
287 Function *PSADBWFn = Intrinsic::getOrInsertDeclaration(M: Op->getModule(), id: IID);
288
289 if (NumElts < 16) {
290 // Pad input with zeroes.
291 SmallVector<int, 32> ConcatMask(16);
292 for (unsigned i = 0; i != NumElts; ++i)
293 ConcatMask[i] = i;
294 for (unsigned i = NumElts; i != 16; ++i)
295 ConcatMask[i] = (i % NumElts) + NumElts;
296
297 Value *Zero = Constant::getNullValue(Ty: Op0->getType());
298 Op0 = Builder.CreateShuffleVector(V1: Op0, V2: Zero, Mask: ConcatMask);
299 Op1 = Builder.CreateShuffleVector(V1: Op1, V2: Zero, Mask: ConcatMask);
300 NumElts = 16;
301 }
302
303 // Intrinsics produce vXi64 and need to be casted to vXi32.
304 auto *I32Ty =
305 FixedVectorType::get(ElementType: Builder.getInt32Ty(), NumElts: IntrinsicNumElts / 4);
306
307 assert(NumElts % IntrinsicNumElts == 0 && "Unexpected number of elements!");
308 unsigned NumSplits = NumElts / IntrinsicNumElts;
309
310 // First collect the pieces we need.
311 SmallVector<Value *, 4> Ops(NumSplits);
312 for (unsigned i = 0; i != NumSplits; ++i) {
313 SmallVector<int, 64> ExtractMask(IntrinsicNumElts);
314 std::iota(first: ExtractMask.begin(), last: ExtractMask.end(), value: i * IntrinsicNumElts);
315 Value *ExtractOp0 = Builder.CreateShuffleVector(V1: Op0, V2: Op0, Mask: ExtractMask);
316 Value *ExtractOp1 = Builder.CreateShuffleVector(V1: Op1, V2: Op0, Mask: ExtractMask);
317 Ops[i] = Builder.CreateCall(Callee: PSADBWFn, Args: {ExtractOp0, ExtractOp1});
318 Ops[i] = Builder.CreateBitCast(V: Ops[i], DestTy: I32Ty);
319 }
320
321 assert(isPowerOf2_32(NumSplits) && "Expected power of 2 splits");
322 unsigned Stages = Log2_32(Value: NumSplits);
323 for (unsigned s = Stages; s > 0; --s) {
324 unsigned NumConcatElts =
325 cast<FixedVectorType>(Val: Ops[0]->getType())->getNumElements() * 2;
326 for (unsigned i = 0; i != 1U << (s - 1); ++i) {
327 SmallVector<int, 64> ConcatMask(NumConcatElts);
328 std::iota(first: ConcatMask.begin(), last: ConcatMask.end(), value: 0);
329 Ops[i] = Builder.CreateShuffleVector(V1: Ops[i*2], V2: Ops[i*2+1], Mask: ConcatMask);
330 }
331 }
332
333 // At this point the final value should be in Ops[0]. Now we need to adjust
334 // it to the final original type.
335 NumElts = cast<FixedVectorType>(Val: OpTy)->getNumElements();
336 if (NumElts == 2) {
337 // Extract down to 2 elements.
338 Ops[0] = Builder.CreateShuffleVector(V1: Ops[0], V2: Ops[0], Mask: ArrayRef<int>{0, 1});
339 } else if (NumElts >= 8) {
340 SmallVector<int, 32> ConcatMask(NumElts);
341 unsigned SubElts =
342 cast<FixedVectorType>(Val: Ops[0]->getType())->getNumElements();
343 for (unsigned i = 0; i != SubElts; ++i)
344 ConcatMask[i] = i;
345 for (unsigned i = SubElts; i != NumElts; ++i)
346 ConcatMask[i] = (i % SubElts) + SubElts;
347
348 Value *Zero = Constant::getNullValue(Ty: Ops[0]->getType());
349 Ops[0] = Builder.CreateShuffleVector(V1: Ops[0], V2: Zero, Mask: ConcatMask);
350 }
351
352 Op->replaceAllUsesWith(V: Ops[0]);
353 Op->eraseFromParent();
354
355 return true;
356}
357
358bool X86PartialReduction::tryByteSumReplacement(Instruction *Op) {
359 if (!ST->hasSSE2())
360 return false;
361
362 auto *OpTy = dyn_cast<FixedVectorType>(Val: Op->getType());
363 if (!OpTy)
364 return false;
365 unsigned ElemBits = OpTy->getElementType()->getScalarSizeInBits();
366 if (ElemBits != 32 && ElemBits != 64)
367 return false;
368
369 auto *ZExt = dyn_cast<ZExtInst>(Val: Op);
370 if (!ZExt)
371 return false;
372
373 auto *SrcTy = dyn_cast<FixedVectorType>(Val: ZExt->getOperand(i_nocapture: 0)->getType());
374 if (!SrcTy || !SrcTy->getElementType()->isIntegerTy(BitWidth: 8))
375 return false;
376
377 unsigned NumElts = OpTy->getNumElements();
378
379 // Below 16 elements, SelectionDAG's SAD matcher handles it.
380 if (NumElts < 16)
381 return false;
382
383 // Select the widest psadbw intrinsic the subtarget supports.
384 unsigned IntrinsicNumElts;
385 Intrinsic::ID IID;
386 if (ST->useBWIRegs() && NumElts >= 64) {
387 IID = Intrinsic::x86_avx512_psad_bw_512;
388 IntrinsicNumElts = 64;
389 } else if (ST->hasAVX2() && NumElts >= 32) {
390 IID = Intrinsic::x86_avx2_psad_bw;
391 IntrinsicNumElts = 32;
392 } else {
393 IID = Intrinsic::x86_sse2_psad_bw;
394 IntrinsicNumElts = 16;
395 }
396
397 if (NumElts % IntrinsicNumElts != 0 ||
398 !isPowerOf2_32(Value: NumElts / IntrinsicNumElts))
399 return false;
400 unsigned NumSplits = NumElts / IntrinsicNumElts;
401
402 IRBuilder<> Builder(Op);
403 Builder.SetCurrentDebugLocation(Op->getDebugLoc());
404
405 Function *PSADBWFn = Intrinsic::getOrInsertDeclaration(M: Op->getModule(), id: IID);
406
407 // psadbw(x, 0) horizontally sums 8 bytes per lane into i64.
408 auto *I8VecTy = FixedVectorType::get(ElementType: Builder.getInt8Ty(), NumElts: IntrinsicNumElts);
409 Value *Zeroes = Constant::getNullValue(Ty: I8VecTy);
410
411 // For i32 accumulators, bitcast each i64 lane to two i32 lanes.
412 // Per-lane sums are at most 8*255 = 2040, so the upper i32 is always zero.
413 FixedVectorType *I32PerSplitTy =
414 ElemBits == 32
415 ? FixedVectorType::get(ElementType: Builder.getInt32Ty(), NumElts: IntrinsicNumElts / 4)
416 : nullptr;
417
418 // Split input into IntrinsicNumElts-byte lanes and compute psadbw per lane.
419 Value *Src = ZExt->getOperand(i_nocapture: 0);
420 SmallVector<Value *, 4> Ops(NumSplits);
421 for (unsigned i = 0; i != NumSplits; ++i) {
422 SmallVector<int, 64> ExtractMask(IntrinsicNumElts);
423 std::iota(first: ExtractMask.begin(), last: ExtractMask.end(), value: i * IntrinsicNumElts);
424 Value *ExtractSrc = Builder.CreateShuffleVector(V1: Src, V2: Src, Mask: ExtractMask);
425 Ops[i] = Builder.CreateCall(Callee: PSADBWFn, Args: {ExtractSrc, Zeroes});
426 if (I32PerSplitTy)
427 Ops[i] = Builder.CreateBitCast(V: Ops[i], DestTy: I32PerSplitTy);
428 }
429
430 // Concat per-split results with a pairwise shuffle tree.
431 unsigned Stages = Log2_32(Value: NumSplits);
432 for (unsigned S = Stages; S > 0; --S) {
433 unsigned NumConcatElts =
434 cast<FixedVectorType>(Val: Ops[0]->getType())->getNumElements() * 2;
435 for (unsigned i = 0; i != 1U << (S - 1); ++i) {
436 SmallVector<int, 64> ConcatMask(NumConcatElts);
437 std::iota(first: ConcatMask.begin(), last: ConcatMask.end(), value: 0);
438 Ops[i] =
439 Builder.CreateShuffleVector(V1: Ops[i * 2], V2: Ops[i * 2 + 1], Mask: ConcatMask);
440 }
441 }
442
443 // Pad with zeros to match the original vector width.
444 SmallVector<int, 32> ConcatMask(NumElts);
445 unsigned SubElts = cast<FixedVectorType>(Val: Ops[0]->getType())->getNumElements();
446 for (unsigned i = 0; i != SubElts; ++i)
447 ConcatMask[i] = i;
448 for (unsigned i = SubElts; i != NumElts; ++i)
449 ConcatMask[i] = (i % SubElts) + SubElts;
450 Value *Zero = Constant::getNullValue(Ty: Ops[0]->getType());
451 Ops[0] = Builder.CreateShuffleVector(V1: Ops[0], V2: Zero, Mask: ConcatMask);
452
453 Op->replaceAllUsesWith(V: Ops[0]);
454 Op->eraseFromParent();
455 return true;
456}
457
458// Walk backwards from the ExtractElementInst and determine if it is the end of
459// a horizontal reduction. Return the input to the reduction if we find one.
460static Value *matchAddReduction(const ExtractElementInst &EE,
461 bool &ReduceInOneBB) {
462 ReduceInOneBB = true;
463 // Make sure we're extracting index 0.
464 auto *Index = dyn_cast<ConstantInt>(Val: EE.getIndexOperand());
465 if (!Index || !Index->isNullValue())
466 return nullptr;
467
468 const auto *BO = dyn_cast<BinaryOperator>(Val: EE.getVectorOperand());
469 if (!BO || BO->getOpcode() != Instruction::Add || !BO->hasOneUse())
470 return nullptr;
471 if (EE.getParent() != BO->getParent())
472 ReduceInOneBB = false;
473
474 unsigned NumElems = cast<FixedVectorType>(Val: BO->getType())->getNumElements();
475 // Ensure the reduction size is a power of 2.
476 if (!isPowerOf2_32(Value: NumElems))
477 return nullptr;
478
479 const Value *Op = BO;
480 unsigned Stages = Log2_32(Value: NumElems);
481 for (unsigned i = 0; i != Stages; ++i) {
482 const auto *BO = dyn_cast<BinaryOperator>(Val: Op);
483 if (!BO || BO->getOpcode() != Instruction::Add)
484 return nullptr;
485 if (EE.getParent() != BO->getParent())
486 ReduceInOneBB = false;
487
488 // If this isn't the first add, then it should only have 2 users, the
489 // shuffle and another add which we checked in the previous iteration.
490 if (i != 0 && !BO->hasNUses(N: 2))
491 return nullptr;
492
493 Value *LHS = BO->getOperand(i_nocapture: 0);
494 Value *RHS = BO->getOperand(i_nocapture: 1);
495
496 auto *Shuffle = dyn_cast<ShuffleVectorInst>(Val: LHS);
497 if (Shuffle) {
498 Op = RHS;
499 } else {
500 Shuffle = dyn_cast<ShuffleVectorInst>(Val: RHS);
501 Op = LHS;
502 }
503
504 // The first operand of the shuffle should be the same as the other operand
505 // of the bin op.
506 if (!Shuffle || Shuffle->getOperand(i_nocapture: 0) != Op)
507 return nullptr;
508
509 // Verify the shuffle has the expected (at this stage of the pyramid) mask.
510 unsigned MaskEnd = 1 << i;
511 for (unsigned Index = 0; Index < MaskEnd; ++Index)
512 if (Shuffle->getMaskValue(Elt: Index) != (int)(MaskEnd + Index))
513 return nullptr;
514 }
515
516 return const_cast<Value *>(Op);
517}
518
519// See if this BO is reachable from this Phi by walking forward through single
520// use BinaryOperators with the same opcode. If we get back then we know we've
521// found a loop and it is safe to step through this Add to find more leaves.
522static bool isReachableFromPHI(PHINode *Phi, BinaryOperator *BO) {
523 // The PHI itself should only have one use.
524 if (!Phi->hasOneUse())
525 return false;
526
527 Instruction *U = cast<Instruction>(Val: *Phi->user_begin());
528 if (U == BO)
529 return true;
530
531 while (U->hasOneUse() && U->getOpcode() == BO->getOpcode())
532 U = cast<Instruction>(Val: *U->user_begin());
533
534 return U == BO;
535}
536
537// Collect all the leaves of the tree of adds that feeds into the horizontal
538// reduction. Root is the Value that is used by the horizontal reduction.
539// We look through single use phis, single use adds, or adds that are used by
540// a phi that forms a loop with the add.
541static void collectLeaves(Value *Root, SmallVectorImpl<Instruction *> &Leaves) {
542 SmallPtrSet<Value *, 8> Visited;
543 SmallVector<Value *, 8> Worklist;
544 Worklist.push_back(Elt: Root);
545
546 while (!Worklist.empty()) {
547 Value *V = Worklist.pop_back_val();
548 if (!Visited.insert(Ptr: V).second)
549 continue;
550
551 if (auto *PN = dyn_cast<PHINode>(Val: V)) {
552 // PHI node should have single use unless it is the root node, then it
553 // has 2 uses.
554 if (!PN->hasNUses(N: PN == Root ? 2 : 1))
555 break;
556
557 // Push incoming values to the worklist.
558 append_range(C&: Worklist, R: PN->incoming_values());
559
560 continue;
561 }
562
563 if (auto *BO = dyn_cast<BinaryOperator>(Val: V)) {
564 if (BO->getOpcode() == Instruction::Add) {
565 // Simple case. Single use, just push its operands to the worklist.
566 if (BO->hasNUses(N: BO == Root ? 2 : 1)) {
567 append_range(C&: Worklist, R: BO->operands());
568 continue;
569 }
570
571 // If there is additional use, make sure it is an unvisited phi that
572 // gets us back to this node.
573 if (BO->hasNUses(N: BO == Root ? 3 : 2)) {
574 PHINode *PN = nullptr;
575 for (auto *U : BO->users())
576 if (auto *P = dyn_cast<PHINode>(Val: U))
577 if (!Visited.count(Ptr: P))
578 PN = P;
579
580 // If we didn't find a 2-input PHI then this isn't a case we can
581 // handle.
582 if (!PN || PN->getNumIncomingValues() != 2)
583 continue;
584
585 // Walk forward from this phi to see if it reaches back to this add.
586 if (!isReachableFromPHI(Phi: PN, BO))
587 continue;
588
589 // The phi forms a loop with this Add, push its operands.
590 append_range(C&: Worklist, R: BO->operands());
591 }
592 }
593 }
594
595 // Not an add or phi, make it a leaf.
596 if (auto *I = dyn_cast<Instruction>(Val: V)) {
597 if (!V->hasNUses(N: I == Root ? 2 : 1))
598 continue;
599
600 // Add this as a leaf.
601 Leaves.push_back(Elt: I);
602 }
603 }
604}
605
606bool X86PartialReduction::run(Function &F) {
607 ST = TM->getSubtargetImpl(F);
608 DL = &F.getDataLayout();
609
610 bool MadeChange = false;
611 for (auto &BB : F) {
612 for (auto &I : BB) {
613 auto *EE = dyn_cast<ExtractElementInst>(Val: &I);
614 if (!EE)
615 continue;
616
617 bool ReduceInOneBB;
618 // First find a reduction tree.
619 // FIXME: Do we need to handle other opcodes than Add?
620 Value *Root = matchAddReduction(EE: *EE, ReduceInOneBB);
621 if (!Root)
622 continue;
623
624 SmallVector<Instruction *, 8> Leaves;
625 collectLeaves(Root, Leaves);
626
627 for (Instruction *I : Leaves) {
628 if (tryMAddReplacement(Op: I, ReduceInOneBB)) {
629 MadeChange = true;
630 continue;
631 }
632
633 // Don't do SAD matching on the root node. SelectionDAG already
634 // has support for that and currently generates better code.
635 if (I != Root && trySADReplacement(Op: I)) {
636 MadeChange = true;
637 continue;
638 }
639
640 // Byte sum via psadbw(x, 0). Same rationale as trySADReplacement:
641 // don't match on the root node because SelectionDAG already handles
642 // small single-vector patterns and generally emits better code for
643 // them. We only help on wider intermediate shapes that reach us
644 // from loop vectorization.
645 if (I != Root && tryByteSumReplacement(Op: I)) {
646 MadeChange = true;
647 continue;
648 }
649 }
650 }
651 }
652
653 return MadeChange;
654}
655
656bool X86PartialReductionLegacy::runOnFunction(Function &F) {
657 if (skipFunction(F))
658 return false;
659
660 auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
661 if (!TPC)
662 return false;
663
664 return X86PartialReduction(&TPC->getTM<X86TargetMachine>()).run(F);
665}
666
667PreservedAnalyses X86PartialReductionPass::run(Function &F,
668 FunctionAnalysisManager &FAM) {
669 bool Changed = X86PartialReduction(TM).run(F);
670 if (!Changed)
671 return PreservedAnalyses::all();
672
673 PreservedAnalyses PA = PreservedAnalyses::none();
674 PA.preserveSet<CFGAnalyses>();
675 return PA;
676}
677