8341127: Extra call to MethodHandle::asType from memory segment var handles fails to inline

Reviewed-by: psandoz, redestad, jvernee

src/java.base/share/classes/java/lang/invoke/MethodHandle.java

@@ -27,6 +27,8 @@ package java.lang.invoke;
 
 
 import jdk.internal.loader.ClassLoaders;
+import jdk.internal.vm.annotation.DontInline;
+import jdk.internal.vm.annotation.ForceInline;
 import jdk.internal.vm.annotation.IntrinsicCandidate;
 
 import java.lang.constant.ClassDesc;
@@ -856,6 +858,7 @@ public abstract sealed class MethodHandle implements Constable
      * @throws WrongMethodTypeException if the conversion cannot be made
      * @see MethodHandles#explicitCastArguments
      */
+    @ForceInline
     public final MethodHandle asType(MethodType newType) {
         // Fast path alternative to a heavyweight {@code asType} call.
         // Return 'this' if the conversion will be a no-op.
@@ -867,7 +870,7 @@ public abstract sealed class MethodHandle implements Constable
         if (at != null) {
             return at;
         }
-        return setAsTypeCache(asTypeUncached(newType));
+        return setAsTypeCache(newType);
     }
 
     private MethodHandle asTypeCached(MethodType newType) {
@@ -885,7 +888,16 @@ public abstract sealed class MethodHandle implements Constable
         return null;
     }
 
-    private MethodHandle setAsTypeCache(MethodHandle at) {
+    /*
+     * We disable inlining here to prevent complex code in the slow path
+     * of MethodHandle::asType from being inlined into that method.
+     * Excessive inlining into MethodHandle::asType can cause that method
+     * to become too big, which will then cause performance issues during
+     * var handle and method handle calls.
+     */
+    @DontInline
+    private MethodHandle setAsTypeCache(MethodType newType) {
+        MethodHandle at = asTypeUncached(newType);
         // Don't introduce a strong reference in the cache if newType depends on any class loader other than
         // current method handle already does to avoid class loader leaks.
         if (isSafeToCache(at.type)) {

test/micro/org/openjdk/bench/java/lang/foreign/LoopOverNonConstantAsType.java

@@ -0,0 +1,138 @@
+/*
+ * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package org.openjdk.bench.java.lang.foreign;
+
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.CompilerControl;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import sun.misc.Unsafe;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.reflect.Proxy;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
+
+import static java.lang.foreign.ValueLayout.*;
+
+@BenchmarkMode(Mode.AverageTime)
+@Warmup(iterations = 5, time = 500, timeUnit = TimeUnit.MILLISECONDS)
+@Measurement(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS)
+@State(org.openjdk.jmh.annotations.Scope.Thread)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Fork(value = 3, jvmArgsAppend = { "-XX:-TieredCompilation" })
+public class LoopOverNonConstantAsType extends JavaLayouts {
+
+    static final Unsafe unsafe = Utils.unsafe;
+
+    static final int ELEM_SIZE = 1_000_000;
+    static final int CARRIER_SIZE = (int)JAVA_LONG.byteSize();
+    static final int ALLOC_SIZE = ELEM_SIZE * CARRIER_SIZE;
+
+    @Param({"false", "true"})
+    public boolean asTypeCompiled;
+
+    Arena arena;
+    MemorySegment segment;
+    long unsafe_addr;
+
+    @Setup
+    public void setup() {
+        unsafe_addr = unsafe.allocateMemory(ALLOC_SIZE);
+        for (int i = 0; i < ELEM_SIZE; i++) {
+            unsafe.putInt(unsafe_addr + (i * CARRIER_SIZE) , i);
+        }
+        arena = Arena.ofConfined();
+        segment = arena.allocate(ALLOC_SIZE, 1);
+        for (int i = 0; i < ELEM_SIZE; i++) {
+            VH_INT.set(segment, (long) i, i);
+        }
+        if (asTypeCompiled) {
+            compileAsType();
+        }
+    }
+
+    public interface T { }
+
+    static final int TYPE_SIZE = 100;
+    static final Class<?>[] types;
+
+    static {
+        types = new Class<?>[TYPE_SIZE];
+        ClassLoader customLoader = new URLClassLoader(new URL[0], LoopOverNonConstantAsType.class.getClassLoader());
+        for (int i = 0 ; i < TYPE_SIZE ; i++) {
+            types[i] = Proxy.newProxyInstance(customLoader,
+                    new Class<?>[] { T.class }, (_, _, _) -> null).getClass();
+        }
+    }
+
+    void compileAsType() {
+        for (Class<?> type : types) {
+            MethodHandle handle = MethodHandles.zero(Object.class);
+            Class<?>[] args = new Class<?>[254];
+            Arrays.fill(args, Object.class);
+            handle = MethodHandles.dropArguments(handle, 0, args);
+            for (int j = 0; j < args.length ; j++) {
+                handle = handle.asType(handle.type().changeParameterType(j, type));
+            }
+        }
+    }
+
+    @TearDown
+    public void tearDown() {
+        arena.close();
+        unsafe.freeMemory(unsafe_addr);
+    }
+
+    @Benchmark
+    public long unsafe_loop() {
+        long res = 0;
+        for (int i = 0; i < ELEM_SIZE; i ++) {
+            res += unsafe.getLong(unsafe_addr + (i * CARRIER_SIZE));
+        }
+        return res;
+    }
+
+    @Benchmark
+    public long segment_loop() {
+        long sum = 0;
+        for (int i = 0; i < ELEM_SIZE; i++) {
+            sum += segment.get(JAVA_LONG, i * CARRIER_SIZE);
+        }
+        return sum;
+    }
+}