JNI: Fast-path for decoding returned jobject.

Results for the timeGetBytesAscii#EMPTY benchmark from the
libcore's StringToBytesBenchmark suite on blueline-userdebug
with the cpu frequencies fixed at 1420800 (cpus 0-3; little)
and 1459200 (cpus 4-7; big):
  32-bit little: ~415 -> ~390
  64-bit little: ~415 -> ~390
  32-bit big:    ~180 -> ~170
  64-bit big:    ~180 -> ~170

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing --debug --ndebug
Test: run-gtests.sh
Test: testrunner.py --target --optimizing --debug --ndebug
Bug: 172332525
Change-Id: I0e19d583e5141e99a8b8c6fd9ae125fe7c9e02e7
diff --git a/compiler/jni/jni_compiler_test.cc b/compiler/jni/jni_compiler_test.cc
index d1d3190..397db25 100644
--- a/compiler/jni/jni_compiler_test.cc
+++ b/compiler/jni/jni_compiler_test.cc
@@ -28,6 +28,7 @@
 #include "common_compiler_test.h"
 #include "compiler.h"
 #include "dex/dex_file.h"
+#include "driver/compiler_options.h"
 #include "entrypoints/entrypoint_utils-inl.h"
 #include "gtest/gtest.h"
 #include "indirect_reference_table.h"
