[code-reflection] RFR: Integrate Java Triton example with Intel Triton Backend [v2]

[Paul Sandoz]

On Wed, 6 Nov 2024 16:14:21 GMT, hanklo6 <duke at openjdk.org> wrote:

>> 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](https://github.com/intel/intel-xpu-backend-for-triton)
>> - [intel-extension-for-pytorch](https://github.com/intel/intel-extension-for-pytorch)
>> - [Intel oneAPI base Toolkit](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html)
>> - [Jextract](https://github.com/openjdk/jextract)
>> - [Level Zero loader](https://github.com/oneapi-src/level-zero)
>> - [compute-runtime](https://github.com/intel/compute-runtime/releases)
>> - [JSON-java](https://github.com/stleary/JSON-java)
>
> hanklo6 has updated the pull request incrementally with one additional commit since the last revision:
> 
>   Add comments

cr-examples/triton/src/main/java/oracle/code/triton/TritonTransformer.java line 1061:

> 1059:             } catch (Throwable e) {
> 1060:                 throw new RuntimeException(e);
> 1061:             }

Suggestion:

            Object zero = 0.0;

cr-examples/triton/src/test/java/oracle/code/triton/TritonTestExtension.java line 114:

> 112:             });
> 113: 
> 114:             String mlirText = MLIRGenerator.transform(actualTritonKernel);

We should write the files to the maven target directory. Add the following to the surefire plugin configuration:

                        <systemPropertyVariables>
                            <project.build.directory>${project.build.directory}</project.build.directory>
                        </systemPropertyVariables>

Then replace the code below with:

            Path buildDir = Path.of(System.getProperty("project.build.directory", ""));
            Path mlirDir = buildDir.resolve("mlir");
            try {
                Files.createDirectories(mlirDir);
                Files.writeString(mlirDir.resolve(javaKernelName + ".mlir"), mlirText, StandardOpenOption.CREATE);
            } catch (IOException e) {
                e.printStackTrace();
            }

-------------

PR Review Comment: https://git.openjdk.org/babylon/pull/241#discussion_r1831864406
PR Review Comment: https://git.openjdk.org/babylon/pull/241#discussion_r1831899439