[Proton][Dialect] Add Initial Frontend and Target Backend Infrastructure For Proton Dialect #5506

CRobeck · 2024-12-28T03:32:18Z

Implement initial basic infrastructure for the Proton Dialect added in #5119

This PR extends the initial boilerplate MLIR Dialect code to the Triton frontend and target backends - currently just lowered to a no-op.

CRobeck · 2024-12-28T15:07:14Z

cc: @fywkevin

include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h

lib/Conversion/TritonGPUToLLVM/CMakeLists.txt

lib/Conversion/TritonToTritonGPU/CMakeLists.txt

Jokeren · 2024-12-28T15:40:29Z

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

+void populateProtonPatterns(TritonGPUTypeConverter &typeConverter,
+                            RewritePatternSet &patterns) {
+  MLIRContext *context = patterns.getContext();
+  patterns.add<GenericOpPattern<triton::proton::RecordOp>>(typeConverter,


This can be taken out if we don't consider the frontend changes for now

Maybe we don't need any conversion here anyway because proton's inputs are scalars but not tensors. Could you please double check? @CRobeck

They are scalars but could envision a case where we pass in tensors in the future - if we want to pass in a Triton tensor object for example for object specific tracing. I don't think the conversion hurts anything here but acknowledge it could potentially be misleading.

You're right though that I can remove this entire function (populateProtonPatterns) and the test still passes.

Jokeren · 2024-12-28T15:40:37Z

python/src/ir.cc

@@ -1603,6 +1606,10 @@ void init_triton_ir(py::module &&m) {
                                               llvm::StringRef(prefix));
             self.create<PrintOp>(prefixAttr, hex, values, isSigned);
           })
+      .def("create_proton_record",


Let's take the python api changes out of this PR. It's not determined yet and probably not very important in the short term. Users, like @pawelszczerbuk, can modify the GPU IR to instrument record ops

This is mostly just to have a template in place for the future Proton front end development to interact with the upper level Triton Python kernel code - we can always modify the interface later but how the Proton Python frontend fit together with the Triton IR builder wasn't super intuitive at first so having this here I think is helpful for the moment even if we ultimately change the API in the future. It also gives us some more testing avenues. What do you think?

The concern I have is that we don't have any real functionalities associated with these record ops anyway. Usually, we add a frontend op when there's at least some functionalities implemented. For example, the most recent tl.gather op. Though its initial commit doesn't address all problems and concerns, frontend users are able to use the op in their kernels but not just placeholders. So IMHO it makes more sense to add the op, and mark it as experimental, after the lowering has been implemented.

What about putting things in a dev module with some warnings about them being under development and not intentioned for use outside the development team:

https://github.com/triton-lang/triton/pull/5506/files#diff-054c0ddf64263bf99d77a237a61deeb79c0cd3e4289a03eba948ce833f1cdce0R17

https://github.com/triton-lang/triton/pull/5506/files#diff-bc216e7a68b5a4545f1ff69bfbf5d0fa64adb6b1acaca6ea0f366f47e93942ebR25

Then that gives us an easy method to develop/test in a way that turning them "on" is just a matter of moving them from the dev module to the experimental namespace?

It's much better now!

Let's just move it under proton.language without using the dev class.

Also I think create_proton_record can be moved to triton_proton.cc

Updated to be under proton.language.

Also I think create_proton_record can be moved to triton_proton.cc

I think this would cause a lot of code disruption. TritonOpBuilder is a unique ptr wrapper around the MLIR OpBuilder. So I think no matter what we'd have to find way to pass that object to triton_proton.cc either through moving the definition of TritonOpBuilder out of ir.cc into a header file or possibly through a helper function defined in ir.cc and called from triton_proton.cc which seems like an equal or more amount of complexity to the existing code.

Maybe we can just invoke the generic MLIR OpBuilder from triton_proton.cc? But then I think we'd be mixing and matching builder objects in the code generator?

I get it now. Let's move it to the last line of all ops and leave a comment:

// proton ops .def("create_proton_record", ...

It's a problem for all third party backends that they cannot define custom ops. We will have to think about a better solution.

Updated.

It's a problem for all third party backends that they cannot define custom ops. We will have to think about a better solution.

Right, there has to be an upper level Triton Op that the backends can specialize but I could not find a way to have a wholly custom third party backend Op. Given that TritonOpBuilder is defined in ir.cc the only solution I see is to pull that class/object out into a standalone module that can be passed to each backend so that they all operate on the same MLIR builder instance.

third_party/proton/dialect/lib/TritonProtonToLLVM/CMakeLists.txt

Jokeren · 2024-12-30T14:46:05Z

python/src/ir.cc

@@ -1603,6 +1606,10 @@ void init_triton_ir(py::module &&m) {
                                               llvm::StringRef(prefix));
             self.create<PrintOp>(prefixAttr, hex, values, isSigned);
           })
