#include "llvm/Transforms/Scalar/NaryReassociate.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/GetElementPtrTypeIterator.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Operator.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"
#include "llvm/IR/ValueHandle.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
#include <cassert>
#include <cstdint>
using namespace llvm;
using namespace PatternMatch;
#define DEBUG_TYPE "nary-reassociate"
namespace {
class NaryReassociateLegacyPass : public FunctionPass {
public:
static char ID;
NaryReassociateLegacyPass() : FunctionPass(ID) {
initializeNaryReassociateLegacyPassPass(*PassRegistry::getPassRegistry());
}
bool doInitialization(Module &M) override {
return false;
}
bool runOnFunction(Function &F) override;
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.addPreserved<DominatorTreeWrapperPass>();
AU.addPreserved<ScalarEvolutionWrapperPass>();
AU.addPreserved<TargetLibraryInfoWrapperPass>();
AU.addRequired<AssumptionCacheTracker>();
AU.addRequired<DominatorTreeWrapperPass>();
AU.addRequired<ScalarEvolutionWrapperPass>();
AU.addRequired<TargetLibraryInfoWrapperPass>();
AU.addRequired<TargetTransformInfoWrapperPass>();
AU.setPreservesCFG();
}
private:
NaryReassociatePass Impl;
};
}
char NaryReassociateLegacyPass::ID = 0;
INITIALIZE_PASS_BEGIN(NaryReassociateLegacyPass, "nary-reassociate",
"Nary reassociation", false, false)
INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
INITIALIZE_PASS_END(NaryReassociateLegacyPass, "nary-reassociate",
"Nary reassociation", false, false)
FunctionPass *llvm::createNaryReassociatePass() {
return new NaryReassociateLegacyPass();
}
bool NaryReassociateLegacyPass::runOnFunction(Function &F) {
if (skipFunction(F))
return false;
auto *AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
auto *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
auto *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
auto *TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
auto *TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
return Impl.runImpl(F, AC, DT, SE, TLI, TTI);
}
PreservedAnalyses NaryReassociatePass::run(Function &F,
FunctionAnalysisManager &AM) {
auto *AC = &AM.getResult<AssumptionAnalysis>(F);
auto *DT = &AM.getResult<DominatorTreeAnalysis>(F);
auto *SE = &AM.getResult<ScalarEvolutionAnalysis>(F);
auto *TLI = &AM.getResult<TargetLibraryAnalysis>(F);
auto *TTI = &AM.getResult<TargetIRAnalysis>(F);
if (!runImpl(F, AC, DT, SE, TLI, TTI))
return PreservedAnalyses::all();
PreservedAnalyses PA;
PA.preserveSet<CFGAnalyses>();
PA.preserve<ScalarEvolutionAnalysis>();
return PA;
}
bool NaryReassociatePass::runImpl(Function &F, AssumptionCache *AC_,
DominatorTree *DT_, ScalarEvolution *SE_,
TargetLibraryInfo *TLI_,
TargetTransformInfo *TTI_) {
AC = AC_;
DT = DT_;
SE = SE_;
TLI = TLI_;
TTI = TTI_;
DL = &F.getParent()->getDataLayout();
bool Changed = false, ChangedInThisIteration;
do {
ChangedInThisIteration = doOneIteration(F);
Changed |= ChangedInThisIteration;
} while (ChangedInThisIteration);
return Changed;
}
bool NaryReassociatePass::doOneIteration(Function &F) {
bool Changed = false;
SeenExprs.clear();
SmallVector<WeakTrackingVH, 16> DeadInsts;
for (const auto Node : depth_first(DT)) {
BasicBlock *BB = Node->getBlock();
for (Instruction &OrigI : *BB) {
const SCEV *OrigSCEV = nullptr;
if (Instruction *NewI = tryReassociate(&OrigI, OrigSCEV)) {
Changed = true;
OrigI.replaceAllUsesWith(NewI);
DeadInsts.push_back(WeakTrackingVH(&OrigI));
const SCEV *NewSCEV = SE->getSCEV(NewI);
SeenExprs[NewSCEV].push_back(WeakTrackingVH(NewI));
if (NewSCEV != OrigSCEV)
SeenExprs[OrigSCEV].push_back(WeakTrackingVH(NewI));
} else if (OrigSCEV)
SeenExprs[OrigSCEV].push_back(WeakTrackingVH(&OrigI));
}
}
RecursivelyDeleteTriviallyDeadInstructionsPermissive(
DeadInsts, TLI, nullptr, [this](Value *V) { SE->forgetValue(V); });
return Changed;
}
template <typename PredT>
Instruction *
NaryReassociatePass::matchAndReassociateMinOrMax(Instruction *I,
const SCEV *&OrigSCEV) {
Value *LHS = nullptr;
Value *RHS = nullptr;
auto MinMaxMatcher =
MaxMin_match<ICmpInst, bind_ty<Value>, bind_ty<Value>, PredT>(
m_Value(LHS), m_Value(RHS));
if (match(I, MinMaxMatcher)) {
OrigSCEV = SE->getSCEV(I);
if (auto *NewMinMax = dyn_cast_or_null<Instruction>(
tryReassociateMinOrMax(I, MinMaxMatcher, LHS, RHS)))
return NewMinMax;
if (auto *NewMinMax = dyn_cast_or_null<Instruction>(
tryReassociateMinOrMax(I, MinMaxMatcher, RHS, LHS)))
return NewMinMax;
}
return nullptr;
}
Instruction *NaryReassociatePass::tryReassociate(Instruction * I,
const SCEV *&OrigSCEV) {
if (!SE->isSCEVable(I->getType()))
return nullptr;
switch (I->getOpcode()) {
case Instruction::Add:
case Instruction::Mul:
OrigSCEV = SE->getSCEV(I);
return tryReassociateBinaryOp(cast<BinaryOperator>(I));
case Instruction::GetElementPtr:
OrigSCEV = SE->getSCEV(I);
return tryReassociateGEP(cast<GetElementPtrInst>(I));
default:
break;
}
Instruction *ResI = nullptr;
if (I->getType()->isIntegerTy())
if ((ResI = matchAndReassociateMinOrMax<umin_pred_ty>(I, OrigSCEV)) ||
(ResI = matchAndReassociateMinOrMax<smin_pred_ty>(I, OrigSCEV)) ||
(ResI = matchAndReassociateMinOrMax<umax_pred_ty>(I, OrigSCEV)) ||
(ResI = matchAndReassociateMinOrMax<smax_pred_ty>(I, OrigSCEV)))
return ResI;
return nullptr;
}
static bool isGEPFoldable(GetElementPtrInst *GEP,
const TargetTransformInfo *TTI) {
SmallVector<const Value *, 4> Indices(GEP->indices());
return TTI->getGEPCost(GEP->getSourceElementType(), GEP->getPointerOperand(),
Indices) == TargetTransformInfo::TCC_Free;
}
Instruction *NaryReassociatePass::tryReassociateGEP(GetElementPtrInst *GEP) {
if (isGEPFoldable(GEP, TTI))
return nullptr;
gep_type_iterator GTI = gep_type_begin(*GEP);
for (unsigned I = 1, E = GEP->getNumOperands(); I != E; ++I, ++GTI) {
if (GTI.isSequential()) {
if (auto *NewGEP = tryReassociateGEPAtIndex(GEP, I - 1,
GTI.getIndexedType())) {
return NewGEP;
}
}
}
return nullptr;
}
bool NaryReassociatePass::requiresSignExtension(Value *Index,
GetElementPtrInst *GEP) {
unsigned PointerSizeInBits =
DL->getPointerSizeInBits(GEP->getType()->getPointerAddressSpace());
return cast<IntegerType>(Index->getType())->getBitWidth() < PointerSizeInBits;
}
GetElementPtrInst *
NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
unsigned I, Type *IndexedType) {
Value *IndexToSplit = GEP->getOperand(I + 1);
if (SExtInst *SExt = dyn_cast<SExtInst>(IndexToSplit)) {
IndexToSplit = SExt->getOperand(0);
} else if (ZExtInst *ZExt = dyn_cast<ZExtInst>(IndexToSplit)) {
if (isKnownNonNegative(ZExt->getOperand(0), *DL, 0, AC, GEP, DT))
IndexToSplit = ZExt->getOperand(0);
}
if (AddOperator *AO = dyn_cast<AddOperator>(IndexToSplit)) {
if (requiresSignExtension(IndexToSplit, GEP) &&
computeOverflowForSignedAdd(AO, *DL, AC, GEP, DT) !=
OverflowResult::NeverOverflows)
return nullptr;
Value *LHS = AO->getOperand(0), *RHS = AO->getOperand(1);
if (auto *NewGEP = tryReassociateGEPAtIndex(GEP, I, LHS, RHS, IndexedType))
return NewGEP;
if (LHS != RHS) {
if (auto *NewGEP =
tryReassociateGEPAtIndex(GEP, I, RHS, LHS, IndexedType))
return NewGEP;
}
}
return nullptr;
}
GetElementPtrInst *
NaryReassociatePass::tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
unsigned I, Value *LHS,
Value *RHS, Type *IndexedType) {
SmallVector<const SCEV *, 4> IndexExprs;
for (Use &Index : GEP->indices())
IndexExprs.push_back(SE->getSCEV(Index));
IndexExprs[I] = SE->getSCEV(LHS);
if (isKnownNonNegative(LHS, *DL, 0, AC, GEP, DT) &&
DL->getTypeSizeInBits(LHS->getType()).getFixedSize() <
DL->getTypeSizeInBits(GEP->getOperand(I)->getType()).getFixedSize()) {
IndexExprs[I] =
SE->getZeroExtendExpr(IndexExprs[I], GEP->getOperand(I)->getType());
}
const SCEV *CandidateExpr = SE->getGEPExpr(cast<GEPOperator>(GEP),
IndexExprs);
Value *Candidate = findClosestMatchingDominator(CandidateExpr, GEP);
if (Candidate == nullptr)
return nullptr;
IRBuilder<> Builder(GEP);
Candidate = Builder.CreateBitOrPointerCast(Candidate, GEP->getType());
assert(Candidate->getType() == GEP->getType());
uint64_t IndexedSize = DL->getTypeAllocSize(IndexedType);
Type *ElementType = GEP->getResultElementType();
uint64_t ElementSize = DL->getTypeAllocSize(ElementType);
if (IndexedSize % ElementSize != 0)
return nullptr;
Type *IntPtrTy = DL->getIntPtrType(GEP->getType());
if (RHS->getType() != IntPtrTy)
RHS = Builder.CreateSExtOrTrunc(RHS, IntPtrTy);
if (IndexedSize != ElementSize) {
RHS = Builder.CreateMul(
RHS, ConstantInt::get(IntPtrTy, IndexedSize / ElementSize));
}
GetElementPtrInst *NewGEP = cast<GetElementPtrInst>(
Builder.CreateGEP(GEP->getResultElementType(), Candidate, RHS));
NewGEP->setIsInBounds(GEP->isInBounds());
NewGEP->takeName(GEP);
return NewGEP;
}
Instruction *NaryReassociatePass::tryReassociateBinaryOp(BinaryOperator *I) {
Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
if (SE->getSCEV(I)->isZero())
return nullptr;
if (auto *NewI = tryReassociateBinaryOp(LHS, RHS, I))
return NewI;
if (auto *NewI = tryReassociateBinaryOp(RHS, LHS, I))
return NewI;
return nullptr;
}
Instruction *NaryReassociatePass::tryReassociateBinaryOp(Value *LHS, Value *RHS,
BinaryOperator *I) {
Value *A = nullptr, *B = nullptr;
if (LHS->hasOneUse() && matchTernaryOp(I, LHS, A, B)) {
const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
const SCEV *RHSExpr = SE->getSCEV(RHS);
if (BExpr != RHSExpr) {
if (auto *NewI =
tryReassociatedBinaryOp(getBinarySCEV(I, AExpr, RHSExpr), B, I))
return NewI;
}
if (AExpr != RHSExpr) {
if (auto *NewI =
tryReassociatedBinaryOp(getBinarySCEV(I, BExpr, RHSExpr), A, I))
return NewI;
}
}
return nullptr;
}
Instruction *NaryReassociatePass::tryReassociatedBinaryOp(const SCEV *LHSExpr,
Value *RHS,
BinaryOperator *I) {
auto *LHS = findClosestMatchingDominator(LHSExpr, I);
if (LHS == nullptr)
return nullptr;
Instruction *NewI = nullptr;
switch (I->getOpcode()) {
case Instruction::Add:
NewI = BinaryOperator::CreateAdd(LHS, RHS, "", I);
break;
case Instruction::Mul:
NewI = BinaryOperator::CreateMul(LHS, RHS, "", I);
break;
default:
llvm_unreachable("Unexpected instruction.");
}
NewI->takeName(I);
return NewI;
}
bool NaryReassociatePass::matchTernaryOp(BinaryOperator *I, Value *V,
Value *&Op1, Value *&Op2) {
switch (I->getOpcode()) {
case Instruction::Add:
return match(V, m_Add(m_Value(Op1), m_Value(Op2)));
case Instruction::Mul:
return match(V, m_Mul(m_Value(Op1), m_Value(Op2)));
default:
llvm_unreachable("Unexpected instruction.");
}
return false;
}
const SCEV *NaryReassociatePass::getBinarySCEV(BinaryOperator *I,
const SCEV *LHS,
const SCEV *RHS) {
switch (I->getOpcode()) {
case Instruction::Add:
return SE->getAddExpr(LHS, RHS);
case Instruction::Mul:
return SE->getMulExpr(LHS, RHS);
default:
llvm_unreachable("Unexpected instruction.");
}
return nullptr;
}
Instruction *
NaryReassociatePass::findClosestMatchingDominator(const SCEV *CandidateExpr,
Instruction *Dominatee) {
auto Pos = SeenExprs.find(CandidateExpr);
if (Pos == SeenExprs.end())
return nullptr;
auto &Candidates = Pos->second;
while (!Candidates.empty()) {
if (Value *Candidate = Candidates.back()) {
Instruction *CandidateInstruction = cast<Instruction>(Candidate);
if (DT->dominates(CandidateInstruction, Dominatee))
return CandidateInstruction;
}
Candidates.pop_back();
}
return nullptr;
}
template <typename MaxMinT> static SCEVTypes convertToSCEVype(MaxMinT &MM) {
if (std::is_same<smax_pred_ty, typename MaxMinT::PredType>::value)
return scSMaxExpr;
else if (std::is_same<umax_pred_ty, typename MaxMinT::PredType>::value)
return scUMaxExpr;
else if (std::is_same<smin_pred_ty, typename MaxMinT::PredType>::value)
return scSMinExpr;
else if (std::is_same<umin_pred_ty, typename MaxMinT::PredType>::value)
return scUMinExpr;
llvm_unreachable("Can't convert MinMax pattern to SCEV type");
return scUnknown;
}
template <typename MaxMinT>
Value *NaryReassociatePass::tryReassociateMinOrMax(Instruction *I,
MaxMinT MaxMinMatch,
Value *LHS, Value *RHS) {
Value *A = nullptr, *B = nullptr;
MaxMinT m_MaxMin(m_Value(A), m_Value(B));
if (LHS->hasNUsesOrMore(3) ||
llvm::any_of(LHS->users(),
[&](auto *U) {
return U != I &&
!(U->hasOneUser() && *U->users().begin() == I);
}) ||
!match(LHS, m_MaxMin))
return nullptr;
auto tryCombination = [&](Value *A, const SCEV *AExpr, Value *B,
const SCEV *BExpr, Value *C,
const SCEV *CExpr) -> Value * {
SmallVector<const SCEV *, 2> Ops1{BExpr, AExpr};
const SCEVTypes SCEVType = convertToSCEVype(m_MaxMin);
const SCEV *R1Expr = SE->getMinMaxExpr(SCEVType, Ops1);
Instruction *R1MinMax = findClosestMatchingDominator(R1Expr, I);
if (!R1MinMax)
return nullptr;
LLVM_DEBUG(dbgs() << "NARY: Found common sub-expr: " << *R1MinMax << "\n");
SmallVector<const SCEV *, 2> Ops2{SE->getUnknown(C),
SE->getUnknown(R1MinMax)};
const SCEV *R2Expr = SE->getMinMaxExpr(SCEVType, Ops2);
SCEVExpander Expander(*SE, *DL, "nary-reassociate");
Value *NewMinMax = Expander.expandCodeFor(R2Expr, I->getType(), I);
NewMinMax->setName(Twine(I->getName()).concat(".nary"));
LLVM_DEBUG(dbgs() << "NARY: Deleting: " << *I << "\n"
<< "NARY: Inserting: " << *NewMinMax << "\n");
return NewMinMax;
};
const SCEV *AExpr = SE->getSCEV(A);
const SCEV *BExpr = SE->getSCEV(B);
const SCEV *RHSExpr = SE->getSCEV(RHS);
if (BExpr != RHSExpr) {
if (auto *NewMinMax = tryCombination(A, AExpr, RHS, RHSExpr, B, BExpr))
return NewMinMax;
}
if (AExpr != RHSExpr) {
if (auto *NewMinMax = tryCombination(RHS, RHSExpr, B, BExpr, A, AExpr))
return NewMinMax;
}
return nullptr;
}