[ImplicitNullCheck] Handle Nonzero faulting pages and complex addressing

[anna]

Summary

This patch:

1. introduces a new TII API for calculating the address of a memory operation in complex addressing mode. This is a general function, which checks for both bas and index registers. This is used in isSuitableMemoryOp in ImplicitNullChecks and all the validity checks are done in this caller function. Earlier, we used to bail out if the memory operation had an index register or a scale != 1 (i.e. we supported offsets with simple addressing mode).
2. handles memory accesses (and makes the null checks implicit) where the faulting page is a non-zero one as

specified through a named module level metadata startaddress_faulting_pages introduced in this patch. The optional metadata is added by front-ends.

See #2 in action with added MIR tests for positive and negative case. #1
allows us to handle a case which was not handled earlier (see
implicit-null-checks.ll).

[anna]

While working on a follow-up change, I realized we don't need to separate out whether we're handling simple or complex addressing mode based on whether we fault in zero page or not. Both are orthogonal issues. The reason I coupled them is because we still need a way to get the "total constant address" *if it exists* and make a distinction between when we succeeded to get one versus when there was no such constant address apart from the offset.

[anna]

handles complex addressing mode using base and index operands. All test cases pass.
I plan to add more counter examples once the basic idea is agreed upon.

[anna]

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

132

I added these to show that none of the negative test cases start failing once we add the metadata with "other faulting pages".

[dantrushin]

llvm/include/llvm/CodeGen/TargetInstrInfo.h

1287

These names are confusing. I've been scrolling back and forth for several minutes :)
Why not to pass and name them separately?

llvm/lib/CodeGen/ImplicitNullChecks.cpp

414

Do you really need to check that BaseRegUsedInConstAddr.first is not zero?
Would simple BaseRegUsedInConstAddr.first != PointerReg work here?

I also found positive names more readable, e.g. bool BaseRegIsNullChecked = BaseRegUsedInConstAddr.first == PointerReg

I think it is reasonable to assert Scale !=0 || IndexRegUsedInConstAddr.first == 0 and remove Scale check.

llvm/lib/Target/X86/X86InstrInfo.cpp

3680

Any reason to use auto * instead of auto & ?
If I read it correctly, BaseOp is used to get base register only. Using reference you'd have less characters to type :)

3693

ditto. auto &IndexOp = MI.getOperand() ?

[anna]

addressed review comments. Added more testcases (positive and countercases).

[anna]

llvm/lib/CodeGen/ImplicitNullChecks.cpp

414

Changed the names and hopefully more readable.
We need to check that the register exists before checking if it is equal to the null checked register because we don't want to fail when we don't have the register (in that case, we wont be using it in the addres calculation). Here we will state that the memory op is unsuitable and then fail. I realized this through some failing test case imp_null_check_load in implicit-null-checks.ll. There, we do not use the index register, so it is zero.

[dantrushin]

Thanks, it reads much better for me now.
Looks good.

[reames]

Must fix:

• Your change does not include a LangRef update for the new metadata.
• The name chosen for the metadata is ambiguous. I'd also suggest phrasing it in terms of ranges, not page starts. Maybe: guaranteed_faulting_ranges, or implicit_check_faulting_ranges?
• As you note in the review, this is mixing two sets of changes in a way which makes it hard to reason about either in isolation. Please split. Either order is fine as you should be able to test either piece in isolation without the other.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

1283

I find the terminology used here to be very confusing. A constant address to me is either a literal constant or a constant expression. You appear to be referring to complex addressing instead. I believe that your ConstantAddr is the Displacement field in a x86 complex addressing, is that true?

If so, rename required. If not, please explain.

llvm/lib/CodeGen/ImplicitNullChecks.cpp

449

From your comment, it sounds like the existing code handles two built in ranges, not one.
0, PageSize
0x100..00, 0x111..11

(i.e. the zeroth page and kernel address range for Linux)

I'd suggest adjusting comments and code structure to reflect same, and to make range handling (not page starts) explicit.

llvm/lib/Target/X86/X86InstrInfo.cpp