@@ -1553,6 +1554,10 @@
 }
 
 void JniCompilerTest::UpcallReturnTypeChecking_InstanceImpl() {
+  // Set debuggable so that the JNI compiler does not emit a fast-path that would skip the
+  // runtime call where we do these checks. Note that while normal gtests use the debug build
+  // which disables the fast path, `art_standalone_compiler_tests` run in the release build.
+  compiler_options_->SetDebuggable(true);
   SetUpForTest(false, "instanceMethodThatShouldReturnClass", "()Ljava/lang/Class;",
                CURRENT_JNI_WRAPPER(Java_MyClassNatives_instanceMethodThatShouldReturnClass));
 
@@ -1580,6 +1585,10 @@
 JNI_TEST(UpcallReturnTypeChecking_Instance)
 
 void JniCompilerTest::UpcallReturnTypeChecking_StaticImpl() {
+  // Set debuggable so that the JNI compiler does not emit a fast-path that would skip the
+  // runtime call where we do these checks. Note that while normal gtests use the debug build
+  // which disables the fast path, `art_standalone_compiler_tests` run in the release build.
+  compiler_options_->SetDebuggable(true);
   SetUpForTest(true, "staticMethodThatShouldReturnClass", "()Ljava/lang/Class;",
                CURRENT_JNI_WRAPPER(Java_MyClassNatives_staticMethodThatShouldReturnClass));
 
diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index 58d11ae..c60d974 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -70,6 +70,12 @@
                                ManagedRegister in_reg);
 
 template <PointerSize kPointerSize>
+static void CallDecodeReferenceResult(JNIMacroAssembler<kPointerSize>* jni_asm,
+                                      JniCallingConvention* jni_conv,
+                                      ManagedRegister mr_return_reg,
+                                      size_t main_out_arg_size);
+
+template <PointerSize kPointerSize>
 static std::unique_ptr<JNIMacroAssembler<kPointerSize>> GetMacroAssembler(
     ArenaAllocator* allocator, InstructionSet isa, const InstructionSetFeatures* features) {
   return JNIMacroAssembler<kPointerSize>::Create(allocator, isa, features);
@@ -103,13 +109,17 @@
   // i.e. if the method was annotated with @CriticalNative
   const bool is_critical_native = (access_flags & kAccCriticalNative) != 0u;
 
-  bool needs_entry_exit_hooks =
-      compiler_options.GetDebuggable() && compiler_options.IsJitCompiler();
+  bool is_debuggable = compiler_options.GetDebuggable();
+  bool needs_entry_exit_hooks = is_debuggable && compiler_options.IsJitCompiler();
   // We don't support JITing stubs for critical native methods in debuggable runtimes yet.
   // TODO(mythria): Add support required for calling method entry / exit hooks from critical native
   // methods.
   DCHECK_IMPLIES(needs_entry_exit_hooks, !is_critical_native);
 
+  // The fast-path for decoding a reference skips CheckJNI checks, so we do not inline the
+  // decoding in debug build or for debuggable apps (both cases enable CheckJNI by default).
+  bool inline_decode_reference = !kIsDebugBuild && !is_debuggable;
+
   // When  walking the stack the top frame doesn't have a pc associated with it. We then depend on
   // the invariant that we don't have JITed code when AOT code is available. In debuggable runtimes
   // this invariant doesn't hold. So we tag the SP for JITed code to indentify if we are executing
@@ -473,8 +483,7 @@
     __ Bind(transition_to_runnable_resume.get());
   }
 
-  // 5.2. For methods that return a reference, do an early exception check so that the
-  //      `JniDecodeReferenceResult()` in the main path does not need to check for exceptions.
+  // 5.2. For methods that return a reference, do an exception check before decoding the reference.
   std::unique_ptr<JNIMacroLabel> exception_slow_path =
       LIKELY(!is_critical_native) ? __ CreateLabel() : nullptr;
   if (reference_return) {
@@ -493,23 +502,23 @@
     __ Bind(suspend_check_resume.get());
   }
 
-  // 5.4 For methods with reference return, decode the `jobject` with `JniDecodeReferenceResult()`.
+  // 5.4 For methods with reference return, decode the `jobject`, either directly
+  //     or with a call to `JniDecodeReferenceResult()`.
+  std::unique_ptr<JNIMacroLabel> decode_reference_slow_path;
+  std::unique_ptr<JNIMacroLabel> decode_reference_resume;
   if (reference_return) {
     DCHECK(!is_critical_native);
-    // We abuse the JNI calling convention here, that is guaranteed to support passing
-    // two pointer arguments, `JNIEnv*` and `jclass`/`jobject`.
-    main_jni_conv->ResetIterator(FrameOffset(main_out_arg_size));
-    ThreadOffset<kPointerSize> jni_decode_reference_result =
-        QUICK_ENTRYPOINT_OFFSET(kPointerSize, pJniDecodeReferenceResult);
-    // Pass result.
-    SetNativeParameter(jni_asm.get(), main_jni_conv.get(), mr_conv->ReturnRegister());
-    main_jni_conv->Next();
-    if (main_jni_conv->IsCurrentParamInRegister()) {
-      __ GetCurrentThread(main_jni_conv->CurrentParamRegister());
-      __ Call(main_jni_conv->CurrentParamRegister(), Offset(jni_decode_reference_result));
+    if (inline_decode_reference) {
+      // Decode local and JNI transition references in the main path.
+      decode_reference_slow_path = __ CreateLabel();
+      decode_reference_resume = __ CreateLabel();
+      __ DecodeJNITransitionOrLocalJObject(mr_conv->ReturnRegister(),
+                                           decode_reference_slow_path.get(),
+                                           decode_reference_resume.get());
+      __ Bind(decode_reference_resume.get());
     } else {
-      __ GetCurrentThread(main_jni_conv->CurrentParamStackOffset());
-      __ CallFromThread(jni_decode_reference_result);
+      CallDecodeReferenceResult<kPointerSize>(
+          jni_asm.get(), main_jni_conv.get(), mr_conv->ReturnRegister(), main_out_arg_size);
     }
   }  // if (!is_critical_native)
 
@@ -639,27 +648,7 @@
     __ Jump(transition_to_runnable_resume.get());
   }
 
-  // 8.4. Suspend check slow path.
-  if (UNLIKELY(is_fast_native)) {
-    __ Bind(suspend_check_slow_path.get());
-    if (reference_return && main_out_arg_size != 0) {
-      jni_asm->cfi().AdjustCFAOffset(main_out_arg_size);
-      __ DecreaseFrameSize(main_out_arg_size);
-    }
-    __ CallFromThread(QUICK_ENTRYPOINT_OFFSET(kPointerSize, pTestSuspend));
-    if (reference_return) {
-      // Suspend check entry point overwrites top of managed stack and leaves it clobbered.
-      // We need to restore the top for subsequent runtime call to `JniDecodeReferenceResult()`.
-      __ StoreStackPointerToThread(Thread::TopOfManagedStackOffset<kPointerSize>(), should_tag_sp);
-    }
-    if (reference_return && main_out_arg_size != 0) {
-      __ IncreaseFrameSize(main_out_arg_size);
-      jni_asm->cfi().AdjustCFAOffset(-main_out_arg_size);
-    }
-    __ Jump(suspend_check_resume.get());
-  }
-
-  // 8.5. Exception poll slow path(s).
+  // 8.4. Exception poll slow path(s).
   if (LIKELY(!is_critical_native)) {
     __ Bind(exception_slow_path.get());
     if (reference_return) {
@@ -675,7 +664,43 @@
     __ DeliverPendingException();
   }
 
-  // 8.6. Method entry / exit hooks slow paths.
+  // 8.5 Slow path for decoding the `jobject`.
+  if (reference_return && inline_decode_reference) {
+    __ Bind(decode_reference_slow_path.get());
+    if (main_out_arg_size != 0) {
+      jni_asm->cfi().AdjustCFAOffset(main_out_arg_size);
+    }
+    CallDecodeReferenceResult<kPointerSize>(
+        jni_asm.get(), main_jni_conv.get(), mr_conv->ReturnRegister(), main_out_arg_size);
+    __ Jump(decode_reference_resume.get());
+    if (main_out_arg_size != 0) {
+      jni_asm->cfi().AdjustCFAOffset(-main_out_arg_size);
+    }
+  }
+
+  // 8.6. Suspend check slow path.
+  if (UNLIKELY(is_fast_native)) {
+    __ Bind(suspend_check_slow_path.get());
+    if (reference_return && main_out_arg_size != 0) {
+      jni_asm->cfi().AdjustCFAOffset(main_out_arg_size);
+      __ DecreaseFrameSize(main_out_arg_size);
+    }
+    __ CallFromThread(QUICK_ENTRYPOINT_OFFSET(kPointerSize, pTestSuspend));
+    if (reference_return) {
+      // Suspend check entry point overwrites top of managed stack and leaves it clobbered.
+      // We need to restore the top for subsequent runtime call to `JniDecodeReferenceResult()`.
+      __ StoreStackPointerToThread(Thread::TopOfManagedStackOffset<kPointerSize>(), should_tag_sp);
+    }
+    if (reference_return && main_out_arg_size != 0) {
+      __ IncreaseFrameSize(main_out_arg_size);
+    }
+    __ Jump(suspend_check_resume.get());
+    if (reference_return && main_out_arg_size != 0) {
+      jni_asm->cfi().AdjustCFAOffset(-main_out_arg_size);
+    }
+  }
+
+  // 8.7. Method entry / exit hooks slow paths.
   if (UNLIKELY(needs_entry_exit_hooks)) {
     __ Bind(method_entry_hook_slow_path.get());
     // Use Jni specific method entry hook that saves all the arguments. We have only saved the
@@ -757,6 +782,31 @@
   }
 }
 
+template <PointerSize kPointerSize>
+static void CallDecodeReferenceResult(JNIMacroAssembler<kPointerSize>* jni_asm,
+                                      JniCallingConvention* jni_conv,
+                                      ManagedRegister mr_return_reg,
+                                      size_t main_out_arg_size) {
+  // We abuse the JNI calling convention here, that is guaranteed to support passing
+  // two pointer arguments, `JNIEnv*` and `jclass`/`jobject`.
+  jni_conv->ResetIterator(FrameOffset(main_out_arg_size));
+  ThreadOffset<kPointerSize> jni_decode_reference_result =
+      QUICK_ENTRYPOINT_OFFSET(kPointerSize, pJniDecodeReferenceResult);
+  // Pass result.
+  SetNativeParameter(jni_asm, jni_conv, mr_return_reg);
+  jni_conv->Next();
+  if (jni_conv->IsCurrentParamInRegister()) {
+    __ GetCurrentThread(jni_conv->CurrentParamRegister());
+    __ Call(jni_conv->CurrentParamRegister(), Offset(jni_decode_reference_result));
+  } else {
+    __ GetCurrentThread(jni_conv->CurrentParamStackOffset());
+    __ CallFromThread(jni_decode_reference_result);
+  }
+  // Note: If the native ABI returns the pointer in a register different from
+  // `mr_return_register`, the `JniDecodeReferenceResult` entrypoint must be
+  // a stub that moves the result to `mr_return_register`.
+}
+
 JniCompiledMethod ArtQuickJniCompileMethod(const CompilerOptions& compiler_options,
                                            uint32_t access_flags,
                                            uint32_t method_idx,
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
index 394575c..5487345 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
@@ -20,6 +20,7 @@
 #include <type_traits>
 
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "indirect_reference_table.h"
 #include "lock_word.h"
 #include "thread.h"
 
@@ -845,6 +846,21 @@
   }
 }
 
