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

 //===- Reassociate.cpp - Reassociate binary expressions -------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass reassociates commutative expressions in an order that is designed
 // to promote better constant propagation, GCSE, LICM, PRE, etc.
 //
 // For example: 4 + (x + 5) -> x + (4 + 5)
 //
-// In the implementation of this algorithm, constants are assigned rank = 0,
-// function arguments are rank = 1, and other values are assigned ranks
-// 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.
+// In the implementation of this algorithm, constants are assigned
+// rank = 1 << 16, function arguments are distanct rank start from
+// (1 << 16) + 3, and other unmovable values are assigned ranks corresponding to
+// the reverse post order traversal of current function (starting at
+// (BlockRank << 16) + 1), which effectively gives values in deep loops higher
+// rank than values not in loops. Rank between 0 and (1 << 16) are reserved for
+// pairing previous associated value.
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
+    const unsigned ConstantRank = (1 << 16);
     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
static bool isUnmovableInstruction(Instruction *I) {
   case Instruction::Call:
     return !isa<DbgInfoIntrinsic>(I);
   default:
     return false;
   }
 }
 
 void Reassociate::BuildRankMap(Function &F) {
-  unsigned i = 2;
+  unsigned i = ConstantRank + 2;
 
-  // Assign distinct ranks to function arguments.
+  // Assign distinct ranks to function arguments whose rank will be larger than
+  // constants.
   for (Function::arg_iterator I = F.arg_begin(), E = F.arg_end(); I != E; ++I) {
     ValueRankMap[&*I] = ++i;
     DEBUG(dbgs() << "Calculated Rank[" << I->getName() << "] = " << i << "\n");
   }
 
+  i = 1;
   ReversePostOrderTraversal<Function*> RPOT(&F);
   for (ReversePostOrderTraversal<Function*>::rpo_iterator I = RPOT.begin(),
          E = RPOT.end(); I != E; ++I) {
     BasicBlock *BB = *I;
     unsigned BBRank = RankMap[BB] = ++i << 16;
 
     // Walk the basic block, adding precomputed ranks for any instructions that
     // we cannot move.  This ensures that the ranks for these instructions are
     // all different in the block.
     for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I)
       if (isUnmovableInstruction(I))
         ValueRankMap[&*I] = ++BBRank;
   }
 }
 
 unsigned Reassociate::getRank(Value *V) {
   Instruction *I = dyn_cast<Instruction>(V);
   if (!I) {
     if (isa<Argument>(V)) return ValueRankMap[V];   // Function argument.
-    return 0;  // Otherwise it's a global or constant, rank 0.
+    return ConstantRank;  // Otherwise it's a global or constant.
   }
 
   if (unsigned Rank = ValueRankMap[I])
     return Rank;    // Rank already known?
 
   // If this is an expression, return the 1+MAX(rank(LHS), rank(RHS)) so that
   // we can reassociate expressions for code motion!  Since we do not recurse
   // for PHI nodes, we cannot have infinite recursion here, because there

 /// 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));
   }
 
if (Ops.empty())
     return Cst;
 
   // Put the combined constant back at the end of the operand list, except if
   // there is no point.  For example, an add of 0 gets dropped here, while a
   // multiplication by zero turns the whole expression into zero.
   if (Cst && Cst != ConstantExpr::getBinOpIdentity(Opcode, I->getType())) {
     if (Cst == ConstantExpr::getBinOpAbsorber(Opcode, I->getType()))
       return Cst;
-    Ops.push_back(ValueEntry(0, Cst));
+    Ops.push_back(ValueEntry(ConstantRank, Cst));
   }
 
   if (Ops.size() == 1) return Ops[0].Op;
 
   // Handle destructive annihilation due to identities between elements in the
   // argument list here.
   unsigned NumOps = Ops.size();
   switch (Opcode) {

 
 /// 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
+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;
+  unsigned PairedRank = 0;
+  // 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");
+    // Assign paired operands low rank
+    // TODO: RewriteExprTree can only guarantee one pair is put together, if
+    // situation changed, we change more pairs rank and will need to check if
+    // an Op is not paired before.
+    Ops[i].Rank = PairedRank;
+    Ops[PairedOpIdx].Rank = PairedRank++;
+    std::stable_sort(Ops.begin(), Ops.end());
+    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
+}