3701

This really doesn't seem like it belongs here. I think this could be cleanly moved into the caller, simplifying the behaviour of this routine greatly.

[anna]

addressed review comments about splitting out patches. This patch handles complex addressing mode.

[reames]

Looks close to ready. Please address review comments, and we'll probably be ready for an LGTM.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

971

"is an instruction"

llvm/lib/CodeGen/ImplicitNullChecks.cpp

429

These should be impossible, use asserts.

447

There's a potential overflow case here which needs consideration. Consider 0xF...F as constant and scale as 2.

You probably need to return some information about the register width/wrapping of the addressing mode to do this in a target independent way.

451

Please return true or false, then set the per register flag. This is too much coupling.

[anna]

llvm/lib/CodeGen/ImplicitNullChecks.cpp

429

Will use assert for the multiplier and add a comment why RegUsedInAddr can be zero reg, i.e. something like: movq 8(,%rdi,8), %rax here the BaseReg is X86::NoRegister and ScaleReg is rdi.

447

good catch! will take a look.

451

yes, agreed.

[anna]

addressed review comments. Bails out in case of overflow (added testcase for multiply overflow).

[reames]

LGTM, but with required changes.

llvm/lib/CodeGen/ImplicitNullChecks.cpp

421

Please add a check here that the size of BaseReg, ScaledReg, and PointerReg are all equal.

I'm concerned about implicit sign extension of registers, and the interaction with the displacement code below. The code you have may be correct, but I'd prefer a explicit bail out just to make it easier to tell.

428

You could if you wished do Displacement mod 2^64, but I don't see any point in that.

434

Add comment here: "If a register used in the address is constant, fold it's effect into the displacement for ease of analysis."

455

Replace IsSignedImmVal with constant true.

This isn't "is negative", this is "interpret this as a signed number instead of unsigned".

457

You should be able to just define the APInt in the definition below. The bitwidth will come from the RHS.

464

e,g, APInt Product =

Also, remember to constant fold isSignedImmVal

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

625

You don't appear to actually have a test here which uses a displacement large enough to need a base register, please add one.

