Use register pair loads/stores in JNI stubs...

... on arm/arm64 for local reference frame manipulation.

Golem results for art-opt-cc (higher is better):
linux-armv7 (Odroid-C2)          before after
NativeDowncallStaticFast         21.622 21.622 (unchanged)
NativeDowncallStaticFast6        18.500 18.491 (-0.0500%)
NativeDowncallStaticFastRefs6    15.354 15.354 (unchanged)
NativeDowncallVirtualFast        21.027 20.741 (-1.361%)
NativeDowncallVirtualFast6       18.953 18.953 (unchanged)
NativeDowncallVirtualFastRefs6   15.504 15.504 (unchanged)
NativeDowncallStaticNormal       14.620 14.620 (unchanged)
NativeDowncallStaticNormal6      13.120 13.120 (unchanged)
NativeDowncallStaticNormalRefs6  11.454 11.454 (unchanged)
NativeDowncallVirtualNormal      14.342 14.216 (-0.8823%)
NativeDowncallVirtualNormal6     13.347 13.347 (unchanged)
NativeDowncallVirtualNormalRefs6 11.538 11.544 (+0.0481%)
linux-armv7 (Raspberry Pi 4)     before after
NativeDowncallStaticFast         41.937 41.564 (-0.8906%)
NativeDowncallStaticFast6        33.234 35.608 (+7.144%)
NativeDowncallStaticFastRefs6    30.527 31.469 (+3.085%)
NativeDowncallVirtualFast        37.531 35.429 (-5.600%)
NativeDowncallVirtualFast6       32.803 34.125 (+4.028%)
NativeDowncallVirtualFastRefs6   30.500 31.500 (+3.279%)
NativeDowncallStaticNormal       13.599 14.112 (+3.773%)
NativeDowncallStaticNormal6      13.599 13.599 (unchanged)
NativeDowncallStaticNormalRefs6  12.358 12.677 (+2.580%)
NativeDowncallVirtualNormal      13.473 13.848 (+2.781%)
NativeDowncallVirtualNormal6     13.235 13.242 (+0.0495%)
NativeDowncallVirtualNormalRefs6 12.165 12.364 (+1.632%)
linux-armv8 (Odroid-C2)          before after
NativeDowncallStaticFast         23.988 24.765 (+3.238%)
NativeDowncallStaticFast6        21.923 22.571 (+2.955%)
NativeDowncallStaticFastRefs6    18.719 19.183 (+2.480%)
NativeDowncallVirtualFast        21.027 21.622 (+2.828%)
NativeDowncallVirtualFast6       20.267 21.319 (+5.190%)
NativeDowncallVirtualFastRefs6   16.683 17.448 (+4.583%)
NativeDowncallStaticNormal       16.683 17.057 (+2.239%)
NativeDowncallStaticNormal6      15.656 15.992 (+2.149%)
NativeDowncallStaticNormalRefs6  13.958 14.085 (+0.9054)
NativeDowncallVirtualNormal      15.196 15.504 (+2.026%)
NativeDowncallVirtualNormal6     15.049 15.347 (+1.980%)
NativeDowncallVirtualNormalRefs6 13.006 13.466 (+3.541%)
linux-armv8 (Raspberry Pi 4)     before after
NativeDowncallStaticFast         36.482 38.366 (+5.164%)
NativeDowncallStaticFast6        37.406 38.366 (+2.564%)
NativeDowncallStaticFastRefs6    28.770 31.652 (+10.02%)
NativeDowncallVirtualFast        34.000 35.201 (+3.532%)
NativeDowncallVirtualFast6       33.251 34.000 (+2.254%)
NativeDowncallVirtualFastRefs6   26.474 27.201 (+2.747%)
NativeDowncallStaticNormal       14.237 14.606 (+2.592%)
NativeDowncallStaticNormal6      14.244 14.948 (+4.942%)
NativeDowncallStaticNormalRefs6  13.012 14.181 (+8.983%)
NativeDowncallVirtualNormal      14.105 14.663 (+3.954%)
NativeDowncallVirtualNormal6     13.979 14.735 (+5.406%)
NativeDowncallVirtualNormalRefs6 13.725 14.244 (+3.775%)

The Odroid-C2 results appear essentially unchanged for armv7
(with some minor regressions within noise) and only little
better for armv8 (but still likely within noise). On the
Raspberry Pi 4, there appears to be some improvement for
armv7 and a decent improvement for armv8 but there is higher
level of noise than on Odroid-C2. Results from this single
run are not very conclusive but we expect to see a clear
trend in the data after submission.

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

