NFC: Migrate PartialInlining to work on InstructionCost

[sdesmalen]

Summary

This patch migrates cost values and arithmetic to work on InstructionCost.
When the interfaces to TargetTransformInfo are changed, any InstructionCost
state will propagate naturally.

See this patch for the introduction of the type: https://reviews.llvm.org/D91174
See this thread for context: http://lists.llvm.org/pipermail/llvm-dev/2020-November/146408.html

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

116

why is this necessary?

[sdesmalen]

llvm/include/llvm/Support/InstructionCost.h

116

I wasn't really sure what type to choose for RHS (float or double), so I thought I'd just template the type for any floating point type.

The operator is needed for a multiplication in PartialInlining with the options MinRegionSizeRatio and ColdBranchRatio which in this case are both of type float.

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

116

It looks like in all cases where those two options are used, they are immediately static_cast<int>'ed before being assigned to anything. Could you just keep the static cast and not have this operator overload?

My fear is that we're going to end up with a billion operator overloads that are only used in one or two places.

[sdesmalen]

llvm/include/llvm/Support/InstructionCost.h

116

Just to make sure we're on the same page, the instance for which I added this operator is:

InstructionCost MinOutlineRegionCost =
    OverallFunctionCost * MinRegionSizeRatio.getValue();

where OverallFunctionCost is now an InstructionCost.

