[X86] Reduce complexity of the LEA optimization pass

[aturetsk]

Summary

In the OptimizeLEA pass keep instructions' positions in the basic block saved and use them for calculation of the distance between two instructions instead of std::distance. This reduces complexity of the pass from O(n^3) to O(n^2) and thus the compile time.

[aturetsk]

Here are the results.

Just -Os (LEA pass is disabled):

$ time ./bin/clang++ -std=c++11 -S a.ll -Os

real    0m8.328s
user    0m8.282s
sys     0m0.041s

-Os with the old LEA pass:

$ time ./bin/clang++ -std=c++11 -S a.ll -Os -mllvm -enable-x86-lea-opt

real    0m11.653s
user    0m11.591s
sys     0m0.059s

-Os with the new LEA pass:

$ time ./bin/clang++ -std=c++11 -S a.ll -Os -mllvm -enable-x86-lea-opt

real    0m8.446s
user    0m8.380s
sys     0m0.064s

a.ll is taken from the example from https://llvm.org/bugs/show_bug.cgi?id=25843 and was generated this way:

$ python a.py 5000 > a.cc
$ ./bin/clang++ -std=c++11 -S -emit-llvm -Os a.cc
(this took about 1.5 hours)

[george.burgess.iv]

Thanks for this! Just one tiny nit.

lib/Target/X86/X86OptimizeLEAs.cpp

93

Is there a reason that we can't use DenseMap here instead?

[aturetsk]

Replaced unordered_map with DenseMap.

[aturetsk]

Hello George,
Thanks for the review.

lib/Target/X86/X86OptimizeLEAs.cpp

93

Fixed.

lib/Target/X86/X86OptimizeLEAs.cpp

 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetInstrInfo.h"
 
+#include <unordered_map>
+
 using namespace llvm;
 
 #define DEBUG_TYPE "x86-optimize-LEAs"
 
 static cl::opt<bool> EnableX86LEAOpt("enable-x86-lea-opt", cl::Hidden,
                                      cl::desc("X86: Enable LEA optimizations."),
                                      cl::init(false));
 
private:
   /// \brief Returns true if two instructions have memory operands that only
   /// differ by displacement. The numbers of the first memory operands for both
   /// instructions are specified through \p N1 and \p N2. The address
   /// displacement is returned through AddrDispShift.
   bool isSimilarMemOp(const MachineInstr &MI1, unsigned N1,
                       const MachineInstr &MI2, unsigned N2,
                       int64_t &AddrDispShift);
 
-  /// \brief Find all LEA instructions in the basic block.
+  /// \brief Find all LEA instructions in the basic block. Also, assign position
+  /// numbers to all instructions in the basic block to speed up calculation of
+  /// distance between them.
   void findLEAs(const MachineBasicBlock &MBB,
                 SmallVectorImpl<MachineInstr *> &List);
 
   /// \brief Removes redundant address calculations.
   bool removeRedundantAddrCalc(const SmallVectorImpl<MachineInstr *> &List);
 
+  std::unordered_map<const MachineInstr *, unsigned> InstrPos;
+
   MachineRegisterInfo *MRI;
   const X86InstrInfo *TII;
   const X86RegisterInfo *TRI;
 
   static char ID;
 };
 char OptimizeLEAPass::ID = 0;
 }
 
 FunctionPass *llvm::createX86OptimizeLEAs() { return new OptimizeLEAPass(); }
 
 int OptimizeLEAPass::calcInstrDist(const MachineInstr &First,
                                    const MachineInstr &Last) {
-  const MachineBasicBlock *MBB = First.getParent();
-
-  // Both instructions must be in the same basic block.
-  assert(Last.getParent() == MBB &&
+  // Both instructions must be in the same basic block and they must be
+  // presented in InstrPos.
+  assert(Last.getParent() == First.getParent() &&
          "Instructions are in different basic blocks");
+  assert(InstrPos.find(&First) != InstrPos.end() &&
+         InstrPos.find(&Last) != InstrPos.end() &&
+         "Instructions' positions are undefined");
 
-  return std::distance(MBB->begin(), MachineBasicBlock::const_iterator(&Last)) -
-         std::distance(MBB->begin(), MachineBasicBlock::const_iterator(&First));
+  return InstrPos[&Last] - InstrPos[&First];
 }
 
 // Find the best LEA instruction in the List to replace address recalculation in
 // MI. Such LEA must meet these requirements:
 // 1) The address calculated by the LEA differs only by the displacement from
 //    the address used in MI.
 // 2) The register class of the definition of the LEA is compatible with the
 //    register class of the address base register of MI.
else
       return false;
   }
 
   return true;
 }
 
 void OptimizeLEAPass::findLEAs(const MachineBasicBlock &MBB,
                                SmallVectorImpl<MachineInstr *> &List) {
+  unsigned Pos = 0;
   for (auto &MI : MBB) {
+    // Assign the position number to the instruction. Note that we are going to
+    // move some instructions during the optimization however there will never
+    // be a need to move two instructions before any selected instruction. So to
+    // avoid multiple positions' updates during moves we just increase position
+    // counter by two leaving a free space for instructions which will be moved.
+    InstrPos[&MI] = Pos += 2;
+
     if (isLEA(MI))
       List.push_back(const_cast<MachineInstr *>(&MI));
   }
 }
 
 // Try to find load and store instructions which recalculate addresses already
 // calculated by some LEA and replace their memory operands with its def
 // register.
for (auto I = MBB->begin(), E = MBB->end(); I != E;) {
     // the instruction. If LEA occurs after, we can lift LEA above the
     // instruction and this way to be able to replace it. Since LEA and the
     // instruction have similar memory operands (thus, the same def
     // instructions for these operands), we can always do that, without
     // worries of using registers before their defs.
     if (Dist < 0) {
       DefMI->removeFromParent();
       MBB->insert(MachineBasicBlock::iterator(&MI), DefMI);
+      InstrPos[DefMI] = InstrPos[&MI] - 1;
+
+      // Make sure the instructions' position numbers are sane.
+      assert(((InstrPos[DefMI] == 1 && DefMI == MBB->begin()) ||
+              InstrPos[DefMI] >
+                  InstrPos[std::prev(MachineBasicBlock::iterator(DefMI))]) &&
+             "Instruction positioning is broken");
     }
 
     // Since we can possibly extend register lifetime, clear kill flags.
     MRI->clearKillFlags(DefMI->getOperand(0).getReg());
 
     ++NumSubstLEAs;
     DEBUG(dbgs() << "OptimizeLEAs: Candidate to replace: "; MI.dump(););
 
bool OptimizeLEAPass::runOnMachineFunction(MachineFunction &MF) {
 
   MRI = &MF.getRegInfo();
   TII = MF.getSubtarget<X86Subtarget>().getInstrInfo();
   TRI = MF.getSubtarget<X86Subtarget>().getRegisterInfo();
 
   // Process all basic blocks.
   for (auto &MBB : MF) {
     SmallVector<MachineInstr *, 16> LEAs;
+    InstrPos.clear();
 
     // Find all LEA instructions in basic block.
     findLEAs(MBB, LEAs);
 
     // If current basic block has no LEAs, move on to the next one.
     if (LEAs.empty())
       continue;
 
     // Remove redundant address calculations.
     Changed |= removeRedundantAddrCalc(LEAs);
   }
 
   return Changed;
 }