[GR-52576] Introduce specialized upcalls for direct method handles.

PullRequest: graal/19448

Author: fangerer

docs/reference-manual/native-image/ForeignInterface.md

@@ -30,14 +30,18 @@ These two kinds of calls are referred to as "downcalls" and "upcalls" respective
 ### Looking Up Native Functions
 
 The FFM API provides the `SymbolLookup` interface to find functions in native libraries by name.
-`SymbolLookup.loaderLookup()` is currently the only supported kind of `SymbolLookup`.
+Native image supports all available symbol lookup methods, i.e., `SymbolLookup.loaderLookup()`, `SymbolLookup.libraryLookup()`, and `Linker.defaultLookup()`.
 
 ### Registering Foreign Calls
 
 In order to perform calls to native code at run time, supporting code must be generated at image build time.
-Therefore, the `native-image` tool must be provided with descriptors that characterize the functions to which downcalls may be performed at run time.
+Therefore, the `native-image` tool must be provided with descriptors that characterize the functions with which downcalls or upcalls can be performed at runtime.
 
-These descriptors can be registered using a custom `Feature`, for example:
+For upcalls, it is recommended to register a specific static method as an upcall target by providing its declaring class and the method name.
+This allows `native-image` to create specialized upcall code that can be orders of magnitude faster than a upcall registered only by function descriptor.
+Whenever possible, this should be the preferred way to register upcalls.
+
+Descriptors and target methods can be registered using a custom `Feature`, for example:
 ```java
 import static java.lang.foreign.ValueLayout.*;
 
@@ -50,6 +54,9 @@ class ForeignRegistrationFeature implements Feature {
     RuntimeForeignAccess.registerForUpcall(FunctionDescriptor.of(JAVA_INT, JAVA_INT, JAVA_INT));
     RuntimeForeignAccess.registerForDowncall(FunctionDescriptor.of(ADDRESS, JAVA_INT, JAVA_INT), Linker.Option.firstVariadicArg(1));
     RuntimeForeignAccess.registerForDowncall(FunctionDescriptor.ofVoid(JAVA_INT), Linker.Option.captureCallState("errno"));
+
+    MethodHandle target = MethodHandles.lookup().findStatic(UserClass.class, "aStaticMethod", MethodType.of(int.class, int.class, int.class));
+    RuntimeForeignAccess.registerForUpcall(target, FunctionDescriptor.of(JAVA_INT, JAVA_INT, JAVA_INT));
   }
 }
 ```

sdk/src/org.graalvm.nativeimage/snapshot.sigtest

@@ -1088,6 +1088,7 @@ meth public !varargs static void initializeAtRunTime(java.lang.String[])
 supr java.lang.Object
 
 CLSS public final org.graalvm.nativeimage.hosted.RuntimeForeignAccess
+meth public !varargs static void registerForDirectUpcall(java.lang.invoke.MethodHandle,java.lang.Object,java.lang.Object[])
 meth public !varargs static void registerForDowncall(java.lang.Object,java.lang.Object[])
 meth public !varargs static void registerForUpcall(java.lang.Object,java.lang.Object[])
 supr java.lang.Object

sdk/src/org.graalvm.nativeimage/src/org/graalvm/nativeimage/hosted/RuntimeForeignAccess.java

@@ -40,6 +40,9 @@
  */
 package org.graalvm.nativeimage.hosted;
 
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodType;
+
 import org.graalvm.nativeimage.ImageSingletons;
 import org.graalvm.nativeimage.Platform;
 import org.graalvm.nativeimage.Platforms;
@@ -86,6 +89,34 @@ public static void registerForUpcall(Object desc, Object... options) {
         ImageSingletons.lookup(RuntimeForeignAccessSupport.class).registerForUpcall(ConfigurationCondition.alwaysTrue(), desc, options);
     }
 