12 files changed, 190 insertions, 79 deletions

diff --git a/compiler/jni/quick/calling_convention.h b/compiler/jni/quick/calling_convention.h
index 2657e943e6..21336edb44 100644
--- a/compiler/jni/quick/calling_convention.h
+++ b/compiler/jni/quick/calling_convention.h
@@ -356,10 +356,6 @@ class JniCallingConvention : public CallingConvention {
 
   virtual ~JniCallingConvention() {}
 
-  static constexpr size_t SavedLocalReferenceCookieSize() {
-    return 4u;
-  }
-
   bool IsFastNative() const {
     return is_fast_native_;
   }
diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index 3357a5f8d7..1c077a1497 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -37,7 +37,7 @@
 #include "driver/compiler_options.h"
 #include "entrypoints/quick/quick_entrypoints.h"
 #include "instrumentation.h"
-#include "jni/jni_env_ext.h"
+#include "jni/local_reference_table.h"
 #include "runtime.h"
 #include "thread.h"
 #include "utils/arm/managed_register_arm.h"
@@ -51,19 +51,6 @@
 
 namespace art HIDDEN {
 
-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 SetNativeParameter(JNIMacroAssembler<kPointerSize>* jni_asm,
                                JniCallingConvention* jni_conv,
@@ -313,15 +300,15 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     // as well as the cookie, so we preserve them across calls in callee-save registers.
     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);
+    constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+    saved_cookie_reg = __ CoreRegisterWithSize(callee_save_scratch_regs[1], kLRTSegmentStateSize);
+    callee_save_temp = __ CoreRegisterWithSize(callee_save_scratch_regs[2], kLRTSegmentStateSize);
 
     // 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);
+    __ PushLocalReferenceFrame(jni_env_reg, saved_cookie_reg, callee_save_temp);
   }
 
   // 4. Make the main native call.
@@ -525,8 +512,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
 
   // 6. Pop local reference frame.
   if (LIKELY(!is_critical_native)) {
-    PopLocalReferenceFrame<kPointerSize>(
-        jni_asm.get(), jni_env_reg, saved_cookie_reg, callee_save_temp);
+    __ PopLocalReferenceFrame(jni_env_reg, saved_cookie_reg, callee_save_temp);
   }
 
   // 7. Return from the JNI stub.
@@ -658,8 +644,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
         jni_asm->cfi().AdjustCFAOffset(main_out_arg_size);
         __ DecreaseFrameSize(main_out_arg_size);
       }
-      PopLocalReferenceFrame<kPointerSize>(
-          jni_asm.get(), jni_env_reg, saved_cookie_reg, callee_save_temp);
+      __ PopLocalReferenceFrame(jni_env_reg, saved_cookie_reg, callee_save_temp);
     }
     DCHECK_EQ(jni_asm->cfi().GetCurrentCFAOffset(), static_cast<int>(current_frame_size));
     __ DeliverPendingException();
@@ -736,40 +721,6 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
 }
 
 template <PointerSize kPointerSize>