(i.e. modify this test to add not 3526, but 0xFFFFF0000 (or something similarly large)

[anna]

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

625

I have a test in the negative test cases file above where the displacement is large and outside the faulting page. I'll modify it to a very large value so that it is in a register and not an immediate.

[anna]

addressed review comments, rebased and updated for AArch64 (this patch is NFC in behaviour for AArch64).

[GMNGeoffrey]

Hey JFYI this created some unused variables. I fixed with https://reviews.llvm.org/D89022

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   /// immediate value and a physical register, and stores the result in
   /// the given physical register \c Reg, return a pair of the source
   /// register and the offset which has been added.
   virtual Optional<RegImmPair> isAddImmediate(const MachineInstr &MI,
                                               Register Reg) const {
     return None;
   }
 
   /// Store the specified register of the given register class to the specified
   /// stack frame index. The store instruction is to be added to the given
   /// machine basic block before the specified machine instruction. If isKill
   /// is true, the register operand is the last use and must be marked kill.
   virtual void storeRegToStackSlot(MachineBasicBlock &MBB,
                                    MachineBasicBlock::iterator MI,
                                    Register SrcReg, bool isKill, int FrameIndex,
                                    const TargetRegisterClass *RC,
                                    const TargetRegisterInfo *TRI) const {
public:
   /// true, the function also sets the operand position in the instruction
   /// for the base register and offset.
   virtual bool getBaseAndOffsetPosition(const MachineInstr &MI,
                                         unsigned &BasePos,
                                         unsigned &OffsetPos) const {
     return false;
   }
 
+  /// Return true if MI is a memory op and it contains base operand, optional
+  /// index operand, scale. This function also calculates the total constant
+  /// address using both the base and the (optional) index operands and
+  /// recorded in ConstantAddr. The base and index registers are recorded in
+  /// BaseReg and IndexReg respectively.
+  /// We also record if the constant address is calculated from the BaseReg in
+  /// \p BaseRegUsedInConstAddr and if the constant address is calculated from
+  /// the IndexReg in \p IndexRegUsedInConstAddr. \p IgnoreInstInAddrCalc
+  /// contains the list of instructions which do not contribute to the address
+  /// calculation (even though they modify the base and index operands).
+  virtual bool getConstantAddrFromMemoryOp(
+      const MachineInstr &MI, int64_t &ConstantAddr, int64_t &Scale,
+      Register &BaseReg, bool &isBaseRegUsedInConstAddr, Register &IndexReg,
+      bool &isIndexRegUsedInConstAddr, const TargetRegisterInfo *TRI,
+      const DenseSet<MachineInstr *> &IgnoreInstInAddrCalc) const {
+    return false;
+  }
+
   /// Returns true if MI's Def is NullValueReg, and the MI
   /// does not change the Zero value. i.e. cases such as rax = shr rax, X where
   /// NullValueReg = rax. Note that if the NullValueReg is non-zero, this
   /// function can return true even if becomes zero. Specifically cases such as
   /// NullValueReg = shl NullValueReg, 63.
   virtual bool preservesZeroValueInReg(const MachineInstr *MI,
                                        const Register NullValueReg,
                                        const TargetRegisterInfo *TRI) const {

llvm/lib/CodeGen/ImplicitNullChecks.cpp

public:
 
     MachineBasicBlock *getNotNullSucc() const { return NotNullSucc; }
 
     MachineBasicBlock *getNullSucc() const { return NullSucc; }
 
     MachineInstr *getOnlyDependency() const { return OnlyDependency; }
   };
 
+  // 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;
+
   const TargetInstrInfo *TII = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
   AliasAnalysis *AA = nullptr;
   MachineFrameInfo *MFI = nullptr;
 
+  // Returns true if some non-zero page is protected.
+  bool canFaultInNonZeroPage() {
+    return StartAddressesOfFaultingPage.size() > 1 ||
+           StartAddressesOfFaultingPage.front() != 0;
+  }
+
   bool analyzeBlockForNullChecks(MachineBasicBlock &MBB,
                                  SmallVectorImpl<NullCheck> &NullCheckList);
   MachineInstr *insertFaultingInstr(MachineInstr *MI, MachineBasicBlock *MBB,
                                     MachineBasicBlock *HandlerMBB);
   void rewriteNullChecks(ArrayRef<NullCheck> NullCheckList);
 
   enum AliasResult {
     AR_NoAlias,
class ImplicitNullChecks : public MachineFunctionPass {
   };
 
   /// Return SR_Suitable if \p MI a memory operation that can be used to
   /// implicitly null check the value in \p PointerReg, SR_Unsuitable if
   /// \p MI cannot be used to null check and SR_Impossible if there is
   /// no sense to continue lookup due to any other instruction will not be able
   /// to be used. \p PrevInsts is the set of instruction seen since
   /// the explicit null check on \p PointerReg.
-  SuitabilityResult isSuitableMemoryOp(const MachineInstr &MI,
-                                       unsigned PointerReg,
-                                       ArrayRef<MachineInstr *> PrevInsts);
+  SuitabilityResult
+  isSuitableMemoryOp(const MachineInstr &MI, unsigned PointerReg,
+                     ArrayRef<MachineInstr *> PrevInsts,
+                     const DenseSet<MachineInstr *> &InstrPreservesZeroReg);
 
   /// Returns true if \p DependenceMI can clobber the liveIns in NullSucc block
   /// 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,
                                            unsigned PointerReg);
 

 }
 
 bool ImplicitNullChecks::runOnMachineFunction(MachineFunction &MF) {
   TII = MF.getSubtarget().getInstrInfo();
   TRI = MF.getRegInfo().getTargetRegisterInfo();
   MFI = &MF.getFrameInfo();
   AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
 
+  // Populate the list of faulting page start addresses.
+  // If the module metadata is attached, populate using that, else consider only
+  // zeroth page to be the faulting page.
+  Module *M = MF.getFunction().getParent();
+  StartAddressesOfFaultingPage.clear();
+  if (auto *SAFP = M->getNamedMetadata("startaddress_faulting_pages")) {
+    for (MDNode *SANode : SAFP->operands()) {
+      uint64_t SA =
+          mdconst::extract<ConstantInt>(SANode->getOperand(0))->getZExtValue();
+      assert(SA >= 0 && "expected positive starting address!");
+      StartAddressesOfFaultingPage.push_back(SA);
+    }
+  } else
+    StartAddressesOfFaultingPage.push_back(0);
+
   SmallVector<NullCheck, 16> NullCheckList;
 
   for (auto &MBB : MF)
     analyzeBlockForNullChecks(MBB, NullCheckList);
 
   if (!NullCheckList.empty())
     rewriteNullChecks(NullCheckList);
 
for (MachineMemOperand *MMO2 : PrevMI->memoperands()) {
                                    MMO2->getAAInfo()));
       if (AAResult != NoAlias)
         return AR_MayAlias;
     }
   }
   return AR_NoAlias;
 }
 
