Do not refresh the Marking Register in CriticalNative methods.

CriticalNative methods shall not be suspended and hence do not
require MR to be refreshed in compiled JNI code.

This change is for ARM and ARM64 only.

Impact on Critical Native benchmarks times (median of 10 runs,
lower is better):

* angler-userdebug - ARMv7

** All cores

   NativeDowncallStaticCritical   -2.78%
   NativeDowncallStaticCritical6  -1.79%

** Little cores only

   NativeDowncallStaticCritical   -1.66%
   NativeDowncallStaticCritical6  -1.27%

** Big cores only

   NativeDowncallStaticCritical   -2.66%
   NativeDowncallStaticCritical6  -1.70%

* angler-userdebug - ARMv8

** All cores

   NativeDowncallStaticCritical   -3.52%
   NativeDowncallStaticCritical6  -1.79%

** Little cores only

   NativeDowncallStaticCritical   -1.63%
   NativeDowncallStaticCritical6  -1.27%

** Big cores only

   NativeDowncallStaticCritical   -3.87%
   NativeDowncallStaticCritical6  -1.75%

Test: m test-art-target
Test: m test-art-target with tree built with ART_USE_READ_BARRIER=false
Test: m test-art-host-gtest
Test: ARM64 device boot test
Test: ARM device boot test
Bug: b/37707231
Change-Id: I95d61b9ecde0afffdd5fd44763b19caa06025ec8
diff --git a/compiler/jni/jni_cfi_test.cc b/compiler/jni/jni_cfi_test.cc
index 347f4ea..28709a1 100644
--- a/compiler/jni/jni_cfi_test.cc
+++ b/compiler/jni/jni_cfi_test.cc
@@ -84,7 +84,7 @@
                         callee_save_regs, mr_conv->EntrySpills());
     jni_asm->IncreaseFrameSize(32);
     jni_asm->DecreaseFrameSize(32);
-    jni_asm->RemoveFrame(frame_size, callee_save_regs);
+    jni_asm->RemoveFrame(frame_size, callee_save_regs, /* may_suspend */ true);
     jni_asm->FinalizeCode();
     std::vector<uint8_t> actual_asm(jni_asm->CodeSize());
     MemoryRegion code(&actual_asm[0], actual_asm.size());
diff --git a/compiler/jni/quick/arm64/calling_convention_arm64.cc b/compiler/jni/quick/arm64/calling_convention_arm64.cc
index 292ce10..3afd701 100644
--- a/compiler/jni/quick/arm64/calling_convention_arm64.cc
+++ b/compiler/jni/quick/arm64/calling_convention_arm64.cc
@@ -110,23 +110,31 @@
 // Calling convention
 ManagedRegister Arm64ManagedRuntimeCallingConvention::InterproceduralScratchRegister() {
   // X20 is safe to use as a scratch register:
-  // - with Baker read barriers, it is reserved as Marking Register,
-  //   and thus does not actually need to be saved/restored; it is
-  //   refreshed on exit (see Arm64JNIMacroAssembler::RemoveFrame);
+  // - with Baker read barriers (in the case of a non-critical native
+  //   method), it is reserved as Marking Register, and thus does not
+  //   actually need to be saved/restored; it is refreshed on exit
+  //   (see Arm64JNIMacroAssembler::RemoveFrame);
   // - in other cases, it is saved on entry (in
   //   Arm64JNIMacroAssembler::BuildFrame) and restored on exit (in
-  //   Arm64JNIMacroAssembler::RemoveFrame).
+  //   Arm64JNIMacroAssembler::RemoveFrame). This is also expected in
+  //   the case of a critical native method in the Baker read barrier
+  //   configuration, where the value of MR must be preserved across
+  //   the JNI call (as there is no MR refresh in that case).
   return Arm64ManagedRegister::FromXRegister(X20);
 }
 
 ManagedRegister Arm64JniCallingConvention::InterproceduralScratchRegister() {
   // X20 is safe to use as a scratch register:
-  // - with Baker read barriers, it is reserved as Marking Register,
-  //   and thus does not actually need to be saved/restored; it is
-  //   refreshed on exit (see Arm64JNIMacroAssembler::RemoveFrame);
+  // - with Baker read barriers (in the case of a non-critical native
+  //   method), it is reserved as Marking Register, and thus does not
+  //   actually need to be saved/restored; it is refreshed on exit
+  //   (see Arm64JNIMacroAssembler::RemoveFrame);
   // - in other cases, it is saved on entry (in
   //   Arm64JNIMacroAssembler::BuildFrame) and restored on exit (in
-  //   Arm64JNIMacroAssembler::RemoveFrame).
+  //   Arm64JNIMacroAssembler::RemoveFrame). This is also expected in
+  //   the case of a critical native method in the Baker read barrier
+  //   configuration, where the value of MR must be preserved across
+  //   the JNI call (as there is no MR refresh in that case).
   return Arm64ManagedRegister::FromXRegister(X20);
 }
 
diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index c66a2a6..e7a5935 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -646,7 +646,10 @@
   // 16. Remove activation - need to restore callee save registers since the GC may have changed
   //     them.
   DCHECK_EQ(jni_asm->cfi().GetCurrentCFAOffset(), static_cast<int>(frame_size));
-  __ RemoveFrame(frame_size, callee_save_regs);
+  // We expect the compiled method to possibly be suspended during its
+  // execution, except in the case of a CriticalNative method.
+  bool may_suspend = !is_critical_native;
+  __ RemoveFrame(frame_size, callee_save_regs, may_suspend);
   DCHECK_EQ(jni_asm->cfi().GetCurrentCFAOffset(), static_cast<int>(frame_size));
 
   // 17. Finalize code generation
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
index ed57ca6..edb3292 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
@@ -117,7 +117,8 @@
 }
 
 void ArmVIXLJNIMacroAssembler::RemoveFrame(size_t frame_size,
-                                           ArrayRef<const ManagedRegister> callee_save_regs) {
+                                           ArrayRef<const ManagedRegister> callee_save_regs,
+                                           bool may_suspend) {
   CHECK_ALIGNED(frame_size, kStackAlignment);
   cfi().RememberState();
 
@@ -152,9 +153,18 @@
   ___ Pop(RegisterList(core_spill_mask));
 
   if (kEmitCompilerReadBarrier && kUseBakerReadBarrier) {
-    // Refresh Mark Register.
-    // TODO: Refresh MR only if suspend is taken.
-    ___ Ldr(mr, MemOperand(tr, Thread::IsGcMarkingOffset<kArmPointerSize>().Int32Value()));
+    if (may_suspend) {
+      // The method may be suspended; refresh the Marking Register.
+      ___ Ldr(mr, MemOperand(tr, Thread::IsGcMarkingOffset<kArmPointerSize>().Int32Value()));
+    } else {
+      // The method shall not be suspended; no need to refresh the Marking Register.
+
+      // Check that the Marking Register is a callee-save register,
+      // and thus has been preserved by native code following the
+      // AAPCS calling convention.
+      DCHECK_NE(core_spill_mask & (1 << MR), 0)
+          << "core_spill_mask should contain Marking Register R" << MR;
+    }
   }
 
   // Return to LR.
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
index f3baf1f..13b52e5 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
@@ -54,7 +54,8 @@
 
   // Emit code that will remove an activation from the stack.
   void RemoveFrame(size_t frame_size,
-                   ArrayRef<const ManagedRegister> callee_save_regs) OVERRIDE;
+                   ArrayRef<const ManagedRegister> callee_save_regs,
+                   bool may_suspend) OVERRIDE;
 
   void IncreaseFrameSize(size_t adjust) OVERRIDE;
   void DecreaseFrameSize(size_t adjust) OVERRIDE;
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.cc b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
index 9732b76..43c0eff 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.cc
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
@@ -743,7 +743,8 @@
 }
 
 void Arm64JNIMacroAssembler::RemoveFrame(size_t frame_size,
-                                         ArrayRef<const ManagedRegister> callee_save_regs) {
+                                         ArrayRef<const ManagedRegister> callee_save_regs,
+                                         bool may_suspend) {
   // Setup VIXL CPURegList for callee-saves.
   CPURegList core_reg_list(CPURegister::kRegister, kXRegSize, 0);
   CPURegList fp_reg_list(CPURegister::kFPRegister, kDRegSize, 0);
@@ -773,10 +774,21 @@
   asm_.UnspillRegisters(fp_reg_list, frame_size - core_reg_size - fp_reg_size);
 
   if (kEmitCompilerReadBarrier && kUseBakerReadBarrier) {
-    // Refresh Mark Register.
-    // TODO: Refresh MR only if suspend is taken.
-    ___ Ldr(reg_w(MR),
-            MemOperand(reg_x(TR), Thread::IsGcMarkingOffset<kArm64PointerSize>().Int32Value()));
+    vixl::aarch64::Register mr = reg_x(MR);  // Marking Register.
+    vixl::aarch64::Register tr = reg_x(TR);  // Thread Register.
+
+    if (may_suspend) {
+      // The method may be suspended; refresh the Marking Register.
+      ___ Ldr(mr.W(), MemOperand(tr, Thread::IsGcMarkingOffset<kArm64PointerSize>().Int32Value()));
+    } else {
+      // The method shall not be suspended; no need to refresh the Marking Register.
+
+      // Check that the Marking Register is a callee-save register,
+      // and thus has been preserved by native code following the
+      // AAPCS64 calling convention.
+      DCHECK(core_reg_list.IncludesAliasOf(mr))
+          << "core_reg_list should contain Marking Register X" << mr.GetCode();
+    }
   }
 
   // Decrease frame size to start of callee saved regs.
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.h b/compiler/utils/arm64/jni_macro_assembler_arm64.h
index baf0434..c993bbf 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.h
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.h
@@ -56,8 +56,9 @@
                   const ManagedRegisterEntrySpills& entry_spills) OVERRIDE;
 
   // Emit code that will remove an activation from the stack.