-static void PushLocalReferenceFrame(JNIMacroAssembler<kPointerSize>* jni_asm,
-                                    ManagedRegister jni_env_reg,
-                                    ManagedRegister saved_cookie_reg,
-                                    ManagedRegister temp_reg) {
-  const size_t kRawPointerSize = static_cast<size_t>(kPointerSize);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kRawPointerSize);
-  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kRawPointerSize);
-
-  // 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 kRawPointerSize = static_cast<size_t>(kPointerSize);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kRawPointerSize);
-  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kRawPointerSize);
-
-  // 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);
-}
-
-template <PointerSize kPointerSize>
 static void SetNativeParameter(JNIMacroAssembler<kPointerSize>* jni_asm,
                                JniCallingConvention* jni_conv,
                                ManagedRegister in_reg) {
diff --git a/compiler/trampolines/trampoline_compiler.cc b/compiler/trampolines/trampoline_compiler.cc
index d9f56629ef..b3fc688cc9 100644
--- a/compiler/trampolines/trampoline_compiler.cc
+++ b/compiler/trampolines/trampoline_compiler.cc
@@ -67,7 +67,7 @@ static std::unique_ptr<const std::vector<uint8_t>> CreateTrampoline(
 
       // VIXL will use the destination as a scratch register if
       // the offset is not encodable as an immediate operand.
-      ___ Ldr(temp_reg, MemOperand(r0, JNIEnvExt::SelfOffset(4).Int32Value()));
+      ___ Ldr(temp_reg, MemOperand(r0, JNIEnvExt::SelfOffset(kArmPointerSize).Int32Value()));
       ___ Ldr(pc, MemOperand(temp_reg, offset.Int32Value()));
       break;
     }
@@ -99,7 +99,7 @@ static std::unique_ptr<const std::vector<uint8_t>> CreateTrampoline(
     case kJniAbi:  // Load via Thread* held in JNIEnv* in first argument (X0).
       __ LoadRawPtr(Arm64ManagedRegister::FromXRegister(IP1),
                       Arm64ManagedRegister::FromXRegister(X0),
-                      Offset(JNIEnvExt::SelfOffset(8).Int32Value()));
+                      Offset(JNIEnvExt::SelfOffset(kArm64PointerSize).Int32Value()));
 
       __ JumpTo(Arm64ManagedRegister::FromXRegister(IP1), Offset(offset.Int32Value()),
                 Arm64ManagedRegister::FromXRegister(IP0));
@@ -134,9 +134,7 @@ static std::unique_ptr<const std::vector<uint8_t>> CreateTrampoline(ArenaAllocat
 
   switch (abi) {
     case kJniAbi:  // Load via Thread* held in JNIEnv* in first argument (A0).
-      __ Loadd(tmp,
-               A0,
-               JNIEnvExt::SelfOffset(static_cast<size_t>(kRiscv64PointerSize)).Int32Value());
+      __ Loadd(tmp, A0, JNIEnvExt::SelfOffset(kRiscv64PointerSize).Int32Value());
       __ Loadd(tmp, tmp, offset.Int32Value());
       __ Jr(tmp);
       break;
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
index 9c2589138c..851d33acfd 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
@@ -21,6 +21,8 @@
 
 #include "entrypoints/quick/quick_entrypoints.h"
 #include "indirect_reference_table.h"
+#include "jni/jni_env_ext.h"
+#include "jni/local_reference_table.h"
 #include "lock_word.h"
 #include "thread.h"
 
@@ -1108,5 +1110,43 @@ void ArmVIXLJNIMacroAssembler::Load(ArmManagedRegister dest,
   }
 }
 
+void ArmVIXLJNIMacroAssembler::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                       ManagedRegister saved_cookie_reg,
+                                                       ManagedRegister temp_reg) {
+  constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+  DCHECK_EQ(kLRTSegmentStateSize, kRegSizeInBytes);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArmPointerSize);
+  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kArmPointerSize);
+  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
+            jni_env_segment_state_offset.SizeValue());
+
+  // Load the old cookie that we shall need to restore together with the current segment state.
+  ___ Ldrd(AsVIXLRegister(saved_cookie_reg.AsArm()),
+           AsVIXLRegister(temp_reg.AsArm()),
+           MemOperand(AsVIXLRegister(jni_env_reg.AsArm()), jni_env_cookie_offset.Int32Value()));
+
+  // Set the cookie in JNI environment to the current segment state.
+  Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kLRTSegmentStateSize);
+}
+
+void ArmVIXLJNIMacroAssembler::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                      ManagedRegister saved_cookie_reg,
+                                                      ManagedRegister temp_reg) {
+  constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+  DCHECK_EQ(kLRTSegmentStateSize, kRegSizeInBytes);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArmPointerSize);
+  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kArmPointerSize);
+  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
+            jni_env_segment_state_offset.SizeValue());
+
+  // Load the current cookie.
+  Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
+
+  // Set the current segment state together with restoring the cookie.
+  ___ Strd(AsVIXLRegister(saved_cookie_reg.AsArm()),
+           AsVIXLRegister(temp_reg.AsArm()),
+           MemOperand(AsVIXLRegister(jni_env_reg.AsArm()), jni_env_cookie_offset.Int32Value()));
+}
+
 }  // namespace arm
 }  // namespace art
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
index f6df7f2c53..9d6b092105 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
@@ -91,6 +91,14 @@ class ArmVIXLJNIMacroAssembler final
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
+  // Manipulating local reference frames.
+  void PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                               ManagedRegister saved_cookie_reg,
+                               ManagedRegister temp_reg) override;
+  void PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                              ManagedRegister saved_cookie_reg,
+                              ManagedRegister temp_reg) override;
+
   // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
   void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
                                          JNIMacroLabel* slow_path,
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.cc b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
index 8ce44b6c63..3e7091b60b 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.cc
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
@@ -18,6 +18,8 @@
 
 #include "entrypoints/quick/quick_entrypoints.h"
 #include "indirect_reference_table.h"
