[ImplicitNullCheck] Hoisting multiple dependencies

[anna]

Summary

This patch introduces a new dependence tracking mechanism for the
faulting operation, wherein multiple dependencies can be tracked. This
allows us to hoist a list of dependencies along with the faulting memory
operation.
We now catch some of the existing testcases with multiple
dependencies as well as some newly added testcases.
The key point to note is that by the time we reach dependency tracking,
we have already completed two other important sections for safety
analysis, namely isSuitableMemoryOp and whether unsafe instructions
modify PointerReg (i.e. those that do not preserve zero value of the
register).

[dantrushin]

llvm/lib/CodeGen/ImplicitNullChecks.cpp

465

Would be great to have a brief comment about what's 'safe' dependence is.

473

You say about explicit uses, but do not have corresponding check.
MI->uses() includes everything except explicit defs. Do you want to use MI->explicit_uses() here?

504

Move this check to the top of loop? (Before canReorder() check)

525

Er, are you checking defs of CurrMI? Would CurrMI->defs() work here (if you're checking explicit defs only)?
Or loop over TrackedRegisterUSes using CurrMI->modifiesRegister() as predicate?

[anna]

may revisit later.

llvm/lib/CodeGen/ImplicitNullChecks.cpp

 #define DEBUG_TYPE "implicit-null-checks"
 
 STATISTIC(NumImplicitNullChecks,
           "Number of explicit null checks made implicit");
 
 namespace {
 
 class ImplicitNullChecks : public MachineFunctionPass {
-  /// Return true if \c computeDependence can process \p MI.
+  /// Return true if this satisfies basic checks that MI does not clobber the
+  /// faulting instruction..
   static bool canHandle(const MachineInstr *MI);
 
-  /// Helper function for \c computeDependence.  Return true if \p A
+  /// Return true if \p A
   /// and \p B do not have any dependences between them, and can be
   /// re-ordered without changing program semantics.
   bool canReorder(const MachineInstr *A, const MachineInstr *B);
 
-  /// A data type for representing the result computed by \c
-  /// computeDependence.  States whether it is okay to reorder the
-  /// instruction passed to \c computeDependence with at most one
-  /// dependency.
-  struct DependenceResult {
-    /// Can we actually re-order \p MI with \p Insts (see \c
-    /// computeDependence).
-    bool CanReorder;
-
-    /// If non-None, then an instruction in \p Insts that also must be
-    /// hoisted.
-    Optional<ArrayRef<MachineInstr *>::iterator> PotentialDependence;
-
-    /*implicit*/ DependenceResult(
-        bool CanReorder,
-        Optional<ArrayRef<MachineInstr *>::iterator> PotentialDependence)
-        : CanReorder(CanReorder), PotentialDependence(PotentialDependence) {
-      assert((!PotentialDependence || CanReorder) &&
-             "!CanReorder && PotentialDependence.hasValue() not allowed!");
-    }
-  };
-
-  /// Compute a result for the following question: can \p MI be
-  /// re-ordered from after \p Insts to before it.
-  ///
-  /// \c canHandle should return true for all instructions in \p
-  /// Insts.
-  DependenceResult computeDependence(const MachineInstr *MI,
-                                     ArrayRef<MachineInstr *> Block);
 
   /// Represents one null check that can be made implicit.
   class NullCheck {
     // The memory operation the null check can be folded into.
     MachineInstr *MemOperation;
 
     // The instruction actually doing the null check (Ptr != 0).
     MachineInstr *CheckOperation;
 
     // The block the check resides in.
     MachineBasicBlock *CheckBlock;
 
     // The block branched to if the pointer is non-null.
     MachineBasicBlock *NotNullSucc;
 
     // The block branched to if the pointer is null.
     MachineBasicBlock *NullSucc;
 
-    // If this is non-null, then MemOperation has a dependency on this
-    // instruction; and it needs to be hoisted to execute before MemOperation.
-    MachineInstr *OnlyDependency;
+    // If this is non-empty, then MemOperation has a dependency on these
+    // instruction; and they need to be hoisted to execute before MemOperation.
+    SmallVector<MachineInstr *, 4> SafeDependenceList;
 
   public:
     explicit NullCheck(MachineInstr *memOperation, MachineInstr *checkOperation,
                        MachineBasicBlock *checkBlock,
                        MachineBasicBlock *notNullSucc,
                        MachineBasicBlock *nullSucc,
-                       MachineInstr *onlyDependency)
+                       SmallVectorImpl<MachineInstr *> &safeDependenceList)
         : MemOperation(memOperation), CheckOperation(checkOperation),
           CheckBlock(checkBlock), NotNullSucc(notNullSucc), NullSucc(nullSucc),
-          OnlyDependency(onlyDependency) {}
+          SafeDependenceList(safeDependenceList.begin(),
+                             safeDependenceList.end()) {}
 
     MachineInstr *getMemOperation() const { return MemOperation; }
 
     MachineInstr *getCheckOperation() const { return CheckOperation; }
 
     MachineBasicBlock *getCheckBlock() const { return CheckBlock; }
 
     MachineBasicBlock *getNotNullSucc() const { return NotNullSucc; }
 
     MachineBasicBlock *getNullSucc() const { return NullSucc; }
 
-    MachineInstr *getOnlyDependency() const { return OnlyDependency; }
+    // This is a backward-list, so we will need to iterate in reverse and place
+    // it in correct MBB when hoisting entire chunk.
+    const SmallVectorImpl<MachineInstr *> &getSafeDependenceList() const {
+      return SafeDependenceList;
+    }
   };
 
   // We can have multiple starting addresses for faulting pages, i.e. page zero
   // need not be the only faulting page defined by the runtime. This is
   // collected through an optional module level metadatade provided by
   // front-ends. The values are expected to be positive or zero.
   SmallVector<uint64_t, 4> StartAddressesOfFaultingPage;
 
class ImplicitNullChecks : public MachineFunctionPass {
   /// if it was hoisted to the NullCheck block. This is used by caller
   /// canHoistInst to decide if DependenceMI can be hoisted safely.
   bool canDependenceHoistingClobberLiveIns(MachineInstr *DependenceMI,
                                            MachineBasicBlock *NullSucc);
 
   /// Return true if \p FaultingMI can be hoisted from after the
   /// instructions in \p InstsSeenSoFar to before them.  Set \p Dependence to a
   /// non-null value if we also need to (and legally can) hoist a depedency.
-  bool canHoistInst(MachineInstr *FaultingMI,
-                    ArrayRef<MachineInstr *> InstsSeenSoFar,
-                    MachineBasicBlock *NullSucc, MachineInstr *&Dependence);
+  bool canHoistInst(MachineInstr *FaultingMI, MachineBasicBlock *NullSucc,
+                    SmallVectorImpl<MachineInstr *> &DependenceList);
 
 public:
   static char ID;
 
   ImplicitNullChecks() : MachineFunctionPass(ID) {
     initializeImplicitNullChecksPass(*PassRegistry::getPassRegistry());
   }
 
bool ImplicitNullChecks::canHandle(const MachineInstr *MI) {
 
   assert(!llvm::any_of(MI->operands(), IsRegMask) &&
          "Calls were filtered out above!");
 
   auto IsUnordered = [](MachineMemOperand *MMO) { return MMO->isUnordered(); };
   return llvm::all_of(MI->memoperands(), IsUnordered);
 }
 
-ImplicitNullChecks::DependenceResult
-ImplicitNullChecks::computeDependence(const MachineInstr *MI,
-                                      ArrayRef<MachineInstr *> Block) {
-  assert(llvm::all_of(Block, canHandle) && "Check this first!");
-  assert(!is_contained(Block, MI) && "Block must be exclusive of MI!");
-
-  Optional<ArrayRef<MachineInstr *>::iterator> Dep;
-
-  for (auto I = Block.begin(), E = Block.end(); I != E; ++I) {
-    if (canReorder(*I, MI))
-      continue;
-
-    if (Dep == None) {
-      // Found one possible dependency, keep track of it.
-      Dep = I;
-    } else {
-      // We found two dependencies, so bail out.
-      return {false, None};
-    }
-  }
-
-  return {true, Dep};
-}
-
 bool ImplicitNullChecks::canReorder(const MachineInstr *A,
                                     const MachineInstr *B) {
   assert(canHandle(A) && canHandle(B) && "Precondition!");
 
   // canHandle makes sure that we _can_ correctly analyze the dependencies
   // between A and B here -- for instance, we should not be dealing with heap
   // load-store dependencies here.
 
if (AnyAliasLiveIn(TRI, NullSucc, DependenceMO.getReg()))
       return true;
 
   }
 
   // The dependence does not clobber live-ins in NullSucc block.
   return false;
 }
 
-bool ImplicitNullChecks::canHoistInst(MachineInstr *FaultingMI,
-                                      ArrayRef<MachineInstr *> InstsSeenSoFar,
-                                      MachineBasicBlock *NullSucc,
-                                      MachineInstr *&Dependence) {
-  auto DepResult = computeDependence(FaultingMI, InstsSeenSoFar);
-  if (!DepResult.CanReorder)
-    return false;
-
-  if (!DepResult.PotentialDependence) {
-    Dependence = nullptr;
-    return true;
+bool ImplicitNullChecks::canHoistInst(
+    MachineInstr *FaultingMI, MachineBasicBlock *NullSucc,
+    SmallVectorImpl<MachineInstr *> &DependenceList) {
+
+  // Dependence state while processing the instructions in the FaultingMI's
+  // basic block. We start with unknown and at any point we should have either
+  // no dependence or a safe dependence. They are mutually exclusive, i.e. we
+  // cannot go from DS_None to DS_Safe and vice-versa.
+  // While handling both cases are possible, the algorithm below becomes
+  // exponential (since we need to check that *each dependency* can be reordered
+  // over
+  // all instructions that were independent of the FaultingMI).
+  enum DependenceState { DS_Unknown, DS_None, DS_Safe };
+
+  DependenceState CurrDepState = DS_Unknown;
+
+  SmallDenseSet<Register, 8> TrackedRegisterUses;
+  // Explicit uses are recorded.
+  auto RecordExplicitUses = [&TrackedRegisterUses](MachineInstr *MI) {
+    for (auto &MO : MI->uses()) {
+      if (MO.isReg() && MO.isUse())
+        TrackedRegisterUses.insert(MO.getReg());
     }
+  };
 
-  auto DependenceItr = *DepResult.PotentialDependence;
-  auto *DependenceMI = *DependenceItr;
-
-  // We don't want to reason about speculating loads.  Note -- at this point
-  // we should have already filtered out all of the other non-speculatable
-  // things, like calls and stores.
-  // We also do not want to hoist stores because it might change the memory
-  // while the FaultingMI may result in faulting.
-  assert(canHandle(DependenceMI) && "Should never have reached here!");
-  if (DependenceMI->mayLoadOrStore())
+  // We perform dependence analysis by a backward traversal starting from
+  // faultingMI
+  // until the start of the basic block and record the dependencies in the
+  // DependenceList. We do not care about dependencies past
+  // the basic block, since we are hoisting the faultingMI and the
+  // DependenceList into the nullcheck basic block.
+  // We first record the uses in faultingMI.
+  RecordExplicitUses(FaultingMI);
+  // A backward traversal looking for register defs of the recorded uses works,
+  // because when we see a def through backward traversal, it is effectively the
+  // "correct def" which corresponds to the recorded use.
+  for (auto It = std::next(MachineBasicBlock::reverse_iterator(FaultingMI));
+       It != FaultingMI->getParent()->rend(); ++It) {
+    MachineInstr *CurrMI = &*It;
+    // No read-after-write, write-afer-write, write-after-read dependencies.
+    if (canReorder(CurrMI, FaultingMI)) {
+      // We cannot move from a safe-dependency state to a no-dependency, since
+      // that will require checking that we can reorder CurrMI over all the
+      // dependencies recorded in DependenceList as well.
+      if (CurrDepState == DS_Safe)
         return false;
+      CurrDepState = DS_None;
+      continue;
+    }
 
-  if (canDependenceHoistingClobberLiveIns(DependenceMI, NullSucc))
+    if (CurrMI->mayLoadOrStore())
       return false;
 
-  auto DepDepResult =
-      computeDependence(DependenceMI, {InstsSeenSoFar.begin(), DependenceItr});
+    // We cannot move from non-dependency to a safe dependency.
+    // This effectively means we either have all instructions that are
+    // independent of faulting MI (i.e. DS_None)  or all instructions that
+    // safely depend on faulting MI (either directly or transitively and we are
+    // at DS_Safe state).
+    if (CurrDepState == DS_None)
+      return false;
 
-  if (!DepDepResult.CanReorder || DepDepResult.PotentialDependence)
+    if (canDependenceHoistingClobberLiveIns(CurrMI, NullSucc))
       return false;
 
-  Dependence = DependenceMI;
+    // At this point, we cannot just simply hoist the FaultingMI. This means
+    // unless it is an MI which defines the uses already recorded, we have to
+    // bail out.
+    if (!llvm::any_of(CurrMI->operands(),
+                      [&TrackedRegisterUses](const MachineOperand &MO) {
+                        return MO.isReg() && MO.isDef() &&
+                               (TrackedRegisterUses.count(MO.getReg()) ||
+                                TrackedRegisterUses.count(MO.getSubReg()));
+                      }))
+      return false;
+    // TODO: We need to handle cases where CurrMI defines a subReg of one of
+    // the registers in trackedRegisterUses.
+    // esi = ..
+    // faulting mi mv32rm = .., rsi, ..
+
+    assert(canHandle(CurrMI) && "Should never have reached here!");
+    DependenceList.push_back(CurrMI);
+    RecordExplicitUses(CurrMI);
+    CurrDepState = DS_Safe;
+  }
   return true;
 }
 
 /// 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.
 bool ImplicitNullChecks::analyzeBlockForNullChecks(
     MachineBasicBlock &MBB, SmallVectorImpl<NullCheck> &NullCheckList) {
bool ImplicitNullChecks::analyzeBlockForNullChecks(
   // ptr -- clause (2) above does not apply in this case.  In the above program
   // the safety of ptr->field can be dependent on some_cond; and, for instance,
   // ptr could be some non-null invalid reference that never gets loaded from
   // because some_cond is always true.
 
   SmallVector<MachineInstr *, 8> InstsSeenSoFar;
   DenseSet<MachineInstr *> InstrPreservesZeroReg;
 
+  SmallVector<MachineInstr *, 4> DependenceList;
+
   for (auto &MI : *NotNullSucc) {
     if (!canHandle(&MI) || InstsSeenSoFar.size() >= MaxInstsToConsider)
       return false;
 
     MachineInstr *Dependence;
     SuitabilityResult SR = isSuitableMemoryOp(MI, PointerReg, InstsSeenSoFar,
                                               InstrPreservesZeroReg);
     if (SR == SR_Impossible)
       return false;
-    if (SR == SR_Suitable &&
-        canHoistInst(&MI, InstsSeenSoFar, NullSucc, Dependence)) {
+
+    if (SR == SR_Suitable && canHoistInst(&MI, NullSucc, DependenceList)) {
       NullCheckList.emplace_back(&MI, MBP.ConditionDef, &MBB, NotNullSucc,
-                                 NullSucc, Dependence);
+                                 NullSucc, DependenceList);
       return true;
     }
 
     // If MI re-defines the PointerReg in a way that changes the value of
     // PointerReg if it was null, then we cannot move further.
     if (!TII->preservesZeroValueInReg(&MI, PointerReg, TRI))
       return false;
     if (MI.modifiesRegister(PointerReg, TRI))
void ImplicitNullChecks::rewriteNullChecks(
   DebugLoc DL;
 
   for (auto &NC : NullCheckList) {
     // Remove the conditional branch dependent on the null check.
     unsigned BranchesRemoved = TII->removeBranch(*NC.getCheckBlock());
     (void)BranchesRemoved;
     assert(BranchesRemoved > 0 && "expected at least one branch!");
 
-    if (auto *DepMI = NC.getOnlyDependency()) {
+    const SmallVectorImpl<MachineInstr *> &DepList = NC.getSafeDependenceList();
+    // We need to iterate in reverse to place dependencies in correct order.
+    for (auto *DepMI : reverse(DepList)) {
       DepMI->removeFromParent();
       NC.getCheckBlock()->insert(NC.getCheckBlock()->end(), DepMI);
     }
 
     // Insert a faulting instruction 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 control flow, we've just made it implicit.
     MachineInstr *FaultingInstr = insertFaultingInstr(
         NC.getMemOperation(), NC.getCheckBlock(), NC.getNullSucc());
     // Now the values defined by MemOperation, if any, are live-in of
     // the block of MemOperation.
     // The original operation may define implicit-defs alongside
     // the value.
     MachineBasicBlock *MBB = NC.getMemOperation()->getParent();
     for (const MachineOperand &MO : FaultingInstr->operands()) {
       if (!MO.isReg() || !MO.isDef())
         continue;
       Register Reg = MO.getReg();
       if (!Reg || MBB->isLiveIn(Reg))
         continue;
       MBB->addLiveIn(Reg);
     }
 
-    if (auto *DepMI = NC.getOnlyDependency()) {
+    for (auto *DepMI : DepList) {
       for (auto &MO : DepMI->operands()) {
         if (!MO.isReg() || !MO.getReg() || !MO.isDef() || MO.isDead())
           continue;
         if (!NC.getNotNullSucc()->isLiveIn(MO.getReg()))
           NC.getNotNullSucc()->addLiveIn(MO.getReg());
       }
     }
 

llvm/test/CodeGen/X86/implicit-null-check.ll

; CHECK-NEXT:    retq
    %y = ptrtoint i64* %x to i64
    %shry = shl i64 %y, 3
    %shry.add = add i64 %shry, 35184372088832
    %y.ptr = inttoptr i64 %shry.add to i64*
    %x.loc = getelementptr i64, i64* %y.ptr, i64 1
    %t = load i64, i64* %x.loc
    ret i64 %t
 }
+
 !startaddress_faulting_pages = !{!1, !2}
 !0 = !{}
 !1 = !{i64 0}
 !2 = !{i64 35184372088832}

llvm/test/CodeGen/X86/implicit-null-checks.mir

--- |
 
   ret_100:
     ret i32 100
 
   ret_200:
     ret i32 200
   }
 
+  define i32 @load_from_non_zero_faulting_page_dep2(i32* %ptr, i32 %val) {
+  entry:
+    %ptr_is_null = icmp eq i32* %ptr, null
+    br i1 %ptr_is_null, label %is_null, label %not_null, !make.implicit !0
+
+  is_null:
+    ret i32 42
+
+  not_null:
+    br i1 undef, label %ret_100, label %ret_200
+
+  ret_100:
+    ret i32 100
+
+  ret_200:
+    ret i32 200
+  }
+
   attributes #0 = { "target-features"="+bmi,+bmi2" }
 
   !startaddress_faulting_pages = !{!1, !2}
   !0 = !{}
   !1 = !{i64 0}
   !2 = !{i64 35184372088832}
   attributes #0 = { "target-features"="+bmi,+bmi2" }
 

 name:            imp_null_check_with_bitwise_op_2
 # CHECK-LABEL: name:            imp_null_check_with_bitwise_op_2
 alignment:       16
 tracksRegLiveness: true
 liveins:
   - { reg: '$rdi' }
   - { reg: '$esi' }
 # CHECK:  bb.0.entry:
-# CHECK:    TEST64rr $rdi, $rdi, implicit-def $eflags
-# CHECK-NEXT:    JCC_1 %bb.3, 4, implicit $eflags
+# CHECK:     $eax = MOV32ri 2200000
+# CHECK:     $eax = ADD32ri killed $eax, 100, implicit-def dead $eflags
+# CHECK:     $eax = FAULTING_OP 1, %bb.3, 444, $eax, $rdi, 1, $noreg, 0, $noreg, implicit-def $eflags
+# CHECK:     JMP_1 %bb.1
 
 body:             |
   bb.0.entry:
     liveins: $esi, $rdi
 
     TEST64rr $rdi, $rdi, implicit-def $eflags
     JCC_1 %bb.3, 4, implicit $eflags
 
body:             |
   bb.2.is_null:
     RETQ
 
 ...
 ---
 name:            inc_store_with_two_dep
 # CHECK-LABEL: inc_store_with_two_dep
 # CHECK: bb.0.entry:
-# CHECK:    TEST64rr $rdi, $rdi, implicit-def $eflags
-# CHECK-NEXT:    JCC_1 %bb.2, 4, implicit killed $eflags
-# CHECK: bb.1.not_null
+# CHECK:   $esi = ADD32rr killed $esi, killed $esi, implicit-def dead $eflags
+# CHECK:   $esi = ADD32ri killed $esi, 15, implicit-def dead $eflags
+# CHECK:   $noreg = FAULTING_OP 3, %bb.2, 1716, $rdi, 1, $noreg, 16, $noreg, $esi
+# CHECK:   JMP_1 %bb.1
 
 alignment:       16
 tracksRegLiveness: true
 liveins:
   - { reg: '$rdi' }
   - { reg: '$rsi' }
 body:             |
   bb.0.entry:
body:             |
 
   bb.3.is_null:
     $eax = MOV32ri 42
     RETQ $eax
 
   bb.4.ret_100:
     $eax = MOV32ri 100
     RETQ $eax
+
+---
+name:          load_from_non_zero_faulting_page_dep2
+# This is a positive test where we have 2 dependencies on the
+# memory op.
+
+# CHECK-LABEL: name: load_from_non_zero_faulting_page_dep2
+# CHECK: bb.0.entry:
+# CHECK:   renamable $rdx = MOV64ri 35184372088832
+# CHECK:   renamable $eax = MOV32rr renamable $eax, implicit killed $rax, implicit-def $rax
+# CHECK:   $ecx = FAULTING_OP 1, %bb.3, 1724, renamable $rdx, 8, renamable $rax, 4, $noreg
+liveins:
+  - { reg: '$rax' }
+  - { reg: '$rbx' }
+
+body:             |
+  bb.0.entry:
+    liveins: $rax, $rbx
+    TEST64rr renamable $rax, renamable $rax, implicit-def $eflags
+    JCC_1 %bb.3, 4, implicit $eflags
+
+  bb.1.not_null:
+    liveins: $rax, $rbx
+    renamable $rdx = MOV64ri 35184372088832
+    renamable $eax = MOV32rr renamable $eax, implicit killed $rax, implicit-def $rax
+    renamable $ecx = MOV32rm renamable $rdx, 8, renamable $rax, 4, $noreg
+    CMP32ri killed renamable $ecx, 1195, implicit-def $eflags
+    JCC_1 %bb.4, 4, implicit $eflags
+
+  bb.2.ret_200:
+    $eax = MOV32ri 200
+    RETQ $eax
+
+  bb.3.is_null:
+    $eax = MOV32ri 42
+    RETQ $eax
+
+  bb.4.ret_100:
+    $eax = MOV32ri 100
+    RETQ $eax