Java Platform Features for AI and HPC(High Performance Computing)

Several proposed Java platform features are designed to enhance AI capabilities and performance by addressing specific pain points in AI and HPC-like solutions, such as optimal utilization of hardware and efficient use of memory. While these features are general and broadly useful across the platform, they significantly benefit AI applications.

The four key features are:

Foreign Function and Memory API (FFM API)

This API, part of Project Panama and final since Java 22, allows Java code to interoperate optimally with off-heap memory and native libraries. This is crucial for AI because it enables zero-copy operations between the JDK and external code, which is extremely important in some cases. It provides a faster, safer, and more productive way to access native APIs and libraries compared to the existing JNI method, which is described as difficult to work with. Using FFM API, developers can leverage highly optimized existing native code, such as linear algebra libraries like Bliss, without needing to reimplement them in Java. A developer productivity tool called Jextract can generate pure Java bindings from C header files, speeding up the mechanical part of interfacing with native libraries. Examples of its use include Llama model inference, an anomaly detection example (MSET 2), and Onyx script for Java. It allows calling out to native libraries like Bliss and efficiently using their matrix multiply algorithms with zero-copy memory access.

Vector API

Also part of Project Panama, the Vector API provides the ability to explicitly express SIMD (Single Instruction, Multiple Data) computations in Java, leading to optimal CPU utilization by leveraging vector hardware for parallel execution. This feature is described as enabling better number crunching, which is a fundamental part of machine learning. It allows for the implementation of high-performance data parallel algorithms in Java, such as dot product and matrix multiply, which are core to machine learning algorithms. Although it is currently incubating, it is already used in some production scenarios. Improvements in performance have been demonstrated by switching to the Vector API, such as a significant boost (from 1.3 to effectively 15 tokens per second) in a Llama model inference example by using vectorized dot product instead of scalar code. The API is being explored for implementing Bliss-like matrix kernels in pure Java experiments. It can be used for efficient O(N) operations directly on the CPU.

Value Classes and Objects (Valhalla)

This project aims to provide optimal utilization of memory on-heap. Crucially for AI, it will enable more kinds of numbers to be easily integrated into the platform, such as float16 and bfloat16. This allows external libraries or developers to add new floating-point types that behave like existing numbers in the system. While Valhalla is not yet ready, work is underway to support and optimize new numeric types like float16 in the runtime (Hotspot) for use with the Vector API and autovectorization. The Vector API is waiting for Valhalla to exit incubation because it can properly represent vectors as values once Valhalla is ready. These value classes could potentially improve memory handling and computation in algorithms and data structures like records.

Code Reflection (Project Babylon)

This feature provides interoperation with foreign programming models that are not part of the standard Java programming model. Examples include programming models for GPUs (like CUDA), ONNX machine learning models, or auto-differentiation. Code reflection enhances Java's existing reflection capabilities by providing reflective access to the code within lambda or method bodies at runtime. It uses a symbolic, in-memory representation of Java code optimized for analysis and transformation. This allows libraries to handle the complexities of these foreign models and their mappings from Java code, enabling capabilities like symbolic analysis, automatic differentiation, and generating machine learning models (such as ONNX models) directly from Java code. The process involves a library extracting a Java code model and transforming it into a foreign code model before generating the final foreign format. It could potentially be used in the future to automatically transform scalar Java code into vectorized code for efficient execution, potentially removing the need for developers to write directly to the Vector API for some tasks. It is seen as a way to leverage diverse hardware and improve number crunching.

Combined, these features are intended to create a more powerful platform for building AI libraries and applications by addressing core areas like memory efficiency, hardware utilization, and interoperation with native or foreign systems. It is suggested that while AI development is currently dominated by Python, these improvements could lead to AI shops increasingly adopting Java. Although building AI solutions in Java is possible without these features, it is described as harder to design, develop, maintain, and evolve. These features aim to make it easier to address the pain points and enable rapid prototyping and development of AI solutions in modern Java.