+void ArmVIXLJNIMacroAssembler::DecodeJNITransitionOrLocalJObject(ManagedRegister mreg,
+                                                                 JNIMacroLabel* slow_path,
+                                                                 JNIMacroLabel* resume) {
+  constexpr uint32_t kGlobalOrWeakGlobalMask =
+      dchecked_integral_cast<uint32_t>(IndirectReferenceTable::GetGlobalOrWeakGlobalMask());
+  constexpr uint32_t kIndirectRefKindMask =
+      dchecked_integral_cast<uint32_t>(IndirectReferenceTable::GetIndirectRefKindMask());
+  vixl32::Register reg = AsVIXLRegister(mreg.AsArm());
+  ___ Tst(reg, kGlobalOrWeakGlobalMask);
+  ___ B(ne, ArmVIXLJNIMacroLabel::Cast(slow_path)->AsArm());
+  ___ Bics(reg, reg, kIndirectRefKindMask);
+  ___ B(eq, ArmVIXLJNIMacroLabel::Cast(resume)->AsArm());  // Skip load for null.
+  ___ Ldr(reg, MemOperand(reg));
+}
+
 void ArmVIXLJNIMacroAssembler::VerifyObject(ManagedRegister src ATTRIBUTE_UNUSED,
                                             bool could_be_null ATTRIBUTE_UNUSED) {
   // TODO: not validating references.
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
index 7f8fd0a..f6df7f2 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
@@ -91,6 +91,11 @@
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
+  // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
+  void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                         JNIMacroLabel* slow_path,
+                                         JNIMacroLabel* resume) override;
+
   // Heap::VerifyObject on src. In some cases (such as a reference to this) we
   // know that src may not be null.
   void VerifyObject(ManagedRegister src, bool could_be_null) override;
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.cc b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
index 807f493..9e9f122 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.cc
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
@@ -17,6 +17,7 @@
 #include "jni_macro_assembler_arm64.h"
 
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "indirect_reference_table.h"
 #include "lock_word.h"
 #include "managed_register_arm64.h"
 #include "offsets.h"
