Add reassociation profitability check to TTI

[dmitry]

Summary

Reassociation for large vector types increases register pressure on some targets.

[mcrosier]

Can you add the functional part for the target of interest along with a test case?

[dmitry]

@mcrosier, I can't. The main motivation is an out-of-tree target, but that it might be beneficial for other targets as well.

[mcrosier]

In D32290#733521, @dmitry wrote:

@mcrosier, I can't. The main motivation is an out-of-tree target, but that it might be beneficial for other targets as well.

I appreciate you posting the patch, but until such a feature is used by an in-tree target it's difficult to justify placing the maintenance burden on the community.

include/llvm/Analysis/TargetTransformInfo.h

public:
 
   /// \brief Return true if it is profitable to hoist instruction in the
   /// then/else to before if.
   bool isProfitableToHoist(Instruction *I) const;
 
   /// \brief Return true if this type is legal.
   bool isTypeLegal(Type *Ty) const;
 
+  /// \brief Return true if reassociation of expressions of the specified type
+  /// is profitable.
+  bool isProfitableToReassociate(Type *Ty) const;
+
   /// \brief Returns the target's jmp_buf alignment in bytes.
   unsigned getJumpBufAlignment() const;
 
   /// \brief Returns the target's jmp_buf size in bytes.
   unsigned getJumpBufSize() const;
 
   /// \brief Return true if switches should be turned into lookup tables for the
   /// target.
public:
   virtual bool isLegalMaskedGather(Type *DataType) = 0;
   virtual int getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
                                    int64_t BaseOffset, bool HasBaseReg,
                                    int64_t Scale, unsigned AddrSpace) = 0;
   virtual bool isFoldableMemAccessOffset(Instruction *I, int64_t Offset) = 0;
   virtual bool isTruncateFree(Type *Ty1, Type *Ty2) = 0;
   virtual bool isProfitableToHoist(Instruction *I) = 0;
   virtual bool isTypeLegal(Type *Ty) = 0;
+  virtual bool isProfitableToReassociate(Type *Ty) = 0;
   virtual unsigned getJumpBufAlignment() = 0;
   virtual unsigned getJumpBufSize() = 0;
   virtual bool shouldBuildLookupTables() = 0;
   virtual bool shouldBuildLookupTablesForConstant(Constant *C) = 0;
   virtual unsigned
   getScalarizationOverhead(Type *Ty, bool Insert, bool Extract) = 0;
   virtual unsigned getOperandsScalarizationOverhead(ArrayRef<const Value *> Args,
                                                     unsigned VF) = 0;
public:
   }
   bool isTruncateFree(Type *Ty1, Type *Ty2) override {
     return Impl.isTruncateFree(Ty1, Ty2);
   }
   bool isProfitableToHoist(Instruction *I) override {
     return Impl.isProfitableToHoist(I);
   }
   bool isTypeLegal(Type *Ty) override { return Impl.isTypeLegal(Ty); }
+  bool isProfitableToReassociate(Type *Ty) override {
+    return Impl.isProfitableToReassociate(Ty);
+  }
   unsigned getJumpBufAlignment() override { return Impl.getJumpBufAlignment(); }
   unsigned getJumpBufSize() override { return Impl.getJumpBufSize(); }
   bool shouldBuildLookupTables() override {
     return Impl.shouldBuildLookupTables();
   }
   bool shouldBuildLookupTablesForConstant(Constant *C) override {
     return Impl.shouldBuildLookupTablesForConstant(C);
   }

include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   bool isFoldableMemAccessOffset(Instruction *I, int64_t Offset) { return true; }
 
   bool isTruncateFree(Type *Ty1, Type *Ty2) { return false; }
 
   bool isProfitableToHoist(Instruction *I) { return true; }
 
   bool isTypeLegal(Type *Ty) { return false; }
 
