[ImplicitNullCheck] Handle instructions that preserve zero value

This is the first in a series of patches to make implicit null checks
more general. This patch identifies instructions that preserves zero
value of a register and considers that as a valid instruction to hoist
along with the faulting load. See added testcases.

Reviewed-By: reames, dantrushin

Differential Revision: https://reviews.llvm.org/D87108

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -1270,6 +1270,17 @@ public:
     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 {
+    return false;
+  }
+
   /// If the instruction is an increment of a constant value, return the amount.
   virtual bool getIncrementValue(const MachineInstr &MI, int &Value) const {
     return false;

llvm/lib/CodeGen/ImplicitNullChecks.cpp

@@ -435,12 +435,6 @@ bool ImplicitNullChecks::canDependenceHoistingClobberLiveIns(
     if (AnyAliasLiveIn(TRI, NullSucc, DependenceMO.getReg()))
       return true;
 
-    // The Dependency can't be re-defining the base register -- then we won't
-    // get the memory operation on the address we want.  This is already
-    // checked in \c IsSuitableMemoryOp.
-    assert(!(DependenceMO.isDef() &&
-             TRI->regsOverlap(DependenceMO.getReg(), PointerReg)) &&
-           "Should have been checked before!");
   }
 
   // The dependence does not clobber live-ins in NullSucc block.
@@ -628,11 +622,9 @@ bool ImplicitNullChecks::analyzeBlockForNullChecks(
       return true;
     }
 
-    // If MI re-defines the PointerReg then we cannot move further.
-    if (llvm::any_of(MI.operands(), [&](MachineOperand &MO) {
-          return MO.isReg() && MO.getReg() && MO.isDef() &&
-                 TRI->regsOverlap(MO.getReg(), PointerReg);
-        }))
+    // 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;
     InstsSeenSoFar.push_back(&MI);
   }

llvm/lib/Target/X86/X86InstrInfo.cpp

@@ -3663,6 +3663,34 @@ static unsigned getLoadStoreRegOpcode(unsigned Reg,
   }
 }
 
+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.
+  case X86::SHR64ri:
+  case X86::SHR32ri:
+  case X86::SHL64ri:
+  case X86::SHL32ri:
+    assert(MI->getOperand(0).isDef() && MI->getOperand(1).isUse() &&
+           "expected for shift opcode!");
+    return MI->getOperand(0).getReg() == NullValueReg &&
+           MI->getOperand(1).getReg() == NullValueReg;
+  // Zero extend of a sub-reg of NullValueReg into itself does not change the
+  // null value.
+  case X86::MOV32rr:
+    return llvm::all_of(MI->operands(), [&](const MachineOperand &MO) {
+      return TRI->isSubRegisterEq(NullValueReg, MO.getReg());
+    });
+  default:
+    return false;
+  }
+  llvm_unreachable("Should be handled above!");
+}
+
 bool X86InstrInfo::getMemOperandsWithOffsetWidth(
     const MachineInstr &MemOp, SmallVectorImpl<const MachineOperand *> &BaseOps,
     int64_t &Offset, bool &OffsetIsScalable, unsigned &Width,

llvm/lib/Target/X86/X86InstrInfo.h

@@ -317,6 +317,10 @@ public:
                      SmallVectorImpl<MachineOperand> &Cond,
                      bool AllowModify) 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,

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

@@ -109,4 +109,24 @@ define i32 @imp_null_check_add_result(i32* %x, i32* %y) {
   ret i32 %p
 }
 
+; This redefines the null check reg by doing a zero-extend, a shift on
+; itself and then an add.
+; Cannot be converted to implicit check since the zero reg is no longer zero.
+define i64 @imp_null_check_load_shift_add_addr(i64* %x, i64 %r) {
+  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, 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
+}
 !0 = !{}

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

@@ -48,6 +48,8 @@ define i32 @imp_null_check_unordered_load(i32* %x) {
   ret i32 %t
 }
 
+
+; TODO: Can be converted into implicit check.
 ;; Probably could be implicit, but we're conservative for now
 define i32 @imp_null_check_seq_cst_load(i32* %x) {
 ; CHECK-LABEL: imp_null_check_seq_cst_load:
@@ -557,4 +559,66 @@ define i32 @imp_null_check_neg_gep_load(i32* %x) {
   ret i32 %t
 }
 
+; This redefines the null check reg by doing a zero-extend and a shift on
+; itself.
+; Converted into implicit null check since both of these operations do not
+; change the nullness of %x (i.e. if it is null, it remains null).
+define i64 @imp_null_check_load_shift_addr(i64* %x) {
+; CHECK-LABEL: imp_null_check_load_shift_addr:
+; CHECK:       ## %bb.0: ## %entry
+; CHECK-NEXT:    shlq $6, %rdi
+; CHECK-NEXT:  Ltmp17:
+; CHECK-NEXT:    movq 8(%rdi), %rax ## on-fault: LBB21_1
+; CHECK-NEXT:  ## %bb.2: ## %not_null
+; CHECK-NEXT:    retq
+; CHECK-NEXT:  LBB21_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, 6
+   %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:  ## %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
+}
 !0 = !{}