@@ -690,6 +691,19 @@
   ___ Str(scratch, MEM_OP(reg_x(SP), out_off.Int32Value()));
 }
 
+void Arm64JNIMacroAssembler::DecodeJNITransitionOrLocalJObject(ManagedRegister m_reg,
+                                                               JNIMacroLabel* slow_path,
+                                                               JNIMacroLabel* resume) {
+  constexpr uint64_t kGlobalOrWeakGlobalMask = IndirectReferenceTable::GetGlobalOrWeakGlobalMask();
+  constexpr uint64_t kIndirectRefKindMask = IndirectReferenceTable::GetIndirectRefKindMask();
+  constexpr size_t kGlobalOrWeakGlobalBit = WhichPowerOf2(kGlobalOrWeakGlobalMask);
+  Register reg = reg_w(m_reg.AsArm64().AsWRegister());
+  ___ Tbnz(reg.X(), kGlobalOrWeakGlobalBit, Arm64JNIMacroLabel::Cast(slow_path)->AsArm64());
+  ___ And(reg.X(), reg.X(), ~kIndirectRefKindMask);
+  ___ Cbz(reg.X(), Arm64JNIMacroLabel::Cast(resume)->AsArm64());  // Skip load for null.
+  ___ Ldr(reg, MEM_OP(reg.X()));
+}
+
 void Arm64JNIMacroAssembler::TryToTransitionFromRunnableToNative(
     JNIMacroLabel* label, ArrayRef<const ManagedRegister> scratch_regs ATTRIBUTE_UNUSED) {
   constexpr uint32_t kNativeStateValue = Thread::StoredThreadStateValue(ThreadState::kNative);
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.h b/compiler/utils/arm64/jni_macro_assembler_arm64.h
index 3e6a23d..2836e09 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.h
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.h
@@ -93,6 +93,11 @@
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
+  // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
+  void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                         JNIMacroLabel* slow_path,
+                                         JNIMacroLabel* resume) override;
+
   // Heap::VerifyObject on src. In some cases (such as a reference to this) we
   // know that src may not be null.
   void VerifyObject(ManagedRegister src, bool could_be_null) override;
