Skip to content

Commit 0b53670

Browse files
ChrisHegartyuschindlerrmuir
authored
[Backport] Integrate the Incubating Panama Vector API #12311 (#12327)
Leverage accelerated vector hardware instructions in Vector Search. Lucene already has a mechanism that enables the use of non-final JDK APIs, currently used for the Previewing Pamana Foreign API. This change expands this mechanism to include the Incubating Pamana Vector API. When the jdk.incubator.vector module is present at run time the Panamaized version of the low-level primitives used by Vector Search is enabled. If not present, the default scalar version of these low-level primitives is used (as it was previously). Currently, we're only targeting support for JDK 20. A subsequent PR should evaluate JDK 21. --------- Co-authored-by: Uwe Schindler <[email protected]> Co-authored-by: Robert Muir <[email protected]>
1 parent 7db2c12 commit 0b53670

File tree

16 files changed

+1142
-288
lines changed

16 files changed

+1142
-288
lines changed

build.gradle

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ apply from: file('gradle/ide/eclipse.gradle')
119119
// (java, tests)
120120
apply from: file('gradle/java/folder-layout.gradle')
121121
apply from: file('gradle/java/javac.gradle')
122-
apply from: file('gradle/java/memorysegment-mrjar.gradle')
122+
apply from: file('gradle/java/core-mrjar.gradle')
123123
apply from: file('gradle/testing/defaults-tests.gradle')
124124
apply from: file('gradle/testing/randomization.gradle')
125125
apply from: file('gradle/testing/fail-on-no-tests.gradle')
@@ -158,7 +158,7 @@ apply from: file('gradle/generation/javacc.gradle')
158158
apply from: file('gradle/generation/forUtil.gradle')
159159
apply from: file('gradle/generation/antlr.gradle')
160160
apply from: file('gradle/generation/unicode-test-classes.gradle')
161-
apply from: file('gradle/generation/panama-foreign.gradle')
161+
apply from: file('gradle/generation/extract-jdk-apis.gradle')
162162

163163
apply from: file('gradle/datasets/external-datasets.gradle')
164164

gradle/generation/panama-foreign.gradle renamed to gradle/generation/extract-jdk-apis.gradle

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,17 @@
1717

1818
def resources = scriptResources(buildscript)
1919

