-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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
1 parent
7db2c12
commit 0b53670
Showing
16 changed files
with
1,142 additions
and
288 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) { | ||
} | ||
|
||
} | ||
|
||
} |
Oops, something went wrong.