[ImplicitNullChecks] Check for rewrite of register used in 'test' instruction

[jskumari]

Summary

The following code pattern:

mov %rax, %rcx
test %rax, %rax
%rax = ....
je  throw_npe
mov(%rcx), %r9
mov(%rax), %r10

gets transformed into the following incorrect code after implicit null check pass:

 mov %rax, %rcx
%rax = ....
faulting_load_op("movl (%rax), %r10", throw_npe)
mov(%rcx), %r9

For implicit null check pass, if the register that is checked for null value (ie, the register used in the 'test' instruction) is written into before the condition jump, we should avoid doing the optimization.

[skatkov]

lib/CodeGen/ImplicitNullChecks.cpp

503 ↗

(On Diff #153032)

This form of comment is useful for understanding the reasoning of the bug but after it is landed it will be difficult to read this comment..
why it is converted in this way?!

I would suggest to re-phrase it. Something like, to prevent the invalid transformation ... we must ensure that ...

512 ↗

(On Diff #153032)

There is also the same assignment later:
const unsigned PointerReg.

514 ↗

(On Diff #153032)

I would suggest to move this check after "if (NotNullSucc->pred_size() != 1)" because the later seems simpler.

514 ↗

(On Diff #153032)

I would add an assert that MBP.ConditionDef is in MBB.

[skatkov]

Test, actually can be simplified bit if you do not mind.

test/CodeGen/X86/implicit-null-chk-reg-rewrite.mir

49 ↗

(On Diff #153032)

Why do you need this?

51 ↗

(On Diff #153032)

I guess you do not need this CMP for reproducing the failure? just return result of the and will be enough?
Do I miss anything?

[jskumari]

Incorporated review comments.

llvm/trunk/lib/CodeGen/ImplicitNullChecks.cpp

if (MBP.Predicate == MachineBranchPredicate::PRED_NE) {
     NullSucc = MBP.TrueDest;
   }
 
   // We handle the simplest case for now.  We can potentially do better by using
   // the machine dominator tree.
   if (NotNullSucc->pred_size() != 1)
     return false;
 
+  // To prevent the invalid transformation of the following code:
+  //
+  //   mov %rax, %rcx
+  //   test %rax, %rax
+  //   %rax = ...
+  //   je throw_npe
+  //   mov(%rcx), %r9
+  //   mov(%rax), %r10
+  //
+  // into:
+  //
+  //   mov %rax, %rcx
+  //   %rax = ....
+  //   faulting_load_op("movl (%rax), %r10", throw_npe)
+  //   mov(%rcx), %r9
+  //
+  // we must ensure that there are no instructions between the 'test' and
+  // conditional jump that modify %rax.
+  const unsigned PointerReg = MBP.LHS.getReg();
+
+  assert(MBP.ConditionDef->getParent() ==  &MBB && "Should be in basic block");
+
+  for (auto I = MBB.rbegin(); MBP.ConditionDef != &*I; ++I)
+    if (I->modifiesRegister(PointerReg, TRI))
+      return false;
+
   // Starting with a code fragment like:
   //
   //   test %rax, %rax
   //   jne LblNotNull
   //
   //  LblNull:
   //   callq throw_NullPointerException
   //
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.
 
-  const unsigned PointerReg = MBP.LHS.getReg();
-
   SmallVector<MachineInstr *, 8> InstsSeenSoFar;
 
   for (auto &MI : *NotNullSucc) {
     if (!canHandle(&MI) || InstsSeenSoFar.size() >= MaxInstsToConsider)
       return false;
 
     MachineInstr *Dependence;
     SuitabilityResult SR = isSuitableMemoryOp(MI, PointerReg, InstsSeenSoFar);

llvm/trunk/test/CodeGen/X86/implicit-null-chk-reg-rewrite.mir

+# RUN: llc -mtriple=x86_64 -run-pass=implicit-null-checks %s -o - | FileCheck %s
+--- |
+
+  define i32 @reg-rewrite(i32* %x) {
+  entry:
+    br i1 undef, label %is_null, label %not_null, !make.implicit !0
+
+  is_null:
+    ret i32 42
+
+  not_null:
+    ret i32 100
+  }
+
+  !0 = !{}
+
+...
+---
+# Check that the TEST instruction is replaced with 
+# FAULTING_OP only if there are no instructions
+# between the TEST and conditional jump
+# that clobber the register used in TEST.
+name:            reg-rewrite
+
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$rdi' }
+
+body:             |
+  bb.0.entry:
+    liveins: $rdi
+
+    TEST64rr $rdi, $rdi, implicit-def $eflags
+    ; CHECK-LABEL: bb.0.entry
+    ; CHECK-NOT: FAULTING_OP
+    renamable $rdi = MOV64ri 5000
+    JE_1 %bb.2, implicit $eflags
+
+  bb.1.not_null:
+    liveins: $rdi, $rsi
+    
+    $rax = MOV64rm renamable $rdi, 1, $noreg, 4, $noreg
+    RETQ $eax
+
+  bb.2.is_null:
+    $eax = MOV32ri 200
+    RETQ $eax
+...