[ImplicitNullChecks] Hoist trivial depdendencies if possible

When trying to convert a loading instruction into a FAULTING_LOAD, we
sometimes face code like this:

  if %R10 is not null:
    %R9<def> = MOV32ri Immediate
    %R9<def, tied> = AND32rm %R9, 0x20(%R10)
  else:
    goto TRAP

In these cases we would like to use the AND32rm instruction as the
faulting operation by hoisting the "depedency" def-ing %R9 also above
the control flow, transforming the program into:

  %R9<def> = MOV32ri Immediate
  %R9<def, tied> = FAULTING_LOAD_OP(AND32rm %R9, 0x20(%R10), FailPath: TRAP)

This change teaches ImplicitNullChecks to do the above, when safe.

llvm-svn: 273501

1 files changed, 159 insertions, 20 deletions

diff --git a/llvm/lib/CodeGen/ImplicitNullChecks.cpp b/llvm/lib/CodeGen/ImplicitNullChecks.cpp
index ada66daac4c..11df2de216d 100644
--- a/llvm/lib/CodeGen/ImplicitNullChecks.cpp
+++ b/llvm/lib/CodeGen/ImplicitNullChecks.cpp
@@ -28,6 +28,7 @@
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
@@ -75,14 +76,19 @@ class ImplicitNullChecks : public MachineFunctionPass {
     // The block branched to if the pointer is null.
     MachineBasicBlock *NullSucc;
 
+    // If this is non-null, then MemOperation has a dependency on on this
+    // instruction; and it needs to be hoisted to execute before MemOperation.
+    MachineInstr *OnlyDependency;
+
   public:
     explicit NullCheck(MachineInstr *memOperation, MachineInstr *checkOperation,
                        MachineBasicBlock *checkBlock,
                        MachineBasicBlock *notNullSucc,
-                       MachineBasicBlock *nullSucc)
+                       MachineBasicBlock *nullSucc,
+                       MachineInstr *onlyDependency)
         : MemOperation(memOperation), CheckOperation(checkOperation),
-          CheckBlock(checkBlock), NotNullSucc(notNullSucc), NullSucc(nullSucc) {
-    }
+          CheckBlock(checkBlock), NotNullSucc(notNullSucc), NullSucc(nullSucc),
+          OnlyDependency(onlyDependency) {}
 
     MachineInstr *getMemOperation() const { return MemOperation; }
 
@@ -93,10 +99,13 @@ class ImplicitNullChecks : public MachineFunctionPass {
     MachineBasicBlock *getNotNullSucc() const { return NotNullSucc; }
 
     MachineBasicBlock *getNullSucc() const { return NullSucc; }
+
+    MachineInstr *getOnlyDependency() const { return OnlyDependency; }
   };
 
   const TargetInstrInfo *TII = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