+#include "jni/jni_env_ext.h"
+#include "jni/local_reference_table.h"
 #include "lock_word.h"
 #include "managed_register_arm64.h"
 #include "offsets.h"
@@ -979,6 +981,48 @@ void Arm64JNIMacroAssembler::RemoveFrame(size_t frame_size,
   cfi().DefCFAOffset(frame_size);
 }
 
+void Arm64JNIMacroAssembler::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                     ManagedRegister saved_cookie_reg,
+                                                     ManagedRegister temp_reg) {
+  constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+  DCHECK_EQ(kLRTSegmentStateSize, kWRegSizeInBytes);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArm64PointerSize);
+  const MemberOffset jni_env_segment_state_offset =
+      JNIEnvExt::SegmentStateOffset(kArm64PointerSize);
+  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
+            jni_env_segment_state_offset.SizeValue());
+
+  // Load the old cookie that we shall need to restore together with the current segment state.
+  ___ Ldp(
+      reg_w(saved_cookie_reg.AsArm64().AsWRegister()),
+      reg_w(temp_reg.AsArm64().AsWRegister()),
+      MemOperand(reg_x(jni_env_reg.AsArm64().AsXRegister()), jni_env_cookie_offset.Int32Value()));
+
+  // Set the cookie in JNI environment to the current segment state.
+  Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kLRTSegmentStateSize);
+}
+
+void Arm64JNIMacroAssembler::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                    ManagedRegister saved_cookie_reg,
+                                                    ManagedRegister temp_reg) {
+  constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+  DCHECK_EQ(kLRTSegmentStateSize, kWRegSizeInBytes);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArm64PointerSize);
+  const MemberOffset jni_env_segment_state_offset =
+      JNIEnvExt::SegmentStateOffset(kArm64PointerSize);
+  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
+            jni_env_segment_state_offset.SizeValue());
+
+  // Load the current cookie.
+  Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
+
+  // Set the current segment state together with restoring the cookie.
+  ___ Stp(
+      reg_w(saved_cookie_reg.AsArm64().AsWRegister()),
+      reg_w(temp_reg.AsArm64().AsWRegister()),
+      MemOperand(reg_x(jni_env_reg.AsArm64().AsXRegister()), jni_env_cookie_offset.Int32Value()));
+}
+
 #undef ___
 
 }  // namespace arm64
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.h b/compiler/utils/arm64/jni_macro_assembler_arm64.h
index 2836e0947d..1482012688 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.h
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.h
@@ -93,6 +93,14 @@ class Arm64JNIMacroAssembler final : public JNIMacroAssemblerFwd<Arm64Assembler,
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
+  // Manipulating local reference frames.
+  void PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                               ManagedRegister saved_cookie_reg,
+                               ManagedRegister temp_reg) override;
+  void PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                              ManagedRegister saved_cookie_reg,
+                              ManagedRegister temp_reg) override;
+
   // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
   void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
                                          JNIMacroLabel* slow_path,
diff --git a/compiler/utils/jni_macro_assembler.cc b/compiler/utils/jni_macro_assembler.cc
index 7a90a46f51..e17d915ad7 100644
--- a/compiler/utils/jni_macro_assembler.cc
+++ b/compiler/utils/jni_macro_assembler.cc
@@ -38,6 +38,8 @@
 #include "base/globals.h"
 #include "base/memory_region.h"
 #include "gc_root.h"
