[Reassociate] recognize more than one pairs for later CSE - NFC

[shchenz]

Summary

Motivated case:

void foo(int a, int b, int c, int d, int* res1, int* res2, int* res3)
{
  *res1 = a + b + c + d;
  *res2 = b + c;
  *res3 = a + d;
}

There are two common expressions a+d & c+d, currently Reassociate can recognize one of them.

It gets:

%add = add i32 %b, %c
%add.1 = add i32 %add, %a
%add.2 = add i32 %add.1, %d
store i32 %add, i32*res2, align 4, !tbaa !2
store i32 %add.2, i32* %res1, align 4, !tbaa !2
%add.3 = add nsw i32 %a, %d
store i32 %add.3, i32* %res3, align 4, !tbaa !2
ret void

But we expect:

%add1 = add i32 %d, %a
%add = add i32 %c, %b
%add2 = add i32 %add, %add1
store i32 %add2, i32* %res1, align 4, !tbaa !2
store i32 %add, i32* %res2, align 4, !tbaa !2
store i32 %add1, i32* %res3, align 4, !tbaa !2
ret void

This is first part to recognize more than one pattern. This is NFC.
Changes about rewriting these common expressions will be posted in following patches.

[lebedev.ri]

I'm not familiar with this pass, but this looks a bit too large to be NFC.

[shchenz]

@lebedev.ri Thanks response for this patch, Roman. Do you happen to know who is familiar with this pass? I add the reviewer by searching the commit history.

[fhahn]

I think it would make the review much easier if the description would include the rational behind the changes, i.e how the patch works, why the changes are necessary in terms of implementation and what issue the patch concretely addressed.

Also you mention follow-on changes that build on this one. Sharing them would probably also help with reviewing this change, to see how things fit together in the end. Phabricator lets you link related/dependent patches.

[shchenz]

In D76057#1921261, @fhahn wrote:

I think it would make the review much easier if the description would include the rational behind the changes, i.e how the patch works, why the changes are necessary in terms of implementation and what issue the patch concretely addressed.

Also you mention follow-on changes that build on this one. Sharing them would probably also help with reviewing this change, to see how things fit together in the end. Phabricator lets you link related/dependent patches.

Make sense. I will ask for reviewing when follow-on patch is ready. Thanks for your comment @fhahn

[shchenz]

https://reviews.llvm.org/D76284 is posted as child of this nfc patch.

[shchenz]

gentle ping

[shchenz]

gentle ping

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

llvm/include/llvm/Transforms/Scalar/Reassociate.h

 /// 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) {}
+  bool operator ==(const ValueEntry &Other) {
+    return Op == Other.Op && Rank == Other.Rank;
+  }
 };
 
 inline bool operator<(const ValueEntry &LHS, const ValueEntry &RHS) {
   return LHS.Rank > RHS.Rank; // Sort so that highest rank goes to start.
 }
 
 /// Utility class representing a base and exponent pair which form one
 /// factor of some product.

llvm/lib/Transforms/Scalar/Reassociate.cpp

if (Ops.size() == 1) {
     // eliminate it.
     I->replaceAllUsesWith(Ops[0].Op);
     if (Instruction *OI = dyn_cast<Instruction>(Ops[0].Op))
       OI->setDebugLoc(I->getDebugLoc());
     RedoInsts.insert(I);
     return;
   }
 