-ImplicitNullChecks::SuitabilityResult
-ImplicitNullChecks::isSuitableMemoryOp(const MachineInstr &MI,
-                                       unsigned PointerReg,
-                                       ArrayRef<MachineInstr *> PrevInsts) {
-  int64_t Offset;
-  bool OffsetIsScalable;
-  const MachineOperand *BaseOp;
+ImplicitNullChecks::SuitabilityResult ImplicitNullChecks::isSuitableMemoryOp(
+    const MachineInstr &MI, unsigned PointerReg,
+    ArrayRef<MachineInstr *> PrevInsts,
+    const DenseSet<MachineInstr *> &InstrPreservesZeroReg) {
+  int64_t ConstantAddr;
+  Register BaseReg, IndexReg;
+  bool BaseRegUsedInConstAddr, IndexRegUsedInConstAddr;
+  int64_t Scale = 0;
 
+  if (!MI.mayLoadOrStore() || MI.isPredicable())
+    return SR_Unsuitable;
 
-  // FIXME: This handles only simple addressing mode.
-  if (!TII->getMemOperandWithOffset(MI, BaseOp, Offset, OffsetIsScalable, TRI))
+  // Check whether this is a valid memory op and get the entire constant address
+  // (i.e. not just offset) for the memory op if we can.
+  if (!TII->getConstantAddrFromMemoryOp(
+          MI, ConstantAddr, Scale, BaseReg, BaseRegUsedInConstAddr, IndexReg,
+          IndexRegUsedInConstAddr, TRI, InstrPreservesZeroReg))
     return SR_Unsuitable;
 
-  // We need the base of the memory instruction to be same as the register
-  // where the null check is performed (i.e. PointerReg).
-  if (!BaseOp->isReg() || BaseOp->getReg() != PointerReg)
+  // We have a degenerate case where we have an index register, but the scale is
+  // zero, i.e. the index register is not used for the memory op.
+  // Just bail out here since it simplifies the checks below (we can handle
+  // these cases, but it does not seem worthwhile).
+  if (IndexReg && !Scale)
     return SR_Unsuitable;
 
-  // Scalable offsets are a part of scalable vectors (SVE for AArch64). That
-  // target is in-practice unsupported for ImplicitNullChecks.
-  if (OffsetIsScalable)
+  // At least one of base or index reg should have been null checked.
+  if (BaseReg != PointerReg && IndexReg != PointerReg)
     return SR_Unsuitable;
 
-  if (!MI.mayLoadOrStore() || MI.isPredicable())
+  // The register which is not null checked should be part of the constant
+  // address calculation, otherwise we do not know whether the ConstantAddress
+  // is just Offset + 0 or Offset + some symbolic address.
+  // This matters because we do not want to incorrectly assume that load from
+  // "Offset + some symbolic address" falls in the zeroth faulting page in the
+  // "sane offset check" below.
+  //
+  // Also, we do not care about getting the constant address through the
+  // register that is null checked. There are two reasons for this:
+  // 1. If the constant address calculation was done before null check, then the
+  // null check will not fail for that register (since it won't be null)
+  // 2. If *some modification* of the constant address should
+  // have occurred in the not_null block, we would have caught it earlier
+  // in the checks at analyzeBlockForNullChecks (see preservesZeroValueInReg).
+  if ((BaseReg && BaseReg != PointerReg && !BaseRegUsedInConstAddr) ||
+      (IndexReg && IndexReg != PointerReg && !IndexRegUsedInConstAddr))
     return SR_Unsuitable;
 
   // We want the mem access to be issued at a sane offset from PointerReg,
   // so that if PointerReg is null then the access reliably page faults.
-  if (!(-PageSize < Offset && Offset < PageSize))
+  if (!llvm::any_of(StartAddressesOfFaultingPage,
+                    [ConstantAddr](uint64_t StartAddress) {
+                      // Zeroth page is special cased when we have a negative
+                      // offset (greater than -PageSize), then the access will
+                      // fault.
+                      return (StartAddress == 0 && ConstantAddr < 0 &&
+                              ConstantAddr > -PageSize) ||
+                             (ConstantAddr >= StartAddress &&
+                              ConstantAddr < StartAddress + PageSize);
+                    }))
     return SR_Unsuitable;
 
   // Finally, check whether the current memory access aliases with previous one.
   for (auto *PrevMI : PrevInsts) {
     AliasResult AR = areMemoryOpsAliased(MI, PrevMI);
     if (AR == AR_WillAliasEverything)
       return SR_Impossible;
     if (AR == AR_MayAlias)
       return SR_Unsuitable;
   }
   return SR_Suitable;
 }
 
 bool ImplicitNullChecks::canDependenceHoistingClobberLiveIns(
     MachineInstr *DependenceMI, MachineBasicBlock *NullSucc,
     unsigned PointerReg) {
bool ImplicitNullChecks::analyzeBlockForNullChecks(
   //
   // we cannot (without code duplication) use the load marked "LD" to null check
   // 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;
 
   for (auto &MI : *NotNullSucc) {
     if (!canHandle(&MI) || InstsSeenSoFar.size() >= MaxInstsToConsider)
       return false;
 
     MachineInstr *Dependence;
-    SuitabilityResult SR = isSuitableMemoryOp(MI, PointerReg, InstsSeenSoFar);
+    SuitabilityResult SR = isSuitableMemoryOp(MI, PointerReg, InstsSeenSoFar,
+                                              InstrPreservesZeroReg);
     if (SR == SR_Impossible)
       return false;
     if (SR == SR_Suitable &&
         canHoistInst(&MI, PointerReg, InstsSeenSoFar, NullSucc, Dependence)) {
       NullCheckList.emplace_back(&MI, MBP.ConditionDef, &MBB, NotNullSucc,
                                  NullSucc, Dependence);
       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))
+      InstrPreservesZeroReg.insert(&MI);
     InstsSeenSoFar.push_back(&MI);
   }
 
   return false;
 }
 
 /// Wrap a machine instruction, MI, into a FAULTING machine instruction.
 /// The FAULTING instruction does the same load/store as MI

llvm/lib/Target/X86/X86InstrInfo.h

void replaceBranchWithTailCall(MachineBasicBlock &MBB,
                                  SmallVectorImpl<MachineOperand> &Cond,
                                  const MachineInstr &TailCall) const override;
 
   bool analyzeBranch(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                      MachineBasicBlock *&FBB,
                      SmallVectorImpl<MachineOperand> &Cond,
                      bool AllowModify) const override;
 
+  bool getConstantAddrFromMemoryOp(
+      const MachineInstr &MI, int64_t &ConstantAddr, int64_t &Scale,
+      Register &BaseReg, bool &BaseRegUsedInConstAddr, Register &IndexReg,
+      bool &IndexRegUsedInConstAddr, const TargetRegisterInfo *TRI,
+      const DenseSet<MachineInstr *> &IgnoreInstInAddrCalc) const override;
+
   bool preservesZeroValueInReg(const MachineInstr *MI,
                                const Register NullValueReg,
                                const TargetRegisterInfo *TRI) const override;
 
   bool getMemOperandsWithOffsetWidth(
       const MachineInstr &LdSt,
       SmallVectorImpl<const MachineOperand *> &BaseOps, int64_t &Offset,
       bool &OffsetIsScalable, unsigned &Width,

llvm/lib/Target/X86/X86InstrInfo.cpp

case 64:
     assert(STI.hasAVX512() && "Using 512-bit register requires AVX512");
     if (isStackAligned)
       return load ? X86::VMOVAPSZrm : X86::VMOVAPSZmr;
     else
       return load ? X86::VMOVUPSZrm : X86::VMOVUPSZmr;
   }
 }
 
