Skip to content

Commit

Permalink
[Backport] Integrate the Incubating Panama Vector API #12311 (#12327)
Browse files Browse the repository at this point in the history
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]>
  • Loading branch information
3 people authored May 25, 2023
1 parent 7db2c12 commit 0b53670
Show file tree
Hide file tree
Showing 16 changed files with 1,142 additions and 288 deletions.
4 changes: 2 additions & 2 deletions build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ apply from: file('gradle/ide/eclipse.gradle')
// (java, tests)
apply from: file('gradle/java/folder-layout.gradle')
apply from: file('gradle/java/javac.gradle')
apply from: file('gradle/java/memorysegment-mrjar.gradle')
apply from: file('gradle/java/core-mrjar.gradle')
apply from: file('gradle/testing/defaults-tests.gradle')
apply from: file('gradle/testing/randomization.gradle')
apply from: file('gradle/testing/fail-on-no-tests.gradle')
Expand Down Expand Up @@ -158,7 +158,7 @@ apply from: file('gradle/generation/javacc.gradle')
apply from: file('gradle/generation/forUtil.gradle')
apply from: file('gradle/generation/antlr.gradle')
apply from: file('gradle/generation/unicode-test-classes.gradle')
apply from: file('gradle/generation/panama-foreign.gradle')
apply from: file('gradle/generation/extract-jdk-apis.gradle')

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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,17 @@

def resources = scriptResources(buildscript)

configure(rootProject) {
ext {
// also change this in extractor tool: ExtractForeignAPI
vectorIncubatorJavaVersions = [ JavaVersion.VERSION_20 ] as Set
}
}

configure(project(":lucene:core")) {
ext {
apijars = file('src/generated/jdk');
panamaJavaVersions = [ 19, 20 ]
mrjarJavaVersions = [ 19, 20 ]
}

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

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

javaLauncher = javaToolchains.launcherFor {
Expand All @@ -45,21 +52,21 @@ configure(project(":lucene:core")) {
javaLauncher.get()
return true
} catch (Exception e) {
logger.warn('Launcher for Java {} is not available; skipping regeneration of Panama Foreign API JAR.', jdkVersion)
logger.warn('Launcher for Java {} is not available; skipping regeneration of Panama Foreign & Vector API JAR.', jdkVersion)
logger.warn('Error: {}', e.cause?.message)
logger.warn("Please make sure to point env 'JAVA{}_HOME' to exactly JDK version {} or enable Gradle toolchain auto-download.", jdkVersion, jdkVersion)
return false
}
}

classpath = configurations.apiextractor
mainClass = file("${resources}/ExtractForeignAPI.java") as String
mainClass = file("${resources}/ExtractJdkApis.java") as String
systemProperties = [
'user.timezone': 'UTC'
]
args = [
jdkVersion,
new File(apijars, "panama-foreign-jdk${jdkVersion}.apijar"),
new File(apijars, "jdk${jdkVersion}.apijar"),
]
}

Expand Down
196 changes: 196 additions & 0 deletions gradle/generation/extract-jdk-apis/ExtractJdkApis.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import java.io.IOException;
import java.net.URI;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.PathMatcher;
import java.nio.file.Paths;
import java.nio.file.attribute.FileTime;
import java.time.Instant;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.TreeMap;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.FieldVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;