-  void RemoveFrame(size_t frame_size, ArrayRef<const ManagedRegister> callee_save_regs)
-      OVERRIDE;
+  void RemoveFrame(size_t frame_size,
+                   ArrayRef<const ManagedRegister> callee_save_regs,
+                   bool may_suspend) OVERRIDE;
 
   void IncreaseFrameSize(size_t adjust) OVERRIDE;
   void DecreaseFrameSize(size_t adjust) OVERRIDE;
diff --git a/compiler/utils/assembler_thumb_test.cc b/compiler/utils/assembler_thumb_test.cc
index 4dbe71b..5622f89 100644
--- a/compiler/utils/assembler_thumb_test.cc
+++ b/compiler/utils/assembler_thumb_test.cc
@@ -285,7 +285,7 @@
 
   __ DecreaseFrameSize(4096);
   __ DecreaseFrameSize(32);
-  __ RemoveFrame(frame_size, callee_save_regs);
+  __ RemoveFrame(frame_size, callee_save_regs, /* may_suspend */ true);
 
   EmitAndCheck(&assembler, "VixlJniHelpers");
 }
diff --git a/compiler/utils/jni_macro_assembler.h b/compiler/utils/jni_macro_assembler.h
index a8ca111..72f1ce0 100644
--- a/compiler/utils/jni_macro_assembler.h
+++ b/compiler/utils/jni_macro_assembler.h
@@ -66,7 +66,13 @@
                           const ManagedRegisterEntrySpills& entry_spills) = 0;
 
   // Emit code that will remove an activation from the stack
-  virtual void RemoveFrame(size_t frame_size, ArrayRef<const ManagedRegister> callee_save_regs) = 0;
+  //
+  // Argument `may_suspend` must be `true` if the compiled method may be
+  // suspended during its execution (otherwise `false`, if it is impossible
+  // to suspend during its execution).
+  virtual void RemoveFrame(size_t frame_size,
+                           ArrayRef<const ManagedRegister> callee_save_regs,
+                           bool may_suspend) = 0;
 
   virtual void IncreaseFrameSize(size_t adjust) = 0;
   virtual void DecreaseFrameSize(size_t adjust) = 0;