+  AliasAnalysis *AA = nullptr;
   MachineModuleInfo *MMI = nullptr;
 
   bool analyzeBlockForNullChecks(MachineBasicBlock &MBB,
@@ -113,6 +122,10 @@ public:
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AAResultsWrapperPass>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
 
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
@@ -126,14 +139,22 @@ public:
 /// machine instruction can be re-ordered from after the machine instructions
 /// seen so far to before them.
 class HazardDetector {
-  DenseSet<unsigned> RegDefs;
+  static MachineInstr *getUnknownMI() {
+    return DenseMapInfo<MachineInstr *>::getTombstoneKey();
+  }
+
+  // Maps physical registers to the instruction defining them.  If there has
+  // been more than one def of an specific register, that register is mapped to
+  // getUnknownMI().
+  DenseMap<unsigned, MachineInstr *> RegDefs;
   DenseSet<unsigned> RegUses;
   const TargetRegisterInfo &TRI;
   bool hasSeenClobber;
+  AliasAnalysis &AA;
 
 public:
-  explicit HazardDetector(const TargetRegisterInfo &TRI) :
-    TRI(TRI), hasSeenClobber(false) {}
+  explicit HazardDetector(const TargetRegisterInfo &TRI, AliasAnalysis &AA)
+      : TRI(TRI), hasSeenClobber(false), AA(AA) {}
 
   /// \brief Make a note of \p MI for later queries to isSafeToHoist.
   ///
@@ -141,8 +162,10 @@ public:
   void rememberInstruction(MachineInstr *MI);
 
   /// \brief Return true if it is safe to hoist \p MI from after all the
-  /// instructions seen so far (via rememberInstruction) to before it.
-  bool isSafeToHoist(MachineInstr *MI);
+  /// instructions seen so far (via rememberInstruction) to before it.  If \p MI
+  /// has one and only one transitive dependency, set \p Dependency to that
+  /// instruction.  If there are more dependencies, return false.
+  bool isSafeToHoist(MachineInstr *MI, MachineInstr *&Dependency);
 
   /// \brief Return true if this instance of HazardDetector has been clobbered
   /// (i.e. has no more useful information).
@@ -181,15 +204,23 @@ void HazardDetector::rememberInstruction(MachineInstr *MI) {
     if (!MO.isReg() || !MO.getReg())
       continue;
 
-    if (MO.isDef())
-      RegDefs.insert(MO.getReg());
-    else
+    if (MO.isDef()) {
+      auto It = RegDefs.find(MO.getReg());
+      if (It == RegDefs.end())
+        RegDefs.insert({MO.getReg(), MI});
+      else {
+        assert(It->second && "Found null MI?");
+        It->second = getUnknownMI();
+      }
+    } else
       RegUses.insert(MO.getReg());
   }
 }
 
-bool HazardDetector::isSafeToHoist(MachineInstr *MI) {
+bool HazardDetector::isSafeToHoist(MachineInstr *MI,
+                                   MachineInstr *&Dependency) {
   assert(!isClobbered() && "isSafeToHoist cannot do anything useful!");
+  Dependency = nullptr;
 
   // Right now we don't want to worry about LLVM's memory model.  This can be
   // made more precise later.
@@ -199,9 +230,54 @@ bool HazardDetector::isSafeToHoist(MachineInstr *MI) {
 
   for (auto &MO : MI->operands()) {
     if (MO.isReg() && MO.getReg()) {
-      for (unsigned Reg : RegDefs)
-        if (TRI.regsOverlap(Reg, MO.getReg()))
-          return false;  // We found a write-after-write or read-after-write
+      for (auto &RegDef : RegDefs) {
+        unsigned Reg = RegDef.first;
+        MachineInstr *MI = RegDef.second;
+        if (!TRI.regsOverlap(Reg, MO.getReg()))
+          continue;
+
+        // We found a write-after-write or read-after-write, see if the
+        // instruction causing this dependency can be hoisted too.
+
+        if (MI == getUnknownMI())
+          // We don't have precise dependency information.
+          return false;
+
+        if (Dependency) {
+          if (Dependency == MI)
+            continue;
+          // We already have one dependency, and we can track only one.
+          return false;
+        }
+
+        // Now check if MI is actually a dependency that can be hoisted.
+
+        // We don't want to track transitive dependencies.  We already know that
+        // MI is the only instruction that defines Reg, but we need to be sure
+        // that it does not use any registers that have been defined (trivially
+        // checked below by ensuring that there are no register uses), and that
+        // it is the only def for every register it defines (otherwise we could
+        // violate a write after write hazard).
+        auto IsMIOperandSafe = [&](MachineOperand &MO) {
+          if (!MO.isReg() || !MO.getReg())
+            return true;
+          if (MO.isUse())
+            return false;
+          assert((!MO.isDef() || RegDefs.count(MO.getReg())) &&
+                 "All defs must be tracked in RegDefs by now!");
+          return !MO.isDef() || RegDefs.find(MO.getReg())->second == MI;
+        };
+
+        if (!all_of(MI->operands(), IsMIOperandSafe))
+          return false;
+
+        // Now check for speculation safety:
+        bool SawStore = true;
+        if (!MI->isSafeToMove(&AA, SawStore) || MI->mayLoad())
+          return false;
+
+        Dependency = MI;
+      }
 
       if (MO.isDef())
         for (unsigned Reg : RegUses)
@@ -217,6 +293,7 @@ bool ImplicitNullChecks::runOnMachineFunction(MachineFunction &MF) {
   TII = MF.getSubtarget().getInstrInfo();
   TRI = MF.getRegInfo().getTargetRegisterInfo();
   MMI = &MF.getMMI();
+  AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
   SmallVector<NullCheck, 16> NullCheckList;
 
@@ -229,6 +306,16 @@ bool ImplicitNullChecks::runOnMachineFunction(MachineFunction &MF) {
   return !NullCheckList.empty();
 }
 
+// Return true if any register aliasing \p Reg is live-in into \p MBB.
+static bool AnyAliasLiveIn(const TargetRegisterInfo *TRI,
+                           MachineBasicBlock *MBB, unsigned Reg) {
+  for (MCRegAliasIterator AR(Reg, TRI, /*IncludeSelf*/ true); AR.isValid();
+       ++AR)
+    if (MBB->isLiveIn(*AR))
+      return true;
+  return false;
+}
+
 /// Analyze MBB to check if its terminating branch can be turned into an
 /// implicit null check.  If yes, append a description of the said null check to
 /// NullCheckList and return true, else return false.
@@ -330,20 +417,56 @@ bool ImplicitNullChecks::analyzeBlockForNullChecks(
 
   unsigned PointerReg = MBP.LHS.getReg();
 
-  HazardDetector HD(*TRI);
+  HazardDetector HD(*TRI, *AA);
 
   for (auto MII = NotNullSucc->begin(), MIE = NotNullSucc->end(); MII != MIE;
        ++MII) {
     MachineInstr *MI = &*MII;
     unsigned BaseReg;
     int64_t Offset;
+    MachineInstr *Dependency = nullptr;
     if (TII->getMemOpBaseRegImmOfs(MI, BaseReg, Offset, TRI))
       if (MI->mayLoad() && !MI->isPredicable() && BaseReg == PointerReg &&
           Offset < PageSize && MI->getDesc().getNumDefs() <= 1 &&
-          HD.isSafeToHoist(MI)) {
-        NullCheckList.emplace_back(MI, MBP.ConditionDef, &MBB, NotNullSucc,
-                                   NullSucc);
-        return true;
+          HD.isSafeToHoist(MI, Dependency)) {
+
+        auto DependencyOperandIsOk = [&](MachineOperand &MO) {
+          assert(!(MO.isReg() && MO.isUse()) &&
+                 "No transitive dependendencies please!");
+          if (!MO.isReg() || !MO.getReg() || !MO.isDef())
+            return true;
+
+          // Make sure that we won't clobber any live ins to the sibling block
+          // by hoisting Dependency.  For instance, we can't hoist INST to
+          // before the null check (even if it safe, and does not violate any
+          // dependencies in the non_null_block) if %rdx is live in to
+          // _null_block.
+          //
+          //    test %rcx, %rcx
+          //    je _null_block
+          //  _non_null_block:
+          //    %rdx<def> = INST
+          //    ...
+          if (AnyAliasLiveIn(TRI, NullSucc, MO.getReg()))
+            return false;
+
+          // Make sure Dependency isn't re-defining the base register.  Then we
+          // won't get the memory operation on the address we want.
+          if (TRI->regsOverlap(MO.getReg(), BaseReg))
+            return false;
+
+          return true;
+        };
+
+        bool DependencyOperandsAreOk =
+            !Dependency ||
+            all_of(Dependency->operands(), DependencyOperandIsOk);
+
+        if (DependencyOperandsAreOk) {
+          NullCheckList.emplace_back(MI, MBP.ConditionDef, &MBB, NotNullSucc,
+                                     NullSucc, Dependency);
+          return true;
+        }
       }
 
     HD.rememberInstruction(MI);
@@ -399,6 +522,11 @@ void ImplicitNullChecks::rewriteNullChecks(
     (void)BranchesRemoved;
     assert(BranchesRemoved > 0 && "expected at least one branch!");
 
+    if (auto *DepMI = NC.getOnlyDependency()) {
+      DepMI->removeFromParent();
+      NC.getCheckBlock()->insert(NC.getCheckBlock()->end(), DepMI);
+    }
+
     // Insert a faulting load where the conditional branch was originally.  We
     // check earlier ensures that this bit of code motion is legal.  We do not
     // touch the successors list for any basic block since we haven't changed
@@ -418,6 +546,16 @@ void ImplicitNullChecks::rewriteNullChecks(
         continue;
       MBB->addLiveIn(Reg);
     }
+
+    if (auto *DepMI = NC.getOnlyDependency()) {
+      for (auto &MO : DepMI->operands()) {
+        if (!MO.isReg() || !MO.getReg() || !MO.isDef())
+          continue;
+        if (!NC.getNotNullSucc()->isLiveIn(MO.getReg()))
+          NC.getNotNullSucc()->addLiveIn(MO.getReg());
+      }
+    }
+
     NC.getMemOperation()->eraseFromParent();
     NC.getCheckOperation()->eraseFromParent();
 
@@ -433,5 +571,6 @@ char ImplicitNullChecks::ID = 0;
 char &llvm::ImplicitNullChecksID = ImplicitNullChecks::ID;
 INITIALIZE_PASS_BEGIN(ImplicitNullChecks, "implicit-null-checks",
                       "Implicit null checks", false, false)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_END(ImplicitNullChecks, "implicit-null-checks",
                     "Implicit null checks", false, false)