diff --git a/compiler/utils/jni_macro_assembler.h b/compiler/utils/jni_macro_assembler.h
index 15a4c3f..0c72970 100644
--- a/compiler/utils/jni_macro_assembler.h
+++ b/compiler/utils/jni_macro_assembler.h
@@ -158,6 +158,11 @@
   virtual void GetCurrentThread(ManagedRegister dest) = 0;
   virtual void GetCurrentThread(FrameOffset dest_offset) = 0;
 
+  // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
+  virtual void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                                 JNIMacroLabel* slow_path,
+                                                 JNIMacroLabel* resume) = 0;
+
   // Heap::VerifyObject on src. In some cases (such as a reference to this) we
   // know that src may not be null.
   virtual void VerifyObject(ManagedRegister src, bool could_be_null) = 0;
diff --git a/compiler/utils/x86/jni_macro_assembler_x86.cc b/compiler/utils/x86/jni_macro_assembler_x86.cc
index 40fdc50..154e50b 100644
--- a/compiler/utils/x86/jni_macro_assembler_x86.cc
+++ b/compiler/utils/x86/jni_macro_assembler_x86.cc
@@ -18,6 +18,7 @@
 
 #include "base/casts.h"
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "indirect_reference_table.h"
 #include "lock_word.h"
 #include "thread.h"
 #include "utils/assembler.h"
@@ -391,6 +392,20 @@
   __ movl(Address(ESP, out_off), scratch);
 }
 
+void X86JNIMacroAssembler::DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                                             JNIMacroLabel* slow_path,
+                                                             JNIMacroLabel* resume) {
+  constexpr uint32_t kGlobalOrWeakGlobalMask =
+      dchecked_integral_cast<uint32_t>(IndirectReferenceTable::GetGlobalOrWeakGlobalMask());
+  constexpr uint32_t kIndirectRefKindMask =
+      dchecked_integral_cast<uint32_t>(IndirectReferenceTable::GetIndirectRefKindMask());
+  __ testl(reg.AsX86().AsCpuRegister(), Immediate(kGlobalOrWeakGlobalMask));
+  __ j(kNotZero, X86JNIMacroLabel::Cast(slow_path)->AsX86());
+  __ andl(reg.AsX86().AsCpuRegister(), Immediate(~kIndirectRefKindMask));
+  __ j(kZero, X86JNIMacroLabel::Cast(resume)->AsX86());  // Skip load for null.
+  __ movl(reg.AsX86().AsCpuRegister(), Address(reg.AsX86().AsCpuRegister(), /*disp=*/ 0));
+}
+
 void X86JNIMacroAssembler::VerifyObject(ManagedRegister /*src*/, bool /*could_be_null*/) {
   // TODO: not validating references
 }
