Inline IRT frame push/pop into JNI stubs.

Golem results for art-opt-cc (higher is better):
linux-ia32                       before after
NativeDowncallStaticNormal       25.704 26.839 (+4.414%)
NativeDowncallStaticNormal6      23.857 25.086 (+5.152%)
NativeDowncallStaticNormalRefs6  23.704 25.248 (+6.513%)
NativeDowncallVirtualNormal      25.578 27.000 (+5.560%)
NativeDowncallVirtualNormal6     23.704 24.925 (+5.153%)
NativeDowncallVirtualNormalRefs6 23.704 25.074 (+5.870%)
NativeDowncallStaticFast         100.65 149.13 (+48.17%)
NativeDowncallStaticFast6        78.304 107.39 (+37.71%)
NativeDowncallStaticFastRefs6    76.962 104.45 (+35.71%)
NativeDowncallVirtualFast        100.40 147.28 (+46.69%)
NativeDowncallVirtualFast6       79.302 106.34 (+34.10%)
NativeDowncallVirtualFastRef26   76.617 103.29 (+34.82%)
linux-x64                        before after
NativeDowncallStaticNormal       26.083 26.987 (+3.465%)
NativeDowncallStaticNormal6      24.606 25.411 (+3.271%)
NativeDowncallStaticNormalRefs6  24.150 25.086 (+3.877%)
NativeDowncallVirtualNormal      25.743 26.812 (+4.156%)
NativeDowncallVirtualNormal6     24.294 25.248 (+3.927%)
NativeDowncallVirtualNormalRefs6 23.857 25.086 (+5.152%)
NativeDowncallStaticFast         109.95 133.10 (+21.06%)
NativeDowncallStaticFast6        90.274 109.12 (+20.87%)
NativeDowncallStaticFastRefs6    87.282 105.29 (+20.63%)
NativeDowncallVirtualFast        104.00 127.55 (+22.65%)
NativeDowncallVirtualFast6       88.191 106.73 (+21.02%)
NativeDowncallVirtualFastRef26   85.530 102.09 (+19.36%)
linux-armv7                      before after
NativeDowncallStaticNormal       6.1148 6.3694 (+4.316%)
NativeDowncallStaticNormal6      5.6845 5.9026 (+3.837%)
NativeDowncallStaticNormalRefs6  5.4054 5.6022 (+3.641%)
NativeDowncallVirtualNormal      5.4726 5.7088 (+4.316%)
NativeDowncallVirtualNormal6     5.1789 5.3685 (+3.660%)
NativeDowncallVirtualNormalRefs6 4.9140 5.0902 (+3.586%)
NativeDowncallStaticFast         16.683 18.058 (+8.239%)
NativeDowncallStaticFast6        13.951 14.896 (+6.770%)
NativeDowncallStaticFastRefs6    12.279 13.006 (+5.919%)
NativeDowncallVirtualFast        16.161 17.848 (+10.44%)
NativeDowncallVirtualFast6       14.085 15.196 (+7.892%)
NativeDowncallVirtualFastRef26   12.089 12.897 (+6.683%)
linux-armv8                      before after
NativeDowncallStaticNormal       6.0663 6.4229 (+5.879%)
NativeDowncallStaticNormal6      5.7252 6.0437 (+5.563%)
NativeDowncallStaticNormalRefs6  5.3114 5.5814 (+5.082%)
NativeDowncallVirtualNormal      5.8795 6.2651 (+6.558%)
NativeDowncallVirtualNormal6     5.6232 5.9494 (+5.801%)
NativeDowncallVirtualNormalRefs6 5.1862 5.4429 (+4.948%)
NativeDowncallStaticFast         17.638 19.183 (+8.760%)
NativeDowncallStaticFast6        14.903 16.161 (+8.438%)
NativeDowncallStaticFastRefs6    12.475 13.235 (+6.094%)
NativeDowncallVirtualFast        15.826 17.848 (+12.78%)
NativeDowncallVirtualFast6       14.064 15.504 (+10.24%)
NativeDowncallVirtualFastRef26   11.628 12.475 (+7.285%)

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Test: run-gtests.sh
Test: testrunner.py --target --optimizing
Bug: 172332525
Change-Id: I5ecfa7a661f08ab63dd2a75d666e1c1b9121935f

1 files changed, 101 insertions, 43 deletions

diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index cdd0263729..25eb919c53 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -50,6 +50,19 @@
 
 namespace art {
 
+constexpr size_t kIRTCookieSize = JniCallingConvention::SavedLocalReferenceCookieSize();
+
+template <PointerSize kPointerSize>
+static void PushLocalReferenceFrame(JNIMacroAssembler<kPointerSize>* jni_asm,
+                                    ManagedRegister jni_env_reg,
+                                    ManagedRegister saved_cookie_reg,
+                                    ManagedRegister temp_reg);
+template <PointerSize kPointerSize>
+static void PopLocalReferenceFrame(JNIMacroAssembler<kPointerSize>* jni_asm,
+                                   ManagedRegister jni_env_reg,
+                                   ManagedRegister saved_cookie_reg,
+                                   ManagedRegister temp_reg);
+
 template <PointerSize kPointerSize>
 static void CopyParameter(JNIMacroAssembler<kPointerSize>* jni_asm,
                           ManagedRuntimeCallingConvention* mr_conv,
@@ -194,11 +207,9 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
   //     method and the current thread.
   const char* jni_end_shorty;
   if (reference_return && is_synchronized) {
-    jni_end_shorty = "ILL";
-  } else if (reference_return) {
     jni_end_shorty = "IL";
-  } else if (is_synchronized) {
-    jni_end_shorty = "VL";
+  } else if (reference_return) {
+    jni_end_shorty = "I";
   } else {
     jni_end_shorty = "V";
   }
@@ -275,18 +286,15 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
   }
 
   // 5. Call into appropriate JniMethodStart passing Thread* so that transition out of Runnable
-  //    can occur. The result is the saved JNI local state that is restored by the exit call. We
-  //    abuse the JNI calling convention here, that is guaranteed to support passing 2 pointer
-  //    arguments.
-  constexpr size_t cookie_size = JniCallingConvention::SavedLocalReferenceCookieSize();
-  ManagedRegister saved_cookie_register = ManagedRegister::NoRegister();
+  //    can occur. We abuse the JNI calling convention here, that is guaranteed to support passing
+  //    two pointer arguments.
   if (LIKELY(!is_critical_native)) {
     // Skip this for @CriticalNative methods. They do not call JniMethodStart.
-    ThreadOffset<kPointerSize> jni_start(
+    ThreadOffset<kPointerSize> jni_start =
         GetJniEntrypointThreadOffset<kPointerSize>(JniEntrypoint::kStart,
                                                    reference_return,
                                                    is_synchronized,
-                                                   is_fast_native).SizeValue());
+                                                   is_fast_native);
     main_jni_conv->ResetIterator(FrameOffset(main_out_arg_size));
     if (is_synchronized) {
       // Pass object for locking.
@@ -322,13 +330,33 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     if (is_synchronized) {  // Check for exceptions from monitor enter.
       __ ExceptionPoll(main_out_arg_size);
     }
+  }
 
-    // Store into stack_frame[saved_cookie_offset] the return value of JniMethodStart.
-    saved_cookie_register = main_jni_conv->SavedLocalReferenceCookieRegister();
-    __ Move(saved_cookie_register, main_jni_conv->IntReturnRegister(), cookie_size);
+  // 6. Push local reference frame.
+  // Skip this for @CriticalNative methods, they cannot use any references.
+  ManagedRegister jni_env_reg = ManagedRegister::NoRegister();
+  ManagedRegister saved_cookie_reg = ManagedRegister::NoRegister();
+  ManagedRegister callee_save_temp = ManagedRegister::NoRegister();
+  if (LIKELY(!is_critical_native)) {
+    // To pop the local reference frame later, we shall need the JNI environment pointer
+    // as well as the cookie, so we preserve them across calls in callee-save registers.
+    // Managed callee-saves were already saved, so these registers are now available.
+    ArrayRef<const ManagedRegister> callee_save_scratch_regs =
+        main_jni_conv->CalleeSaveScratchRegisters();
+    CHECK_GE(callee_save_scratch_regs.size(), 3u);  // At least 3 for each supported architecture.
+    jni_env_reg = callee_save_scratch_regs[0];
+    saved_cookie_reg = __ CoreRegisterWithSize(callee_save_scratch_regs[1], kIRTCookieSize);
+    callee_save_temp = __ CoreRegisterWithSize(callee_save_scratch_regs[2], kIRTCookieSize);
+
+    // Load the JNI environment pointer.
+    __ LoadRawPtrFromThread(jni_env_reg, Thread::JniEnvOffset<kPointerSize>());
+
+    // Push the local reference frame.
+    PushLocalReferenceFrame<kPointerSize>(
+        jni_asm.get(), jni_env_reg, saved_cookie_reg, callee_save_temp);
   }
 
-  // 6. Fill arguments.
+  // 7. Fill arguments.
   if (UNLIKELY(is_critical_native)) {
     ArenaVector<ArgumentLocation> src_args(allocator.Adapter());
     ArenaVector<ArgumentLocation> dest_args(allocator.Adapter());
@@ -388,7 +416,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
       CopyParameter(jni_asm.get(), mr_conv.get(), main_jni_conv.get());
     }
 
-    // 7. For static method, create jclass argument as a pointer to the method's declaring class.
+    // 8. For static method, create jclass argument as a pointer to the method's declaring class.
     if (is_static) {
       main_jni_conv->ResetIterator(FrameOffset(main_out_arg_size));
       main_jni_conv->Next();  // Skip JNIEnv*
@@ -413,18 +441,17 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     // Set the iterator back to the incoming Method*.
     main_jni_conv->ResetIterator(FrameOffset(main_out_arg_size));
 
-    // 8. Create 1st argument, the JNI environment ptr.
-    // Register that will hold local indirect reference table
+    // 9. Create 1st argument, the JNI environment ptr.
     if (main_jni_conv->IsCurrentParamInRegister()) {
-      ManagedRegister jni_env = main_jni_conv->CurrentParamRegister();
-      __ LoadRawPtrFromThread(jni_env, Thread::JniEnvOffset<kPointerSize>());
+      ManagedRegister jni_env_arg = main_jni_conv->CurrentParamRegister();
+      __ Move(jni_env_arg, jni_env_reg, static_cast<size_t>(kPointerSize));
     } else {
-      FrameOffset jni_env = main_jni_conv->CurrentParamStackOffset();
-      __ CopyRawPtrFromThread(jni_env, Thread::JniEnvOffset<kPointerSize>());
+      FrameOffset jni_env_arg_offset = main_jni_conv->CurrentParamStackOffset();
+      __ Store(jni_env_arg_offset, jni_env_reg, static_cast<size_t>(kPointerSize));
     }
   }
 
-  // 9. Plant call to native code associated with method.
+  // 10. Plant call to native code associated with method.
   MemberOffset jni_entrypoint_offset =
       ArtMethod::EntryPointFromJniOffset(InstructionSetPointerSize(instruction_set));
   if (UNLIKELY(is_critical_native)) {
@@ -442,7 +469,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     }
   }
 