+#include "jni/jni_env_ext.h"
+#include "jni/local_reference_table.h"
 #include "stack_reference.h"
 
 namespace art HIDDEN {
@@ -129,4 +131,54 @@ template
 void JNIMacroAssembler<PointerSize::k64>::LoadStackReference(ManagedRegister dest,
                                                              FrameOffset offs);
 
+template <PointerSize kPointerSize>
+void JNIMacroAssembler<kPointerSize>::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                              ManagedRegister saved_cookie_reg,
+                                                              ManagedRegister temp_reg) {
+  constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kPointerSize);
+  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kPointerSize);
+
+  // Load the old cookie that we shall need to restore.
+  Load(saved_cookie_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
+
+  // Set the cookie to the current segment state.
+  Load(temp_reg, jni_env_reg, jni_env_segment_state_offset, kLRTSegmentStateSize);
+  Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kLRTSegmentStateSize);
+}
+
+template
+void JNIMacroAssembler<PointerSize::k32>::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                                  ManagedRegister saved_cookie_reg,
+                                                                  ManagedRegister temp_reg);
+template
+void JNIMacroAssembler<PointerSize::k64>::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                                  ManagedRegister saved_cookie_reg,
+                                                                  ManagedRegister temp_reg);
+
+template <PointerSize kPointerSize>
+void JNIMacroAssembler<kPointerSize>::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                             ManagedRegister saved_cookie_reg,
+                                                             ManagedRegister temp_reg) {
+  constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
+  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kPointerSize);
+  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kPointerSize);
+
+  // Set the current segment state to the current cookie.
+  Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
+  Store(jni_env_reg, jni_env_segment_state_offset, temp_reg, kLRTSegmentStateSize);
+
+  // Restore the cookie to the saved value.
+  Store(jni_env_reg, jni_env_cookie_offset, saved_cookie_reg, kLRTSegmentStateSize);
+}
+
+template
+void JNIMacroAssembler<PointerSize::k32>::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                                 ManagedRegister saved_cookie_reg,
+                                                                 ManagedRegister temp_reg);
+template
+void JNIMacroAssembler<PointerSize::k64>::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                                                 ManagedRegister saved_cookie_reg,
+                                                                 ManagedRegister temp_reg);
+
 }  // namespace art
diff --git a/compiler/utils/jni_macro_assembler.h b/compiler/utils/jni_macro_assembler.h
index 2d52eada08..a415141fbe 100644
--- a/compiler/utils/jni_macro_assembler.h
+++ b/compiler/utils/jni_macro_assembler.h
@@ -167,6 +167,17 @@ class JNIMacroAssembler : public DeletableArenaObject<kArenaAllocAssembler> {
   virtual void GetCurrentThread(ManagedRegister dest) = 0;
   virtual void GetCurrentThread(FrameOffset dest_offset) = 0;
 
+  // Manipulating local reference frames.
+  //
+  // These have a default implementation but they can be overridden to use register pair
+  // load/store instructions on architectures that support them (arm, arm64).
+  virtual void PushLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                       ManagedRegister saved_cookie_reg,
+                                       ManagedRegister temp_reg);
+  virtual void PopLocalReferenceFrame(ManagedRegister jni_env_reg,
+                                      ManagedRegister saved_cookie_reg,
+                                      ManagedRegister temp_reg);
+
   // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
   virtual void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
                                                  JNIMacroLabel* slow_path,
