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

• [x] 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

Webrev

Link to Webrev Comment

[bridgekeeper[bot]]

:wave: 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[bot]]

@hanklo6 This change now passes all automated pre-integration checks.

ℹ️ This project also has non-automated pre-integration requirements. Please see the file CONTRIBUTING.md for details.

After integration, the commit message for the final commit will be:

Integrate Java Triton example with Intel Triton Backend

Reviewed-by: psandoz

You can use pull request commands such as /summary, /contributor and /issue to adjust it as needed.

At the time when this comment was updated there had been no new commits pushed to the code-reflection branch. If another commit should be pushed before you perform the /integrate command, your PR will be automatically rebased. If you prefer to avoid any potential automatic rebasing, please check the documentation for the /integrate command for further details.

As you do not have Committer status in this project an existing Committer must agree to sponsor your change. Possible candidates are the reviewers of this PR (@PaulSandoz) but any other Committer may sponsor as well.

➡️ To flag this PR as ready for integration with the above commit message, type /integrate in a new comment. (Afterwards, your sponsor types /sponsor in a new comment to perform the integration).

[mlbridge[bot]]

Webrevs

• 04: Full (c92a8adc)
• 03: Full - Incremental (f585076b)
• 02: Full - Incremental (a91434d4)
• 01: Full - Incremental (2dbe2df9)
• 00: Full (79696b4e)

[bridgekeeper[bot]]

@hanklo6 This pull request has been inactive for more than 4 weeks and will be automatically closed if another 4 weeks passes without any activity. To avoid this, simply add a new comment to the pull request. Feel free to ask for assistance if you need help with progressing this pull request towards integration!

[PaulSandoz]

Keep alive

[bridgekeeper[bot]]

@hanklo6 This pull request has been inactive for more than 4 weeks and will be automatically closed if another 4 weeks passes without any activity. To avoid this, simply add a new comment to the pull request. Feel free to ask for assistance if you need help with progressing this pull request towards integration!

[PaulSandoz]

Keep alive, need to give it another quick review.

[openjdk[bot]]

@hanklo6 this pull request can not be integrated into code-reflection due to one or more merge conflicts. To resolve these merge conflicts and update this pull request you can run the following commands in the local repository for your personal fork:

git checkout triton
git fetch https://git.openjdk.org/babylon.git code-reflection
git merge FETCH_HEAD
# resolve conflicts and follow the instructions given by git merge
git commit -m "Merge code-reflection"
git push

[hanklo6]

@PaulSandoz Thank you for the review. I've merged the code.

[PaulSandoz]

@PaulSandoz Thank you for the review. I've merged the code.

Thanks. Type "/integrate" and then i can sponsor it.

[hanklo6]

/integrate

[openjdk[bot]]

@hanklo6 Your change (at version c92a8adc23f4ae264043ec4278eb4a5186355d63) is now ready to be sponsored by a Committer.

[PaulSandoz]

/sponsor

[openjdk[bot]]

Going to push as commit c0c70a431bd35542bbef8ba94b0c419a4371f441.

[openjdk[bot]]

@PaulSandoz @hanklo6 Pushed as commit c0c70a431bd35542bbef8ba94b0c419a4371f441.

:bulb: You may see a message that your pull request was closed with unmerged commits. This can be safely ignored.