public final class ExtractJdkApis {

private static final FileTime FIXED_FILEDATE = FileTime.from(Instant.parse("2022-01-01T00:00:00Z"));

private static final String PATTERN_PANAMA_FOREIGN = "java.base/java/{lang/foreign/*,nio/channels/FileChannel,util/Objects}";
private static final String PATTERN_VECTOR_INCUBATOR = "jdk.incubator.vector/jdk/incubator/vector/*";
private static final String PATTERN_VECTOR_VM_INTERNALS = "java.base/jdk/internal/vm/vector/VectorSupport{,$Vector,$VectorMask,$VectorPayload,$VectorShuffle}";

static final Map<Integer,List<String>> CLASSFILE_PATTERNS = Map.of(
19, List.of(PATTERN_PANAMA_FOREIGN),
20, List.of(PATTERN_PANAMA_FOREIGN, PATTERN_VECTOR_VM_INTERNALS, PATTERN_VECTOR_INCUBATOR)
);

public static void main(String... args) throws IOException {
if (args.length != 2) {
throw new IllegalArgumentException("Need two parameters: java version, output file");
}
Integer jdk = Integer.valueOf(args[0]);
if (jdk.intValue() != Runtime.version().feature()) {
throw new IllegalStateException("Incorrect java version: " + Runtime.version().feature());
}
if (!CLASSFILE_PATTERNS.containsKey(jdk)) {
throw new IllegalArgumentException("No support to extract stubs from java version: " + jdk);
}
var outputPath = Paths.get(args[1]);

// create JRT filesystem and build a combined FileMatcher:
var jrtPath = Paths.get(URI.create("jrt:/")).toRealPath();
var patterns = CLASSFILE_PATTERNS.get(jdk).stream()
.map(pattern -> jrtPath.getFileSystem().getPathMatcher("glob:" + pattern + ".class"))
.toArray(PathMatcher[]::new);
PathMatcher pattern = p -> Arrays.stream(patterns).anyMatch(matcher -> matcher.matches(p));

// Collect all files to process:
final List<Path> filesToExtract;
try (var stream = Files.walk(jrtPath)) {
filesToExtract = stream.filter(p -> pattern.matches(jrtPath.relativize(p))).collect(Collectors.toList());
}

// Process all class files:
try (var out = new ZipOutputStream(Files.newOutputStream(outputPath))) {
process(filesToExtract, out);
}
}

private static void process(List<Path> filesToExtract, ZipOutputStream out) throws IOException {
var classesToInclude = new HashSet<String>();
var references = new HashMap<String, String[]>();
var processed = new TreeMap<String, byte[]>();
System.out.println("Transforming " + filesToExtract.size() + " class files...");
for (Path p : filesToExtract) {
try (var in = Files.newInputStream(p)) {
var reader = new ClassReader(in);
var cw = new ClassWriter(0);
var cleaner = new Cleaner(cw, classesToInclude, references);
reader.accept(cleaner, ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
processed.put(reader.getClassName(), cw.toByteArray());
}
}
// recursively add all superclasses / interfaces of visible classes to classesToInclude:
for (Set<String> a = classesToInclude; !a.isEmpty();) {
a = a.stream().map(references::get).filter(Objects::nonNull).flatMap(Arrays::stream).collect(Collectors.toSet());
classesToInclude.addAll(a);
}
// remove all non-visible or not referenced classes:
processed.keySet().removeIf(Predicate.not(classesToInclude::contains));
System.out.println("Writing " + processed.size() + " visible classes...");
for (var cls : processed.entrySet()) {
String cn = cls.getKey();
System.out.println("Writing stub for class: " + cn);
out.putNextEntry(new ZipEntry(cn.concat(".class")).setLastModifiedTime(FIXED_FILEDATE));
out.write(cls.getValue());
out.closeEntry();
}
classesToInclude.removeIf(processed.keySet()::contains);
System.out.println("Referenced classes not included: " + classesToInclude);
}

static boolean isVisible(int access) {
return (access & (Opcodes.ACC_PROTECTED | Opcodes.ACC_PUBLIC)) != 0;
}

static class Cleaner extends ClassVisitor {
private static final String PREVIEW_ANN = "jdk/internal/javac/PreviewFeature";
private static final String PREVIEW_ANN_DESCR = Type.getObjectType(PREVIEW_ANN).getDescriptor();

private final Set<String> classesToInclude;
private final Map<String, String[]> references;

Cleaner(ClassWriter out, Set<String> classesToInclude, Map<String, String[]> references) {
super(Opcodes.ASM9, out);
this.classesToInclude = classesToInclude;
this.references = references;
}

@Override
public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
super.visit(Opcodes.V11, access, name, signature, superName, interfaces);
if (isVisible(access)) {
classesToInclude.add(name);
}
references.put(name, Stream.concat(Stream.of(superName), Arrays.stream(interfaces)).toArray(String[]::new));
}

@Override
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
return Objects.equals(descriptor, PREVIEW_ANN_DESCR) ? null : super.visitAnnotation(descriptor, visible);
}

@Override
public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
if (!isVisible(access)) {
return null;
}
return new FieldVisitor(Opcodes.ASM9, super.visitField(access, name, descriptor, signature, value)) {
@Override
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
return Objects.equals(descriptor, PREVIEW_ANN_DESCR) ? null : super.visitAnnotation(descriptor, visible);
}
};
}

@Override
public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
if (!isVisible(access)) {
return null;
}
return new MethodVisitor(Opcodes.ASM9, super.visitMethod(access, name, descriptor, signature, exceptions)) {
@Override
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
return Objects.equals(descriptor, PREVIEW_ANN_DESCR) ? null : super.visitAnnotation(descriptor, visible);
}
};
}

@Override
public void visitInnerClass(String name, String outerName, String innerName, int access) {
if (!Objects.equals(outerName, PREVIEW_ANN)) {
super.visitInnerClass(name, outerName, innerName, access);
}
}

@Override
public void visitPermittedSubclass​(String c) {
}

}

}
Loading

0 comments on commit 0b53670

Please sign in to comment.