I am trying to read and understand the internals of openAI’s triton’s IR lowering, from tritonIR to tritonGPUIR to LLVMIR. I had added quite a lot of print, either by using module.dump() or other print methods provided by LLVM, it helps me to understand the purpose of some passes.
But the problem is, just when everything is on the edge to get converted to LLVM dialect finally, now a lot of pass pipelines are added because it has to handle the logic of converting every op ( triton ops, tritonGPU ops, scf/cf ops, arith ops, etc.) to LLVM op. For each specific op, when converting to LLVM ops, it has to extract the information from the attribute of the op (layout), and emit many LLVM IRs for each thread to execute. So, without further optimizations( CSE, DCE, canonicalizer stuff), the generated LLVM dialect IRs are really really long. If I dump the whole moduleOp out, it will be quite difficult to spot like which part of this IRs are generated from which triton op lowering. So, I tried to dump the moduleOp only in specific triton op’s matchAndRewrite method. I dump the moduleOp when entering this method, and just before exiting this method. But still, the dumped Module-level IR is huge (10s of thousands of lines or more), it takes quite a while just dumping it out when executing the code. I could then extract the dumped module IR, and use tools to compare and figure out the difference of the two IRs, but it’s still not very convenient. So, I’m wondering is there any other more cononical way to checkout what some matchAndRewrite method does to the IR, instead of dumping the IR? Another problem is when I’m dumping out the IR inside specific op’s “matchAndRewrite” method, since the whole moduleOp has several ops with the same type. So eventually I dumped many module-level IRs.
Have you tried the debug option? It should be quite verbose and describe every pattern application.
There is also --log-actions-to that may log patterns as well.
One more thing: try adding the snapshot-op-locations pass to the pipeline right before the lowering: it’ll annotate every operation with a location matching the current IR, if you print the IR before and after it’ll help directly tracking down where is an operation post-lowering coming from.
Interesting, I have not followed this line of work.
Would it be possible to put some full 1-liners that demonstrate this?
Oh you should definitely watch the recording (or skim through) Open MLIR Meeting 2/23/2023: MLIR Actions: Tracing and Debugging MLIR-based Compilers - #2 by mehdi_amini
I have some cool demo there!
The doc: Action: Tracing and Debugging MLIR-based Compilers - MLIR
We should add more Actions to the codebase also (hook in the transformation dialect?)!
I think right now --log-actions-to wouldn’t be enough because it won’t print the IR after conversion, only before.
In addition to what Mehdi suggested, try -debug-only=dialect-conversion, which provides information related to the dialect conversion infrastructure specifically rather than dumping a lot of debug information (provided triton uses the dialect conversion infrastructure).
Also, it is possible to dump any operation, not just the top-level module.