Accessible Dynamic SPIR-V Code Generation from Java

At The University of Manchester, we have developed a Java library, called Beehive SPIR-V Toolkit, for generating SPIR-V binary modules that can be dispatched on supported devices, for example, on GPUs. In this post, I will explain why we need this, and how developers can use it to build their own SPIR-V modules from Java.

You might think, well, it is just a code generator, what’s special about it? The Beehive SPIR-V Toolkit is a programming library to build SPIR-V binary modules that is also autogenerated. I hope you are still with me. In this post, I will explain why we followed this design, and show a few examples. Note, this post is a summary of a recent academic publication at VMIL’23.

Why did we do this? A bit of history

To understand the main reasons why we developed this library, let’s start with a bit of history. Back in 2020, we initiated a collaboration with Intel to work on the TornadoVM JIT compiler and the runtime system and try out the brand new (back then) Level Zero API. The Level Zero API is a low-level bare-metal API to program hardware accelerators such as Integrated GPUs, CPUs and FPGAs. Note that at the time of writing this post, Level Zero API is only available for programming Intel GPUs (discrete and integrated GPUs), but the programming model is generic enough to also accommodate another type of hardware, such as multi-core CPUs and FPGAs.

This looked promising because this API allows us to have more control at the runtime system level compared to other solutions, such as OpenCL, and it matches very well on Intel GPU architectures. For example, having features such as explicit types for shared memory and low-latency command queues are very appealing for a parallel programming framework like TornadoVM. BUT, Level Zero requires the dispatch of SPIR-V binary Compute Kernels, and at the time we initiated this collaboration, TornadoVM did not have such a backend.

A bit of context

Just to add a bit of context in case you are not familiar with TornadoVM. TornadoVM is a Java parallel programming framework to offload Java programs to be accelerated on modern hardware, such as GPUs. TornadoVM itself is also written in Java, and this includes the Just In Time (JIT) compiler, and the runtime system. Thus, what we wanted to do is a tool to generate SPIR-V code from the TornadoVM internal IR, and we did not want to use any external LLVM tool to do so.

SPIR-V Toolkit as a Standalone Application

Hence, we planned to start adding a new backend to support SPIR-V devices. SPIR-V is a Standard Portable Intermediate Representation (IR) in binary format for hardware accelerators. Kernels written in SPIR-V can be dispatched through supported APIs (e.g., OpenCL and Level Zero). But, generating SPIR-V is not a simple task. The main complication is that the SPIR-V is in binary format, and this makes the initial development and debugging process very difficult. In our case, it was not only to develop a code generator, we also needed an optimizing compiler (transforming Java bytecode to an internal IR, called Tornado IR), optimizing the program in this IR format, and then performing the code generation from the optimized TornadoVM IR to a SPIR-V binary.

Thus, to tackle this process more efficiently, we decided to fully decouple the optimization process from the code generation process by having three different projects that can be combined:

1. The actual optimization pipeline (TornadoVM JIT Compiler for SPIR-V).
2. The SPIR-V binary module generation, and this is what we called Beehive SPIR-V Toolkit.
3. The SPIR-V code dispatcher through the Level Zero APIs from Java.

All of the projects were released as open-source projects.

Beehive SPIR-V Toolkit for Java

The Beehive SPIR-V Toolkit has three main components:

1. A library generator: a system engine that generates a Java library to compose standard SPIR-V binary modules. The engine, called Template System Engine (TSE) takes, as inputs, a JSON file that specifies the standard SPIR-V grammar and a set of Java templates. Using both the JSON files and the templates, the TSE generates a set of Java classes for generating SPIR-V instructions and the SPIR-V operands for every new SPIR-V version release.
2. The SPIR-V Library: this is the result of the TSE, and, once compiled, the SPIR-V library is ready to be consumed by client applications (e.g., an optimizing compiler, or a runtime system).
3. A SPIR-V command line utility for assembly and disassembly of SPIR-V code. This is useful for debugging.

Why is the library also generated?

The library is generated from the files published by the Khronos Group (a non-profit consortium of organizations developing standards, and, among them, the SPIR-V standard) on their GitHub repository. These JSON files represent a specific version of the SPIR-V standard (e.g., SPIR-V 1.2). At the time of writing this post, the tool generates all Java classes to compose SPIR-V modules compliant with SPIR-V 1.2, which is the version required by the Intel drivers to work with the Level Zero API.

However, there are new standards, such as 1.5 and Unified. To be able to quickly comply with the latest standard version, we automatize the process of library generation based on these file descriptions, and in fact, we also have tested other versions of the standard.

How the API looks like?

If we know what SPIR-V instructions to generate, it is straightforward to use the Beehive SPIR-V Toolkit API. Besides, all instructions in the API follow the same order as the SPIR-V specification. Thus, it is easy to look up the instructions and reason about the type of operands needed for each instruction. Let’s take a look at a few examples:

1. Creating a SPIR-V Module Header

We create a new SPIR-V module by creating an object of type SPIRVModule. This class receives an instance object of type SPIRVHeader:

SPIRVModule module = new SPIRVModule(
  new SPIRVHeader(
    1,     // Major Version
    2,     // Minor Version
    32,   //  ID-Generator 
    0,     // Bounds : set to 0 initially
    0));   // schema 

This will generate the following SPIR-V code:

; SPIR-V​
; Version: 1.2​
; Generator: Khronos; 32​
; Bound: 77​
; Schema: 0​

2. Adding Capabilities

We can also add SPIR-V capabilities such as using 64-bit integers, Kernel mode, etc. The following code snippet shows some of the most common options in compute kernels:

module.add(new SPIRVOpCapability(SPIRVCapability.Addresses()));     // Uses physical addressing, non-logical addressing modes.
module.add(new SPIRVOpCapability(SPIRVCapability.Linkage()));       // Uses partially linked modules and libraries. (e.g., OpenCL)
module.add(new SPIRVOpCapability(SPIRVCapability.Kernel()));        // Uses the Kernel Execution Model.
module.add(new SPIRVOpCapability(SPIRVCapability.Int64()));         // Uses OpTypeInt to declare 64-bit integer types

3. Performing an Integer-Addition

To perform an addition, we need the resulting type, an ID where to store the result, and the two operands. The following code snippet shows an example.

SPIRVId add = module.getNextId();
blockScope.add(new SPIRVOpIAdd(
  ulong,      // resulting type
  add,        // result id (just created in the previous line)
  id28,       // operand 1
  ulongConstant24   // operand-2
));

This just shows a sneak peak of some of the SPIR-V constructs that can be built with the SPIR-V Toolkit. You can review full examples from OpenCL semantics down to SPIR-V in the project main repository:

More Information

Our VMIL’23 paper also contains more details about each component, and shows a performance evaluation with TornadovM comparing the quality of its OpenCL backend versus the SPIR-V backend.

Lastly, we gave a presentation at VMIL’23 at the SPLASH 2023 venue. You can check the video if this presentation in the following link:

Follow-up and discussions

If you get this far, and you’re still interested in this work, you can check the paper, and/or drop me an email for more discussions. You can also leave your comments on GitHub, or contact me via Twitter.