In my previous comment I suggested ColdBranchRatio was also used in a computation with InstructionCost, but that was not correct (that is the one being static_cast'ed to int).

The alternative way to write this is to check for validity and ask the for the value, multiply with the FP value, and return a new InstructionCost.

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

116

Oh, I see the issue. Instead of providing overloads for float, you could add a map function like Optional has:

// F is a function from CostType to CostType
template <class Function>
InstructionCost map(const Function &F) const {
  if (isValid()) 
    return F(*getValue());
  return getInvalid();
}

Then you could have:

InstructionCost MinOutlineRegionCost = OverallFunctionCost.map(
        [&](auto OFC) { return OFC * MinRegionSizeRatio; });

(I guess you could delegate to Optional<InstructionCost::CostType>::map, but I don't think that would save any keystrokes)

[sdesmalen]

llvm/include/llvm/Support/InstructionCost.h

116

That's an interesting way to write it. Maybe the use of auto here would be a bit confusing, as this is a InstructionCost::CostType.

Would writing:

InstructionCost MinOutlineRegionCost = OverallFunctionCost * MinRegionSizeRatio;

not be more intuitive though?

I guess I'm trying to understand your reasoning for not doing the operator overload. Is that because you don't like the fact that it now needs to instantiate two functions (one for float and one for double, depending on it's use), or are you concerned about the truncation that happens when the result is converted back is unexpected when working on an InstructionCost? If it is the former, then I think in practice there will be at most two overloads that need instantiating, whereas for the latter, I'm not really concerned because all existing code that would use this already assumes truncation.

The getValue() part of MinRegionSizeRatio.getValue() is needed because MinRegionSizeRatio is a cl::opt, hence:

InstructionCost MinOutlineRegionCost = OverallFunctionCost * MinRegionSizeRatio.getValue();

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

116

I don't think two template instantiations is that much overhead. I'm trying to avoid going down the slippery slope of having InstructionCost.h be filled with thousands of operator overloads. I see a dystopian future where we have overloads of some of the operators for floating point types, cl::opt of floats and ints, Optional of floats and ints, APFloat, APInt, Constant * floats and ints, and gosh knows what else, all only used in one or two places. In this world, every time somebody hits an edge case, they go and add an operator overload to InstructionCost, and maybe they do it right, and maybe they don't.

I would think that 99% of the usages of the operators are between int and InstructionCost, or between InstructionCost and InstructionCost. If we provide a map function, we can give people a way to handle these edge cases that isn't too much more verbose but keeps the interface clean. I want it to be the case where InstructionCost.h hits a steady state where it just provides all the tools anybody could ever need. Realistically, InstructionCost is just Optional<int> with convenient operator overloads for arithmetic. map is the only function (and getValueOr, but that's just gravy), from the interface of Optional<int> that we don't have.

If we must go down this route, I'd like for you to add all the operators all at once. But I don't think we need that; we really shouldn't be doing lots of floating point math anyways.

[sdesmalen]

Replaced overloaded operator with InstructionCost::map() function.

[sdesmalen]

llvm/include/llvm/Support/InstructionCost.h

116

You were right about this being an edge case, this was actually the only instance that broke when I removed the operator.
I agree with you now. Removing these operators means that all operations on InstructionCost are integer-based, and if someone needs to handle an edge case like this FP -> Int conversion, map provides the functionality. I've updated the patch to reflect this. Thank you for the suggestion!

[ctetreau]

Aside from the function signature for InstructionCost::map, this patch looks good to me.

llvm/include/llvm/Support/InstructionCost.h

21

if map takes a template parameter for the function type, then this should be unnecessary.

204–208

Having the function type be a template parameter (as it is defined in Optional.h) is preferred because lambda types can be directly used without the indirection of assigning them to a std::function. It also avoids having to #include <functional>

204–208

also, I believe map can be const.

[sdesmalen]

Use template cost function for InstructionCost::map.

[sdesmalen]

llvm/include/llvm/Support/InstructionCost.h

204–208

Having the function type be a template parameter (as it is defined in Optional.h) is preferred because lambda types can be directly used without the indirection of assigning them to a std::function. It also avoids having to #include <functional>

Okay, I wasn't sure if it was better to limit the signature of Function by specifying it would take/return a CostType by using std::function. Any issues resulting from the lambda return the wrong value type probably show up when calling the constructor of InstructionCost.

[ctetreau]

llvm/include/llvm/Support/InstructionCost.h

204–208

Yeah, the function object has to take one argument that coerces to CostType, and return some type that coerces to InstructionCost, so it should be safe. I guess it would be nice to constrain the signature to prevent people from thinking they need to return InstructionCost and invoking the copy constructor for no reason, but this version is safer to call in a hot inner loop which I think most C++ programmers would prefer.

llvm/lib/Transforms/IPO/PartialInlining.cpp

1363

this looks like a behavioral change. Please just assert that the costs are valid.

[sdesmalen]

Removed functional change and replaced it by an assert.

[paulwalker-arm]

llvm/lib/Transforms/IPO/PartialInlining.cpp

1361

This comment doesn't make sense anymore given the most recent change. Probably easier to just remove it?

[sdesmalen]

Removed stale comment.

[sdesmalen]

@ctetreau are you also happy with the patch after the latest changes?

llvm/include/llvm/Support/InstructionCost.h

 /// unknown.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_INSTRUCTIONCOST_H
 #define LLVM_SUPPORT_INSTRUCTIONCOST_H
 
 #include "llvm/ADT/Optional.h"
 
 namespace llvm {
 
 class raw_ostream;
 
 class InstructionCost {
 public:
   using CostType = int;
 
public:
   }
 
   InstructionCost &operator*=(const CostType RHS) {
     InstructionCost RHS2(RHS);
     *this *= RHS2;
     return *this;
   }
 
   InstructionCost &operator/=(const InstructionCost &RHS) {
     propagateState(RHS);
     Value /= RHS.Value;
     return *this;
   }
 
   InstructionCost &operator/=(const CostType RHS) {
     InstructionCost RHS2(RHS);
     *this /= RHS2;
public:
   }
 
   bool operator>=(const CostType RHS) const {
     InstructionCost RHS2(RHS);
     return *this >= RHS2;
   }
 
   void print(raw_ostream &OS) const;
+
+  template <class Function>
+  auto map(const Function &F) const -> InstructionCost {
+    if (isValid())
+      return F(*getValue());
+    return getInvalid();
+  }
 };
 
 inline InstructionCost operator+(const InstructionCost &LHS,
                                  const InstructionCost &RHS) {
   InstructionCost LHS2(LHS);
   LHS2 += RHS;
   return LHS2;
 }

llvm/lib/Transforms/IPO/PartialInlining.cpp

struct FunctionCloner {
     typedef std::pair<Function *, BasicBlock *> FuncBodyCallerPair;
     // Keep track of Outlined Functions and the basic block they're called from.
     SmallVector<FuncBodyCallerPair, 4> OutlinedFunctions;
 
     // ClonedFunc is inlined in one of its callers after function
     // outlining.
     bool IsFunctionInlined = false;
     // The cost of the region to be outlined.
-    int OutlinedRegionCost = 0;
+    InstructionCost OutlinedRegionCost = 0;
     // ClonedOI is specific to outlining non-early return blocks.
     std::unique_ptr<FunctionOutliningInfo> ClonedOI = nullptr;
     // ClonedOMRI is specific to outlining cold regions.
     std::unique_ptr<FunctionOutliningMultiRegionInfo> ClonedOMRI = nullptr;
     std::unique_ptr<BlockFrequencyInfo> ClonedFuncBFI = nullptr;
     OptimizationRemarkEmitter &ORE;
     function_ref<AssumptionCache *(Function &)> LookupAC;
     function_ref<TargetTransformInfo &(Function &)> GetTTI;
private:
   }
 
   // Returns the costs associated with function outlining:
   // - The first value is the non-weighted runtime cost for making the call
   //   to the outlined function, including the addtional  setup cost in the
   //    outlined function itself;
   // - The second value is the estimated size of the new call sequence in
   //   basic block Cloner.OutliningCallBB;
-  std::tuple<int, int> computeOutliningCosts(FunctionCloner &Cloner) const;
+  std::tuple<InstructionCost, InstructionCost>
+  computeOutliningCosts(FunctionCloner &Cloner) const;
 
   // Compute the 'InlineCost' of block BB. InlineCost is a proxy used to
   // approximate both the size and runtime cost (Note that in the current
   // inline cost analysis, there is no clear distinction there either).
-  static int computeBBInlineCost(BasicBlock *BB, TargetTransformInfo *TTI);
+  static InstructionCost computeBBInlineCost(BasicBlock *BB,
+                                             TargetTransformInfo *TTI);
 
   std::unique_ptr<FunctionOutliningInfo>
   computeOutliningInfo(Function &F) const;
 
   std::unique_ptr<FunctionOutliningMultiRegionInfo>
   computeOutliningColdRegionsInfo(Function &F,
                                   OptimizationRemarkEmitter &ORE) const;
 };
auto BBProfileCount = [BFI](BasicBlock *BB) {
     return BFI->getBlockProfileCount(BB)
                ? BFI->getBlockProfileCount(BB).getValue()
                : 0;
   };
 
   // Use the same computeBBInlineCost function to compute the cost savings of
   // the outlining the candidate region.
   TargetTransformInfo *FTTI = &GetTTI(F);
-  int OverallFunctionCost = 0;
+  InstructionCost OverallFunctionCost = 0;
   for (auto &BB : F)
     OverallFunctionCost += computeBBInlineCost(&BB, FTTI);
 
   LLVM_DEBUG(dbgs() << "OverallFunctionCost = " << OverallFunctionCost
                     << "\n";);
 
-  int MinOutlineRegionCost =
-      static_cast<int>(OverallFunctionCost * MinRegionSizeRatio);
+  InstructionCost MinOutlineRegionCost = OverallFunctionCost.map(
+      [&](auto Cost) { return Cost * MinRegionSizeRatio; });
+
   BranchProbability MinBranchProbability(
       static_cast<int>(ColdBranchRatio * MinBlockCounterExecution),
       MinBlockCounterExecution);
   bool ColdCandidateFound = false;
   BasicBlock *CurrEntry = EntryBlock;
   std::vector<BasicBlock *> DFS;
   DenseMap<BasicBlock *, bool> VisitedMap;
   DFS.push_back(CurrEntry);
for (auto SI = succ_begin(ThisBB); SI != succ_end(ThisBB); ++SI) {
       BasicBlock *ExitBlock = nullptr;
       // We can only outline single exit regions (for now).
       if (!(ExitBlock = IsSingleExit(DominateVector))) {
         LLVM_DEBUG(dbgs() << "ABORT: Block " << SI->getName()
                           << " doesn't have a unique successor\n";);
         continue;
       }
 
-      int OutlineRegionCost = 0;
+      InstructionCost OutlineRegionCost = 0;
       for (auto *BB : DominateVector)
         OutlineRegionCost += computeBBInlineCost(BB, &GetTTI(*BB->getParent()));
 
       LLVM_DEBUG(dbgs() << "OutlineRegionCost = " << OutlineRegionCost
                         << "\n";);
 
       if (!SkipCostAnalysis && OutlineRegionCost < MinOutlineRegionCost) {
         ORE.emit([&]() {
return OptimizationRemarkAnalysis(DEBUG_TYPE, "CanBePartiallyInlined", &CB)
            << NV("Threshold", IC.getCostDelta() + IC.getCost()) << ")";
   });
   return true;
 }
 
 // TODO: Ideally  we should share Inliner's InlineCost Analysis code.
 // For now use a simplified version. The returned 'InlineCost' will be used
 // to esimate the size cost as well as runtime cost of the BB.
-int PartialInlinerImpl::computeBBInlineCost(BasicBlock *BB,
-                                            TargetTransformInfo *TTI) {
-  int InlineCost = 0;
+InstructionCost
+PartialInlinerImpl::computeBBInlineCost(BasicBlock *BB,
+  InstructionCost InlineCost = 0;
   const DataLayout &DL = BB->getParent()->getParent()->getDataLayout();
   for (Instruction &I : BB->instructionsWithoutDebug()) {
     // Skip free instructions.
     switch (I.getOpcode()) {
     case Instruction::BitCast:
     case Instruction::PtrToInt:
     case Instruction::IntToPtr:
     case Instruction::Alloca:
for (Instruction &I : BB->instructionsWithoutDebug()) {
     }
 
     if (SwitchInst *SI = dyn_cast<SwitchInst>(&I)) {
       InlineCost += (SI->getNumCases() + 1) * InlineConstants::InstrCost;
       continue;
     }
     InlineCost += InlineConstants::InstrCost;
   }
+
   return InlineCost;
 }
 
-std::tuple<int, int>
+std::tuple<InstructionCost, InstructionCost>
 PartialInlinerImpl::computeOutliningCosts(FunctionCloner &Cloner) const {
-  int OutliningFuncCallCost = 0, OutlinedFunctionCost = 0;
+  InstructionCost OutliningFuncCallCost = 0, OutlinedFunctionCost = 0;
   for (auto FuncBBPair : Cloner.OutlinedFunctions) {
     Function *OutlinedFunc = FuncBBPair.first;
     BasicBlock* OutliningCallBB = FuncBBPair.second;
     // Now compute the cost of the call sequence to the outlined function
     // 'OutlinedFunction' in BB 'OutliningCallBB':
     auto *OutlinedFuncTTI = &GetTTI(*OutlinedFunc);
     OutliningFuncCallCost +=
         computeBBInlineCost(OutliningCallBB, OutlinedFuncTTI);
 
     // Now compute the cost of the extracted/outlined function itself:
     for (BasicBlock &BB : *OutlinedFunc)
       OutlinedFunctionCost += computeBBInlineCost(&BB, OutlinedFuncTTI);
   }
   assert(OutlinedFunctionCost >= Cloner.OutlinedRegionCost &&
          "Outlined function cost should be no less than the outlined region");
 
   // The code extractor introduces a new root and exit stub blocks with
   // additional unconditional branches. Those branches will be eliminated
   // later with bb layout. The cost should be adjusted accordingly:
   OutlinedFunctionCost -=
       2 * InlineConstants::InstrCost * Cloner.OutlinedFunctions.size();
 
-  int OutliningRuntimeOverhead =
+  InstructionCost OutliningRuntimeOverhead =
       OutliningFuncCallCost +
       (OutlinedFunctionCost - Cloner.OutlinedRegionCost) +
-      ExtraOutliningPenalty;
+      ExtraOutliningPenalty.getValue();
 
   return std::make_tuple(OutliningFuncCallCost, OutliningRuntimeOverhead);
 }
 
 // Create the callsite to profile count map which is
 // used to update the original function's entry count,
 // after the function is partially inlined into the callsite.
 void PartialInlinerImpl::computeCallsiteToProfCountMap(
for (auto *DP : DeadPhis)
     DP->eraseFromParent();
 
   for (auto *E : ClonedOI->ReturnBlockPreds)
     E->getTerminator()->replaceUsesOfWith(PreReturn, ClonedOI->ReturnBlock);
 }
 
 bool PartialInlinerImpl::FunctionCloner::doMultiRegionFunctionOutlining() {
 
-  auto ComputeRegionCost = [&](SmallVectorImpl<BasicBlock *> &Region) {
-    int Cost = 0;
+  auto ComputeRegionCost =
+      [&](SmallVectorImpl<BasicBlock *> &Region) -> InstructionCost {
+    InstructionCost Cost = 0;
     for (BasicBlock* BB : Region)
       Cost += computeBBInlineCost(BB, &GetTTI(*BB->getParent()));
     return Cost;
   };
 
   assert(ClonedOMRI && "Expecting OutlineInfo for multi region outline");
 
   if (ClonedOMRI->ORI.empty())
bool PartialInlinerImpl::FunctionCloner::doMultiRegionFunctionOutlining() {
   ClonedFuncBFI.reset(new BlockFrequencyInfo(*ClonedFunc, BPI, LI));
 
   // Cache and recycle the CodeExtractor analysis to avoid O(n^2) compile-time.
   CodeExtractorAnalysisCache CEAC(*ClonedFunc);
 
   SetVector<Value *> Inputs, Outputs, Sinks;
   for (FunctionOutliningMultiRegionInfo::OutlineRegionInfo RegionInfo :
        ClonedOMRI->ORI) {
-    int CurrentOutlinedRegionCost = ComputeRegionCost(RegionInfo.Region);
+    InstructionCost CurrentOutlinedRegionCost =
+        ComputeRegionCost(RegionInfo.Region);
 
     CodeExtractor CE(RegionInfo.Region, &DT, /*AggregateArgs*/ false,
                      ClonedFuncBFI.get(), &BPI,
                      LookupAC(*RegionInfo.EntryBlock->getParent()),
                      /* AllowVarargs */ false);
 
     CE.findInputsOutputs(Inputs, Outputs, Sinks);
 

 
 bool PartialInlinerImpl::tryPartialInline(FunctionCloner &Cloner) {
   if (Cloner.OutlinedFunctions.empty())
     return false;
 
   int SizeCost = 0;
   BlockFrequency WeightedRcost;
   int NonWeightedRcost;
-  std::tie(SizeCost, NonWeightedRcost) = computeOutliningCosts(Cloner);
+
+  auto OutliningCosts = computeOutliningCosts(Cloner);
+  assert(std::get<0>(OutliningCosts).isValid() &&
+         std::get<1>(OutliningCosts).isValid() && "Expected valid costs");
+
+  SizeCost = *std::get<0>(OutliningCosts).getValue();
+  NonWeightedRcost = *std::get<1>(OutliningCosts).getValue();
 
   // Only calculate RelativeToEntryFreq when we are doing single region
   // outlining.
   BranchProbability RelativeToEntryFreq;
   if (Cloner.ClonedOI)
     RelativeToEntryFreq = getOutliningCallBBRelativeFreq(Cloner);
   else
     // RelativeToEntryFreq doesn't make sense when we have more than one