20+
configure(rootProject) {
21+
ext {
22+
// also change this in extractor tool: ExtractForeignAPI
23+
vectorIncubatorJavaVersions = [ JavaVersion.VERSION_20 ] as Set
24+
}
25+
}
26+
2027
configure(project(":lucene:core")) {
2128
ext {
2229
apijars = file('src/generated/jdk');
23-
panamaJavaVersions = [ 19, 20 ]
30+
mrjarJavaVersions = [ 19, 20 ]
2431
}
2532

2633
configurations {
@@ -31,9 +38,9 @@ configure(project(":lucene:core")) {
3138
apiextractor "org.ow2.asm:asm:${scriptDepVersions['asm']}"
3239
}
3340

34-
for (jdkVersion : panamaJavaVersions) {
35-
def task = tasks.create(name: "generatePanamaForeignApiJar${jdkVersion}", type: JavaExec) {
36-
description "Regenerate the API-only JAR file with public Panama Foreign API from JDK ${jdkVersion}"
41+
for (jdkVersion : mrjarJavaVersions) {
42+
def task = tasks.create(name: "generateJdkApiJar${jdkVersion}", type: JavaExec) {
43+
description "Regenerate the API-only JAR file with public Panama Foreign & Vector API from JDK ${jdkVersion}"
3744
group "generation"
3845

3946
javaLauncher = javaToolchains.launcherFor {
@@ -45,21 +52,21 @@ configure(project(":lucene:core")) {
4552
javaLauncher.get()
4653
return true
4754
} catch (Exception e) {
48-
logger.warn('Launcher for Java {} is not available; skipping regeneration of Panama Foreign API JAR.', jdkVersion)
55+
logger.warn('Launcher for Java {} is not available; skipping regeneration of Panama Foreign & Vector API JAR.', jdkVersion)
4956
logger.warn('Error: {}', e.cause?.message)
5057
logger.warn("Please make sure to point env 'JAVA{}_HOME' to exactly JDK version {} or enable Gradle toolchain auto-download.", jdkVersion, jdkVersion)
5158
return false
5259
}
5360
}
54-
61+
5562
classpath = configurations.apiextractor
56-
mainClass = file("${resources}/ExtractForeignAPI.java") as String
63+
mainClass = file("${resources}/ExtractJdkApis.java") as String
5764
systemProperties = [
5865
'user.timezone': 'UTC'
5966
]
6067
args = [
6168
jdkVersion,
62-
new File(apijars, "panama-foreign-jdk${jdkVersion}.apijar"),
69+
new File(apijars, "jdk${jdkVersion}.apijar"),
6370
]
6471
}
6572

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
import java.io.IOException;
18+
import java.net.URI;
19+
import java.nio.file.Files;
20+
import java.nio.file.Path;
21+
import java.nio.file.PathMatcher;
22+
import java.nio.file.Paths;
23+
import java.nio.file.attribute.FileTime;
24+
import java.time.Instant;
25+
import java.util.Arrays;
26+
import java.util.HashMap;
27+
import java.util.HashSet;
28+
import java.util.List;
29+
import java.util.Map;
30+
import java.util.Objects;
31+
import java.util.Set;
32+
import java.util.TreeMap;
33+
import java.util.function.Predicate;
34+
import java.util.stream.Collectors;
35+
import java.util.stream.Stream;
36+
import java.util.zip.ZipEntry;
37+
import java.util.zip.ZipOutputStream;
38+
39+
import org.objectweb.asm.AnnotationVisitor;
40+
import org.objectweb.asm.ClassReader;
41+
import org.objectweb.asm.ClassVisitor;
42+
import org.objectweb.asm.ClassWriter;
43+
import org.objectweb.asm.FieldVisitor;
44+
import org.objectweb.asm.MethodVisitor;
45+
import org.objectweb.asm.Opcodes;
46+
import org.objectweb.asm.Type;
47+
48+
public final class ExtractJdkApis {
49+
50+
private static final FileTime FIXED_FILEDATE = FileTime.from(Instant.parse("2022-01-01T00:00:00Z"));
51+
52+
private static final String PATTERN_PANAMA_FOREIGN = "java.base/java/{lang/foreign/*,nio/channels/FileChannel,util/Objects}";
53+
private static final String PATTERN_VECTOR_INCUBATOR = "jdk.incubator.vector/jdk/incubator/vector/*";
54+
private static final String PATTERN_VECTOR_VM_INTERNALS = "java.base/jdk/internal/vm/vector/VectorSupport{,$Vector,$VectorMask,$VectorPayload,$VectorShuffle}";
55+
56+
static final Map<Integer,List<String>> CLASSFILE_PATTERNS = Map.of(
57+
19, List.of(PATTERN_PANAMA_FOREIGN),
58+
20, List.of(PATTERN_PANAMA_FOREIGN, PATTERN_VECTOR_VM_INTERNALS, PATTERN_VECTOR_INCUBATOR)
59+
);
60+
61+
public static void main(String... args) throws IOException {
62+
if (args.length != 2) {
63+
throw new IllegalArgumentException("Need two parameters: java version, output file");
64+
}
65+
Integer jdk = Integer.valueOf(args[0]);
66+
if (jdk.intValue() != Runtime.version().feature()) {
67+
throw new IllegalStateException("Incorrect java version: " + Runtime.version().feature());
68+
}
69+
if (!CLASSFILE_PATTERNS.containsKey(jdk)) {
70+
throw new IllegalArgumentException("No support to extract stubs from java version: " + jdk);
71+
}
72+
var outputPath = Paths.get(args[1]);
73+
74+
// create JRT filesystem and build a combined FileMatcher:
75+
var jrtPath = Paths.get(URI.create("jrt:/")).toRealPath();
76+
var patterns = CLASSFILE_PATTERNS.get(jdk).stream()
77+
.map(pattern -> jrtPath.getFileSystem().getPathMatcher("glob:" + pattern + ".class"))
78+
.toArray(PathMatcher[]::new);
79+
PathMatcher pattern = p -> Arrays.stream(patterns).anyMatch(matcher -> matcher.matches(p));
80+
81+
// Collect all files to process:
82+
final List<Path> filesToExtract;
83+
try (var stream = Files.walk(jrtPath)) {
84+
filesToExtract = stream.filter(p -> pattern.matches(jrtPath.relativize(p))).collect(Collectors.toList());
85+
}
86+
87+
// Process all class files:
88+
try (var out = new ZipOutputStream(Files.newOutputStream(outputPath))) {
89+
process(filesToExtract, out);
90+
}
91+
}
92+
93+
private static void process(List<Path> filesToExtract, ZipOutputStream out) throws IOException {
94+
var classesToInclude = new HashSet<String>();
95+
var references = new HashMap<String, String[]>();
96+
var processed = new TreeMap<String, byte[]>();
97+
System.out.println("Transforming " + filesToExtract.size() + " class files...");
98+
for (Path p : filesToExtract) {
99+
try (var in = Files.newInputStream(p)) {
100+
var reader = new ClassReader(in);
101+
var cw = new ClassWriter(0);
102+
var cleaner = new Cleaner(cw, classesToInclude, references);
103+
reader.accept(cleaner, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
104+
processed.put(reader.getClassName(), cw.toByteArray());
105+
}
106+
}
107+
// recursively add all superclasses / interfaces of visible classes to classesToInclude:
108+
for (Set<String> a = classesToInclude; !a.isEmpty();) {
109+
a = a.stream().map(references::get).filter(Objects::nonNull).flatMap(Arrays::stream).collect(Collectors.toSet());
110+
classesToInclude.addAll(a);
111+
}
112+
// remove all non-visible or not referenced classes:
113+
processed.keySet().removeIf(Predicate.not(classesToInclude::contains));
114+
System.out.println("Writing " + processed.size() + " visible classes...");
115+
for (var cls : processed.entrySet()) {
116+
String cn = cls.getKey();
117+
System.out.println("Writing stub for class: " + cn);
118+
out.putNextEntry(new ZipEntry(cn.concat(".class")).setLastModifiedTime(FIXED_FILEDATE));
119+
out.write(cls.getValue());
120+
out.closeEntry();
121+
}
122+
classesToInclude.removeIf(processed.keySet()::contains);
123+
System.out.println("Referenced classes not included: " + classesToInclude);
124+
}
125+
126+
static boolean isVisible(int access) {
127+
return (access & (Opcodes.ACC_PROTECTED | Opcodes.ACC_PUBLIC)) != 0;
128+
}
129+
130+
static class Cleaner extends ClassVisitor {
131+
private static final String PREVIEW_ANN = "jdk/internal/javac/PreviewFeature";
132+
private static final String PREVIEW_ANN_DESCR = Type.getObjectType(PREVIEW_ANN).getDescriptor();
133+
134+
private final Set<String> classesToInclude;
135+
private final Map<String, String[]> references;
136+
137+
Cleaner(ClassWriter out, Set<String> classesToInclude, Map<String, String[]> references) {
138+
super(Opcodes.ASM9, out);
139+
this.classesToInclude = classesToInclude;
140+
this.references = references;
141+
}
142+
143+
@Override
144+
public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
145+
super.visit(Opcodes.V11, access, name, signature, superName, interfaces);
146+
if (isVisible(access)) {
147+
classesToInclude.add(name);
148+
}
149+
references.put(name, Stream.concat(Stream.of(superName), Arrays.stream(interfaces)).toArray(String[]::new));
150+
}
151+
152+
@Override
153+
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
154+
return Objects.equals(descriptor, PREVIEW_ANN_DESCR) ? null : super.visitAnnotation(descriptor, visible);
155+
}
156+
157+
@Override
158+
public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
159+
if (!isVisible(access)) {
160+
return null;
161+
}
162+
return new FieldVisitor(Opcodes.ASM9, super.visitField(access, name, descriptor, signature, value)) {
163+
@Override
164+
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
165+
return Objects.equals(descriptor, PREVIEW_ANN_DESCR) ? null : super.visitAnnotation(descriptor, visible);
166+
}
167+
};
168+
}
169+
170+
@Override
171+
public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
172+
if (!isVisible(access)) {
173+
return null;
174+
}
175+
return new MethodVisitor(Opcodes.ASM9, super.visitMethod(access, name, descriptor, signature, exceptions)) {
176+
@Override
177+
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
178+
return Objects.equals(descriptor, PREVIEW_ANN_DESCR) ? null : super.visitAnnotation(descriptor, visible);
179+
}
180+
};
181+
}
182+
183+
@Override
184+
public void visitInnerClass(String name, String outerName, String innerName, int access) {
185+
if (!Objects.equals(outerName, PREVIEW_ANN)) {
186+
super.visitInnerClass(name, outerName, innerName, access);
187+
}
188+
}
189+
190+
@Override
191+
public void visitPermittedSubclass​(String c) {
192+
}
193+
194+
}
195+
196+
}

0 commit comments

Comments
 (0)