diff --git a/compiler/utils/x86/jni_macro_assembler_x86.h b/compiler/utils/x86/jni_macro_assembler_x86.h
index c5e8ad5..6b177f5 100644
--- a/compiler/utils/x86/jni_macro_assembler_x86.h
+++ b/compiler/utils/x86/jni_macro_assembler_x86.h
@@ -88,6 +88,11 @@
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
+  // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
+  void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                         JNIMacroLabel* slow_path,
+                                         JNIMacroLabel* resume) override;
+
   // Heap::VerifyObject on src. In some cases (such as a reference to this) we
   // know that src may not be null.
   void VerifyObject(ManagedRegister src, bool could_be_null) override;
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 e552d29..3888457 100644
--- a/compiler/utils/x86_64/jni_macro_assembler_x86_64.cc
+++ b/compiler/utils/x86_64/jni_macro_assembler_x86_64.cc
@@ -19,6 +19,7 @@
 #include "base/casts.h"
 #include "base/memory_region.h"
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "indirect_reference_table.h"
 #include "lock_word.h"
 #include "thread.h"
 
@@ -464,6 +465,19 @@
   __ movq(Address(CpuRegister(RSP), out_off), scratch);
 }
 
+void X86_64JNIMacroAssembler::DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                                                JNIMacroLabel* slow_path,
+                                                                JNIMacroLabel* resume) {
+  constexpr uint64_t kGlobalOrWeakGlobalMask = IndirectReferenceTable::GetGlobalOrWeakGlobalMask();
+  constexpr uint64_t kIndirectRefKindMask = IndirectReferenceTable::GetIndirectRefKindMask();
+  // TODO: Add `testq()` with `imm32` to assembler to avoid using 64-bit pointer as 32-bit value.
+  __ testl(reg.AsX86_64().AsCpuRegister(), Immediate(kGlobalOrWeakGlobalMask));
+  __ j(kNotZero, X86_64JNIMacroLabel::Cast(slow_path)->AsX86_64());
+  __ andq(reg.AsX86_64().AsCpuRegister(), Immediate(~kIndirectRefKindMask));
+  __ j(kZero, X86_64JNIMacroLabel::Cast(resume)->AsX86_64());  // Skip load for null.
+  __ movl(reg.AsX86_64().AsCpuRegister(), Address(reg.AsX86_64().AsCpuRegister(), /*disp=*/ 0));
+}
+
 void X86_64JNIMacroAssembler::VerifyObject(ManagedRegister /*src*/, bool /*could_be_null*/) {
   // TODO: not validating references
 }
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 2c1fc35..da0aef9 100644
--- a/compiler/utils/x86_64/jni_macro_assembler_x86_64.h
+++ b/compiler/utils/x86_64/jni_macro_assembler_x86_64.h
@@ -89,6 +89,11 @@
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
+  // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
+  void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
+                                         JNIMacroLabel* slow_path,
+                                         JNIMacroLabel* resume) override;
+
   // Heap::VerifyObject on src. In some cases (such as a reference to this) we
   // know that src may not be null.
   void VerifyObject(ManagedRegister src, bool could_be_null) override;
diff --git a/runtime/indirect_reference_table.h b/runtime/indirect_reference_table.h
index 9773f15..ded6bcd 100644
--- a/runtime/indirect_reference_table.h
+++ b/runtime/indirect_reference_table.h
@@ -243,6 +243,10 @@
         reinterpret_cast<uintptr_t>(iref) & ~static_cast<uintptr_t>(kKindMask));
   }
 
+  static constexpr uintptr_t GetIndirectRefKindMask() {
+    return kKindMask;
+  }
+
   /* Reference validation for CheckJNI. */
   bool IsValidReference(IndirectRef, /*out*/std::string* error_msg) const
       REQUIRES_SHARED(Locks::mutator_lock_);