Integrate Java Triton example with Intel Triton Backend #241

hanklo6 · 2024-09-25T15:32:53Z

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

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

Using Webrev

Link to Webrev Comment

bridgekeeper · 2024-09-25T15:33:50Z

👋 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 · 2024-09-25T15:34:58Z

@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

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 1203 new commits pushed to the code-reflection branch:

86dfc36: Aggregated work on Verifier and BytecodeLift deconstruction
a40e0f2: Simplify Interpreter invocation
1d23ceb: Support conversions of primitives in testing context
fa4dd44: Simplified bld script and added stricter file types
13b180f: restoring older versions of CoreOp and Interpreter
1fb552a: ANF Experimental Code Model
a9fbdee: Experiment: Alternative SSA construction algorithm
f69e732: Dont use hashbang for launcher scripts
cb52322: Fixed env.bash issue found on Linux (Mac ok?)
94b5149: bld script and bldr tools cleanup
... and 1193 more: https://git.openjdk.org/babylon/compare/7746e41b9fc35a3765ec2e662e8d87cf68edd5a1...code-reflection

As there are no conflicts, your changes will automatically be rebased on top of these commits when integrating. If you prefer to avoid this 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 · 2024-10-09T21:44:24Z

Webrevs

bridgekeeper · 2024-10-31T21:11:24Z

@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 · 2024-11-04T20:47:04Z

Keep alive

PaulSandoz

Thank you for this effort. It's not easy to hook this up and you did it without any help!

Initially i have just focused on the code in the triton package.

cr-examples/triton/src/main/java/oracle/code/triton/ArithMathOps.java

cr-examples/triton/src/test/java/oracle/code/triton/TestMatrix.java

PaulSandoz · 2024-11-05T22:41:38Z

cr-examples/triton/src/test/java/oracle/code/triton/TestMatrix.java

+            int stride_am, @Constant int stride_ak,
+            int stride_bk, @Constant int stride_bn,
+            int stride_cm, @Constant int stride_cn,


Why the marking as constants?

Triton will mark these three stride values as constants if they are equal to one at compile time. We mark them as constants directly as our code is not compiled by Triton.

Oh i see, i don't see those equivalent parameters marked as tl.constexpr in the Python matmul example. I think we need the concept of optional constants e.g., a boolean on the @Constant annotation. An optional constant parameter can be associated with type or a constant type.
(Alternatively we could drop the constant annotation, but it seems useful to distinguish various parameters)

PaulSandoz · 2024-11-05T22:47:34Z

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

@@ -1018,8 +1051,15 @@ public Value dot(TensorType rType, Op.Result r,
                         TypeElement bType, Value b) {
            a = block.context().getValue(a);
            b = block.context().getValue(b);
-
-            return block.op(TritonOps.dot(rType, a, b));
+            Object zero;


Add a comment such as

// Computed result is tensor of floats, regardless of inputs

like when we compute the type. Since the result is hard coded we don't need to use reflection and can directly use the constant expression "0.0".

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

cr-examples/triton/src/test/java/oracle/code/triton/TestMatrix.java

PaulSandoz · 2024-11-06T23:55:52Z

cr-examples/triton/src/test/java/oracle/code/triton/TestMatrix.java

+            int stride_am, @Constant int stride_ak,
+            int stride_bk, @Constant int stride_bn,
+            int stride_cm, @Constant int stride_cn,


Oh i see, i don't see those equivalent parameters marked as tl.constexpr in the Python matmul example. I think we need the concept of optional constants e.g., a boolean on the @Constant annotation. An optional constant parameter can be associated with type or a constant type.
(Alternatively we could drop the constant annotation, but it seems useful to distinguish various parameters)

PaulSandoz · 2024-11-07T00:39:42Z

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

@@ -107,6 +111,17 @@ void test(CoreOp.FuncOp javaKernel,
                return TritonTransformer.tritonModule(javaKernel, JavaType.VOID, argTypes);
            });

+            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(); }

Thanks, fixed.

hanklo6 added 8 commits September 24, 2024 11:26

Add attributes name support and enhance dot and load op

166f5a3

Add MLIR generator

b2b8f91

Update test cases

17f7689

Add triton translator

33bbdcf

Update the generated file names

a5d6800

Add compile patch

9a4af49

Update translate.py

9f9e387

Add launcher code and translate code

1a91283

hanklo6 added 2 commits September 25, 2024 08:48

Use system home in the path

72695bb

Fix trailing and leading whitespaces

79696b4

hanklo6 changed the title ~~Enhance Triton MLIR for Triton Backend~~ Integrate Java Triton example with Intel Triton Backend Sep 25, 2024

hanklo6 marked this pull request as ready for review October 9, 2024 21:38

openjdk bot added ready Pull request is ready to be integrated rfr Pull request is ready for review labels Oct 9, 2024

PaulSandoz reviewed Nov 5, 2024

View reviewed changes

Add comments

2dbe2df

PaulSandoz reviewed Nov 7, 2024

View reviewed changes

hanklo6 added 4 commits November 7, 2024 15:24

Update attribute name

66ba69c

Update MLIR output paths

335e207

Simplify zero initialization in dot method

a91434d

Add the test for matrix multiplication with float16

f585076

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Integrate Java Triton example with Intel Triton Backend #241

Integrate Java Triton example with Intel Triton Backend #241

hanklo6 commented Sep 25, 2024 •

edited by openjdk bot

Loading

bridgekeeper bot commented Sep 25, 2024

openjdk bot commented Sep 25, 2024 •

edited

Loading

mlbridge bot commented Oct 9, 2024 •

edited

Loading

bridgekeeper bot commented Oct 31, 2024

PaulSandoz commented Nov 4, 2024

PaulSandoz left a comment

PaulSandoz Nov 5, 2024

hanklo6 Nov 6, 2024

PaulSandoz Nov 6, 2024

PaulSandoz Nov 5, 2024

hanklo6 Nov 6, 2024

PaulSandoz Nov 6, 2024

PaulSandoz Nov 7, 2024

hanklo6 Nov 8, 2024

Integrate Java Triton example with Intel Triton Backend #241

Are you sure you want to change the base?

Integrate Java Triton example with Intel Triton Backend #241

Conversation

hanklo6 commented Sep 25, 2024 • edited by openjdk bot Loading

Usage

Dependencies

Progress

Reviewing

bridgekeeper bot commented Sep 25, 2024

openjdk bot commented Sep 25, 2024 • edited Loading

mlbridge bot commented Oct 9, 2024 • edited Loading

Webrevs

bridgekeeper bot commented Oct 31, 2024

PaulSandoz commented Nov 4, 2024

PaulSandoz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hanklo6 commented Sep 25, 2024 •

edited by openjdk bot

Loading

openjdk bot commented Sep 25, 2024 •

edited

Loading

mlbridge bot commented Oct 9, 2024 •

edited

Loading