+    /**
+     * Registers a specific static method (denoted by a method handle) as a fast upcall target. This
+     * will create a specialized upcall stub that will invoke only the specified method, which is
+     * much faster than using {@link #registerForUpcall(Object, Object...)}).
+     * <p>
+     * The provided method handle must be a direct method handle. Those are most commonly created
+     * using {@link java.lang.invoke.MethodHandles.Lookup#findStatic(Class, String, MethodType)}.
+     * However, a strict requirement is that it must be possible to create a non-empty descriptor
+     * for the method handle using {@link MethodHandle#describeConstable()}. The denoted static
+     * method will also be registered for reflective access since run-time code will also create a
+     * method handle to denoted static method.
+     * </p>
+     * <p>
+     * Even though this method is weakly typed for compatibility reasons, runtime checks will be
+     * performed to ensure that the arguments have the expected type. It will be deprecated in favor
+     * of strongly typed variant as soon as possible.
+     * </p>
+     *
+     * @param target A direct method handle denoting a static method.
+     * @param desc A {@link java.lang.foreign.FunctionDescriptor} to register for upcalls.
+     * @param options An array of {@link java.lang.foreign.Linker.Option} used for the upcalls.
+     *
+     * @since 24.2
+     */
+    public static void registerForDirectUpcall(MethodHandle target, Object desc, Object... options) {
+        ImageSingletons.lookup(RuntimeForeignAccessSupport.class).registerForDirectUpcall(ConfigurationCondition.alwaysTrue(), target, desc, options);
+    }
+
     private RuntimeForeignAccess() {
     }
 }

sdk/src/org.graalvm.nativeimage/src/org/graalvm/nativeimage/impl/RuntimeForeignAccessSupport.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2023, 2024, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * The Universal Permissive License (UPL), Version 1.0
@@ -40,8 +40,12 @@
  */
 package org.graalvm.nativeimage.impl;
 
+import java.lang.invoke.MethodHandle;
+
 public interface RuntimeForeignAccessSupport {
     void registerForDowncall(ConfigurationCondition condition, Object desc, Object... options);
 
     void registerForUpcall(ConfigurationCondition condition, Object desc, Object... options);
+
+    void registerForDirectUpcall(ConfigurationCondition condition, MethodHandle target, Object desc, Object... options);
 }

substratevm/CHANGELOG.md

@@ -10,6 +10,7 @@ At runtime, premain runtime options are set along with main class' arguments in
   The warning is planned to be replaced by an error in GraalVM for JDK 25.
 * (GR-48384) Added a GDB Python script (`gdb-debughelpers.py`) to improve the Native Image debugging experience.
 * (GR-49517) Add support for emitting Windows x64 unwind info. This enables stack walking in native tooling such as debuggers and profilers.
+* (GR-52576) Optimize FFM API upcalls for specifiable static upcall target methods.
 * (GR-56599) Update native image debuginfo from DWARF4 to DWARF5 and store type information for debugging in DWARF type units.
 * (GR-56601) Together with Red Hat, we added experimental support for `jcmd` on Linux and macOS. Add `--enable-monitoring=jcmd` to your build arguments to try it out.
 * (GR-57384) Preserve the origin of a resource included in a native image. The information is included in the report produced by -H:+GenerateEmbeddedResourcesFile.

substratevm/mx.substratevm/suite.py

