Reassociate in favor of grouping previously paired operands

[jingyue]

Summary

This patch enhances Reassociate to group operands that are paired
previously. With this patch, Reassociate can expose in general more
opportunities for later redundancy elimination. For example,

foo(a + c);
foo(a + (b + c));

Previously,

t1 = a + c;
foo(t);
t2 = b + c;
t3 = a + t2;
foo(t3); // using 5 instructions

With this patch, since a, c are paired together in t1, later instruction
will try to pair a, c together so it can be optimized to

t1 = a + c;
foo(t1);
t2 = t1 + b;
foo(t2); // using 4 instructions

We implement this idea leveraging a data structure called
LastOperandPairsMap that stores the most recently paired operands. It is
indexed by a Value and its corresponding Opcode in this pair.

For example, LastOperandPairsMap[x][op] = <y, I> means there is an instruction I
in form of (op x, y) or or (op y, x) in the current block.

Each time when function ReassocaiteExpression sorts the operands of an
expression, it searches the map for any pair of leaves that have been
reassociated recently, and moves such pair (if any) to the end of the
operand list so that they are guaranteed to be grouped together by function
RewriteExprTree later.

Note that this optimization overlaps quite a bit with NaryReassociate.
However, since NaryReassociate requires ScalarEvolution and is a
iterative algorithm, NaryReassociate is much more expensive than the
algorithm introduced in this patch and thus cannot be run as pervasively
as Reassociate. This optimization allows Reassociate to cheaply handle
many cases NaryReassociate is designed for.

[• dberlin]

For reference, can you explain the difference between this and the
last patch I posted that does the same thing?

[wengxt]

Related discussion and original patch on llvmdev.

http://lists.cs.uiuc.edu/pipermail/llvmdev/2015-May/085246.html

There are some differences between the original patch and this version:

1. Because OptimizeExpression has some assumptions about the order of linearized ops, in this version the order of operands is adjusted after calling function OptimizeExpression.
2. Because RewriteExprTree can only pair the last two operands in Ops, in this patch the order of operands is not adjusted by rank but by simply moving one pair of operands to the end of Ops.
3. Paired operands are indexed with their corresponding Opcode and Instruction. So (add a, b) will not make a, b together in a multiplication expression.
4. The pair will be discarded correctly in following case:
5. The corresponding instruction of the pair is removed.
6. The corresponding instruction of the pair is being linearized
7. Last pairing map is cleared after processing every basic block.

[wengxt]

Updated based on Daniel's comment.

This patch shift the rank range by (1 << 16) to have some reserved
space for pairing operands. Ranks are now adjusted after function
OptimizeExpresssion and resort again.

[jingyue]

Pressed the wrong button. Meant to say "needs revision". Feel free to reclaim it.

lib/Transforms/Scalar/Reassociate.cpp

 // corresponding to the reverse post order traversal of current function
 // (starting at 2), which effectively gives values in deep loops higher rank
 // than values not in loops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/IndexedMap.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.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"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
+#include <unordered_map>
 using namespace llvm;
 
 #define DEBUG_TYPE "reassociate"
 
 STATISTIC(NumChanged, "Number of insts reassociated");
 STATISTIC(NumAnnihil, "Number of expr tree annihilated");
 STATISTIC(NumFactor , "Number of multiplies factored");
 
