[ImplicitNullCheck] Handle instructions that do not modify null behaviour of null checked reg

[anna]

Summary

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

[anna]

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

58

FYI - all changes below are cosmetic ones from auto generation of check lines with llc. I initially ran this with just the --function option but wasn't sure if it's worth the bother, given it's been a while this test file is updated.

[dantrushin]

llvm/lib/Target/X86/X86InstrInfo.cpp

3670

Nit: Since you check specific opcodes anyway, that check is unnecessary, IMHO

3677

Again, since you're checking specific opcodes, you know that opnd 0 and 1 are registers,
so you can do something like

return MI->getOperand(0).getReg() == NullValueReg && MI->getOperand(1).getReg() == NullValueReg;

I.e., no need to check MO.isReg and direct register comparision might be more readable that isTied

[anna]

llvm/lib/Target/X86/X86InstrInfo.cpp

3670

We need this check because if you see the caller in ImplicitNullChecks, it was previously checking if If MI re-defines the PointerReg (i.e. the NullValueReg). We now have both checks within isNullBehaviourUnchanged. I originally thought of having the modifiesRegister in the caller itself and leaving an assert here instead of the check, but I felt this is better.

3677

agreed.

[dantrushin]

Looks good for me.
The only thing that worries me is shouldn't we have broader audience since we're changing 'global' APIs?

llvm/lib/Target/X86/X86InstrInfo.cpp

3670

Yeah, my bad, I've misread the code. :) If it is not modifies register, then it's safe.

[anna]

addressed review comments around using simpler form rather than isTied. Whitespace testcase changes avoided. NFC w.r.t. prev diff.

[anna]

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

58

ignore comment. was w.r.t prev diff. all whitespace changes avoided now,

[anna]

thanks Denis for your comments. I'll let others have a chance in review as well. Regarding broader audience today, implicit null checks works only with X86 (or rather that's what we test with). Also, none of the global APIs are changed, precisely because they are used in other parts of the code. Hence a new API is introduced and used currently only for implicit null checks. It is broad enough that we can reuse it in other passes if we wish to do so.

[reames]

Approach generally looks reasonable. Please address nitpicks and then I plan to LGTM.

llvm/include/llvm/CodeGen/TargetInstrInfo.h

1276

minor name change: preservesZeroReg(...)

Null is simply zero in MI land, and we don't want to imply there's anything special about the pointer interpretation here.

I'd also suggest being really explicit in the comment that the register can become null. This function only returns whether a zero is guaranteed to remain a zero.

llvm/lib/Target/X86/X86InstrInfo.cpp

3673

Minor: Use a switch instead of a series of if clauses please.

Also, I believe you can handle all the rr variants here.

Also, please assert that operand 0 is a def and operand 1 is a use. Mostly just to make it easier for someone with minimal context to understand form.

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

595

not an LEA - I think you meant complex addressing? Or maybe the LEA gets folded late?

Regardless, the comment is out of sync with tested output.

623

Please add a couple of counter examples. Particularly:
%rax = shift %rdx, %rax (i.e. wrong operand order)
%rax = add64ri %rax, 1 (i.e. a different instruction)

[anna]

llvm/include/llvm/CodeGen/TargetInstrInfo.h

1276

I'd keep the name as preservesZeroValueInReg because "zero reg" could be confusing to mean "NoRegister" which is 0 for all targets.

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

595

ah yes. I meant complex addressing :)

[anna]

llvm/lib/Target/X86/X86InstrInfo.cpp

3673

I tried writing up the register variant and got compile errors that such opcodes don't exist (i.e. SHL64rr etc..). Checked the intel guide and indeed we only support ri and mi variants, i.e. the MI.getOperand(2) is guaranteed to be an immediate.

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

623

wrong operand order for shift doesn't work (see above comment about rr variant). will add the different instruction case.

[anna]

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

623

I might have to play around with the IR to get the %rax "ri" form working the way we want, i.e. don't convert to implicit null check. Today the "add ri" form doesn't exist before implicit null check. X86 does a pretty decent job and patterns such as "add immediate OR shift, add immediate" is converted into complex addressing mode load.
For example IR such as:

not_null:
+    %y = ptrtoint i64* %x to i64 
+    %shry = shl i64 %y, 6
+    %shry.add = add i64 %shry, 1
+    %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

becomes this before implicit-null-checks:

liveins: $rdi
 renamable $rdi = SHL64ri killed renamable $rdi(tied-def 0), 6, implicit-def dead $eflags
 renamable $rax = MOV64rm killed renamable $rdi, 1, $noreg, 9, $noreg :: (load 8 from %ir.x.loc) <-- rax = rdi + 1 * 0 + 9 
 RETQ $rax

where rdi is the ZeroValueReg.

So, I'm adding an "rr" form and that works as a counter example.

[anna]

addressed review comments.

[anna]

handle clang-tidy comment.

[reames]

LGTM

llvm/include/llvm/CodeGen/TargetInstrInfo.h

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;
   }
 
+  /// Returns true if MI's Def is NullValueReg, and the MI
+  /// does not change the NullValue. i.e. cases such as rax = shr rax, X where
+  /// NullValueReg = rax.
+  virtual bool isNullBehaviourUnchanged(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;
   }
 
   /// Returns true if the two given memory operations should be scheduled
   /// adjacent. Note that you have to add:
   ///   DAG->addMutation(createLoadClusterDAGMutation(DAG->TII, DAG->TRI));

llvm/lib/CodeGen/ImplicitNullChecks.cpp

for (auto &DependenceMO : DependenceMI->operands()) {
     // case the pointer loaded from is in the null page, the load will not
     // semantically execute, and affect machine state.  That is, if the load
     // was loading into %rax and it faults, the value of %rax should stay the
     // same as it would have been had the load not have executed and we'd have
     // branched to NullSucc directly.
     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.
   return false;
 }
 
 bool ImplicitNullChecks::canHoistInst(MachineInstr *FaultingMI,
                                       unsigned PointerReg,
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 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->isNullBehaviourUnchanged(&MI, PointerReg, TRI))
       return false;
     InstsSeenSoFar.push_back(&MI);
   }
 
   return false;
 }
 
 /// Wrap a machine instruction, MI, into a FAULTING machine instruction.

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 isNullBehaviourUnchanged(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,
       const TargetRegisterInfo *TRI) const override;
   bool analyzeBranchPredicate(MachineBasicBlock &MBB,
                               TargetInstrInfo::MachineBranchPredicate &MBP,
                               bool AllowModify = false) const override;

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::isNullBehaviourUnchanged(
+    const MachineInstr *MI, const Register NullValueReg,
+    const TargetRegisterInfo *TRI) const {
+  if (!MI->modifiesRegister(NullValueReg, TRI))
+    return true;
+  // Shift right/left of a null unto itself is still a null, i.e. rax = shl rax
+  // X.
+  if (MI->getOpcode() == X86::SHR64ri || MI->getOpcode() == X86::SHR32ri ||
+      MI->getOpcode() == X86::SHL64ri || MI->getOpcode() == X86::SHR64ri)
+    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.
+  if (MI->getOpcode() == X86::MOV32rr)
+    if (llvm::all_of(MI->operands(), [&](const MachineOperand &MO) {
+          return TRI->isSubRegisterEq(NullValueReg, MO.getReg());
+        }))
+      return true;
+
+  return false;
+}
+
 bool X86InstrInfo::getMemOperandsWithOffsetWidth(
     const MachineInstr &MemOp, SmallVectorImpl<const MachineOperand *> &BaseOps,
     int64_t &Offset, bool &OffsetIsScalable, unsigned &Width,
     const TargetRegisterInfo *TRI) const {
   const MCInstrDesc &Desc = MemOp.getDesc();
   int MemRefBegin = X86II::getMemoryOperandNo(Desc.TSFlags);
   if (MemRefBegin < 0)
     return false;

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

  is_null:
   ret i32 42
 
  not_null:
   %t = load atomic i32, i32* %x unordered, align 4
   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:
 ; CHECK:       ## %bb.0: ## %entry
 ; CHECK-NEXT:    testq %rdi, %rdi
 ; CHECK-NEXT:    je LBB2_1
 ; CHECK-NEXT:  ## %bb.2: ## %not_null
 ; CHECK-NEXT:    movl (%rdi), %eax
 ; CHECK-NEXT:    retq
 ; CHECK-NEXT:  LBB2_1: ## %is_null
 ; CHECK-NEXT:    movl $42, %eax
 ; CHECK-NEXT:    retq
 
  entry:
is_null:
   ret i32 42
 
  not_null:
   %x.gep = getelementptr i32, i32* %x, i32 -32
   %t = load i32, i32* %x.gep
   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 an LEA.
+; 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 = !{}