-  // 10. Fix differences in result widths.
+  // 11. Fix differences in result widths.
   if (main_jni_conv->RequiresSmallResultTypeExtension()) {
     DCHECK(main_jni_conv->HasSmallReturnType());
     CHECK(!is_critical_native || !main_jni_conv->UseTailCall());
@@ -458,7 +485,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     }
   }
 
-  // 11. Process return value
+  // 12. Process return value
   bool spill_return_value = main_jni_conv->SpillsReturnValue();
   FrameOffset return_save_location =
       spill_return_value ? main_jni_conv->ReturnValueSaveLocation() : FrameOffset(0);
@@ -504,26 +531,17 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     }
     end_jni_conv->ResetIterator(FrameOffset(end_out_arg_size));
 
-    // 12. Call JniMethodEnd
-    ThreadOffset<kPointerSize> jni_end(
+    // 13. Call JniMethodEnd
+    ThreadOffset<kPointerSize> jni_end =
         GetJniEntrypointThreadOffset<kPointerSize>(JniEntrypoint::kEnd,
                                                    reference_return,
                                                    is_synchronized,
-                                                   is_fast_native).SizeValue());
+                                                   is_fast_native);
     if (reference_return) {
       // Pass result.
       SetNativeParameter(jni_asm.get(), end_jni_conv.get(), end_jni_conv->ReturnRegister());
       end_jni_conv->Next();
     }
-    // Pass saved local reference state.
-    if (end_jni_conv->IsCurrentParamOnStack()) {
-      FrameOffset out_off = end_jni_conv->CurrentParamStackOffset();
-      __ Store(out_off, saved_cookie_register, cookie_size);
-    } else {
-      ManagedRegister out_reg = end_jni_conv->CurrentParamRegister();
-      __ Move(out_reg, saved_cookie_register, cookie_size);
-    }
-    end_jni_conv->Next();
     if (is_synchronized) {
       // Pass object for unlocking.
       if (is_static) {
@@ -563,26 +581,32 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
       __ CallFromThread(jni_end);
     }
 
-    // 13. Reload return value
+    // 14. Reload return value
     if (spill_return_value) {
       __ Load(mr_conv->ReturnRegister(), return_save_location, mr_conv->SizeOfReturnValue());
     }
   }  // if (!is_critical_native)
 
