Add inductor provenance tracking higlighter

[yushangdi]

Add the inductor provenance tracking highlighter to the output of tlparse. To enable, use flag inductor-provenance.
Note that to produce the artifacts required for the highlighter to work, you need the following flags when running your program: TORCH_TRACE=~/my_trace_log_dir TORCH_LOGS="+inductor" TORCH_COMPILE_DEBUG=1

demo.mov

It works for both JIT inductor and AOT inductor. For JIT inductor, the third panel is the generated python code. For AOT inductor, the third panel is the generated cpp code.

See example outputs with:

tlparse tests/inputs/inductor_provenance_aot_log.txt --inductor-provenance
tlparse tests/inputs/inductor_provenance_jit_log.txt --inductor-provenance

[jamesjwu]

What does it look like when there's more than one compile id? Is there one provenance tracking view per compile ID?

[yushangdi]

What does it look like when there's more than one compile id? Is there one provenance tracking view per compile ID?

@jamesjwu For now it only supports one compile id and a single graph. If there're more than one compile id, it might just show one of the graphs from a random compile id. This was originally designed for the inference use case, where there's usually only one graph.

I can make it support multiple graphs in a followup PR.

[jamesjwu]

Hmm, I'm not a fan of having it be broken for multiple compiles like that: do you need it urgently for this use case? If not I'd very much prefer we implement it to be less brittle from the get go, as it will make it harder for us to refactor later. I feel like the output file here should be composed per compile id, not one big link at the bottom of the page. Doing it now would probably be easier than having to rewrite it later.

It shouldn't be too hard to do: each file from get_file_content's path will have the compile id directly attached to the filename (i.e. /0_0/inductor_provenance_tracked_node_mappings.json). So you could organize it so it produces N output html files, one for each inductor_provenance_tracked_node_mappings.json, and insert it properly into CompileDirectory. You could even use CompileDirectory directly to query for the output files maybe.

A better solution to this would may be to write a custom StructuredLogParser (making sure that the parser always runs after the ones that generate existing inductor outputs), something like this one: https://github.com/pytorch/tlparse/blob/main/src/parsers.rs#L346

[jamesjwu]

Instead of having this be a config to pass to tlparse, can we have it so this logic turns on if there's at least one inductor_provenance_tracking_node_mappings? If we already turned it on in our structured logging, there's no reason to also have to flip a config on tlparse side

[yushangdi]

@jamesjwu I added it because it's a new feature and I don't want to break the normal tlparse flow accidentally. I'm thinking we can remove this flag after it has been running smoothly for a while.

[jamesjwu]

Actually, similar question for @angelayi , why did we end up needing this again? Shouldn't we be able to tell just from the structured log input whether there are export related logs?

[yushangdi]

@jamesjwu I added support for multiple graphs now, I also updated the summary with a new demo video.

[jamesjwu]

Nice! Can you file an issue to put these links in the compile directory too?