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

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.

Patch by Surya Kumari Jangala!

Differential Revision: https://reviews.llvm.org/D48627
Reviewed By: skatkov

llvm-svn: 336241

llvm/lib/CodeGen/ImplicitNullChecks.cpp

@@ -496,6 +496,32 @@ bool ImplicitNullChecks::analyzeBlockForNullChecks(
   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
@@ -550,8 +576,6 @@ bool ImplicitNullChecks::analyzeBlockForNullChecks(
   // 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) {

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

@@ -0,0 +1,49 @@
+# 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
+...