@@ -719,13 +719,13 @@
             ],
             "requiresConcealed": {
                 "java.base": [
-                    "jdk.internal.loader",
-                    "jdk.internal.reflect",
                     "jdk.internal.foreign",
                     "jdk.internal.foreign.abi",
                     "jdk.internal.foreign.abi.x64",
                     "jdk.internal.foreign.abi.x64.sysv",
                     "jdk.internal.foreign.abi.x64.windows",
+                    "jdk.internal.loader",
+                    "jdk.internal.reflect",
                 ],
                 "jdk.internal.vm.ci" : [
                     "jdk.vm.ci.amd64",
@@ -758,6 +758,9 @@
                 "java.base": [
                     "jdk.internal.foreign",
                     "jdk.internal.foreign.abi",
+                    "jdk.internal.foreign.abi.x64.windows",
+                    "jdk.internal.foreign.abi.x64.sysv",
+                    "jdk.internal.foreign.layout",
                 ],
                 "jdk.internal.vm.ci" : [
                     "jdk.vm.ci.code",

substratevm/src/com.oracle.svm.core.foreign/src/com/oracle/svm/core/foreign/ForeignFunctionsRuntime.java

@@ -26,9 +26,11 @@
 
 import static jdk.graal.compiler.core.common.spi.ForeignCallDescriptor.CallSideEffect.HAS_SIDE_EFFECT;
 
+import java.lang.constant.DirectMethodHandleDesc;
 import java.lang.invoke.MethodHandle;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.BiConsumer;
 
 import org.graalvm.collections.EconomicMap;
 import org.graalvm.nativeimage.ImageSingletons;
@@ -61,11 +63,15 @@ public static ForeignFunctionsRuntime singleton() {
 
     private final AbiUtils.TrampolineTemplate trampolineTemplate = AbiUtils.singleton().generateTrampolineTemplate();
     private final EconomicMap<NativeEntryPointInfo, FunctionPointerHolder> downcallStubs = EconomicMap.create();
+    private final EconomicMap<DirectMethodHandleDesc, FunctionPointerHolder> directUpcallStubs = EconomicMap.create();
     private final EconomicMap<JavaEntryPointInfo, FunctionPointerHolder> upcallStubs = EconomicMap.create();
 
     private final Map<Long, TrampolineSet> trampolines = new HashMap<>();
     private TrampolineSet currentTrampolineSet;
 
+    // for testing: callback if direct upcall lookup succeeded
+    private BiConsumer<Long, DirectMethodHandleDesc> usingSpecializedUpcallListener;
+
     @Platforms(Platform.HOSTED_ONLY.class)
     public ForeignFunctionsRuntime() {
     }
@@ -82,6 +88,12 @@ public void addUpcallStubPointer(JavaEntryPointInfo jep, CFunctionPointer ptr) {
         VMError.guarantee(upcallStubs.put(jep, new FunctionPointerHolder(ptr)) == null);
     }
 
+    @Platforms(Platform.HOSTED_ONLY.class)
+    public void addDirectUpcallStubPointer(DirectMethodHandleDesc desc, CFunctionPointer ptr) {
+        VMError.guarantee(!directUpcallStubs.containsKey(desc), "Seems like multiple stubs were generated for " + desc);
+        VMError.guarantee(directUpcallStubs.put(desc, new FunctionPointerHolder(ptr)) == null);
+    }
+
     /**
      * We'd rather report the function descriptor than the native method type, but we don't have it
      * available here. One could intercept this exception in
@@ -105,15 +117,61 @@ CFunctionPointer getUpcallStubPointer(JavaEntryPointInfo jep) {
     }
 
     Pointer registerForUpcall(MethodHandle methodHandle, JavaEntryPointInfo jep) {
+        /*
+         * Look up the upcall stub pointer first to avoid unnecessary allocation and synchronization
+         * if it doesn't exist.
+         */
+        CFunctionPointer upcallStubPointer = getUpcallStubPointer(jep);
         synchronized (trampolines) {
             if (currentTrampolineSet == null || !currentTrampolineSet.hasFreeTrampolines()) {
                 currentTrampolineSet = new TrampolineSet(trampolineTemplate);
                 trampolines.put(currentTrampolineSet.base().rawValue(), currentTrampolineSet);
             }
-            return currentTrampolineSet.assignTrampoline(methodHandle, getUpcallStubPointer(jep));
+            return currentTrampolineSet.assignTrampoline(methodHandle, upcallStubPointer);
         }
     }
 
+    /**
+     * Updates the stub address in the upcall trampoline with the address of a direct upcall stub.
+     * The trampoline is identified by the given native address and the direct upcall stub is
+     * identified by the method handle descriptor.
+     *
+     * @param trampolineAddress The address of the upcall trampoline.
+     * @param desc A direct method handle descriptor used to lookup the direct upcall stub.
+     */
+    void patchForDirectUpcall(long trampolineAddress, DirectMethodHandleDesc desc) {
+        FunctionPointerHolder functionPointerHolder = directUpcallStubs.get(desc);
+        if (functionPointerHolder == null) {
+            return;
+        }
+
+        Pointer trampolinePointer = WordFactory.pointer(trampolineAddress);
+        Pointer trampolineSetBase = TrampolineSet.getAllocationBase(trampolinePointer);
+        TrampolineSet trampolineSet = trampolines.get(trampolineSetBase.rawValue());
+        if (trampolineSet == null) {
+            return;
+        }
+        /*
+         * Synchronizing on 'trampolineSet' is not necessary at this point since we are still in the
+         * call context of 'Linker.upcallStub' and the allocated trampoline is owned by the
+         * allocating thread until it returns from the call. Also, the trampoline cannot be free'd
+         * between allocation and patching because the associated arena is still on the stack.
+         */
+        trampolineSet.patchTrampolineForDirectUpcall(trampolinePointer, functionPointerHolder.functionPointer);
+        /*
+         * If we reach this point, everything went fine and the trampoline was patched with the
+         * specialized upcall stub's address. For testing, now report that the lookup and patching
+         * succeeded.
+         */
+        if (usingSpecializedUpcallListener != null) {
+            usingSpecializedUpcallListener.accept(trampolineAddress, desc);
+        }
+    }
+
+    public void setUsingSpecializedUpcallListener(BiConsumer<Long, DirectMethodHandleDesc> listener) {
+        usingSpecializedUpcallListener = listener;
+    }
+
     void freeTrampoline(long addr) {
         synchronized (trampolines) {
             long base = TrampolineSet.getAllocationBase(WordFactory.pointer(addr)).rawValue();

substratevm/src/com.oracle.svm.core.foreign/src/com/oracle/svm/core/foreign/Target_jdk_internal_foreign_abi_AbstractLinker.java

@@ -24,11 +24,25 @@
  */
 package com.oracle.svm.core.foreign;
 
+import java.lang.constant.DirectMethodHandleDesc;
+import java.lang.constant.MethodHandleDesc;
+import java.lang.foreign.Arena;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.MemorySegment;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodType;
+import java.util.Optional;
+
 import com.oracle.svm.core.annotate.Alias;
 import com.oracle.svm.core.annotate.RecomputeFieldValue;
+import com.oracle.svm.core.annotate.Substitute;
 import com.oracle.svm.core.annotate.TargetClass;
 
 import jdk.internal.foreign.abi.AbstractLinker;
+import jdk.internal.foreign.abi.AbstractLinker.UpcallStubFactory;
+import jdk.internal.foreign.abi.LinkerOptions;
+import jdk.internal.foreign.abi.x64.sysv.SysVx64Linker;
+import jdk.internal.foreign.abi.x64.windows.Windowsx64Linker;
 
 @TargetClass(AbstractLinker.class)
 public final class Target_jdk_internal_foreign_abi_AbstractLinker {
@@ -46,3 +60,55 @@ public final class Target_jdk_internal_foreign_abi_AbstractLinker {
 @TargetClass(className = "jdk.internal.foreign.abi.SoftReferenceCache")
 final class Target_jdk_internal_foreign_abi_SoftReferenceCache {
 }
+
+/**
+ * A decorator for jdk.internal.foreign.abi.UpcallStubFactory which intercepts the call to method
+ * 'makeStub'. It will (1) call the original factory to create the upcall, and (2) then use the
+ * method handle's descriptor to lookup if a specialized (direct) upcall stub is available. If so,
+ * the trampoline will be updated with the specialized stub's address.
+ * 
+ * @param delegate The original upcall stub factory as created by JDK's call arranger.
+ */
+record UpcallStubFactoryDecorator(UpcallStubFactory delegate) implements UpcallStubFactory {
+
+    @Override
+    public MemorySegment makeStub(MethodHandle target, Arena arena) {
+        MemorySegment segment = delegate.makeStub(target, arena);
+
+        /*
+         * We cannot do this in 'UpcallLinker.makeUpcallStub' because that one already gets a
+         * different method handle that will handle parameter/return value bindings. Further, method
+         * handles cannot be compared. If the provided method handle is a DirectMethodHandle, we use
+         * the MH descriptor to check if there is a registered direct upcall stub. Then, we will
+         * patch the already allocated trampoline with a different upcall stub pointer.
+         */
+        Optional<MethodHandleDesc> methodHandleDesc = target.describeConstable();
+        if (methodHandleDesc.isPresent() && methodHandleDesc.get() instanceof DirectMethodHandleDesc desc) {
+            ForeignFunctionsRuntime.singleton().patchForDirectUpcall(segment.address(), desc);
+        }
+        return segment;
+    }
+}
+
+@TargetClass(value = SysVx64Linker.class, onlyWith = ForeignFunctionsEnabled.class)
+final class Target_jdk_internal_foreign_abi_x64_sysv_SysVx64Linker {
+
+    @Substitute
+    UpcallStubFactory arrangeUpcall(MethodType targetType, FunctionDescriptor function, LinkerOptions options) {
+        return new UpcallStubFactoryDecorator(jdk.internal.foreign.abi.x64.sysv.CallArranger.arrangeUpcall(targetType, function, options));
+    }
+}
+
+@TargetClass(value = Windowsx64Linker.class, onlyWith = ForeignFunctionsEnabled.class)
+final class Target_jdk_internal_foreign_abi_x64_windows_Windowsx64Linker {
+
+    @Substitute
+    UpcallStubFactory arrangeUpcall(MethodType targetType, FunctionDescriptor function, LinkerOptions options) {
+        return new UpcallStubFactoryDecorator(jdk.internal.foreign.abi.x64.windows.CallArranger.arrangeUpcall(targetType, function, options));
+    }
+}
+
+/*
+ * GR-58659, GR-58660: add substitutions for LinuxAArch64Linker and MacOsAArch64Linker here once we
+ * support them.
+ */