Capture constexpr and delay evaluation

[joker-eph]

I don't know if this was already considered, but I felt I'll ask about directions before just implementing a solution: I'm looking into way to capture in IR a representation that does not expand the constexpr values. The goal is to be able to "AOT" some parametric kernels to this MLIR "Extended" Triton IR representation and "instantiate" them using the Triton compiler later, and only inject the constexpr values at this time. We would be able to ship Triton kernels in this MLIR format and JIT them later, without providing the constexpr while evaluating the Python code.

Any appetite for this? Happy to meet and brainstorm about this!

[ptillet]

Hey! I'll answer briefly because I'm on time off starting today :p Generally, yes, I agree that this would solve some existing problems. Constexprs being resolved by the frontend causes a bunch of issues, as you mentioned. The one thing I am worried about is that we have bigger long-term plans for the frontend involving some more advanced metaprogramming using Python, which would mean that constexprs could potentially be python objects too.

[vinodgro]

we have bigger long-term plans for the frontend involving some more advanced metaprogramming using Python

is that at bytecode level or at python AST manipulation level :) @ptillet

[ptillet]

No concrete design decisions have been made yet, but probably the latter. We have seen cases where users have to generate Triton kernel strings because we don't have much metaprogramming :')

[vinodgro]

higher order AST transformers in trition.lang would be an interesting approach :)