maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity
llvm-project/llvm/lib/Target/X86/X86PartialReduction.cpp

464 lines
15 KiB
C++
Raw Blame History

//===-- X86PartialReduction.cpp -------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This pass looks for add instructions used by a horizontal reduction to see
// if we might be able to use pmaddwd or psadbw. Some cases of this require
// cross basic block knowledge and can't be done in SelectionDAG.
//
//===----------------------------------------------------------------------===//
#include "X86.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsX86.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Operator.h"
#include "llvm/Pass.h"
#include "X86TargetMachine.h"
using namespace llvm;
#define DEBUG_TYPE "x86-partial-reduction"
namespace {
class X86PartialReduction : public FunctionPass {
const DataLayout *DL;
const X86Subtarget *ST;
public:
static char ID; // Pass identification, replacement for typeid.
X86PartialReduction() : FunctionPass(ID) { }
bool runOnFunction(Function &Fn) override;
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.setPreservesCFG();
}
StringRef getPassName() const override {
return "X86 Partial Reduction";
}
private:
bool tryMAddPattern(BinaryOperator *BO);
bool tryMAddReplacement(Value *Op, BinaryOperator *Add);
bool trySADPattern(BinaryOperator *BO);
bool trySADReplacement(Value *Op, BinaryOperator *Add);
};
}
FunctionPass *llvm::createX86PartialReductionPass() {
return new X86PartialReduction();
}
char X86PartialReduction::ID = 0;
INITIALIZE_PASS(X86PartialReduction, DEBUG_TYPE,
"X86 Partial Reduction", false, false)
static bool isVectorReductionOp(const BinaryOperator &BO) {
if (!BO.getType()->isVectorTy())
return false;
unsigned Opcode = BO.getOpcode();
switch (Opcode) {
case Instruction::Add:
case Instruction::Mul:
case Instruction::And:
case Instruction::Or:
case Instruction::Xor:
break;
case Instruction::FAdd:
case Instruction::FMul:
if (auto *FPOp = dyn_cast<FPMathOperator>(&BO))
if (FPOp->getFastMathFlags().isFast())
break;
LLVM_FALLTHROUGH;
default:
return false;
}
unsigned ElemNum = cast<VectorType>(BO.getType())->getNumElements();
// Ensure the reduction size is a power of 2.
if (!isPowerOf2_32(ElemNum))
return false;
unsigned ElemNumToReduce = ElemNum;
// Do DFS search on the def-use chain from the given instruction. We only
// allow four kinds of operations during the search until we reach the
// instruction that extracts the first element from the vector:
//
// 1. The reduction operation of the same opcode as the given instruction.
//
// 2. PHI node.
//
// 3. ShuffleVector instruction together with a reduction operation that
// does a partial reduction.
//
// 4. ExtractElement that extracts the first element from the vector, and we
// stop searching the def-use chain here.
//
// 3 & 4 above perform a reduction on all elements of the vector. We push defs
// from 1-3 to the stack to continue the DFS. The given instruction is not
// a reduction operation if we meet any other instructions other than those
// listed above.
SmallVector<const User *, 16> UsersToVisit{&BO};
SmallPtrSet<const User *, 16> Visited;
bool ReduxExtracted = false;
while (!UsersToVisit.empty()) {
auto User = UsersToVisit.back();
UsersToVisit.pop_back();
if (!Visited.insert(User).second)
continue;
for (const auto *U : User->users()) {
auto *Inst = dyn_cast<Instruction>(U);
if (!Inst)
return false;
if (Inst->getOpcode() == Opcode || isa<PHINode>(U)) {
if (auto *FPOp = dyn_cast<FPMathOperator>(Inst))
if (!isa<PHINode>(FPOp) && !FPOp->getFastMathFlags().isFast())
return false;
UsersToVisit.push_back(U);
} else if (auto *ShufInst = dyn_cast<ShuffleVectorInst>(U)) {
// Detect the following pattern: A ShuffleVector instruction together
// with a reduction that do partial reduction on the first and second
// ElemNumToReduce / 2 elements, and store the result in
// ElemNumToReduce / 2 elements in another vector.
unsigned ResultElements = ShufInst->getType()->getNumElements();
if (ResultElements < ElemNum)
return false;
if (ElemNumToReduce == 1)
return false;
if (!isa<UndefValue>(U->getOperand(1)))
return false;
for (unsigned i = 0; i < ElemNumToReduce / 2; ++i)
if (ShufInst->getMaskValue(i) != int(i + ElemNumToReduce / 2))
return false;
for (unsigned i = ElemNumToReduce / 2; i < ElemNum; ++i)
if (ShufInst->getMaskValue(i) != -1)
return false;
// There is only one user of this ShuffleVector instruction, which
// must be a reduction operation.
if (!U->hasOneUse())
return false;
auto *U2 = dyn_cast<BinaryOperator>(*U->user_begin());
if (!U2 || U2->getOpcode() != Opcode)
return false;
// Check operands of the reduction operation.
if ((U2->getOperand(0) == U->getOperand(0) && U2->getOperand(1) == U) ||
(U2->getOperand(1) == U->getOperand(0) && U2->getOperand(0) == U)) {
UsersToVisit.push_back(U2);
ElemNumToReduce /= 2;
} else
return false;
} else if (isa<ExtractElementInst>(U)) {
// At this moment we should have reduced all elements in the vector.
if (ElemNumToReduce != 1)
return false;
auto *Val = dyn_cast<ConstantInt>(U->getOperand(1));
if (!Val || !Val->isZero())
return false;
ReduxExtracted = true;
} else
return false;
}
}
return ReduxExtracted;
}
bool X86PartialReduction::tryMAddReplacement(Value *Op, BinaryOperator *Add) {
BasicBlock *BB = Add->getParent();
auto *BO = dyn_cast<BinaryOperator>(Op);
if (!BO || BO->getOpcode() != Instruction::Mul || !BO->hasOneUse() ||
BO->getParent() != BB)
return false;
Value *LHS = BO->getOperand(0);
Value *RHS = BO->getOperand(1);
// LHS and RHS should be only used once or if they are the same then only
// used twice. Only check this when SSE4.1 is enabled and we have zext/sext
// instructions, otherwise we use punpck to emulate zero extend in stages. The
// trunc/ we need to do likely won't introduce new instructions in that case.
if (ST->hasSSE41()) {
if (LHS == RHS) {
if (!isa<Constant>(LHS) && !LHS->hasNUses(2))
return false;
} else {
if (!isa<Constant>(LHS) && !LHS->hasOneUse())
return false;
if (!isa<Constant>(RHS) && !RHS->hasOneUse())
return false;
}
}
auto canShrinkOp = [&](Value *Op) {
if (isa<Constant>(Op) && ComputeNumSignBits(Op, *DL, 0, nullptr, BO) > 16)
return true;
if (auto *Cast = dyn_cast<CastInst>(Op)) {
if (Cast->getParent() == BB &&
(Cast->getOpcode() == Instruction::SExt ||
Cast->getOpcode() == Instruction::ZExt) &&
ComputeNumSignBits(Op, *DL, 0, nullptr, BO) > 16)
return true;
}
return false;
};
// Both Ops need to be shrinkable.
if (!canShrinkOp(LHS) && !canShrinkOp(RHS))
return false;
IRBuilder<> Builder(Add);
auto *MulTy = cast<VectorType>(Op->getType());
unsigned NumElts = MulTy->getNumElements();
// Extract even elements and odd elements and add them together. This will
// be pattern matched by SelectionDAG to pmaddwd. This instruction will be
// half the original width.
SmallVector<int, 16> EvenMask(NumElts / 2);
SmallVector<int, 16> OddMask(NumElts / 2);
for (int i = 0, e = NumElts / 2; i != e; ++i) {
EvenMask[i] = i * 2;
OddMask[i] = i * 2 + 1;
}
Value *EvenElts = Builder.CreateShuffleVector(BO, BO, EvenMask);
Value *OddElts = Builder.CreateShuffleVector(BO, BO, OddMask);
Value *MAdd = Builder.CreateAdd(EvenElts, OddElts);
// Concatenate zeroes to extend back to the original type.
SmallVector<int, 32> ConcatMask(NumElts);
std::iota(ConcatMask.begin(), ConcatMask.end(), 0);
Value *Zero = Constant::getNullValue(MAdd->getType());
Value *Concat = Builder.CreateShuffleVector(MAdd, Zero, ConcatMask);
// Replaces the use of mul in the original Add with the pmaddwd and zeroes.
Add->replaceUsesOfWith(BO, Concat);
Add->setHasNoSignedWrap(false);
Add->setHasNoUnsignedWrap(false);
return true;
}
// Try to replace operans of this add with pmaddwd patterns.
bool X86PartialReduction::tryMAddPattern(BinaryOperator *BO) {
if (!ST->hasSSE2())
return false;
// Need at least 8 elements.
if (cast<VectorType>(BO->getType())->getNumElements() < 8)
return false;
// Element type should be i32.
if (!cast<VectorType>(BO->getType())->getElementType()->isIntegerTy(32))
return false;
bool Changed = false;
Changed |= tryMAddReplacement(BO->getOperand(0), BO);
Changed |= tryMAddReplacement(BO->getOperand(1), BO);
return Changed;
}
bool X86PartialReduction::trySADReplacement(Value *Op, BinaryOperator *Add) {
// Operand should be a select.
auto *SI = dyn_cast<SelectInst>(Op);
if (!SI)
return false;
// Select needs to implement absolute value.
Value *LHS, *RHS;
auto SPR = matchSelectPattern(SI, LHS, RHS);
if (SPR.Flavor != SPF_ABS)
return false;
// Need a subtract of two values.
auto *Sub = dyn_cast<BinaryOperator>(LHS);
if (!Sub || Sub->getOpcode() != Instruction::Sub)
return false;
// Look for zero extend from i8.
auto getZeroExtendedVal = [](Value *Op) -> Value * {
if (auto *ZExt = dyn_cast<ZExtInst>(Op))
if (cast<VectorType>(ZExt->getOperand(0)->getType())
->getElementType()
->isIntegerTy(8))
return ZExt->getOperand(0);
return nullptr;
};
// Both operands of the subtract should be extends from vXi8.
Value *Op0 = getZeroExtendedVal(Sub->getOperand(0));
Value *Op1 = getZeroExtendedVal(Sub->getOperand(1));
if (!Op0 || !Op1)
return false;
IRBuilder<> Builder(Add);
auto *OpTy = cast<VectorType>(Op->getType());
unsigned NumElts = OpTy->getNumElements();
unsigned IntrinsicNumElts;
Intrinsic::ID IID;
if (ST->hasBWI() && NumElts >= 64) {
IID = Intrinsic::x86_avx512_psad_bw_512;
IntrinsicNumElts = 64;
} else if (ST->hasAVX2() && NumElts >= 32) {
IID = Intrinsic::x86_avx2_psad_bw;
IntrinsicNumElts = 32;
} else {
IID = Intrinsic::x86_sse2_psad_bw;
IntrinsicNumElts = 16;
}
Function *PSADBWFn = Intrinsic::getDeclaration(Add->getModule(), IID);
if (NumElts < 16) {
// Pad input with zeroes.
SmallVector<int, 32> ConcatMask(16);
for (unsigned i = 0; i != NumElts; ++i)
ConcatMask[i] = i;
for (unsigned i = NumElts; i != 16; ++i)
ConcatMask[i] = (i % NumElts) + NumElts;
Value *Zero = Constant::getNullValue(Op0->getType());
Op0 = Builder.CreateShuffleVector(Op0, Zero, ConcatMask);
Op1 = Builder.CreateShuffleVector(Op1, Zero, ConcatMask);
NumElts = 16;
}
// Intrinsics produce vXi64 and need to be casted to vXi32.
Type *I32Ty = VectorType::get(Builder.getInt32Ty(), IntrinsicNumElts / 4);
assert(NumElts % IntrinsicNumElts == 0 && "Unexpected number of elements!");
unsigned NumSplits = NumElts / IntrinsicNumElts;
// First collect the pieces we need.
SmallVector<Value *, 4> Ops(NumSplits);
for (unsigned i = 0; i != NumSplits; ++i) {
SmallVector<int, 64> ExtractMask(IntrinsicNumElts);
std::iota(ExtractMask.begin(), ExtractMask.end(), i * IntrinsicNumElts);
Value *ExtractOp0 = Builder.CreateShuffleVector(Op0, Op0, ExtractMask);
Value *ExtractOp1 = Builder.CreateShuffleVector(Op1, Op0, ExtractMask);
Ops[i] = Builder.CreateCall(PSADBWFn, {ExtractOp0, ExtractOp1});
Ops[i] = Builder.CreateBitCast(Ops[i], I32Ty);
}
assert(isPowerOf2_32(NumSplits) && "Expected power of 2 splits");
unsigned Stages = Log2_32(NumSplits);
for (unsigned s = Stages; s > 0; --s) {
unsigned NumConcatElts =
cast<VectorType>(Ops[0]->getType())->getNumElements() * 2;
for (unsigned i = 0; i != 1U << (s - 1); ++i) {
SmallVector<int, 64> ConcatMask(NumConcatElts);
std::iota(ConcatMask.begin(), ConcatMask.end(), 0);
Ops[i] = Builder.CreateShuffleVector(Ops[i*2], Ops[i*2+1], ConcatMask);
}
}
// At this point the final value should be in Ops[0]. Now we need to adjust
// it to the final original type.
NumElts = cast<VectorType>(OpTy)->getNumElements();
if (NumElts == 2) {
// Extract down to 2 elements.
Ops[0] = Builder.CreateShuffleVector(Ops[0], Ops[0], ArrayRef<int>{0, 1});
} else if (NumElts >= 8) {
SmallVector<int, 32> ConcatMask(NumElts);
unsigned SubElts = cast<VectorType>(Ops[0]->getType())->getNumElements();
for (unsigned i = 0; i != SubElts; ++i)
ConcatMask[i] = i;
for (unsigned i = SubElts; i != NumElts; ++i)
ConcatMask[i] = (i % SubElts) + SubElts;
Value *Zero = Constant::getNullValue(Ops[0]->getType());
Ops[0] = Builder.CreateShuffleVector(Ops[0], Zero, ConcatMask);
}
// Replaces the uses of Op in Add with the new sequence.
Add->replaceUsesOfWith(Op, Ops[0]);
Add->setHasNoSignedWrap(false);
Add->setHasNoUnsignedWrap(false);
return false;
}
bool X86PartialReduction::trySADPattern(BinaryOperator *BO) {
if (!ST->hasSSE2())
return false;
// TODO: There's nothing special about i32, any integer type above i16 should
// work just as well.
if (!cast<VectorType>(BO->getType())->getElementType()->isIntegerTy(32))
return false;
bool Changed = false;
Changed |= trySADReplacement(BO->getOperand(0), BO);
Changed |= trySADReplacement(BO->getOperand(1), BO);
return Changed;
}
bool X86PartialReduction::runOnFunction(Function &F) {
if (skipFunction(F))
return false;
auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
if (!TPC)
return false;
auto &TM = TPC->getTM<X86TargetMachine>();
ST = TM.getSubtargetImpl(F);
DL = &F.getParent()->getDataLayout();
bool MadeChange = false;
for (auto &BB : F) {
for (auto &I : BB) {
auto *BO = dyn_cast<BinaryOperator>(&I);
if (!BO)
continue;
if (!isVectorReductionOp(*BO))
continue;
if (BO->getOpcode() == Instruction::Add) {
if (tryMAddPattern(BO)) {
MadeChange = true;
continue;
}
if (trySADPattern(BO)) {
MadeChange = true;
continue;
}
}
}
}
return MadeChange;
}
Reference in New Issue View Git Blame Copy Permalink