+      .def("create_proton_record",


It's much better now!

Let's just move it under proton.language without using the dev class.

Jokeren · 2024-12-30T14:48:15Z

python/src/ir.cc

@@ -1603,6 +1606,10 @@ void init_triton_ir(py::module &&m) {
                                               llvm::StringRef(prefix));
             self.create<PrintOp>(prefixAttr, hex, values, isSigned);
           })
+      .def("create_proton_record",


Also I think create_proton_record can be moved to triton_proton.cc

python/src/ir.cc

Jokeren · 2024-12-30T14:51:12Z

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

+void populateProtonPatterns(TritonGPUTypeConverter &typeConverter,
+                            RewritePatternSet &patterns) {
+  MLIRContext *context = patterns.getContext();
+  patterns.add<GenericOpPattern<triton::proton::RecordOp>>(typeConverter,


Maybe we don't need any conversion here anyway because proton's inputs are scalars but not tensors. Could you please double check? @CRobeck

third_party/proton/proton/profile.py

third_party/proton/proton/language.py

Jokeren · 2024-12-31T01:44:23Z

python/src/ir.cc

@@ -1603,6 +1606,10 @@ void init_triton_ir(py::module &&m) {
                                               llvm::StringRef(prefix));
             self.create<PrintOp>(prefixAttr, hex, values, isSigned);
           })
+      .def("create_proton_record",


I get it now. Let's move it to the last line of all ops and leave a comment:

// proton ops .def("create_proton_record", ...

It's a problem for all third party backends that they cannot define custom ops. We will have to think about a better solution.

third_party/proton/proton/__init__.py

third_party/proton/proton/profile.py

Jokeren · 2024-12-31T03:09:12Z

It looks good to me now mostly. Do you want to postpone merging into main a bit until Yuanwei has taken a look?

CRobeck · 2024-12-31T14:48:52Z

Yes, that's fine. I don't think should be anything too controversial.

CRobeck added 8 commits December 28, 2024 02:01

temp

c27f894

update

7496112

update

e47e4c9

update

435fbe6

update

a33f513

update

f4095cc

update

09b7cf9

update

877e2ea

Jokeren self-assigned this Dec 28, 2024

update

ae0ff69

CRobeck changed the title ~~[Proton][Dialect] Add Initial Frontend Support For Proton Dialect~~ [Proton][Dialect] Add Initial Frontend and Target Backend Support For Proton Dialect Dec 28, 2024

CRobeck added 3 commits December 28, 2024 04:18

update

e4784f4

update

754fdd1

update

d2f9c0b

CRobeck marked this pull request as ready for review December 28, 2024 15:04

CRobeck requested review from antiagainst, zhanglx13 and ptillet as code owners December 28, 2024 15:04

Jokeren reviewed Dec 28, 2024

View reviewed changes

CRobeck added 8 commits December 29, 2024 02:10

move some paths into proton dialect

7b0485d

update

36f1861

move more things into proton dialect

3897f81

update cmake

bf1011e

update cmake

97b92b2

update cmake

6a0dc4c

add dev module

f1911e4

formatting

b085d75

CRobeck changed the title ~~[Proton][Dialect] Add Initial Frontend and Target Backend Support For Proton Dialect~~ [Proton][Dialect] Add Initial Frontend and Target Backend infrastructure For Proton Dialect Dec 30, 2024

CRobeck changed the title ~~[Proton][Dialect] Add Initial Frontend and Target Backend infrastructure For Proton Dialect~~ [Proton][Dialect] Add Initial Frontend and Target Backend Infrastructure For Proton Dialect Dec 30, 2024

update

5cb4ad2

Jokeren reviewed Dec 30, 2024

View reviewed changes

third_party/proton/proton/profile.py Outdated Show resolved Hide resolved

third_party/proton/proton/language.py Outdated Show resolved Hide resolved

CRobeck added 6 commits December 30, 2024 19:39

update

190d26c

update

3c72681

update

d980aea

update

134f102

update

6a64945

update

ca63a07

Jokeren reviewed Dec 31, 2024

View reviewed changes

CRobeck added 2 commits December 31, 2024 02:40

update

3184568

formatting

a61a21e

Merge branch 'main' into proton_record

9e98d14

Jokeren approved these changes Jan 6, 2025

View reviewed changes

zhanglx13 merged commit dcad5ac into triton-lang:main Jan 6, 2025
7 checks passed

CRobeck deleted the proton_record branch January 8, 2025 02:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Proton][Dialect] Add Initial Frontend and Target Backend Infrastructure For Proton Dialect #5506

[Proton][Dialect] Add Initial Frontend and Target Backend Infrastructure For Proton Dialect #5506

CRobeck commented Dec 28, 2024 •

edited

Loading

CRobeck commented Dec 28, 2024

Jokeren Dec 28, 2024

Jokeren Dec 30, 2024 •

edited

Loading

CRobeck Dec 30, 2024 •

edited

Loading

Jokeren Dec 28, 2024

CRobeck Dec 28, 2024 •

edited

Loading

Jokeren Dec 28, 2024

CRobeck Dec 29, 2024 •

edited

Loading

Jokeren Dec 30, 2024

Jokeren Dec 30, 2024

CRobeck Dec 30, 2024 •

edited

Loading

Jokeren Dec 31, 2024

CRobeck Dec 31, 2024 •

edited

Loading

Jokeren Dec 30, 2024

Jokeren Dec 30, 2024

Jokeren Dec 30, 2024 •

edited

Loading

Jokeren Dec 31, 2024

Jokeren commented Dec 31, 2024

CRobeck commented Dec 31, 2024 •

edited

Loading

[Proton][Dialect] Add Initial Frontend and Target Backend Infrastructure For Proton Dialect #5506

[Proton][Dialect] Add Initial Frontend and Target Backend Infrastructure For Proton Dialect #5506

Conversation

CRobeck commented Dec 28, 2024 • edited Loading

CRobeck commented Dec 28, 2024

Choose a reason for hiding this comment

Jokeren Dec 30, 2024 • edited Loading

Choose a reason for hiding this comment

CRobeck Dec 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CRobeck Dec 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CRobeck Dec 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CRobeck Dec 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CRobeck Dec 31, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jokeren Dec 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jokeren commented Dec 31, 2024

CRobeck commented Dec 31, 2024 • edited Loading

CRobeck commented Dec 28, 2024 •

edited

Loading

Jokeren Dec 30, 2024 •

edited

Loading

CRobeck Dec 30, 2024 •

edited

Loading

CRobeck Dec 28, 2024 •

edited

Loading

CRobeck Dec 29, 2024 •

edited

Loading

CRobeck Dec 30, 2024 •

edited

Loading

CRobeck Dec 31, 2024 •

edited

Loading

Jokeren Dec 30, 2024 •

edited

Loading

CRobeck commented Dec 31, 2024 •

edited

Loading