[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

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

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

514

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

514

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

Why do you need this?

51

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.

lib/CodeGen/ImplicitNullChecks.cpp

bool ImplicitNullChecks::analyzeBlockForNullChecks(
   if (MBP.Predicate == MachineBranchPredicate::PRED_NE) {
     NotNullSucc = MBP.TrueDest;
     NullSucc = MBP.FalseDest;
   } else {
     NotNullSucc = MBP.FalseDest;
     NullSucc = MBP.TrueDest;
   }
 
+  // The following code fragment:
+  //
+  //   mov %rax, %rcx
+  //   test %rax, %rax
+  //   %rax = ...
+  //   je throw_npe
+  //   mov(%rcx), %r9
+  //   mov(%rax), %r10
+  //
+  // gets incorrectly converted into:
+  //
+  //   mov %rax, %rcx
+  //   %rax = ....
+  //   faulting_load_op("movl (%rax), %r10", throw_npe)
+  //   mov(%rcx), %r9
+  //
+  // Check if there are any instructions between the 'test' and conditional
+  // jump that modify %rax.
+  unsigned cmpReg = MBP.LHS.getReg();
+
+  for (auto I = MBB.rbegin(); MBP.ConditionDef != &*I; ++I)
+    if (I->modifiesRegister(cmpReg, TRI))
+      return false;
+
   // 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;
 
   // Starting with a code fragment like:
   //
   //   test %rax, %rax

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, i32 %val) {
+  entry:
+    br i1 undef, 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
+  }
+
+  !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' }
+  - { reg: '$rsi' }
+body:             |
+  bb.0.entry:
+    liveins: $rdi, $rsi
+
+    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 = MOV64ri 2200000
+    $rax = AND64rm killed $rax, killed $rdi, 1, $noreg, 0, $noreg, implicit-def dead $eflags
+    CMP64rr killed $rax, killed $rsi, implicit-def $eflags
+    JE_1 %bb.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
+...