diff --git a/compiler/utils/mips/assembler_mips.cc b/compiler/utils/mips/assembler_mips.cc
index b300cc5..b83e3f5 100644
--- a/compiler/utils/mips/assembler_mips.cc
+++ b/compiler/utils/mips/assembler_mips.cc
@@ -5016,7 +5016,8 @@
 }
 
 void MipsAssembler::RemoveFrame(size_t frame_size,
-                                ArrayRef<const ManagedRegister> callee_save_regs) {
+                                ArrayRef<const ManagedRegister> callee_save_regs,
+                                bool may_suspend ATTRIBUTE_UNUSED) {
   CHECK_ALIGNED(frame_size, kStackAlignment);
   DCHECK(!overwriting_);
   cfi_.RememberState();
diff --git a/compiler/utils/mips/assembler_mips.h b/compiler/utils/mips/assembler_mips.h
index 0b4eb9c..e82693a 100644
--- a/compiler/utils/mips/assembler_mips.h
+++ b/compiler/utils/mips/assembler_mips.h
@@ -1090,8 +1090,9 @@
                   const ManagedRegisterEntrySpills& entry_spills) OVERRIDE;
 
   // Emit code that will remove an activation from the stack.
-  void RemoveFrame(size_t frame_size, ArrayRef<const ManagedRegister> callee_save_regs)
-      OVERRIDE;
+  void RemoveFrame(size_t frame_size,
+                   ArrayRef<const ManagedRegister> callee_save_regs,
+                   bool may_suspend) OVERRIDE;
 
   void IncreaseFrameSize(size_t adjust) OVERRIDE;
   void DecreaseFrameSize(size_t adjust) OVERRIDE;
diff --git a/compiler/utils/mips64/assembler_mips64.cc b/compiler/utils/mips64/assembler_mips64.cc
index 183b5e5..606d4c3 100644
--- a/compiler/utils/mips64/assembler_mips64.cc
+++ b/compiler/utils/mips64/assembler_mips64.cc
@@ -3406,7 +3406,8 @@
 }
 
 void Mips64Assembler::RemoveFrame(size_t frame_size,
-                                  ArrayRef<const ManagedRegister> callee_save_regs) {
+                                  ArrayRef<const ManagedRegister> callee_save_regs,
+                                  bool may_suspend ATTRIBUTE_UNUSED) {
   CHECK_ALIGNED(frame_size, kStackAlignment);
   DCHECK(!overwriting_);
   cfi_.RememberState();
diff --git a/compiler/utils/mips64/assembler_mips64.h b/compiler/utils/mips64/assembler_mips64.h
index bb54382..ee78cdb 100644
--- a/compiler/utils/mips64/assembler_mips64.h
+++ b/compiler/utils/mips64/assembler_mips64.h
@@ -1278,7 +1278,9 @@
                   const ManagedRegisterEntrySpills& entry_spills) OVERRIDE;
 
   // Emit code that will remove an activation from the stack.
-  void RemoveFrame(size_t frame_size, ArrayRef<const ManagedRegister> callee_save_regs) OVERRIDE;
+  void RemoveFrame(size_t frame_size,
+                   ArrayRef<const ManagedRegister> callee_save_regs,
+                   bool may_suspend) OVERRIDE;
 
   void IncreaseFrameSize(size_t adjust) OVERRIDE;
   void DecreaseFrameSize(size_t adjust) OVERRIDE;