diff --git a/runtime/jni/jni_env_ext.cc b/runtime/jni/jni_env_ext.cc
index 26b154c89a..30f2c1e604 100644
--- a/runtime/jni/jni_env_ext.cc
+++ b/runtime/jni/jni_env_ext.cc
@@ -142,27 +142,27 @@ void JNIEnvExt::PopFrame() {
 //       are tests in jni_internal_test to match the results against the actual values.
 
 // This is encoding the knowledge of the structure and layout of JNIEnv fields.
-static size_t JNIEnvSize(size_t pointer_size) {
+static size_t JNIEnvSize(PointerSize pointer_size) {
   // A single pointer.
-  return pointer_size;
+  return static_cast<size_t>(pointer_size);
 }
 
-inline MemberOffset JNIEnvExt::LocalReferenceTableOffset(size_t pointer_size) {
+inline MemberOffset JNIEnvExt::LocalReferenceTableOffset(PointerSize pointer_size) {
   return MemberOffset(JNIEnvSize(pointer_size) +
-                      2 * pointer_size);          // Thread* self + JavaVMExt* vm
+                      2 * static_cast<size_t>(pointer_size));  // Thread* self + JavaVMExt* vm
 }
 
-MemberOffset JNIEnvExt::SegmentStateOffset(size_t pointer_size) {
+MemberOffset JNIEnvExt::SegmentStateOffset(PointerSize pointer_size) {
   return MemberOffset(LocalReferenceTableOffset(pointer_size).SizeValue() +
                       jni::LocalReferenceTable::SegmentStateOffset().SizeValue());
 }
 
-MemberOffset JNIEnvExt::LocalRefCookieOffset(size_t pointer_size) {
+MemberOffset JNIEnvExt::LocalRefCookieOffset(PointerSize pointer_size) {
   return MemberOffset(LocalReferenceTableOffset(pointer_size).SizeValue() +
                       jni::LocalReferenceTable::PreviousStateOffset().SizeValue());
 }
 
-MemberOffset JNIEnvExt::SelfOffset(size_t pointer_size) {
+MemberOffset JNIEnvExt::SelfOffset(PointerSize pointer_size) {
   return MemberOffset(JNIEnvSize(pointer_size));
 }
 
diff --git a/runtime/jni/jni_env_ext.h b/runtime/jni/jni_env_ext.h
index 2e8b52e4cc..7172acc7b3 100644
--- a/runtime/jni/jni_env_ext.h
+++ b/runtime/jni/jni_env_ext.h
@@ -42,9 +42,9 @@ class JNIEnvExt : public JNIEnv {
   // Creates a new JNIEnvExt. Returns null on error, in which case error_msg
   // will contain a description of the error.
   static JNIEnvExt* Create(Thread* self, JavaVMExt* vm, std::string* error_msg);
-  static MemberOffset SegmentStateOffset(size_t pointer_size);
-  static MemberOffset LocalRefCookieOffset(size_t pointer_size);
-  static MemberOffset SelfOffset(size_t pointer_size);
+  static MemberOffset SegmentStateOffset(PointerSize pointer_size);
+  static MemberOffset LocalRefCookieOffset(PointerSize pointer_size);
+  static MemberOffset SelfOffset(PointerSize pointer_size);
   static jint GetEnvHandler(JavaVMExt* vm, /*out*/void** out, jint version);
 
   ~JNIEnvExt();
@@ -147,7 +147,7 @@ class JNIEnvExt : public JNIEnv {
       REQUIRES(!Locks::thread_list_lock_, !Locks::jni_function_table_lock_);
 
  private:
-  static MemberOffset LocalReferenceTableOffset(size_t pointer_size);
+  static MemberOffset LocalReferenceTableOffset(PointerSize pointer_size);
 
   // Override of function tables. This applies to both default as well as instrumented (CheckJNI)
   // function tables.
diff --git a/runtime/jni/jni_internal_test.cc b/runtime/jni/jni_internal_test.cc
index 155048a30d..1c878b079a 100644
--- a/runtime/jni/jni_internal_test.cc
+++ b/runtime/jni/jni_internal_test.cc
@@ -2634,19 +2634,22 @@ TEST_F(JniInternalTest, IndirectReferenceTableOffsets) {
 
 // Test the offset computation of JNIEnvExt offsets. b/26071368.
 TEST_F(JniInternalTest, JNIEnvExtOffsets) {
-  EXPECT_EQ(OFFSETOF_MEMBER(JNIEnvExt, self_), JNIEnvExt::SelfOffset(sizeof(void*)).Uint32Value());
+  EXPECT_EQ(OFFSETOF_MEMBER(JNIEnvExt, self_),
+            JNIEnvExt::SelfOffset(kRuntimePointerSize).Uint32Value());
 
   // `previous_state_` amd `segment_state_` are private in the IndirectReferenceTable.
   // So this test isn't as good as we'd hope it to be.
   uint32_t previous_state_now =
       OFFSETOF_MEMBER(JNIEnvExt, locals_) +
       jni::LocalReferenceTable::PreviousStateOffset().Uint32Value();
-  uint32_t previous_state_computed = JNIEnvExt::LocalRefCookieOffset(sizeof(void*)).Uint32Value();
+  uint32_t previous_state_computed =
+      JNIEnvExt::LocalRefCookieOffset(kRuntimePointerSize).Uint32Value();
   EXPECT_EQ(previous_state_now, previous_state_computed);
   uint32_t segment_state_now =
       OFFSETOF_MEMBER(JNIEnvExt, locals_) +
       jni::LocalReferenceTable::SegmentStateOffset().Uint32Value();
-  uint32_t segment_state_computed = JNIEnvExt::SegmentStateOffset(sizeof(void*)).Uint32Value();
+  uint32_t segment_state_computed =
+      JNIEnvExt::SegmentStateOffset(kRuntimePointerSize).Uint32Value();
   EXPECT_EQ(segment_state_now, segment_state_computed);
 }