1 files changed, 49 insertions, 0 deletions
diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index 5e6d6190c63..a10bd7f52e7 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -27,6 +27,7 @@
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -673,6 +674,49 @@ MachineBasicBlock *MachineSinking::FindSuccToSinkTo(MachineInstr *MI,
   return SuccToSinkTo;
 }
 
+/// \brief Return true if MI is likely to be usable as a memory operation by the
+/// implicit null check optimization.
+///
+/// This is a "best effort" heuristic, and should not be relied upon for
+/// correctness.  This returning true does not guarantee that the implicit null
+/// check optimization is legal over MI, and this returning false does not
+/// guarantee MI cannot possibly be used to do a null check.
+static bool SinkingPreventsImplicitNullCheck(MachineInstr *MI,
+                                             const TargetInstrInfo *TII,
+                                             const TargetRegisterInfo *TRI) {
+  typedef TargetInstrInfo::MachineBranchPredicate MachineBranchPredicate;
+
+  auto *MBB = MI->getParent();
+  if (MBB->pred_size() != 1)
+    return false;
+
+  auto *PredMBB = *MBB->pred_begin();
+  auto *PredBB = PredMBB->getBasicBlock();
+
+  // Frontends that don't use implicit null checks have no reason to emit
+  // branches with make.implicit metadata, and this function should always
+  // return false for them.
+  if (!PredBB ||
+      !PredBB->getTerminator()->getMetadata(LLVMContext::MD_make_implicit))
+    return false;
+
+  unsigned BaseReg, Offset;
+  if (!TII->getMemOpBaseRegImmOfs(MI, BaseReg, Offset, TRI))
+    return false;
+
+  if (!(MI->mayLoad() && !MI->isPredicable()))
+    return false;
+
+  MachineBranchPredicate MBP;
+  if (TII->AnalyzeBranchPredicate(*PredMBB, MBP, false))
+    return false;
+
+  return MBP.LHS.isReg() && MBP.RHS.isImm() && MBP.RHS.getImm() == 0 &&
+         (MBP.Predicate == MachineBranchPredicate::PRED_NE ||
+          MBP.Predicate == MachineBranchPredicate::PRED_EQ) &&
+         MBP.LHS.getReg() == BaseReg;
+}
+
 /// SinkInstruction - Determine whether it is safe to sink the specified machine
 /// instruction out of its current block into a successor.
 bool MachineSinking::SinkInstruction(MachineInstr *MI, bool &SawStore,
@@ -691,6 +735,11 @@ bool MachineSinking::SinkInstruction(MachineInstr *MI, bool &SawStore,
   if (MI->isConvergent())
     return false;
 
+  // Don't break implicit null checks.  This is a performance heuristic, and not
+  // required for correctness.
+  if (SinkingPreventsImplicitNullCheck(MI, TII, TRI))
+    return false;
+
   // FIXME: This should include support for sinking instructions within the
   // block they are currently in to shorten the live ranges.  We often get
   // instructions sunk into the top of a large block, but it would be better to