diff --git a/compiler/utils/x86/jni_macro_assembler_x86.cc b/compiler/utils/x86/jni_macro_assembler_x86.cc
index e074346..2b3c65b 100644
--- a/compiler/utils/x86/jni_macro_assembler_x86.cc
+++ b/compiler/utils/x86/jni_macro_assembler_x86.cc
@@ -85,7 +85,8 @@
 }
 
 void X86JNIMacroAssembler::RemoveFrame(size_t frame_size,
-                                       ArrayRef<const ManagedRegister> spill_regs) {
+                                       ArrayRef<const ManagedRegister> spill_regs,
+                                       bool may_suspend ATTRIBUTE_UNUSED) {
   CHECK_ALIGNED(frame_size, kStackAlignment);
   cfi().RememberState();
   // -kFramePointerSize for ArtMethod*.
diff --git a/compiler/utils/x86/jni_macro_assembler_x86.h b/compiler/utils/x86/jni_macro_assembler_x86.h
index 8ffda64..75cdd1e 100644
--- a/compiler/utils/x86/jni_macro_assembler_x86.h
+++ b/compiler/utils/x86/jni_macro_assembler_x86.h
@@ -48,8 +48,9 @@
                   const ManagedRegisterEntrySpills& entry_spills) OVERRIDE;
 
   // Emit code that will remove an activation from the stack
-  void RemoveFrame(size_t frame_size, ArrayRef<const ManagedRegister> callee_save_regs)
-      OVERRIDE;
+  void RemoveFrame(size_t frame_size,
+                   ArrayRef<const ManagedRegister> callee_save_regs,
+                   bool may_suspend) OVERRIDE;
 
   void IncreaseFrameSize(size_t adjust) OVERRIDE;
   void DecreaseFrameSize(size_t adjust) OVERRIDE;
diff --git a/compiler/utils/x86_64/assembler_x86_64_test.cc b/compiler/utils/x86_64/assembler_x86_64_test.cc
index aff8871..b08ba4a 100644
--- a/compiler/utils/x86_64/assembler_x86_64_test.cc
+++ b/compiler/utils/x86_64/assembler_x86_64_test.cc
@@ -2043,7 +2043,7 @@
   ArrayRef<const ManagedRegister> spill_regs(raw_spill_regs);
 
   size_t frame_size = 10 * kStackAlignment;
-  assembler->RemoveFrame(frame_size, spill_regs);
+  assembler->RemoveFrame(frame_size, spill_regs, /* may_suspend */ true);
 
   // Construct assembly text counterpart.
   std::ostringstream str;
diff --git a/compiler/utils/x86_64/jni_macro_assembler_x86_64.cc b/compiler/utils/x86_64/jni_macro_assembler_x86_64.cc
index ec86254..72d9c43 100644
--- a/compiler/utils/x86_64/jni_macro_assembler_x86_64.cc
+++ b/compiler/utils/x86_64/jni_macro_assembler_x86_64.cc
@@ -100,7 +100,8 @@
 }
 
 void X86_64JNIMacroAssembler::RemoveFrame(size_t frame_size,
-                                          ArrayRef<const ManagedRegister> spill_regs) {
+                                          ArrayRef<const ManagedRegister> spill_regs,
+                                          bool may_suspend ATTRIBUTE_UNUSED) {
   CHECK_ALIGNED(frame_size, kStackAlignment);
   cfi().RememberState();
   int gpr_count = 0;
diff --git a/compiler/utils/x86_64/jni_macro_assembler_x86_64.h b/compiler/utils/x86_64/jni_macro_assembler_x86_64.h
index aa058f7..734ed96 100644
--- a/compiler/utils/x86_64/jni_macro_assembler_x86_64.h
+++ b/compiler/utils/x86_64/jni_macro_assembler_x86_64.h
@@ -49,8 +49,9 @@
                   const ManagedRegisterEntrySpills& entry_spills) OVERRIDE;
 
   // Emit code that will remove an activation from the stack
-  void RemoveFrame(size_t frame_size, ArrayRef<const ManagedRegister> callee_save_regs)
-      OVERRIDE;
+  void RemoveFrame(size_t frame_size,
+                   ArrayRef<const ManagedRegister> callee_save_regs,
+                   bool may_suspend) OVERRIDE;
 
   void IncreaseFrameSize(size_t adjust) OVERRIDE;
   void DecreaseFrameSize(size_t adjust) OVERRIDE;