+  bool isProfitableToReassociate(Type *Ty) { return true; }
+
   unsigned getJumpBufAlignment() { return 0; }
 
   unsigned getJumpBufSize() { return 0; }
 
   bool shouldBuildLookupTables() { return true; }
   bool shouldBuildLookupTablesForConstant(Constant *C) { return true; }
 
   unsigned getScalarizationOverhead(Type *Ty, bool Insert, bool Extract) {

include/llvm/CodeGen/BasicTTIImpl.h

bool isProfitableToHoist(Instruction *I) {
     return getTLI()->isProfitableToHoist(I);
   }
 
   bool isTypeLegal(Type *Ty) {
     EVT VT = getTLI()->getValueType(DL, Ty);
     return getTLI()->isTypeLegal(VT);
   }
 
+  bool isProfitableToReassociate(Type *Ty) {
+    return true;
+  }
+
   int getGEPCost(Type *PointeeType, const Value *Ptr,
                  ArrayRef<const Value *> Operands) {
     return BaseT::getGEPCost(PointeeType, Ptr, Operands);
   }
 
   unsigned getIntrinsicCost(Intrinsic::ID IID, Type *RetTy,
                             ArrayRef<const Value *> Arguments) {
     return BaseT::getIntrinsicCost(IID, RetTy, Arguments);

include/llvm/Transforms/Scalar/Reassociate.h

 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
+class TargetTransformInfo;
+
 /// A private "module" namespace for types and utilities used by Reassociate.
 /// These are implementation details and should not be used by clients.
 namespace reassociate {
 struct ValueEntry {
   unsigned Rank;
   Value *Op;
   ValueEntry(unsigned R, Value *O) : Rank(R), Op(O) {}
 };

 /// Reassociate commutative expressions.
 class ReassociatePass : public PassInfoMixin<ReassociatePass> {
   DenseMap<BasicBlock *, unsigned> RankMap;
   DenseMap<AssertingVH<Value>, unsigned> ValueRankMap;
   SetVector<AssertingVH<Instruction>> RedoInsts;
   bool MadeChange;
 
 public:
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  bool runImpl(Function &F, FunctionAnalysisManager &,
+               const TargetTransformInfo &TTI);
 
 private:
   void BuildRankMap(Function &F, ReversePostOrderTraversal<Function *> &RPOT);
   unsigned getRank(Value *V);
   void canonicalizeOperands(Instruction *I);
   void ReassociateExpression(BinaryOperator *I);
   void RewriteExprTree(BinaryOperator *I,
                        SmallVectorImpl<reassociate::ValueEntry> &Ops);
private:
   Value *buildMinimalMultiplyDAG(IRBuilder<> &Builder,
                                  SmallVectorImpl<reassociate::Factor> &Factors);
   Value *OptimizeMul(BinaryOperator *I,
                      SmallVectorImpl<reassociate::ValueEntry> &Ops);
   Value *RemoveFactorFromExpression(Value *V, Value *Factor);
   void EraseInst(Instruction *I);
   void RecursivelyEraseDeadInsts(Instruction *I,
                                  SetVector<AssertingVH<Instruction>> &Insts);
-  void OptimizeInst(Instruction *I);
+  void OptimizeInst(Instruction *I, const TargetTransformInfo &TTI);
   Instruction *canonicalizeNegConstExpr(Instruction *I);
 };
 }
 
 #endif // LLVM_TRANSFORMS_SCALAR_REASSOCIATE_H

lib/Analysis/TargetTransformInfo.cpp

 bool TargetTransformInfo::isProfitableToHoist(Instruction *I) const {
   return TTIImpl->isProfitableToHoist(I);
 }
 
 bool TargetTransformInfo::isTypeLegal(Type *Ty) const {
   return TTIImpl->isTypeLegal(Ty);
 }
 
+bool TargetTransformInfo::isProfitableToReassociate(Type *Ty) const {
+  return TTIImpl->isProfitableToReassociate(Ty);
+}
+
 unsigned TargetTransformInfo::getJumpBufAlignment() const {
   return TTIImpl->getJumpBufAlignment();
 }
 
 unsigned TargetTransformInfo::getJumpBufSize() const {
   return TTIImpl->getJumpBufSize();
 }
 

lib/Transforms/Scalar/Reassociate.cpp

 
 #include "llvm/Transforms/Scalar/Reassociate.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
Instruction *ReassociatePass::canonicalizeNegConstExpr(Instruction *I) {
   NI->setDebugLoc(I->getDebugLoc());
   RedoInsts.insert(I);
   MadeChange = true;
   return NI;
 }
 
 /// Inspect and optimize the given instruction. Note that erasing
 /// instructions is not allowed.
-void ReassociatePass::OptimizeInst(Instruction *I) {
+void ReassociatePass::OptimizeInst(Instruction *I, const TargetTransformInfo &TTI) {
   // Only consider operations that we understand.
   if (!isa<BinaryOperator>(I))
     return;
 
   if (I->getOpcode() == Instruction::Shl && isa<ConstantInt>(I->getOperand(1)))
     // If an operand of this shift is a reassociable multiply, or if the shift
     // is used by a reassociable multiply or add, turn into a multiply.
     if (isReassociableOp(I->getOperand(0), Instruction::Mul) ||
void ReassociatePass::OptimizeInst(Instruction *I, const TargetTransformInfo &TTI) {
   // currently unsupported.
   if (I->getType()->isVectorTy() && I->getOpcode() == Instruction::Xor)
     return;
 
   // Don't optimize floating point instructions that don't have unsafe algebra.
   if (I->getType()->isFPOrFPVectorTy() && !I->hasUnsafeAlgebra())
     return;
 
+  if (!TTI.isProfitableToReassociate(I->getType()))
+    return;
+
   // Do not reassociate boolean (i1) expressions.  We want to preserve the
   // original order of evaluation for short-circuited comparisons that
   // SimplifyCFG has folded to AND/OR expressions.  If the expression
   // is not further optimized, it is likely to be transformed back to a
   // short-circuited form for code gen, and the source order may have been
   // optimized for the most likely conditions.
   if (I->getType()->isIntegerTy(1))
     return;
if (Ops.size() == 1) {
     return;
   }
 
   // Now that we ordered and optimized the expressions, splat them back into
   // the expression tree, removing any unneeded nodes.
   RewriteExprTree(I, Ops);
 }
 
-PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
+PreservedAnalyses ReassociatePass::run(Function &F,
+                                       FunctionAnalysisManager &AM) {
+  auto &TTI = AM.getResult<TargetIRAnalysis>(F);
+  if (runImpl(F, AM, TTI)) {
+    PreservedAnalyses PA;
+    PA.preserveSet<CFGAnalyses>();
+    PA.preserve<GlobalsAA>();
+    return PA;
+  }
+  return PreservedAnalyses::all();
+}
+
+bool ReassociatePass::runImpl(Function &F, FunctionAnalysisManager &,
+                              const TargetTransformInfo &TTI) {
   // Get the functions basic blocks in Reverse Post Order. This order is used by
   // BuildRankMap to pre calculate ranks correctly. It also excludes dead basic
   // blocks (it has been seen that the analysis in this pass could hang when
   // analysing dead basic blocks).
   ReversePostOrderTraversal<Function *> RPOT(&F);
 
   // Calculate the rank map for F.
   BuildRankMap(F, RPOT);
 
   MadeChange = false;
   // Traverse the same blocks that was analysed by BuildRankMap.
   for (BasicBlock *BI : RPOT) {
     assert(RankMap.count(&*BI) && "BB should be ranked.");
     // Optimize every instruction in the basic block.
     for (BasicBlock::iterator II = BI->begin(), IE = BI->end(); II != IE;)
       if (isInstructionTriviallyDead(&*II)) {
         EraseInst(&*II++);
       } else {
-        OptimizeInst(&*II);
+        OptimizeInst(&*II, TTI);
         assert(II->getParent() == &*BI && "Moved to a different block!");
         ++II;
       }
 
     // Make a copy of all the instructions to be redone so we can remove dead
     // instructions.
     SetVector<AssertingVH<Instruction>> ToRedo(RedoInsts);
     // Iterate over all instructions to be reevaluated and remove trivially dead
for (BasicBlock *BI : RPOT) {
 
     // Now that we have removed dead instructions, we can reoptimize the
     // remaining instructions.
     while (!RedoInsts.empty()) {
       Instruction *I = RedoInsts.pop_back_val();
       if (isInstructionTriviallyDead(I))
         EraseInst(I);
       else
-        OptimizeInst(I);
+        OptimizeInst(I, TTI);
     }
   }
 
   // We are done with the rank map.
   RankMap.clear();
   ValueRankMap.clear();
 
-  if (MadeChange) {
-    PreservedAnalyses PA;
-    PA.preserveSet<CFGAnalyses>();
-    PA.preserve<GlobalsAA>();
-    return PA;
-  }
-
-  return PreservedAnalyses::all();
+  return MadeChange;
 }
 
 namespace {
   class ReassociateLegacyPass : public FunctionPass {
     ReassociatePass Impl;
   public:
     static char ID; // Pass identification, replacement for typeid
     ReassociateLegacyPass() : FunctionPass(ID) {
       initializeReassociateLegacyPassPass(*PassRegistry::getPassRegistry());
     }
 
     bool runOnFunction(Function &F) override {
       if (skipFunction(F))
         return false;
 
       FunctionAnalysisManager DummyFAM;
-      auto PA = Impl.run(F, DummyFAM);
-      return !PA.areAllPreserved();
+      const auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+      return Impl.runImpl(F, DummyFAM, TTI);
     }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<TargetTransformInfoWrapperPass>();
       AU.setPreservesCFG();
       AU.addPreserved<GlobalsAAWrapperPass>();
     }
   };
 }
 
 char ReassociateLegacyPass::ID = 0;
-INITIALIZE_PASS(ReassociateLegacyPass, "reassociate",
+INITIALIZE_PASS_BEGIN(ReassociateLegacyPass, "reassociate",
+                      "Reassociate expressions", false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
+INITIALIZE_PASS_END(ReassociateLegacyPass, "reassociate",
                     "Reassociate expressions", false, false)
 
 // Public interface to the Reassociate pass
 FunctionPass *llvm::createReassociatePass() {
   return new ReassociateLegacyPass();
 }