[RFC] Triton IR Printer Decorator

[henryg-d-Matrix]

TL;DR: decorator like jit except it prints Triton IR when kernel is called.

I propose a decorator like jit except the compiled Triton IR is printed when the decorated python function / kernel is called. This would mainly be used for testing and development purposes to simplify getting Triton IR code organically compiled from Triton source code. This would prevent having to handwrite kernels in Triton IR or having to dig through the cache directory trying to find the right cached IR file. It would also enable natural generation of Triton IR on simple development machines without GPUs or perhaps without even CUDA. Hacking the Triton compiler to achieve this means is obviously possible, but a dedicated utility for this purpose would prevent unnecessary disturbances to the actual Triton compiler stack meant for kernel execution.

Though not meant for frontend integration onto other third-party custom compiler stacks, it can definitely be expanded for that purpose. This printer decorator utility can also be expanded to print other forms of IR such as Triton GPU IR, or the converted LinAlg output of triton-to-linalg.

My current work on this has the printer decorator contained in a single file and no changes to other parts of the Triton Python frontend. Another similar proposal would be to add another additional parameter to the jit compiler for such debug compile outputs of intermediate IRs. The latter method has less separation and less room for future expansions.

Hopefully this utility can be useful despite its simplicity.

Related discussions:
A verbose mode printing commands used
Is it possible to get the intermediate ttir and ttgir on a non-CUDA machine?
How do I view Triton IR and LLVM IR generated by the triton compiler?

[joker-eph]

Why do you need a decorator instead of explicitly calling the APIs?

For example:

ret = triton.compile(kernel, signature="*fp32,i32,*fp32,i32", constants={"BLOCK_M": 64, "BLOCK_N": 64})
print(ret.asm["ttgir"])

https://github.com/openai/triton/blob/main/python/examples/copy_strided.py#L7-L18

[joker-eph]

By the way, isn't such decorator something you can implement yourself already? I imagine it just needs to call triton.compile and print? It's sugar.

[henryg-d-Matrix]

I was hoping that Triton could provide something official, but now I see that the triton.compile API fulfills this role as of the moment. I still think that a simplified IR printing interface could be useful in the future, but for now I agree to disagree.

@joker-eph please feel free to close this discussion.

[joker-eph]

To be clear, I'm not a developer of Triton, just a user (and occasional contributor)! So just my opinion here :)

[henryg-d-Matrix]

Your opinion is greatly appreciated! I'm very grateful to have this nice community to bounce my ideas against. Your suggestions of using the triton.compile API has inspired me a lot. I'll develop my idea more based on top of it.

[awgr]

thank you so much for sharing this party of the api here, I'd been reading through the documentation, but somehow missed how to work more directly with the IR, and it's so relaxing now to be able to read it ❤️