private:
     Value *SymbolicPart;
     APInt ConstPart;
     unsigned SymbolicRank;
     bool isOr;
   };
 }
 
 namespace {
+  struct OpcodeToIndex : public std::unary_function<unsigned, unsigned> {
+    unsigned operator()(unsigned op) const {
+      switch(op) {
+      case Instruction::Add: return 0;
+      case Instruction::Mul: return 1;
+      case Instruction::And: return 2;
+      case Instruction::Or: return 3;
+      case Instruction::Xor: return 4;
+      case Instruction::FAdd: return 5;
+      case Instruction::FMul: return 6;
+      // other op code
+      default: return 7;
+      }
+    }
+  };
+
+  // Utility class storing the last paired Operands per basic block
+  // It maps from (Value, Opcode) to (Value, Instruction).
+  // Mapped instruction is where these two value get paired.
+  class LastOperandPairsMap {
+  public:
+    // We don't want to walk through the map to figure out the dangling Value
+    // and instructions, so just use WeakVH here and check the pointer before
+    // use it.
+    typedef std::pair<WeakVH, WeakVH> PairedValue;
+    void Clear() { PairingMap.clear(); }
+    void Add(Value *LHS, Value *RHS, Instruction *I);
+    void Remove(Instruction *I);
+    std::pair<WeakVH, WeakVH> Lookup(Value *V, unsigned Opcode) const;
+    void Cleanup(Instruction *I);
+
+  private:
+    DenseMap<AssertingVH<Value>, IndexedMap<PairedValue, OpcodeToIndex>> PairingMap;
+  };
+}
+
+namespace {
   class Reassociate : public FunctionPass {
     DenseMap<BasicBlock*, unsigned> RankMap;
     DenseMap<AssertingVH<Value>, unsigned> ValueRankMap;
     SetVector<AssertingVH<Instruction> > RedoInsts;
     bool MadeChange;
+    LastOperandPairsMap LastPairingMap;
   public:
     static char ID; // Pass identification, replacement for typeid
     Reassociate() : FunctionPass(ID) {
       initializeReassociatePass(*PassRegistry::getPassRegistry());
     }
 
     bool runOnFunction(Function &F) override;
 
XorOpnd::XorOpnd(Value *V) {
   }
 
   // view the operand as "V | 0"
   SymbolicPart = V;
   ConstPart = APInt::getNullValue(V->getType()->getIntegerBitWidth());
   isOr = true;
 }
 
+void LastOperandPairsMap::Add(Value *LHS, Value *RHS, Instruction *I) {
+  auto Opcode = I->getOpcode();
+  Value *Op[] = {LHS, RHS};
+  for (unsigned i = 0; i != 2; ++i) {
+    auto &Map = PairingMap[Op[i]];
+    Map.grow(Opcode);
+    Map[I->getOpcode()] = std::make_pair(WeakVH(Op[1 - i]), WeakVH(I));
+  }
+}
+
+void LastOperandPairsMap::Remove(Instruction *I) {
+  if (!isa<BinaryOperator>(I))
+    return;
+
+  auto Opcode = I->getOpcode();
+  for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i) {
+    auto Iter = PairingMap.find(I->getOperand(i));
+    if (Iter != PairingMap.end() && Iter->second.inBounds(Opcode) &&
+        Iter->second[Opcode].second == I) {
+      Iter->second[Opcode].first = nullptr;
+      Iter->second[Opcode].second = nullptr;
+    }
+  }
+}
+
+std::pair<WeakVH, WeakVH> LastOperandPairsMap::Lookup(Value *V,
+                                                      unsigned Opcode) const
+{
+  auto Iter = PairingMap.find(V);
+  if (Iter == PairingMap.end()) {
+    return PairedValue();
+  }
+  auto &Map = Iter->second;
+  if (!Map.inBounds(Opcode)) {
+    return PairedValue();
+  }
+
+  return Map[Opcode];
+}
+
+void LastOperandPairsMap::Cleanup(Instruction *I) {
+  Remove(I);
+  PairingMap.erase(I);
+}
+
 char Reassociate::ID = 0;
 INITIALIZE_PASS(Reassociate, "reassociate",
                 "Reassociate expressions", false, false)
 
 // Public interface to the Reassociate pass
 FunctionPass *llvm::createReassociatePass() { return new Reassociate(); }
 
 /// Return true if V is an instruction of the specified opcode and if it

 /// which requires passing through a phi node.
 ///
 /// Note that this routine may also mutate binary operators of the wrong type
 /// that have all uses inside the expression (i.e. only used by non-leaf nodes
 /// of the expression) if it can turn them into binary operators of the right
 /// type and thus make the expression bigger.
 
 static bool LinearizeExprTree(BinaryOperator *I,
-                              SmallVectorImpl<RepeatedValue> &Ops) {
+                              SmallVectorImpl<RepeatedValue> &Ops,
+                              LastOperandPairsMap &OpndPairsMap) {
   DEBUG(dbgs() << "LINEARIZE: " << *I << '\n');
   unsigned Bitwidth = I->getType()->getScalarType()->getPrimitiveSizeInBits();
   unsigned Opcode = I->getOpcode();
   assert(I->isAssociative() && I->isCommutative() &&
          "Expected an associative and commutative operation!");
 
   // Visit all operands of the expression, keeping track of their weight (the
   // number of paths from the expression root to the operand, or if you like
static bool LinearizeExprTree(BinaryOperator *I,
   SmallVector<Value*, 8> LeafOrder; // Ensure deterministic leaf output order.
 
 #ifndef NDEBUG
   SmallPtrSet<Value*, 8> Visited; // For sanity checking the iteration scheme.
 #endif
   while (!Worklist.empty()) {
     std::pair<BinaryOperator*, APInt> P = Worklist.pop_back_val();
     I = P.first; // We examine the operands of this binary operator.
+    // The pairs introduced by inner nodes should be ignored.
+    OpndPairsMap.Remove(I);
 
     for (unsigned OpIdx = 0; OpIdx < 2; ++OpIdx) { // Visit operands.
       Value *Op = I->getOperand(OpIdx);
       APInt Weight = P.second; // Number of paths to this operand.
       DEBUG(dbgs() << "OPERAND: " << *Op << " (" << Weight << ")\n");
       assert(!Op->use_empty() && "No uses, so how did we get to it?!");
 
       // If this is a binary operation of the right kind with only one use then
for (unsigned i = 0; ; ++i) {
     // The last operation (which comes earliest in the IR) is special as both
     // operands will come from Ops, rather than just one with the other being
     // a subexpression.
     if (i+2 == Ops.size()) {
       Value *NewLHS = Ops[i].Op;
       Value *NewRHS = Ops[i+1].Op;
       Value *OldLHS = Op->getOperand(0);
       Value *OldRHS = Op->getOperand(1);
+      LastPairingMap.Add(NewLHS, NewRHS, Op);
 
       if (NewLHS == OldLHS && NewRHS == OldRHS)
         // Nothing changed, leave it alone.
         break;
 
       if (NewLHS == OldRHS && NewRHS == OldLHS) {
         // The order of the operands was reversed.  Swap them.
         DEBUG(dbgs() << "RA: " << *Op << '\n');
if (NewRHS != Op->getOperand(1)) {
       } else {
         // Overwrite with the new right-hand side.
         BinaryOperator *BO = isReassociableOp(Op->getOperand(1), Opcode);
         if (BO && !NotRewritable.count(BO))
           NodesToRewrite.push_back(BO);
         Op->setOperand(1, NewRHS);
         ExpressionChanged = Op;
       }
+      LastPairingMap.Add(Op->getOperand(0), NewRHS, Op);
       DEBUG(dbgs() << "TO: " << *Op << '\n');
       MadeChange = true;
       ++NumChanged;
     }
 
     // Now deal with the left-hand side.  If this is already an operation node
     // from the original expression then just rewrite the rest of the expression
     // into it.

 /// and if this sequence contains a multiply by Factor,
 /// remove Factor from the tree and return the new tree.
 Value *Reassociate::RemoveFactorFromExpression(Value *V, Value *Factor) {
   BinaryOperator *BO = isReassociableOp(V, Instruction::Mul, Instruction::FMul);
   if (!BO)
     return nullptr;
 
   SmallVector<RepeatedValue, 8> Tree;
-  MadeChange |= LinearizeExprTree(BO, Tree);
+  MadeChange |= LinearizeExprTree(BO, Tree, LastPairingMap);
   SmallVector<ValueEntry, 8> Factors;
   Factors.reserve(Tree.size());
   for (unsigned i = 0, e = Tree.size(); i != e; ++i) {
     RepeatedValue E = Tree[i];
     Factors.append(E.second.getZExtValue(),
                    ValueEntry(getRank(E.first), E.first));
   }
 

 
 /// Zap the given instruction, adding interesting operands to the work list.
 void Reassociate::EraseInst(Instruction *I) {
   assert(isInstructionTriviallyDead(I) && "Trivially dead instructions only!");
   SmallVector<Value*, 8> Ops(I->op_begin(), I->op_end());
   // Erase the dead instruction.
   ValueRankMap.erase(I);
   RedoInsts.remove(I);
+  LastPairingMap.Cleanup(I);
+
   I->eraseFromParent();
   // Optimize its operands.
   SmallPtrSet<Instruction *, 8> Visited; // Detect self-referential nodes.
   for (unsigned i = 0, e = Ops.size(); i != e; ++i)
     if (Instruction *Op = dyn_cast<Instruction>(Ops[i])) {
       // If this is a node in an expression tree, climb to the expression root
       // and add that since that's where optimization actually happens.
       unsigned Opcode = Op->getOpcode();
if (BO->hasOneUse() && BO->getOpcode() == Instruction::Add &&
     return;
   if (BO->hasOneUse() && BO->getOpcode() == Instruction::FAdd &&
       cast<Instruction>(BO->user_back())->getOpcode() == Instruction::FSub)
     return;
 
   ReassociateExpression(BO);
 }
 
+static void MoveOperandsToEnd(unsigned OpIdxA, unsigned OpIdxB,
+                              SmallVectorImpl<ValueEntry> &Ops)
+{
+  using std::swap;
+  // If the operands are already in their place, do nothing
+  if ((OpIdxA + 2 == Ops.size() && OpIdxB + 1 == Ops.size()) ||
+      (OpIdxA + 1 == Ops.size() && OpIdxB + 2 == Ops.size())) {
+    return;
+  }
+
+  // First move two Operands to the back of Ops
+  std::move(Ops.begin() + OpIdxA, Ops.begin() + OpIdxA + 1, std::back_inserter(Ops));
+  std::move(Ops.begin() + OpIdxB, Ops.begin() + OpIdxB + 1, std::back_inserter(Ops));
+
+  // make sure i < j
+  if (OpIdxA > OpIdxB)
+    swap(OpIdxA, OpIdxB);
+  // Erase the old element at OpIdxA and OpIdxB, because Ops is stored in a
+  // use std::remove_if instead of calling erase twice to avoid redundant copy.
+  unsigned count = 0;
+  auto Delta = OpIdxB - OpIdxA;
+  Ops.erase(std::remove_if(Ops.begin() + OpIdxA, Ops.end(),
+                           [&count, Delta](const ValueEntry&) -> bool {
+                             auto result = (count == 0 || count == Delta);
+                             count ++;
+                             return result;
+                           }),
+            Ops.end());
+}
+
+static
+std::pair<bool, unsigned> FindPairedOperandIdx(
+  const LastOperandPairsMap& LastPairingMap,
+  unsigned Opcode,
+  const SmallVectorImpl<ValueEntry> &Ops,
+  unsigned Idx,
+  const std::unordered_multimap<Value *, unsigned> &Leaves) {
+    Value* PairedValue;
+    Value* Inst;
+    std::tie(PairedValue, Inst) = LastPairingMap.Lookup(Ops[Idx].Op, Opcode);
+    // Make sure the pair we found is valid, the pair may be invalid when:
+    // 1. The pair doesn't exist.
+    // 2. the instruction contains this pair is erased.
+    // 3. the value in this pair is erased.
+    if (!PairedValue || !Inst) {
+      return std::make_pair(false, 0U);
+    }
+
+    // Search all the paired value appeared in Ops
+    auto IterPair = Leaves.equal_range(PairedValue);
+    // Find a Value which is not itself. A value can be paired with itself.
+    // For example,
+    //   j = i * i
+    //   ...
+    //   k = a * i
+    //   l = k * i
+    // will introduce a pair (i, i), but if we want to pair them together later,
+    // we need to make sure that x pairs with another x but not itself.
+    auto Iter = std::find_if(
+        IterPair.first, IterPair.second,
+        [Idx](const std::pair<Value*, unsigned> &Item) {
+          // check if we don't pair to currrent Op itself
+          return Item.second != Idx;
+        });
+
+    if (Iter == IterPair.second) {
+      return std::make_pair(false, 0U);
+    }
+
+    return std::make_pair(true, Iter->second);
+}
+
+static void PairOperands(const LastOperandPairsMap& LastPairingMap,
+                         unsigned Opcode,
+                         SmallVectorImpl<ValueEntry> &Ops)
+{
+  SmallPtrSet<Value *, 8> Visited;
+  // First put Value and its original position in Ops in to Leaves
+  std::unordered_multimap<Value *, unsigned> Leaves;
+  for (unsigned i = 0, e = Ops.size(); i != e ; ++i) {
+    Leaves.insert(std::make_pair(Ops[i].Op, i));
+  }
+  for (unsigned i = 0, e = Ops.size(); i != e ; ++i) {
+    // Find a paired operand from Ops
+    bool PairFound;
+    unsigned PairedOpIdx;
+    std::tie(PairFound, PairedOpIdx) = FindPairedOperandIdx(LastPairingMap,
+                                                            Opcode, Ops,
+                                                            i, Leaves);
+    if (!PairFound) {
+      continue;
+    }
+
+    DEBUG(dbgs() << "PAIR OPERANDS: [" << *Ops[i].Op << "], ["
+                 << *Ops[PairedOpIdx].Op << "]\n");
+    // Move paired operands to the end of Ops in order to make RewriteExprTree
+    // pair them together.
+    MoveOperandsToEnd(i, PairedOpIdx, Ops);
+    return;
+  }
+}
+
 void Reassociate::ReassociateExpression(BinaryOperator *I) {
   // First, walk the expression tree, linearizing the tree, collecting the
   // operand information.
   SmallVector<RepeatedValue, 8> Tree;
-  MadeChange |= LinearizeExprTree(I, Tree);
+  MadeChange |= LinearizeExprTree(I, Tree, LastPairingMap);
   SmallVector<ValueEntry, 8> Ops;
   Ops.reserve(Tree.size());
   for (unsigned i = 0, e = Tree.size(); i != e; ++i) {
     RepeatedValue E = Tree[i];
     Ops.append(E.second.getZExtValue(),
                ValueEntry(getRank(E.first), E.first));
   }
 
if (Value *V = OptimizeExpression(I, Ops)) {
     I->replaceAllUsesWith(V);
     if (Instruction *VI = dyn_cast<Instruction>(V))
       VI->setDebugLoc(I->getDebugLoc());
     RedoInsts.insert(I);
     ++NumAnnihil;
     return;
   }
 
+  // Make pair before OptimizeExpression may mess up constant merge and
+  // factorize in OptimizeExpression, so we do pair after that.
+  PairOperands(LastPairingMap, I->getOpcode(), Ops);
+
   // We want to sink immediates as deeply as possible except in the case where
   // this is a multiply tree used only by an add, and the immediate is a -1.
   // In this case we reassociate to put the negation on the outside so that we
   // can fold the negation into the add: (-X)*Y + Z -> Z-X*Y
   if (I->hasOneUse()) {
     if (I->getOpcode() == Instruction::Mul &&
         cast<Instruction>(I->user_back())->getOpcode() == Instruction::Add &&
         isa<ConstantInt>(Ops.back().Op) &&
for (Function::iterator BI = F.begin(), BE = F.end(); BI != BE; ++BI) {
     // If this produced extra instructions to optimize, handle them now.
     while (!RedoInsts.empty()) {
       Instruction *I = RedoInsts.pop_back_val();
       if (isInstructionTriviallyDead(I))
         EraseInst(I);
       else
         OptimizeInst(I);
     }
+
+    // Clear LastPairingMap per basic block
+    LastPairingMap.Clear();
   }
 
   // We are done with the rank map.
   RankMap.clear();
   ValueRankMap.clear();
 
   return MadeChange;
 }

test/Transforms/Reassociate/pair.ll

+; RUN: opt < %s -reassociate -early-cse -S | FileCheck %s
+
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+
+declare void @foo(i32)
+
+; foo(a + c);
+; foo((a + (b + c));
+;   =>
+; t = a + c;
+; foo(t);
+; foo(t + b);
+define void @test1(i32 %a, i32 %b, i32 %c) {
+; CHECK-LABEL: @test1(
+  %1 = mul i32 %a, %c
+; CHECK: [[BASE:%[a-zA-Z0-9]+]] = mul i32 %c, %a
+  call void @foo(i32 %1)
+  %2 = mul i32 %b, %c
+  %3 = mul i32 %a, %2
+; CHECK: mul i32 [[BASE]], %b
+  call void @foo(i32 %3)
+  ret void
+}
+
+; Test that we will not pair inner nodes that are processed during linearization
+define i32 @test2(i32 %a, i32 %b, i32 %z) {
+; CHECK-LABEL: test2
+
+  %d = mul i32 %z, 40
+  %c = sub i32 0, %d
+  %e = mul i32 %a, %c
+; CHECK-NOT: %e = mul i32 %z, 40
+  %f = sub i32 0, %e
+  ret i32 %f
+}