Implement JIT optimization for MethodHandle.linkToNative calls

[nbhuiyan]

Note: description currently not up to date with changes in the implementation. This will be updated prior to being marked ready for review.

This commit implements partial support for optimizing MethodHandle.linkToNative invocations by transforming them into direct native function calls, bypassing the J2I transition overhead. Currently, only primitive argument and return types are supported (no structs/arrays).

The linkToNative polymorphic signature method is used by the Foreign Function & Memory API (Project Panama) to invoke native functions through MethodHandles. At the bytecode level, linkToNative receives:

• arg0: native function address as a MemorySegment
• arg1: appendix object from the invoke cache
• args 2..N-1: primitive arguments for the native function
• argN: NativeMethodHandle object

Without this optimization, such invocations would be dispatched through the interpreter, which has a lot more overhead.

Implementation Details:

1. Inlining Support Through InterpreterEmulator Enhancements:

   • Added handling of linkToNative in InterpreterEmulator
   • Extracts native entry point, MemberName, and appendix from the NativeMethodHandle's invoke cache array
   • Reconstructs the operand stack by removing the NativeMethodHandle and inserting the native entry point and appendix in the correct positions
   • Establishes const provenance relationships for known objects

2. MethodHandleTransformer:

   • New method: process_java_lang_invoke_MethodHandle_linkToNative() to handle MethodHandle.linkToNative transformation.
   • Validates that the native function address is a known constant
   • Extracts the raw function pointer from the NativeMemorySegmentImpl
   • Validates and transforms the method signature:
     • Strips the MemorySegment, appendix, and NativeMethodHandle parameters
     • Ensures only primitive types remain (rejects arrays and references) * Constructs a new signature with only primitive parameters and return
   • Creates a new resolved method with the transformed signature
   • Marks the call as a direct native call (canDirectNativeCall=true)
   • Adjusts the call node's children to match the new signature

3. Platform-Specific JNI Linkage Updates:

   • x86-64 (AMD64JNILinkage.cpp): * Disables JNIEnv* passing and reference wrapping * Fixes argument cleanup calculation to account for the base argument size * Adds logging for debugging argument size calculations
   • Z (S390PrivateLinkage.cpp): Disables JNIEnv* passing and JNI reference frame collapsing
   • aarch64 (ARM64JNILinkage.cpp): Disables JNIEnv* passing and reference wrapping for linkToNative calls
   • Power (PPCJNILinkage.cpp): Same adjustments as aarch64
   • i386: as this platform is only used for Java 8, no changes necessary for linkToNative, which was introduced in Java 16

4. CodeGen (common):

   • Modified lowerTreeIfNeeded() to skip dummy null argument insertion for linkToNative calls that have been prepared for direct JNI dispatch
   • Uses isPreparedForDirectJNI() flag to distinguish transformed calls

5. New Recognized Method (and Class):

   • Added java_lang_invoke_NativeMethodHandle_internalNativeEntryPoint to the recognized methods enum (required clang-format assisted reformatting)

6. Method Metadata:

   • Added linkToNativeJNITargetAddress field to TR_ResolvedJ9Method
   • Added setter/getter methods for the native target address
   • Modified startAddressForJNIMethod() to return the stored target address for linkToNative calls
   • TR_ResolvedJ9JITServerMethod's functions isJNINative and startAddressForJNIMethod checks linkToNativeJniTargetAddress field

With the above changes implemented, we achieve the following optimizations for linkToNative calls that consist of primitive-only args and return types:

• Eliminate J2I transition overhead for eligible linkToNative invocations
• Remove JNI frame creation and teardown overhead that's part of the linkToNative INL function in the interpreter
• Bypass JNI exception checking

This is built on top of @jdmpapin's prototype implementation. Currently in WIP state to complete further testing and performance analysis on all platforms, and some refactoring work in the interpreter.

Commit/PR message composed using IBM Bob.

[nbhuiyan]

@0xdaryl @vijaysun-omr @hzongaro @zl-wang FYI

@TobiAjila Please confirm if there are any planned VM refactoring work that will require changes in how the linkToNative optimization is implemented.

[jdmpapin]

Some thoughts. Of course, you might have already considered some or all of this.

It's not obvious to me that we want to refine linkToNative() like that in interpreter emulator. IIRC I did leave behind code that did so, but I think I only wrote it to see what the callee method handle was doing.

AFAICT from a quick look, a case where the inlining kicks in is one where during inlining we could determine the MH but not the native entry point. If we were to leave the linkToNative() call as-is, then it might be possible for the native entry point to be determined later, e.g. by global VP, and if so, then maybe we could do the direct JNI transformation after that. But by inlining the MH, we would preclude that possibility by eliminating the linkToNative() call. (I can't say that I know whether this case is even realistically possible, but if not, then the inlining would be pointless anyway.)

I also worry that inlining is overkill if the goal is just to avoid a J2I transition. Inlining might add a lot of code to the IL for the benefit of a call that's necessarily going to be pretty slow anyway. For the case where we fail to do direct JNI, might it be better to do a direct call transformation in recognized call transformer, like for linkToStatic()? That would also help the case where we fail to determine even the MH, though I don't know whether to expect that case to matter.

[zl-wang]

i haven't read the design yet, but agreed with @jdmpapin from the outset. for the best performance, staying the IL at a level, as long as it allows the codegen realizing (by symRef name recognition or otherwise) to go JNI-Direct. of course, the real callee first needs to pass through some limitation tests (afaik, the only criteria is the number of pass-down arguments doesn't exceed 32 or something like that). if it fails the test, i guessed the existing approach still applies (i.e. through libffi).

one of the future optimizations possible is to be able to query if the callee is suitable for fast-JNI.

if there are upCall(s) involved, the other optimization in the future is to cache upCall thunk (to be used by multiple threads in multiple scopes for example). no need to (re)generate it per thread per call-site per scope (i.e. inefficient).

[zl-wang]

or, ascertaining/investigating: this particular linkToNative situation can always assume to be fast-JNI (i.e. don't build call-out frame, don't drop VMAccess, no GC allowed, don't kill preserved GPRs, don't need to tear down JNI frame -- ref pool left-behind, and no need to check for pending throws).

[tajila]

this particular linkToNative situation can always assume to be fast-JNI (i.e. don't build call-out frame, don't drop VMAccess, no GC allowed, don't kill preserved GPRs, don't need to tear down JNI frame -- ref pool left-behind, and no need to check for pending throws).

We cant assume this, we need to operate with the same JNI assumptions where the JNI frame is built and VM access is dropped.

That being said it is possible to allow fastJNI like behaviour where we dont need to drop VMAccess in some cases, however, this is beyond what we initially discussed.

See #23631