+  SmallVector<std::pair<Value *, Value *>, 4> ChosenPairs;
   if (Ops.size() > 2 && Ops.size() <= GlobalReassociateLimit) {
-    // Find the pair with the highest count in the pairmap and move it to the
-    // back of the list so that it can later be CSE'd.
+    // Find the pairs with more than one use in the pairmap and move them to the
+    // back of the list so that they can later be CSE'd.
     // example:
     //   a*b*c*d*e
-    // if c*e is the most "popular" pair, we can express this as
-    //   (((c*e)*d)*b)*a
+    // if c*e and b*d have more than one use, we can express this as
+    // (((c*e)*(b*d))*a)
+    SmallVector<ValueEntry, 8> OpsWorkList(Ops);
+    while (OpsWorkList.size() >= 2) {
       unsigned Max = 1;
       unsigned BestRank = 0;
       std::pair<unsigned, unsigned> BestPair;
       unsigned Idx = I->getOpcode() - Instruction::BinaryOpsBegin;
-    for (unsigned i = 0; i < Ops.size() - 1; ++i)
-      for (unsigned j = i + 1; j < Ops.size(); ++j) {
+      for (unsigned i = 0; i < OpsWorkList.size() - 1; ++i)
+        for (unsigned j = i + 1; j < OpsWorkList.size(); ++j) {
           unsigned Score = 0;
-        Value *Op0 = Ops[i].Op;
-        Value *Op1 = Ops[j].Op;
+          Value *Op0 = OpsWorkList[i].Op;
+          Value *Op1 = OpsWorkList[j].Op;
           if (std::less<Value *>()(Op1, Op0))
             std::swap(Op0, Op1);
           auto it = PairMap[Idx].find({Op0, Op1});
-        if (it != PairMap[Idx].end()) {
+          if (it == PairMap[Idx].end())
+            continue;
           // Functions like BreakUpSubtract() can erase the Values we're using
           // as keys and create new Values after we built the PairMap. There's a
           // small chance that the new nodes can have the same address as
           // something already in the table. We shouldn't accumulate the stored
           // score in that case as it refers to the wrong Value.
           if (it->second.isValid())
-            Score += it->second.Score;
-        }
+            Score = it->second.Score;
 
-        unsigned MaxRank = std::max(Ops[i].Rank, Ops[j].Rank);
+          unsigned MaxRank = std::max(OpsWorkList[i].Rank, OpsWorkList[j].Rank);
           if (Score > Max || (Score == Max && MaxRank < BestRank)) {
             BestPair = {i, j};
             Max = Score;
             BestRank = MaxRank;
           }
         }
+      if (Max == 1)
+        break;
       if (Max > 1) {
-      auto Op0 = Ops[BestPair.first];
-      auto Op1 = Ops[BestPair.second];
-      Ops.erase(&Ops[BestPair.second]);
-      Ops.erase(&Ops[BestPair.first]);
-      Ops.push_back(Op0);
-      Ops.push_back(Op1);
+        auto Op0InWorkList = OpsWorkList[BestPair.first];
+        auto Op1InWorkList = OpsWorkList[BestPair.second];
+
+        // Update OpsWorkList.
+        OpsWorkList.erase(&OpsWorkList[BestPair.second]);
+        OpsWorkList.erase(&OpsWorkList[BestPair.first]);
+
+        Value *Op0 = Op0InWorkList.Op;
+        Value *Op1 = Op1InWorkList.Op;
+        if (std::less<Value *>()(Op1, Op0))
+          std::swap(Op0, Op1);
+
+        // If current best pair is already chosen, skip.
+        if (std::find(ChosenPairs.begin(), ChosenPairs.end(),
+                      std::make_pair(Op0, Op1)) != ChosenPairs.end())
+          continue;
+
+        // Get one best pair, put it into back of Ops.
+        // ChosenPairs.size() * 2 elements at tail position are selected pairs,
+        // do not touch them.
+        auto Op0InOpsIter = std::find(
+            Ops.begin(), Ops.end() - ChosenPairs.size() * 2, Op0InWorkList);
+        auto Op0InOps = *Op0InOpsIter;
+        assert((Op0InOpsIter != (Ops.end() - ChosenPairs.size() * 2)) &&
+               "ops wrong\n");
+        Ops.erase(Op0InOpsIter);
+
+        auto Op1InOpsIter = std::find(
+            Ops.begin(), Ops.end() - ChosenPairs.size() * 2, Op1InWorkList);
+        assert((Op1InOpsIter != (Ops.end() - ChosenPairs.size() * 2)) &&
+               "ops wrong\n");
+        auto Op1InOps = *Op1InOpsIter;
+        Ops.erase(Op1InOpsIter);
+
+        // FIXME: for now we want a NFC patch which keeps most 'popular' pair at
+        // the very end like before. But we should use push_back instead of
+        // insert bacause all pairs should be equal.
+        Ops.insert(Ops.end() - ChosenPairs.size() * 2, Op0InOps);
+        Ops.insert(Ops.end() - ChosenPairs.size() * 2, Op1InOps);
+
+        // Update ChosenPairs.
+        ChosenPairs.push_back(std::make_pair(Op0, Op1));
+      }
     }
   }
   // Now that we ordered and optimized the expressions, splat them back into
   // the expression tree, removing any unneeded nodes.
   RewriteExprTree(I, Ops);
 }
 
 void