NFC: Migrate SpeculateAroundPHIs 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]

Aside from the one nit, this looks good to me.

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

329

I assume this assert is to keep this patch NFC? If so, please reword the message to imply that it should be possible for this to happen, but it's not implemented yet or add a FIXME comment.

[sdesmalen]

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

329

Not entirely; if a Constant cannot be folded when speculating, it must be possible to materialize it fully, and so we must be able to cost it. Otherwise there is something wrong with the existing IR. I don't think we should ever encounter this case.

[ctetreau]

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

329

If a user writes bad IR, the system shouldn't crash. If this is truly impossible except in cases of bad IR, then we should raise a fatal error.

[sdesmalen]

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

329

Perhaps I should have said: a constant should always be materializable and if it is materializable, it should also be costable. If that is not the case, the cost-model in the compiler is broken and needs fixing, not this pass.

If TotalMatCost is valid, then TotalFoldedCost must be valid as well, by the assert above.

[ctetreau]

Aside from the nit, this looks good to me.

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

329

If TotalMatCost is valid, then TotalFoldedCost must be valid as well, by the assert above.

If TotalFoldedCost and TotalMatCost are both invalid, and both have the same value in the Value field, then an invalid TotalFoldingCost can survive the assert on line 329. It would assert on the next line, so in practice it wouldn't be an issue, but code changes could lead to issues.

Looking at the code, the comparators of InstructionCost seem to be messed up. operator== considers the Value field and the State field in all cases, but operator< does not consider the Value field if either the rhs or lhs are invalid. We should probably fix this.

a constant should always be materializable and if it is materializable, it should also be costable.

Then this assert should be above the other. This one states that constants can be materialized. The next one uses the cost of materialized constants.

Sorry if this all sounds super nitpicky. I'm just trying to understand the situation. I think reordering the asserts would help future readers understand the logic, but this is just a NIT.

[sdesmalen]

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

329

Then this assert should be above the other. This one states that constants can be materialized. The next one uses the cost of materialized constants.

Good point, I'll reorder the two asserts before I commit the patch.

Sorry if this all sounds super nitpicky. I'm just trying to understand the situation.

No problem, I appreciate your attention to detail! I occasionally still try to get my head around what the right approach is, decisions like "assert vs fatal error vs bail out early" aren't always obvious. Any new insights I can use in the other InstructionCost patches.