+bool X86InstrInfo::getConstantAddrFromMemoryOp(
+    const MachineInstr &MemI, int64_t &ConstantAddr, int64_t &Scale,
+    Register &BaseReg, bool &BaseRegUsedInConstAddr, Register &IndexReg,
+    bool &IndexRegUsedInConstAddr, const TargetRegisterInfo *TRI,
+    const DenseSet<MachineInstr *> &IgnoreInstInAddrCalc) const {
+
+  const MCInstrDesc &Desc = MemI.getDesc();
+  int MemRefBegin = X86II::getMemoryOperandNo(Desc.TSFlags);
+  if (MemRefBegin < 0)
+    return false;
+
+  MemRefBegin += X86II::getOperandBias(Desc);
+
+  auto &BaseOp = MemI.getOperand(MemRefBegin + X86::AddrBaseReg);
+  if (!BaseOp.isReg()) // Can be an MO_FrameIndex
+    return false;
+  BaseReg = BaseOp.getReg();
+  BaseRegUsedInConstAddr = false;
+
+  const MachineOperand &DispMO = MemI.getOperand(MemRefBegin + X86::AddrDisp);
+  // Displacement can be symbolic
+  if (!DispMO.isImm())
+    return false;
+
+  // The first part of the constant address is the displacement itself.
+  ConstantAddr = DispMO.getImm();
+  Scale = MemI.getOperand(MemRefBegin + X86::AddrScaleAmt).getImm();
+  IndexReg = MemI.getOperand(MemRefBegin + X86::AddrIndexReg).getReg();
+  IndexRegUsedInConstAddr = false;
+
+  // Check if there are other constant values that make up the address.
+  // TODO: Extend this beyond the current basic block because usually the
+  // immediate is constant hoisted into a dominator MBB and used in a
+  // MOV/OR/ADD.
+  // Returns true if we could safely update ConstantAddr.
+  auto UpdateAddrForComplexMode = [&](Register RegUsedInAddr,
+                                      int64_t multiplier) {
+    if (RegUsedInAddr == X86::NoRegister || !multiplier)
+      return;
+    auto It = std::next(MachineBasicBlock::const_reverse_iterator(&MemI)),
+         End = MemI.getParent()->rend();
+    for (; It != End; It++) {
+      const MachineInstr *CurrMI = &*It;
+      if (CurrMI->modifiesRegister(RegUsedInAddr, TRI) &&
+          !IgnoreInstInAddrCalc.count(CurrMI))
+        break;
+    }
+    // No modifications.
+    if (It == End)
+      return;
+
+    const MachineInstr *CurrMI = &*It;
+    // We have found the most recent def of the baseOp which is not in one
+    // of the forms required to check for immediate.
+    // We only support cases where we have a direct mov since that means all
+    // other previous values for that register has been invalidated.
+    // TODO: We can also support instructions of the form: rax = add rax, rcx
+    // where
+    // the *most recent* def of rcx is a mov and we know rax to be zero at the
+    // point where the add is done. In other words:
+    // mov rcx, 12345 <-- this could be in another dominating block.
+    // rax = add rax, rcx
+    if (CurrMI->getOpcode() != X86::MOV32ri &&
+        CurrMI->getOpcode() != X86::MOV64ri)
+      return;
+    ConstantAddr += CurrMI->getOperand(1).getImm() * multiplier;
+    if (RegUsedInAddr == BaseReg)
+      BaseRegUsedInConstAddr = true;
+    else
+      IndexRegUsedInConstAddr = true;
+  };
+
+  UpdateAddrForComplexMode(BaseReg, 1);
+  UpdateAddrForComplexMode(IndexReg, Scale);
+  return true;
+}
+
 bool X86InstrInfo::preservesZeroValueInReg(
     const MachineInstr *MI, const Register NullValueReg,
     const TargetRegisterInfo *TRI) const {
   if (!MI->modifiesRegister(NullValueReg, TRI))
     return true;
   switch (MI->getOpcode()) {
   // Shift right/left of a null unto itself is still a null, i.e. rax = shl rax
   // X.

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

define i64 @imp_null_check_load_shift_add_addr(i64* %x, i64 %r) {
    %y = ptrtoint i64* %x to i64
    %shry = shl i64 %y, 6
    %shry.add = add i64 %shry, %r
    %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
 }
+
+; the memory op is outside the faulting regions.
+define i64 @imp_null_check_load_addr_outside_faulting_page(i64* %x) {
+  entry:
+   %c = icmp eq i64* %x, null
+   br i1 %c, label %is_null, label %not_null, !make.implicit !0
+
+  is_null:
+   ret i64 42
+
+  not_null:
+   %y = ptrtoint i64* %x to i64
+   %shry = shl i64 %y, 3
+   %shry.add = add i64 %shry, 3518432
+   %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-check.ll

-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -verify-machineinstrs -O3 -mtriple=x86_64-apple-macosx -enable-implicit-null-checks < %s | FileCheck %s
 
 define i32 @imp_null_check_load(i32* %x) {
 ; CHECK-LABEL: imp_null_check_load:
 ; CHECK:       ## %bb.0: ## %entry
 ; CHECK-NEXT:  Ltmp0:
 ; CHECK-NEXT:    movl (%rdi), %eax ## on-fault: LBB0_1
 ; CHECK-NEXT:  ## %bb.2: ## %not_null
; CHECK-NEXT:    retq
    %y.ptr = inttoptr i64 %shry to i64*
    %x.loc = getelementptr i64, i64* %y.ptr, i64 1
    %t = load i64, i64* %x.loc
    ret i64 %t
 }
 
 ; Same as imp_null_check_load_shift_addr but shift is by 3 and this is now
 ; converted into complex addressing.
-; TODO: Can be converted into implicit null check
 define i64 @imp_null_check_load_shift_by_3_addr(i64* %x) {
 ; CHECK-LABEL: imp_null_check_load_shift_by_3_addr:
 ; CHECK:       ## %bb.0: ## %entry
-; CHECK-NEXT:    testq %rdi, %rdi
-; CHECK-NEXT:    je LBB22_1
+; CHECK-NEXT:  Ltmp18:
+; CHECK-NEXT:    movq 8(,%rdi,8), %rax ## on-fault: LBB22_1
 ; CHECK-NEXT:  ## %bb.2: ## %not_null
-; CHECK-NEXT:    movq 8(,%rdi,8), %rax
 ; CHECK-NEXT:    retq
 ; CHECK-NEXT:  LBB22_1: ## %is_null
 ; CHECK-NEXT:    movl $42, %eax
 ; CHECK-NEXT:    retq
 
   entry:
    %c = icmp eq i64* %x, null
    br i1 %c, label %is_null, label %not_null, !make.implicit !0
 
   is_null:
    ret i64 42
 
   not_null:
    %y = ptrtoint i64* %x to i64
    %shry = shl i64 %y, 3
    %y.ptr = inttoptr i64 %shry to i64*
    %x.loc = getelementptr i64, i64* %y.ptr, i64 1
    %t = load i64, i64* %x.loc
    ret i64 %t
 }
+
+define i64 @imp_null_check_load_shift_add_addr(i64* %x) {
+; CHECK-LABEL: imp_null_check_load_shift_add_addr:
+; CHECK:       ## %bb.0: ## %entry
+; CHECK-NEXT:    movabsq $35184372088840, %rax ## imm = 0x200000000008
+; CHECK-NEXT:  Ltmp19:
+; CHECK-NEXT:    movq (%rax,%rdi,8), %rax ## on-fault: LBB23_1
+; CHECK-NEXT:  ## %bb.2: ## %not_null
+; CHECK-NEXT:    retq
+; CHECK-NEXT:  LBB23_1: ## %is_null
+; CHECK-NEXT:    movl $42, %eax
+; CHECK-NEXT:    retq
+
+  entry:
+   %c = icmp eq i64* %x, null
+   br i1 %c, label %is_null, label %not_null, !make.implicit !0
+
+  is_null:
+   ret i64 42
+
+  not_null:
+   %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

 # RUN: llc -run-pass implicit-null-checks -mtriple=x86_64-apple-macosx -o - %s | FileCheck %s
 
 --- |
   target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
   target triple = "x86_64-apple-macosx"
 
+  ;; Just to avoid confusion here - none of these IR actually correspond to MIR test cases below (other than the names and basic blocks).
+  ;; We need those since we rely on the make.implicit metadata in the pass.
   ;; Positive test
   define i32 @imp_null_check_with_bitwise_op_0(i32* %x, i32 %val) {
   entry:
     br i1 undef, label %is_null, label %not_null, !make.implicit !0
 
   is_null:
     ret i32 42
 
--- |
 
   not_null:
     ret i32 undef
 
   is_null:
     ret i32 undef
   }
 
+  define i32 @load_from_non_zero_faulting_page(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
+  }
+
+
+  define i32 @load_from_non_zero_faulting_page_negative_test(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
+  }
+
+  define i32 @load_from_non_zero_faulting_page_neg_test_2(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_0
 # CHECK-LABEL: name:            imp_null_check_with_bitwise_op_0
 alignment:       16
 tracksRegLiveness: true
 liveins:
   - { reg: '$rdi' }
bb.1.not_null:
     $eax = MOV32rr $edi
     RETQ $eax
 
   bb.2.is_null:
     $eax = XOR32rr undef $eax, undef $eax, implicit-def dead $eflags
     RETQ $eax
 
 ...
+---
+name:          load_from_non_zero_faulting_page
+# This is a positive test since the address for the memory load is within
+# the valid regions known to be faulting (according to the
+# startaddress_faulting_pages metadata).
+# CHECK-LABEL: name: load_from_non_zero_faulting_page
+# CHECK: bb.0.entry:
+# CHECK:   renamable $rdx = MOV64ri 35184372088832
+# CHECK:   $ecx = FAULTING_OP 1, %bb.3, 1724, renamable $rdx, 8, renamable $rax, 4, $noreg
+# CHECK:   JMP_1 %bb.1
+# CHECK: bb.1.not_null:
+# CHECK:   CMP32ri killed renamable $ecx, 1195, implicit-def $eflags
+# CHECK:   JCC_1 %bb.4, 4, implicit $eflags
+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 $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
+
+...
+---
+name:          load_from_non_zero_faulting_page_negative_test
+# This is a negative test since the address for the memory load is not within
+# the valid regions known to be faulting (according to the
+# startaddress_faulting_pages metadata).
+
+# CHECK-LABEL: name: load_from_non_zero_faulting_page_negative_test
+# CHECK: bb.0.entry:
+# CHECK:   TEST64rr renamable $rax, renamable $rax, implicit-def $eflags
+# CHECK-NOT: FAULTING_OP
+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 351843720
+    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
+
+---
+name:          load_from_non_zero_faulting_page_neg_test_2
+# This is a negative test since we do not know what makes up the address for
+# memory op and hence we cannot say if it is within faulting region.
+
+# CHECK-LABEL: name: load_from_non_zero_faulting_page_neg_test_2
+# CHECK: bb.0.entry:
+# CHECK:   TEST64rr renamable $rax, renamable $rax, implicit-def $eflags
+# CHECK-NOT: FAULTING_OP
+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 $ecx = MOV32rm renamable $rbx, 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