15ffd83dbSDimitry Andric //===-- X86PartialReduction.cpp -------------------------------------------===//
25ffd83dbSDimitry Andric //
35ffd83dbSDimitry Andric // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
45ffd83dbSDimitry Andric // See https://llvm.org/LICENSE.txt for license information.
55ffd83dbSDimitry Andric // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
65ffd83dbSDimitry Andric //
75ffd83dbSDimitry Andric //===----------------------------------------------------------------------===//
85ffd83dbSDimitry Andric //
95ffd83dbSDimitry Andric // This pass looks for add instructions used by a horizontal reduction to see
105ffd83dbSDimitry Andric // if we might be able to use pmaddwd or psadbw. Some cases of this require
115ffd83dbSDimitry Andric // cross basic block knowledge and can't be done in SelectionDAG.
125ffd83dbSDimitry Andric //
135ffd83dbSDimitry Andric //===----------------------------------------------------------------------===//
145ffd83dbSDimitry Andric
155ffd83dbSDimitry Andric #include "X86.h"
1604eeddc0SDimitry Andric #include "X86TargetMachine.h"
175ffd83dbSDimitry Andric #include "llvm/Analysis/ValueTracking.h"
185ffd83dbSDimitry Andric #include "llvm/CodeGen/TargetPassConfig.h"
195ffd83dbSDimitry Andric #include "llvm/IR/Constants.h"
2004eeddc0SDimitry Andric #include "llvm/IR/IRBuilder.h"
215ffd83dbSDimitry Andric #include "llvm/IR/Instructions.h"
2281ad6265SDimitry Andric #include "llvm/IR/IntrinsicInst.h"
235ffd83dbSDimitry Andric #include "llvm/IR/IntrinsicsX86.h"
245ffd83dbSDimitry Andric #include "llvm/IR/Operator.h"
2581ad6265SDimitry Andric #include "llvm/IR/PatternMatch.h"
265ffd83dbSDimitry Andric #include "llvm/Pass.h"
2704eeddc0SDimitry Andric #include "llvm/Support/KnownBits.h"
285ffd83dbSDimitry Andric
295ffd83dbSDimitry Andric using namespace llvm;
305ffd83dbSDimitry Andric
315ffd83dbSDimitry Andric #define DEBUG_TYPE "x86-partial-reduction"
325ffd83dbSDimitry Andric
335ffd83dbSDimitry Andric namespace {
345ffd83dbSDimitry Andric
355ffd83dbSDimitry Andric class X86PartialReduction : public FunctionPass {
36*06c3fb27SDimitry Andric const DataLayout *DL = nullptr;
37*06c3fb27SDimitry Andric const X86Subtarget *ST = nullptr;
385ffd83dbSDimitry Andric
395ffd83dbSDimitry Andric public:
405ffd83dbSDimitry Andric static char ID; // Pass identification, replacement for typeid.
415ffd83dbSDimitry Andric
X86PartialReduction()425ffd83dbSDimitry Andric X86PartialReduction() : FunctionPass(ID) { }
435ffd83dbSDimitry Andric
445ffd83dbSDimitry Andric bool runOnFunction(Function &Fn) override;
455ffd83dbSDimitry Andric
getAnalysisUsage(AnalysisUsage & AU) const465ffd83dbSDimitry Andric void getAnalysisUsage(AnalysisUsage &AU) const override {
475ffd83dbSDimitry Andric AU.setPreservesCFG();
485ffd83dbSDimitry Andric }
495ffd83dbSDimitry Andric
getPassName() const505ffd83dbSDimitry Andric StringRef getPassName() const override {
515ffd83dbSDimitry Andric return "X86 Partial Reduction";
525ffd83dbSDimitry Andric }
535ffd83dbSDimitry Andric
545ffd83dbSDimitry Andric private:
5504eeddc0SDimitry Andric bool tryMAddReplacement(Instruction *Op, bool ReduceInOneBB);
565ffd83dbSDimitry Andric bool trySADReplacement(Instruction *Op);
575ffd83dbSDimitry Andric };
585ffd83dbSDimitry Andric }
595ffd83dbSDimitry Andric
createX86PartialReductionPass()605ffd83dbSDimitry Andric FunctionPass *llvm::createX86PartialReductionPass() {
615ffd83dbSDimitry Andric return new X86PartialReduction();
625ffd83dbSDimitry Andric }
635ffd83dbSDimitry Andric
645ffd83dbSDimitry Andric char X86PartialReduction::ID = 0;
655ffd83dbSDimitry Andric
665ffd83dbSDimitry Andric INITIALIZE_PASS(X86PartialReduction, DEBUG_TYPE,
675ffd83dbSDimitry Andric "X86 Partial Reduction", false, false)
685ffd83dbSDimitry Andric
6904eeddc0SDimitry Andric // This function should be aligned with detectExtMul() in X86ISelLowering.cpp.
matchVPDPBUSDPattern(const X86Subtarget * ST,BinaryOperator * Mul,const DataLayout * DL)7004eeddc0SDimitry Andric static bool matchVPDPBUSDPattern(const X86Subtarget *ST, BinaryOperator *Mul,
7104eeddc0SDimitry Andric const DataLayout *DL) {
7204eeddc0SDimitry Andric if (!ST->hasVNNI() && !ST->hasAVXVNNI())
7304eeddc0SDimitry Andric return false;
7404eeddc0SDimitry Andric
7504eeddc0SDimitry Andric Value *LHS = Mul->getOperand(0);
7604eeddc0SDimitry Andric Value *RHS = Mul->getOperand(1);
7704eeddc0SDimitry Andric
7804eeddc0SDimitry Andric if (isa<SExtInst>(LHS))
7904eeddc0SDimitry Andric std::swap(LHS, RHS);
8004eeddc0SDimitry Andric
8104eeddc0SDimitry Andric auto IsFreeTruncation = [&](Value *Op) {
8204eeddc0SDimitry Andric if (auto *Cast = dyn_cast<CastInst>(Op)) {
8304eeddc0SDimitry Andric if (Cast->getParent() == Mul->getParent() &&
8404eeddc0SDimitry Andric (Cast->getOpcode() == Instruction::SExt ||
8504eeddc0SDimitry Andric Cast->getOpcode() == Instruction::ZExt) &&
8604eeddc0SDimitry Andric Cast->getOperand(0)->getType()->getScalarSizeInBits() <= 8)
8704eeddc0SDimitry Andric return true;
8804eeddc0SDimitry Andric }
8904eeddc0SDimitry Andric
9004eeddc0SDimitry Andric return isa<Constant>(Op);
9104eeddc0SDimitry Andric };
9204eeddc0SDimitry Andric
9304eeddc0SDimitry Andric // (dpbusd (zext a), (sext, b)). Since the first operand should be unsigned
9404eeddc0SDimitry Andric // value, we need to check LHS is zero extended value. RHS should be signed
9504eeddc0SDimitry Andric // value, so we just check the signed bits.
9604eeddc0SDimitry Andric if ((IsFreeTruncation(LHS) &&
9704eeddc0SDimitry Andric computeKnownBits(LHS, *DL).countMaxActiveBits() <= 8) &&
9804eeddc0SDimitry Andric (IsFreeTruncation(RHS) && ComputeMaxSignificantBits(RHS, *DL) <= 8))
9904eeddc0SDimitry Andric return true;
10004eeddc0SDimitry Andric
10104eeddc0SDimitry Andric return false;
10204eeddc0SDimitry Andric }
10304eeddc0SDimitry Andric
tryMAddReplacement(Instruction * Op,bool ReduceInOneBB)10404eeddc0SDimitry Andric bool X86PartialReduction::tryMAddReplacement(Instruction *Op,
10504eeddc0SDimitry Andric bool ReduceInOneBB) {
1065ffd83dbSDimitry Andric if (!ST->hasSSE2())
1075ffd83dbSDimitry Andric return false;
1085ffd83dbSDimitry Andric
1095ffd83dbSDimitry Andric // Need at least 8 elements.
1105ffd83dbSDimitry Andric if (cast<FixedVectorType>(Op->getType())->getNumElements() < 8)
1115ffd83dbSDimitry Andric return false;
1125ffd83dbSDimitry Andric
1135ffd83dbSDimitry Andric // Element type should be i32.
1145ffd83dbSDimitry Andric if (!cast<VectorType>(Op->getType())->getElementType()->isIntegerTy(32))
1155ffd83dbSDimitry Andric return false;
1165ffd83dbSDimitry Andric
1175ffd83dbSDimitry Andric auto *Mul = dyn_cast<BinaryOperator>(Op);
1185ffd83dbSDimitry Andric if (!Mul || Mul->getOpcode() != Instruction::Mul)
1195ffd83dbSDimitry Andric return false;
1205ffd83dbSDimitry Andric
1215ffd83dbSDimitry Andric Value *LHS = Mul->getOperand(0);
1225ffd83dbSDimitry Andric Value *RHS = Mul->getOperand(1);
1235ffd83dbSDimitry Andric
12404eeddc0SDimitry Andric // If the target support VNNI, leave it to ISel to combine reduce operation
12504eeddc0SDimitry Andric // to VNNI instruction.
12604eeddc0SDimitry Andric // TODO: we can support transforming reduce to VNNI intrinsic for across block
12704eeddc0SDimitry Andric // in this pass.
12804eeddc0SDimitry Andric if (ReduceInOneBB && matchVPDPBUSDPattern(ST, Mul, DL))
12904eeddc0SDimitry Andric return false;
13004eeddc0SDimitry Andric
1315ffd83dbSDimitry Andric // LHS and RHS should be only used once or if they are the same then only
1325ffd83dbSDimitry Andric // used twice. Only check this when SSE4.1 is enabled and we have zext/sext
1335ffd83dbSDimitry Andric // instructions, otherwise we use punpck to emulate zero extend in stages. The
1345ffd83dbSDimitry Andric // trunc/ we need to do likely won't introduce new instructions in that case.
1355ffd83dbSDimitry Andric if (ST->hasSSE41()) {
1365ffd83dbSDimitry Andric if (LHS == RHS) {
1375ffd83dbSDimitry Andric if (!isa<Constant>(LHS) && !LHS->hasNUses(2))
1385ffd83dbSDimitry Andric return false;
1395ffd83dbSDimitry Andric } else {
1405ffd83dbSDimitry Andric if (!isa<Constant>(LHS) && !LHS->hasOneUse())
1415ffd83dbSDimitry Andric return false;
1425ffd83dbSDimitry Andric if (!isa<Constant>(RHS) && !RHS->hasOneUse())
1435ffd83dbSDimitry Andric return false;
1445ffd83dbSDimitry Andric }
1455ffd83dbSDimitry Andric }
1465ffd83dbSDimitry Andric
1475ffd83dbSDimitry Andric auto CanShrinkOp = [&](Value *Op) {
1485ffd83dbSDimitry Andric auto IsFreeTruncation = [&](Value *Op) {
1495ffd83dbSDimitry Andric if (auto *Cast = dyn_cast<CastInst>(Op)) {
1505ffd83dbSDimitry Andric if (Cast->getParent() == Mul->getParent() &&
1515ffd83dbSDimitry Andric (Cast->getOpcode() == Instruction::SExt ||
1525ffd83dbSDimitry Andric Cast->getOpcode() == Instruction::ZExt) &&
1535ffd83dbSDimitry Andric Cast->getOperand(0)->getType()->getScalarSizeInBits() <= 16)
1545ffd83dbSDimitry Andric return true;
1555ffd83dbSDimitry Andric }
1565ffd83dbSDimitry Andric
1575ffd83dbSDimitry Andric return isa<Constant>(Op);
1585ffd83dbSDimitry Andric };
1595ffd83dbSDimitry Andric
1605ffd83dbSDimitry Andric // If the operation can be freely truncated and has enough sign bits we
1615ffd83dbSDimitry Andric // can shrink.
1625ffd83dbSDimitry Andric if (IsFreeTruncation(Op) &&
1635ffd83dbSDimitry Andric ComputeNumSignBits(Op, *DL, 0, nullptr, Mul) > 16)
1645ffd83dbSDimitry Andric return true;
1655ffd83dbSDimitry Andric
1665ffd83dbSDimitry Andric // SelectionDAG has limited support for truncating through an add or sub if
1675ffd83dbSDimitry Andric // the inputs are freely truncatable.
1685ffd83dbSDimitry Andric if (auto *BO = dyn_cast<BinaryOperator>(Op)) {
1695ffd83dbSDimitry Andric if (BO->getParent() == Mul->getParent() &&
1705ffd83dbSDimitry Andric IsFreeTruncation(BO->getOperand(0)) &&
1715ffd83dbSDimitry Andric IsFreeTruncation(BO->getOperand(1)) &&
1725ffd83dbSDimitry Andric ComputeNumSignBits(Op, *DL, 0, nullptr, Mul) > 16)
1735ffd83dbSDimitry Andric return true;
1745ffd83dbSDimitry Andric }
1755ffd83dbSDimitry Andric
1765ffd83dbSDimitry Andric return false;
1775ffd83dbSDimitry Andric };
1785ffd83dbSDimitry Andric
1795ffd83dbSDimitry Andric // Both Ops need to be shrinkable.
1805ffd83dbSDimitry Andric if (!CanShrinkOp(LHS) && !CanShrinkOp(RHS))
1815ffd83dbSDimitry Andric return false;
1825ffd83dbSDimitry Andric
1835ffd83dbSDimitry Andric IRBuilder<> Builder(Mul);
1845ffd83dbSDimitry Andric
1855ffd83dbSDimitry Andric auto *MulTy = cast<FixedVectorType>(Op->getType());
1865ffd83dbSDimitry Andric unsigned NumElts = MulTy->getNumElements();
1875ffd83dbSDimitry Andric
1885ffd83dbSDimitry Andric // Extract even elements and odd elements and add them together. This will
1895ffd83dbSDimitry Andric // be pattern matched by SelectionDAG to pmaddwd. This instruction will be
1905ffd83dbSDimitry Andric // half the original width.
1915ffd83dbSDimitry Andric SmallVector<int, 16> EvenMask(NumElts / 2);
1925ffd83dbSDimitry Andric SmallVector<int, 16> OddMask(NumElts / 2);
1935ffd83dbSDimitry Andric for (int i = 0, e = NumElts / 2; i != e; ++i) {
1945ffd83dbSDimitry Andric EvenMask[i] = i * 2;
1955ffd83dbSDimitry Andric OddMask[i] = i * 2 + 1;
1965ffd83dbSDimitry Andric }
1975ffd83dbSDimitry Andric // Creating a new mul so the replaceAllUsesWith below doesn't replace the
1985ffd83dbSDimitry Andric // uses in the shuffles we're creating.
1995ffd83dbSDimitry Andric Value *NewMul = Builder.CreateMul(Mul->getOperand(0), Mul->getOperand(1));
2005ffd83dbSDimitry Andric Value *EvenElts = Builder.CreateShuffleVector(NewMul, NewMul, EvenMask);
2015ffd83dbSDimitry Andric Value *OddElts = Builder.CreateShuffleVector(NewMul, NewMul, OddMask);
2025ffd83dbSDimitry Andric Value *MAdd = Builder.CreateAdd(EvenElts, OddElts);
2035ffd83dbSDimitry Andric
2045ffd83dbSDimitry Andric // Concatenate zeroes to extend back to the original type.
2055ffd83dbSDimitry Andric SmallVector<int, 32> ConcatMask(NumElts);
2065ffd83dbSDimitry Andric std::iota(ConcatMask.begin(), ConcatMask.end(), 0);
2075ffd83dbSDimitry Andric Value *Zero = Constant::getNullValue(MAdd->getType());
2085ffd83dbSDimitry Andric Value *Concat = Builder.CreateShuffleVector(MAdd, Zero, ConcatMask);
2095ffd83dbSDimitry Andric
2105ffd83dbSDimitry Andric Mul->replaceAllUsesWith(Concat);
2115ffd83dbSDimitry Andric Mul->eraseFromParent();
2125ffd83dbSDimitry Andric
2135ffd83dbSDimitry Andric return true;
2145ffd83dbSDimitry Andric }
2155ffd83dbSDimitry Andric
trySADReplacement(Instruction * Op)2165ffd83dbSDimitry Andric bool X86PartialReduction::trySADReplacement(Instruction *Op) {
2175ffd83dbSDimitry Andric if (!ST->hasSSE2())
2185ffd83dbSDimitry Andric return false;
2195ffd83dbSDimitry Andric
2205ffd83dbSDimitry Andric // TODO: There's nothing special about i32, any integer type above i16 should
2215ffd83dbSDimitry Andric // work just as well.
2225ffd83dbSDimitry Andric if (!cast<VectorType>(Op->getType())->getElementType()->isIntegerTy(32))
2235ffd83dbSDimitry Andric return false;
2245ffd83dbSDimitry Andric
22581ad6265SDimitry Andric Value *LHS;
22681ad6265SDimitry Andric if (match(Op, PatternMatch::m_Intrinsic<Intrinsic::abs>())) {
22781ad6265SDimitry Andric LHS = Op->getOperand(0);
22881ad6265SDimitry Andric } else {
2295ffd83dbSDimitry Andric // Operand should be a select.
2305ffd83dbSDimitry Andric auto *SI = dyn_cast<SelectInst>(Op);
2315ffd83dbSDimitry Andric if (!SI)
2325ffd83dbSDimitry Andric return false;
2335ffd83dbSDimitry Andric
23481ad6265SDimitry Andric Value *RHS;
2355ffd83dbSDimitry Andric // Select needs to implement absolute value.
2365ffd83dbSDimitry Andric auto SPR = matchSelectPattern(SI, LHS, RHS);
2375ffd83dbSDimitry Andric if (SPR.Flavor != SPF_ABS)
2385ffd83dbSDimitry Andric return false;
23981ad6265SDimitry Andric }
2405ffd83dbSDimitry Andric
2415ffd83dbSDimitry Andric // Need a subtract of two values.
2425ffd83dbSDimitry Andric auto *Sub = dyn_cast<BinaryOperator>(LHS);
2435ffd83dbSDimitry Andric if (!Sub || Sub->getOpcode() != Instruction::Sub)
2445ffd83dbSDimitry Andric return false;
2455ffd83dbSDimitry Andric
2465ffd83dbSDimitry Andric // Look for zero extend from i8.
2475ffd83dbSDimitry Andric auto getZeroExtendedVal = [](Value *Op) -> Value * {
2485ffd83dbSDimitry Andric if (auto *ZExt = dyn_cast<ZExtInst>(Op))
2495ffd83dbSDimitry Andric if (cast<VectorType>(ZExt->getOperand(0)->getType())
2505ffd83dbSDimitry Andric ->getElementType()
2515ffd83dbSDimitry Andric ->isIntegerTy(8))
2525ffd83dbSDimitry Andric return ZExt->getOperand(0);
2535ffd83dbSDimitry Andric
2545ffd83dbSDimitry Andric return nullptr;
2555ffd83dbSDimitry Andric };
2565ffd83dbSDimitry Andric
2575ffd83dbSDimitry Andric // Both operands of the subtract should be extends from vXi8.
2585ffd83dbSDimitry Andric Value *Op0 = getZeroExtendedVal(Sub->getOperand(0));
2595ffd83dbSDimitry Andric Value *Op1 = getZeroExtendedVal(Sub->getOperand(1));
2605ffd83dbSDimitry Andric if (!Op0 || !Op1)
2615ffd83dbSDimitry Andric return false;
2625ffd83dbSDimitry Andric
26381ad6265SDimitry Andric IRBuilder<> Builder(Op);
2645ffd83dbSDimitry Andric
2655ffd83dbSDimitry Andric auto *OpTy = cast<FixedVectorType>(Op->getType());
2665ffd83dbSDimitry Andric unsigned NumElts = OpTy->getNumElements();
2675ffd83dbSDimitry Andric
2685ffd83dbSDimitry Andric unsigned IntrinsicNumElts;
2695ffd83dbSDimitry Andric Intrinsic::ID IID;
2705ffd83dbSDimitry Andric if (ST->hasBWI() && NumElts >= 64) {
2715ffd83dbSDimitry Andric IID = Intrinsic::x86_avx512_psad_bw_512;
2725ffd83dbSDimitry Andric IntrinsicNumElts = 64;
2735ffd83dbSDimitry Andric } else if (ST->hasAVX2() && NumElts >= 32) {
2745ffd83dbSDimitry Andric IID = Intrinsic::x86_avx2_psad_bw;
2755ffd83dbSDimitry Andric IntrinsicNumElts = 32;
2765ffd83dbSDimitry Andric } else {
2775ffd83dbSDimitry Andric IID = Intrinsic::x86_sse2_psad_bw;
2785ffd83dbSDimitry Andric IntrinsicNumElts = 16;
2795ffd83dbSDimitry Andric }
2805ffd83dbSDimitry Andric
28181ad6265SDimitry Andric Function *PSADBWFn = Intrinsic::getDeclaration(Op->getModule(), IID);
2825ffd83dbSDimitry Andric
2835ffd83dbSDimitry Andric if (NumElts < 16) {
2845ffd83dbSDimitry Andric // Pad input with zeroes.
2855ffd83dbSDimitry Andric SmallVector<int, 32> ConcatMask(16);
2865ffd83dbSDimitry Andric for (unsigned i = 0; i != NumElts; ++i)
2875ffd83dbSDimitry Andric ConcatMask[i] = i;
2885ffd83dbSDimitry Andric for (unsigned i = NumElts; i != 16; ++i)
2895ffd83dbSDimitry Andric ConcatMask[i] = (i % NumElts) + NumElts;
2905ffd83dbSDimitry Andric
2915ffd83dbSDimitry Andric Value *Zero = Constant::getNullValue(Op0->getType());
2925ffd83dbSDimitry Andric Op0 = Builder.CreateShuffleVector(Op0, Zero, ConcatMask);
2935ffd83dbSDimitry Andric Op1 = Builder.CreateShuffleVector(Op1, Zero, ConcatMask);
2945ffd83dbSDimitry Andric NumElts = 16;
2955ffd83dbSDimitry Andric }
2965ffd83dbSDimitry Andric
2975ffd83dbSDimitry Andric // Intrinsics produce vXi64 and need to be casted to vXi32.
2985ffd83dbSDimitry Andric auto *I32Ty =
2995ffd83dbSDimitry Andric FixedVectorType::get(Builder.getInt32Ty(), IntrinsicNumElts / 4);
3005ffd83dbSDimitry Andric
3015ffd83dbSDimitry Andric assert(NumElts % IntrinsicNumElts == 0 && "Unexpected number of elements!");
3025ffd83dbSDimitry Andric unsigned NumSplits = NumElts / IntrinsicNumElts;
3035ffd83dbSDimitry Andric
3045ffd83dbSDimitry Andric // First collect the pieces we need.
3055ffd83dbSDimitry Andric SmallVector<Value *, 4> Ops(NumSplits);
3065ffd83dbSDimitry Andric for (unsigned i = 0; i != NumSplits; ++i) {
3075ffd83dbSDimitry Andric SmallVector<int, 64> ExtractMask(IntrinsicNumElts);
3085ffd83dbSDimitry Andric std::iota(ExtractMask.begin(), ExtractMask.end(), i * IntrinsicNumElts);
3095ffd83dbSDimitry Andric Value *ExtractOp0 = Builder.CreateShuffleVector(Op0, Op0, ExtractMask);
3105ffd83dbSDimitry Andric Value *ExtractOp1 = Builder.CreateShuffleVector(Op1, Op0, ExtractMask);
3115ffd83dbSDimitry Andric Ops[i] = Builder.CreateCall(PSADBWFn, {ExtractOp0, ExtractOp1});
3125ffd83dbSDimitry Andric Ops[i] = Builder.CreateBitCast(Ops[i], I32Ty);
3135ffd83dbSDimitry Andric }
3145ffd83dbSDimitry Andric
3155ffd83dbSDimitry Andric assert(isPowerOf2_32(NumSplits) && "Expected power of 2 splits");
3165ffd83dbSDimitry Andric unsigned Stages = Log2_32(NumSplits);
3175ffd83dbSDimitry Andric for (unsigned s = Stages; s > 0; --s) {
3185ffd83dbSDimitry Andric unsigned NumConcatElts =
3195ffd83dbSDimitry Andric cast<FixedVectorType>(Ops[0]->getType())->getNumElements() * 2;
3205ffd83dbSDimitry Andric for (unsigned i = 0; i != 1U << (s - 1); ++i) {
3215ffd83dbSDimitry Andric SmallVector<int, 64> ConcatMask(NumConcatElts);
3225ffd83dbSDimitry Andric std::iota(ConcatMask.begin(), ConcatMask.end(), 0);
3235ffd83dbSDimitry Andric Ops[i] = Builder.CreateShuffleVector(Ops[i*2], Ops[i*2+1], ConcatMask);
3245ffd83dbSDimitry Andric }
3255ffd83dbSDimitry Andric }
3265ffd83dbSDimitry Andric
3275ffd83dbSDimitry Andric // At this point the final value should be in Ops[0]. Now we need to adjust
3285ffd83dbSDimitry Andric // it to the final original type.
3295ffd83dbSDimitry Andric NumElts = cast<FixedVectorType>(OpTy)->getNumElements();
3305ffd83dbSDimitry Andric if (NumElts == 2) {
3315ffd83dbSDimitry Andric // Extract down to 2 elements.
3325ffd83dbSDimitry Andric Ops[0] = Builder.CreateShuffleVector(Ops[0], Ops[0], ArrayRef<int>{0, 1});
3335ffd83dbSDimitry Andric } else if (NumElts >= 8) {
3345ffd83dbSDimitry Andric SmallVector<int, 32> ConcatMask(NumElts);
3355ffd83dbSDimitry Andric unsigned SubElts =
3365ffd83dbSDimitry Andric cast<FixedVectorType>(Ops[0]->getType())->getNumElements();
3375ffd83dbSDimitry Andric for (unsigned i = 0; i != SubElts; ++i)
3385ffd83dbSDimitry Andric ConcatMask[i] = i;
3395ffd83dbSDimitry Andric for (unsigned i = SubElts; i != NumElts; ++i)
3405ffd83dbSDimitry Andric ConcatMask[i] = (i % SubElts) + SubElts;
3415ffd83dbSDimitry Andric
3425ffd83dbSDimitry Andric Value *Zero = Constant::getNullValue(Ops[0]->getType());
3435ffd83dbSDimitry Andric Ops[0] = Builder.CreateShuffleVector(Ops[0], Zero, ConcatMask);
3445ffd83dbSDimitry Andric }
3455ffd83dbSDimitry Andric
34681ad6265SDimitry Andric Op->replaceAllUsesWith(Ops[0]);
34781ad6265SDimitry Andric Op->eraseFromParent();
3485ffd83dbSDimitry Andric
3495ffd83dbSDimitry Andric return true;
3505ffd83dbSDimitry Andric }
3515ffd83dbSDimitry Andric
3525ffd83dbSDimitry Andric // Walk backwards from the ExtractElementInst and determine if it is the end of
3535ffd83dbSDimitry Andric // a horizontal reduction. Return the input to the reduction if we find one.
matchAddReduction(const ExtractElementInst & EE,bool & ReduceInOneBB)35404eeddc0SDimitry Andric static Value *matchAddReduction(const ExtractElementInst &EE,
35504eeddc0SDimitry Andric bool &ReduceInOneBB) {
35604eeddc0SDimitry Andric ReduceInOneBB = true;
3575ffd83dbSDimitry Andric // Make sure we're extracting index 0.
3585ffd83dbSDimitry Andric auto *Index = dyn_cast<ConstantInt>(EE.getIndexOperand());
3595ffd83dbSDimitry Andric if (!Index || !Index->isNullValue())
3605ffd83dbSDimitry Andric return nullptr;
3615ffd83dbSDimitry Andric
3625ffd83dbSDimitry Andric const auto *BO = dyn_cast<BinaryOperator>(EE.getVectorOperand());
3635ffd83dbSDimitry Andric if (!BO || BO->getOpcode() != Instruction::Add || !BO->hasOneUse())
3645ffd83dbSDimitry Andric return nullptr;
36504eeddc0SDimitry Andric if (EE.getParent() != BO->getParent())
36604eeddc0SDimitry Andric ReduceInOneBB = false;
3675ffd83dbSDimitry Andric
3685ffd83dbSDimitry Andric unsigned NumElems = cast<FixedVectorType>(BO->getType())->getNumElements();
3695ffd83dbSDimitry Andric // Ensure the reduction size is a power of 2.
3705ffd83dbSDimitry Andric if (!isPowerOf2_32(NumElems))
3715ffd83dbSDimitry Andric return nullptr;
3725ffd83dbSDimitry Andric
3735ffd83dbSDimitry Andric const Value *Op = BO;
3745ffd83dbSDimitry Andric unsigned Stages = Log2_32(NumElems);
3755ffd83dbSDimitry Andric for (unsigned i = 0; i != Stages; ++i) {
3765ffd83dbSDimitry Andric const auto *BO = dyn_cast<BinaryOperator>(Op);
3775ffd83dbSDimitry Andric if (!BO || BO->getOpcode() != Instruction::Add)
3785ffd83dbSDimitry Andric return nullptr;
37904eeddc0SDimitry Andric if (EE.getParent() != BO->getParent())
38004eeddc0SDimitry Andric ReduceInOneBB = false;
3815ffd83dbSDimitry Andric
3825ffd83dbSDimitry Andric // If this isn't the first add, then it should only have 2 users, the
3835ffd83dbSDimitry Andric // shuffle and another add which we checked in the previous iteration.
3845ffd83dbSDimitry Andric if (i != 0 && !BO->hasNUses(2))
3855ffd83dbSDimitry Andric return nullptr;
3865ffd83dbSDimitry Andric
3875ffd83dbSDimitry Andric Value *LHS = BO->getOperand(0);
3885ffd83dbSDimitry Andric Value *RHS = BO->getOperand(1);
3895ffd83dbSDimitry Andric
3905ffd83dbSDimitry Andric auto *Shuffle = dyn_cast<ShuffleVectorInst>(LHS);
3915ffd83dbSDimitry Andric if (Shuffle) {
3925ffd83dbSDimitry Andric Op = RHS;
3935ffd83dbSDimitry Andric } else {
3945ffd83dbSDimitry Andric Shuffle = dyn_cast<ShuffleVectorInst>(RHS);
3955ffd83dbSDimitry Andric Op = LHS;
3965ffd83dbSDimitry Andric }
3975ffd83dbSDimitry Andric
3985ffd83dbSDimitry Andric // The first operand of the shuffle should be the same as the other operand
3995ffd83dbSDimitry Andric // of the bin op.
4005ffd83dbSDimitry Andric if (!Shuffle || Shuffle->getOperand(0) != Op)
4015ffd83dbSDimitry Andric return nullptr;
4025ffd83dbSDimitry Andric
4035ffd83dbSDimitry Andric // Verify the shuffle has the expected (at this stage of the pyramid) mask.
4045ffd83dbSDimitry Andric unsigned MaskEnd = 1 << i;
4055ffd83dbSDimitry Andric for (unsigned Index = 0; Index < MaskEnd; ++Index)
4065ffd83dbSDimitry Andric if (Shuffle->getMaskValue(Index) != (int)(MaskEnd + Index))
4075ffd83dbSDimitry Andric return nullptr;
4085ffd83dbSDimitry Andric }
4095ffd83dbSDimitry Andric
4105ffd83dbSDimitry Andric return const_cast<Value *>(Op);
4115ffd83dbSDimitry Andric }
4125ffd83dbSDimitry Andric
4135ffd83dbSDimitry Andric // See if this BO is reachable from this Phi by walking forward through single
4145ffd83dbSDimitry Andric // use BinaryOperators with the same opcode. If we get back then we know we've
4155ffd83dbSDimitry Andric // found a loop and it is safe to step through this Add to find more leaves.
isReachableFromPHI(PHINode * Phi,BinaryOperator * BO)4165ffd83dbSDimitry Andric static bool isReachableFromPHI(PHINode *Phi, BinaryOperator *BO) {
4175ffd83dbSDimitry Andric // The PHI itself should only have one use.
4185ffd83dbSDimitry Andric if (!Phi->hasOneUse())
4195ffd83dbSDimitry Andric return false;
4205ffd83dbSDimitry Andric
4215ffd83dbSDimitry Andric Instruction *U = cast<Instruction>(*Phi->user_begin());
4225ffd83dbSDimitry Andric if (U == BO)
4235ffd83dbSDimitry Andric return true;
4245ffd83dbSDimitry Andric
4255ffd83dbSDimitry Andric while (U->hasOneUse() && U->getOpcode() == BO->getOpcode())
4265ffd83dbSDimitry Andric U = cast<Instruction>(*U->user_begin());
4275ffd83dbSDimitry Andric
4285ffd83dbSDimitry Andric return U == BO;
4295ffd83dbSDimitry Andric }
4305ffd83dbSDimitry Andric
4315ffd83dbSDimitry Andric // Collect all the leaves of the tree of adds that feeds into the horizontal
4325ffd83dbSDimitry Andric // reduction. Root is the Value that is used by the horizontal reduction.
4335ffd83dbSDimitry Andric // We look through single use phis, single use adds, or adds that are used by
4345ffd83dbSDimitry Andric // a phi that forms a loop with the add.
collectLeaves(Value * Root,SmallVectorImpl<Instruction * > & Leaves)4355ffd83dbSDimitry Andric static void collectLeaves(Value *Root, SmallVectorImpl<Instruction *> &Leaves) {
4365ffd83dbSDimitry Andric SmallPtrSet<Value *, 8> Visited;
4375ffd83dbSDimitry Andric SmallVector<Value *, 8> Worklist;
4385ffd83dbSDimitry Andric Worklist.push_back(Root);
4395ffd83dbSDimitry Andric
4405ffd83dbSDimitry Andric while (!Worklist.empty()) {
4415ffd83dbSDimitry Andric Value *V = Worklist.pop_back_val();
4425ffd83dbSDimitry Andric if (!Visited.insert(V).second)
4435ffd83dbSDimitry Andric continue;
4445ffd83dbSDimitry Andric
4455ffd83dbSDimitry Andric if (auto *PN = dyn_cast<PHINode>(V)) {
4465ffd83dbSDimitry Andric // PHI node should have single use unless it is the root node, then it
4475ffd83dbSDimitry Andric // has 2 uses.
4485ffd83dbSDimitry Andric if (!PN->hasNUses(PN == Root ? 2 : 1))
4495ffd83dbSDimitry Andric break;
4505ffd83dbSDimitry Andric
4515ffd83dbSDimitry Andric // Push incoming values to the worklist.
452e8d8bef9SDimitry Andric append_range(Worklist, PN->incoming_values());
4535ffd83dbSDimitry Andric
4545ffd83dbSDimitry Andric continue;
4555ffd83dbSDimitry Andric }
4565ffd83dbSDimitry Andric
4575ffd83dbSDimitry Andric if (auto *BO = dyn_cast<BinaryOperator>(V)) {
4585ffd83dbSDimitry Andric if (BO->getOpcode() == Instruction::Add) {
4595ffd83dbSDimitry Andric // Simple case. Single use, just push its operands to the worklist.
4605ffd83dbSDimitry Andric if (BO->hasNUses(BO == Root ? 2 : 1)) {
461e8d8bef9SDimitry Andric append_range(Worklist, BO->operands());
4625ffd83dbSDimitry Andric continue;
4635ffd83dbSDimitry Andric }
4645ffd83dbSDimitry Andric
4655ffd83dbSDimitry Andric // If there is additional use, make sure it is an unvisited phi that
4665ffd83dbSDimitry Andric // gets us back to this node.
4675ffd83dbSDimitry Andric if (BO->hasNUses(BO == Root ? 3 : 2)) {
4685ffd83dbSDimitry Andric PHINode *PN = nullptr;
469753f127fSDimitry Andric for (auto *U : BO->users())
4705ffd83dbSDimitry Andric if (auto *P = dyn_cast<PHINode>(U))
4715ffd83dbSDimitry Andric if (!Visited.count(P))
4725ffd83dbSDimitry Andric PN = P;
4735ffd83dbSDimitry Andric
4745ffd83dbSDimitry Andric // If we didn't find a 2-input PHI then this isn't a case we can
4755ffd83dbSDimitry Andric // handle.
4765ffd83dbSDimitry Andric if (!PN || PN->getNumIncomingValues() != 2)
4775ffd83dbSDimitry Andric continue;
4785ffd83dbSDimitry Andric
4795ffd83dbSDimitry Andric // Walk forward from this phi to see if it reaches back to this add.
4805ffd83dbSDimitry Andric if (!isReachableFromPHI(PN, BO))
4815ffd83dbSDimitry Andric continue;
4825ffd83dbSDimitry Andric
4835ffd83dbSDimitry Andric // The phi forms a loop with this Add, push its operands.
484e8d8bef9SDimitry Andric append_range(Worklist, BO->operands());
4855ffd83dbSDimitry Andric }
4865ffd83dbSDimitry Andric }
4875ffd83dbSDimitry Andric }
4885ffd83dbSDimitry Andric
4895ffd83dbSDimitry Andric // Not an add or phi, make it a leaf.
4905ffd83dbSDimitry Andric if (auto *I = dyn_cast<Instruction>(V)) {
4915ffd83dbSDimitry Andric if (!V->hasNUses(I == Root ? 2 : 1))
4925ffd83dbSDimitry Andric continue;
4935ffd83dbSDimitry Andric
4945ffd83dbSDimitry Andric // Add this as a leaf.
4955ffd83dbSDimitry Andric Leaves.push_back(I);
4965ffd83dbSDimitry Andric }
4975ffd83dbSDimitry Andric }
4985ffd83dbSDimitry Andric }
4995ffd83dbSDimitry Andric
runOnFunction(Function & F)5005ffd83dbSDimitry Andric bool X86PartialReduction::runOnFunction(Function &F) {
5015ffd83dbSDimitry Andric if (skipFunction(F))
5025ffd83dbSDimitry Andric return false;
5035ffd83dbSDimitry Andric
5045ffd83dbSDimitry Andric auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
5055ffd83dbSDimitry Andric if (!TPC)
5065ffd83dbSDimitry Andric return false;
5075ffd83dbSDimitry Andric
5085ffd83dbSDimitry Andric auto &TM = TPC->getTM<X86TargetMachine>();
5095ffd83dbSDimitry Andric ST = TM.getSubtargetImpl(F);
5105ffd83dbSDimitry Andric
5115ffd83dbSDimitry Andric DL = &F.getParent()->getDataLayout();
5125ffd83dbSDimitry Andric
5135ffd83dbSDimitry Andric bool MadeChange = false;
5145ffd83dbSDimitry Andric for (auto &BB : F) {
5155ffd83dbSDimitry Andric for (auto &I : BB) {
5165ffd83dbSDimitry Andric auto *EE = dyn_cast<ExtractElementInst>(&I);
5175ffd83dbSDimitry Andric if (!EE)
5185ffd83dbSDimitry Andric continue;
5195ffd83dbSDimitry Andric
52004eeddc0SDimitry Andric bool ReduceInOneBB;
5215ffd83dbSDimitry Andric // First find a reduction tree.
5225ffd83dbSDimitry Andric // FIXME: Do we need to handle other opcodes than Add?
52304eeddc0SDimitry Andric Value *Root = matchAddReduction(*EE, ReduceInOneBB);
5245ffd83dbSDimitry Andric if (!Root)
5255ffd83dbSDimitry Andric continue;
5265ffd83dbSDimitry Andric
5275ffd83dbSDimitry Andric SmallVector<Instruction *, 8> Leaves;
5285ffd83dbSDimitry Andric collectLeaves(Root, Leaves);
5295ffd83dbSDimitry Andric
5305ffd83dbSDimitry Andric for (Instruction *I : Leaves) {
53104eeddc0SDimitry Andric if (tryMAddReplacement(I, ReduceInOneBB)) {
5325ffd83dbSDimitry Andric MadeChange = true;
5335ffd83dbSDimitry Andric continue;
5345ffd83dbSDimitry Andric }
5355ffd83dbSDimitry Andric
5365ffd83dbSDimitry Andric // Don't do SAD matching on the root node. SelectionDAG already
5375ffd83dbSDimitry Andric // has support for that and currently generates better code.
5385ffd83dbSDimitry Andric if (I != Root && trySADReplacement(I))
5395ffd83dbSDimitry Andric MadeChange = true;
5405ffd83dbSDimitry Andric }
5415ffd83dbSDimitry Andric }
5425ffd83dbSDimitry Andric }
5435ffd83dbSDimitry Andric
5445ffd83dbSDimitry Andric return MadeChange;
5455ffd83dbSDimitry Andric }
546