Integrate Java Triton example with Intel Triton Backend

[hanklo6]

Babylon Java Triton example translates Java source code with Java Triton API into code model by code reflection.

In this PR, we traverse the given code model and output Triton MLIR dialect in the generic form, and then inject generated MLIR dialect into the Intel Triton backend. We then utilize Intel Triton backend to compile the Triton MLIR dialect into a SPIR-V module. Use Jextract to create Java binding of Intel Level Zero runtime and launch the given kernel function with it on Intel GPUs.

Usage

Navigate to the cr-example/triton directory and execute mvn clean test. This will generate multiple MLIR files in the result directory ready to be processed by the Triton backend.

Next, modify the compiler.py file within the intel-xpu-triton-backend project by applying the patch git apply add-mlir-insertion.patch. Then run the Triton backend by running python3 translate.py.

The Triton backend will generate SPIR-V files, which will be located under ~/.triton/cache/{hash_value}/{kernel_name}/{kernel_name}.spv.

To create a binding for Level Zero, execute the below commands:

$JEXTRACT_DIR/bin/jextract --output src/gen/java -I /usr/include -t oneapi.levelzero level-zero/include/ze_api.h
$JAVA_HOME/bin/javac -cp target/classes -d target/classes src/gen/java/oneapi/levelzero/*.java
$JAVA_HOME/bin/jar cf levelzero.jar -C target/classes/ .

The will generate levelzero.jar in the current directory.

After getting JAR files for Level Zero and JSON-java, proceed to compile and run the launcher LevelZero.java with the following commands:

babylon/build/linux-x86_64-server-release/jdk/bin/javac -cp .:levelzero.jar:json-java.jar LevelZero.java
babylon/build/linux-x86_64-server-release/jdk/bin/java -ea -cp .:levelzero.jar:json-java.jar LevelZero

Ensure the hash values in~/.triton/cache match those used in the LevelZero.java.

Dependencies

• intel-xpu-backend-for-triton
• intel-extension-for-pytorch
• Intel oneAPI base Toolkit
• Jextract
• Level Zero loader
• compute-runtime
• JSON-java

---

Progress

• Change must not contain extraneous whitespace

Reviewing

Using git

Checkout this PR locally:
$ git fetch https://git.openjdk.org/babylon.git pull/241/head:pull/241
$ git checkout pull/241

Update a local copy of the PR:
$ git checkout pull/241
$ git pull https://git.openjdk.org/babylon.git pull/241/head

Using Skara CLI tools

Checkout this PR locally:
$ git pr checkout 241

View PR using the GUI difftool:
$ git pr show -t 241

Using diff file

Download this PR as a diff file:
https://git.openjdk.org/babylon/pull/241.diff

[bridgekeeper]

👋 Welcome back hanklo6! A progress list of the required criteria for merging this PR into code-reflection will be added to the body of your pull request. There are additional pull request commands available for use with this pull request.

[openjdk]

❗ This change is not yet ready to be integrated.
See the Progress checklist in the description for automated requirements.