llvm/lib/Transforms/Scalar/SpeculateAroundPHIs.cpp

 /// folded into its users rather than materialized. This higher cost could
 /// result in a dynamic *path* that is more expensive even when the total cost
 /// is lower. Currently, all of the interesting cases where this optimization
 /// should fire are ones where it is a no-loss operation in this sense. If we
 /// ever want to be more aggressive here, we would need to balance the
 /// different incoming edges' cost by looking at their respective
 /// probabilities.
 static bool isSafeAndProfitableToSpeculateAroundPHI(
-    PHINode &PN, SmallDenseMap<PHINode *, int, 16> &CostSavingsMap,
+    PHINode &PN, SmallDenseMap<PHINode *, InstructionCost, 16> &CostSavingsMap,
     SmallPtrSetImpl<Instruction *> &PotentialSpecSet,
     SmallPtrSetImpl<Instruction *> &UnsafeSet, DominatorTree &DT,
     TargetTransformInfo &TTI) {
   // First see whether there is any cost savings to speculating around this
   // PHI, and build up a map of the constant inputs to how many times they
   // occur.
   bool NonFreeMat = false;
   struct CostsAndCount {
-    int MatCost = TargetTransformInfo::TCC_Free;
-    int FoldedCost = TargetTransformInfo::TCC_Free;
+    InstructionCost MatCost = TargetTransformInfo::TCC_Free;
+    InstructionCost FoldedCost = TargetTransformInfo::TCC_Free;
     int Count = 0;
   };
   SmallDenseMap<ConstantInt *, CostsAndCount, 16> CostsAndCounts;
   SmallPtrSet<BasicBlock *, 16> IncomingConstantBlocks;
   for (int i : llvm::seq<int>(0, PN.getNumIncomingValues())) {
     auto *IncomingC = dyn_cast<ConstantInt>(PN.getIncomingValue(i));
     if (!IncomingC)
       continue;
 
     // Only visit each incoming edge with a constant input once.
     if (!IncomingConstantBlocks.insert(PN.getIncomingBlock(i)).second)
       continue;
 
     auto InsertResult = CostsAndCounts.insert({IncomingC, {}});
     // Count how many edges share a given incoming costant.
     ++InsertResult.first->second.Count;
     // Only compute the cost the first time we see a particular constant.
     if (!InsertResult.second)
       continue;
 
-    int &MatCost = InsertResult.first->second.MatCost;
+    InstructionCost &MatCost = InsertResult.first->second.MatCost;
     MatCost = TTI.getIntImmCost(IncomingC->getValue(), IncomingC->getType(),
                                 TargetTransformInfo::TCK_SizeAndLatency);
     NonFreeMat |= MatCost != TTI.TCC_Free;
   }
   if (!NonFreeMat) {
     LLVM_DEBUG(dbgs() << "    Free: " << PN << "\n");
     // No profit in free materialization.
     return false;
for (Use &U : PN.uses()) {
 
     // Get the intrinsic ID if this user is an intrinsic.
     Intrinsic::ID IID = Intrinsic::not_intrinsic;
     if (auto *UserII = dyn_cast<IntrinsicInst>(UserI))
       IID = UserII->getIntrinsicID();
 
     for (auto &IncomingConstantAndCostsAndCount : CostsAndCounts) {
       ConstantInt *IncomingC = IncomingConstantAndCostsAndCount.first;
-      int MatCost = IncomingConstantAndCostsAndCount.second.MatCost;
-      int &FoldedCost = IncomingConstantAndCostsAndCount.second.FoldedCost;
+      InstructionCost MatCost = IncomingConstantAndCostsAndCount.second.MatCost;
+      InstructionCost &FoldedCost =
+          IncomingConstantAndCostsAndCount.second.FoldedCost;
       if (IID)
         FoldedCost +=
           TTI.getIntImmCostIntrin(IID, Idx, IncomingC->getValue(),
                                   IncomingC->getType(),
                                   TargetTransformInfo::TCK_SizeAndLatency);
       else
         FoldedCost +=
             TTI.getIntImmCostInst(UserI->getOpcode(), Idx,
for (auto &IncomingConstantAndCostsAndCount : CostsAndCounts) {
                              "    Accumulated folded cost: "
                           << FoldedCost << "\n");
         return false;
       }
     }
   }
 
   // Compute the total cost savings afforded by this PHI node.
-  int TotalMatCost = TTI.TCC_Free, TotalFoldedCost = TTI.TCC_Free;
+  InstructionCost TotalMatCost = TTI.TCC_Free, TotalFoldedCost = TTI.TCC_Free;
   for (auto IncomingConstantAndCostsAndCount : CostsAndCounts) {
-    int MatCost = IncomingConstantAndCostsAndCount.second.MatCost;
-    int FoldedCost = IncomingConstantAndCostsAndCount.second.FoldedCost;
+    InstructionCost MatCost = IncomingConstantAndCostsAndCount.second.MatCost;
+    InstructionCost FoldedCost =
+        IncomingConstantAndCostsAndCount.second.FoldedCost;
     int Count = IncomingConstantAndCostsAndCount.second.Count;
 
     TotalMatCost += MatCost * Count;
     TotalFoldedCost += FoldedCost * Count;
   }
   assert(TotalFoldedCost <= TotalMatCost && "If each constant's folded cost is "
                                             "less that its materialized cost, "
                                             "the sum must be as well.");
-
+  assert(TotalMatCost.isValid() && "Constants must be  materializable");
   LLVM_DEBUG(dbgs() << "    Cost savings " << (TotalMatCost - TotalFoldedCost)
                     << ": " << PN << "\n");
   CostSavingsMap[&PN] = TotalMatCost - TotalFoldedCost;
   return true;
 }
 
 /// Simple helper to walk all the users of a list of phis depth first, and call
 /// a visit function on each one in post-order.

 /// There is an orthogonal extra complexity to all of this: computing the cost
 /// itself can easily become a linear computation making everything again (at
 /// best) quadratic. Using a postorder over the operand graph makes it
 /// particularly easy to avoid this through dynamic programming. As we do the
 /// postorder walk, we build the transitive cost of that subgraph. It is also
 /// straightforward to then update these costs when we mark a PHI for
 /// speculation so that subsequent PHIs don't re-pay the cost of already
 /// speculated instructions.
-static SmallVector<PHINode *, 16>
-findProfitablePHIs(ArrayRef<PHINode *> PNs,
-                   const SmallDenseMap<PHINode *, int, 16> &CostSavingsMap,
-                   const SmallPtrSetImpl<Instruction *> &PotentialSpecSet,
-                   int NumPreds, DominatorTree &DT, TargetTransformInfo &TTI) {
+static SmallVector<PHINode *, 16> findProfitablePHIs(
+    ArrayRef<PHINode *> PNs,
+    const SmallDenseMap<PHINode *, InstructionCost, 16> &CostSavingsMap,
+    const SmallPtrSetImpl<Instruction *> &PotentialSpecSet, int NumPreds,
+    DominatorTree &DT, TargetTransformInfo &TTI) {
   SmallVector<PHINode *, 16> SpecPNs;
 
   // First, establish a reverse mapping from immediate users of the PHI nodes
   // to the nodes themselves, and count how many users each PHI node has in
   // a way we can update while processing them.
   SmallDenseMap<Instruction *, TinyPtrVector<PHINode *>, 16> UserToPNMap;
   SmallDenseMap<PHINode *, int, 16> PNUserCountMap;
   SmallPtrSet<Instruction *, 16> UserSet;
   for (auto *PN : PNs) {
     assert(UserSet.empty() && "Must start with an empty user set!");
     for (Use &U : PN->uses())
       UserSet.insert(cast<Instruction>(U.getUser()));
     PNUserCountMap[PN] = UserSet.size();
     for (auto *UI : UserSet)
       UserToPNMap.insert({UI, {}}).first->second.push_back(PN);
     UserSet.clear();
   }
 
   // Now do a DFS across the operand graph of the users, computing cost as we
   // go and when all costs for a given PHI are known, checking that PHI for
   // profitability.
-  SmallDenseMap<Instruction *, int, 16> SpecCostMap;
+  SmallDenseMap<Instruction *, InstructionCost, 16> SpecCostMap;
   visitPHIUsersAndDepsInPostOrder(
       PNs,
       /*IsVisited*/
       [&](Instruction *I) {
         // We consider anything that isn't potentially speculated to be
         // "visited" as it is already handled. Similarly, anything that *is*
         // potentially speculated but for which we have an entry in our cost
         // map, we're done.
         return !PotentialSpecSet.count(I) || SpecCostMap.count(I);
       },
       /*Visit*/
       [&](Instruction *I) {
         // We've fully visited the operands, so sum their cost with this node
         // and update the cost map.
-        int Cost = TTI.TCC_Free;
+        InstructionCost Cost = TTI.TCC_Free;
         for (Value *OpV : I->operand_values())
           if (auto *OpI = dyn_cast<Instruction>(OpV)) {
             auto CostMapIt = SpecCostMap.find(OpI);
             if (CostMapIt != SpecCostMap.end())
               Cost += CostMapIt->second;
           }
         Cost += TTI.getUserCost(I, TargetTransformInfo::TCK_SizeAndLatency);
         bool Inserted = SpecCostMap.insert({I, Cost}).second;
visitPHIUsersAndDepsInPostOrder(
               --PNUserCount;
               return PNUserCount != 0;
             });
 
         // FIXME: Rather than one at a time, we should sum the savings as the
         // cost will be completely shared.
         SmallVector<Instruction *, 16> SpecWorklist;
         for (auto *PN : llvm::make_range(UserPNsSplitIt, UserPNs.end())) {
-          int SpecCost = TTI.TCC_Free;
+          InstructionCost SpecCost = TTI.TCC_Free;
           for (Use &U : PN->uses())
             SpecCost +=
                 SpecCostMap.find(cast<Instruction>(U.getUser()))->second;
           SpecCost *= (NumPreds - 1);
           // When the user count of a PHI node hits zero, we should check its
           // profitability. If profitable, we should mark it for speculation
           // and zero out the cost of everything it depends on.
-          int CostSavings = CostSavingsMap.find(PN)->second;
+          InstructionCost CostSavings = CostSavingsMap.find(PN)->second;
           if (SpecCost > CostSavings) {
             LLVM_DEBUG(dbgs() << "  Not profitable, speculation cost: " << *PN
                               << "\n"
                                  "    Cost savings:     "
                               << CostSavings
                               << "\n"
                                  "    Speculation cost: "
                               << SpecCost << "\n");

 /// This routine checks whether any of these PHIs are profitable to speculate
 /// users around. If safe and profitable, it does the speculation. It returns
 /// true when at least some speculation occurs.
 static bool tryToSpeculatePHIs(SmallVectorImpl<PHINode *> &PNs,
                                DominatorTree &DT, TargetTransformInfo &TTI) {
   LLVM_DEBUG(dbgs() << "Evaluating phi nodes for speculation:\n");
 
   // Savings in cost from speculating around a PHI node.
-  SmallDenseMap<PHINode *, int, 16> CostSavingsMap;
+  SmallDenseMap<PHINode *, InstructionCost, 16> CostSavingsMap;
 
   // Remember the set of instructions that are candidates for speculation so
   // that we can quickly walk things within that space. This prunes out
   // instructions already available along edges, etc.
   SmallPtrSet<Instruction *, 16> PotentialSpecSet;
 
   // Remember the set of instructions that are (transitively) unsafe to
   // speculate into the incoming edges of this basic block. This avoids