-  // 14. Move frame up now we're done with the out arg space.
+  // 15. Pop local reference frame.
+  if (!is_critical_native) {
+    PopLocalReferenceFrame<kPointerSize>(
+        jni_asm.get(), jni_env_reg, saved_cookie_reg, callee_save_temp);
+  }
+
+  // 16. Move frame up now we're done with the out arg space.
   //     @CriticalNative remove out args together with the frame in RemoveFrame().
   if (LIKELY(!is_critical_native)) {
     __ DecreaseFrameSize(current_out_arg_size);
     current_frame_size -= current_out_arg_size;
   }
 
-  // 15. Process pending exceptions from JNI call or monitor exit.
+  // 17. Process pending exceptions from JNI call or monitor exit.
   //     @CriticalNative methods do not need exception poll in the stub.
   if (LIKELY(!is_critical_native)) {
     __ ExceptionPoll(/* stack_adjust= */ 0);
   }
 
-  // 16. Remove activation - need to restore callee save registers since the GC may have changed
+  // 18. Remove activation - need to restore callee save registers since the GC may have changed
   //     them.
   DCHECK_EQ(jni_asm->cfi().GetCurrentCFAOffset(), static_cast<int>(current_frame_size));
   if (LIKELY(!is_critical_native) || !main_jni_conv->UseTailCall()) {
@@ -593,7 +617,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     DCHECK_EQ(jni_asm->cfi().GetCurrentCFAOffset(), static_cast<int>(current_frame_size));
   }
 
-  // 17. Read barrier slow path for the declaring class in the method for a static call.
+  // 19. Read barrier slow path for the declaring class in the method for a static call.
   //     Skip this for @CriticalNative because we're not passing a `jclass` to the native method.
   if (kUseReadBarrier && is_static && !is_critical_native) {
     __ Bind(jclass_read_barrier_slow_path.get());
@@ -649,7 +673,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     }
   }
 
-  // 18. Finalize code generation
+  // 20. Finalize code generation
   __ FinalizeCode();
   size_t cs = __ CodeSize();
   std::vector<uint8_t> managed_code(cs);
@@ -664,6 +688,40 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
                            ArrayRef<const uint8_t>(*jni_asm->cfi().data()));
 }
 
+template <PointerSize kPointerSize>
+static void PushLocalReferenceFrame(JNIMacroAssembler<kPointerSize>* jni_asm,
+                                    ManagedRegister jni_env_reg,
+                                    ManagedRegister saved_cookie_reg,
+                                    ManagedRegister temp_reg) {
+  const size_t pointer_size = static_cast<size_t>(kPointerSize);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(pointer_size);
+  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(pointer_size);
+
+  // Load the old cookie that we shall need to restore.
+  __ Load(saved_cookie_reg, jni_env_reg, jni_env_cookie_offset, kIRTCookieSize);
+
+  // Set the cookie in JNI environment to the current segment state.
+  __ Load(temp_reg, jni_env_reg, jni_env_segment_state_offset, kIRTCookieSize);
+  __ Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kIRTCookieSize);
+}
+
+template <PointerSize kPointerSize>
+static void PopLocalReferenceFrame(JNIMacroAssembler<kPointerSize>* jni_asm,
+                                   ManagedRegister jni_env_reg,
+                                   ManagedRegister saved_cookie_reg,
+                                   ManagedRegister temp_reg) {
+  const size_t pointer_size = static_cast<size_t>(kPointerSize);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(pointer_size);
+  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(pointer_size);
+
+  // Set the current segment state to the current cookie in JNI environment.
+  __ Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kIRTCookieSize);
+  __ Store(jni_env_reg, jni_env_segment_state_offset, temp_reg, kIRTCookieSize);
+
+  // Restore the cookie in JNI environment to the saved value.
+  __ Store(jni_env_reg, jni_env_cookie_offset, saved_cookie_reg, kIRTCookieSize);
+}
+
 // Copy a single parameter from the managed to the JNI calling convention.
 template <PointerSize kPointerSize>
 static void CopyParameter(JNIMacroAssembler<kPointerSize>* jni_asm,