[Feature] Add launch_diff event

[FindHao]

This pull request introduces a major enhancement to the trace parsing functionality. Instead of treating each event in isolation, it now intelligently groups compilation and launch events by kernel. The key feature is the generation of a launch_diff event, which provides a concise summary of how different launches of the same kernel vary in their parameters. This is extremely useful for debugging, as it highlights what changes between calls. The changes in structured_logging.py support this by ensuring that all necessary metadata is correctly serialized.

Detailed Changes

triton_trace/parser.py

• parse_single_trace_content: The logic is now wrapped in a check for entry.get("event_type") == "compilation". This ensures that parsing logic for IR files and source mappings only runs for compilation events.
• parse_single_file:
  • The function has been significantly refactored to group events by kernel hash.
  • It now differentiates between compilation and launch events.
  • compilation events and their associated launch events are grouped under a kernel hash.
  • A new launch_diff event is generated for each kernel, summarizing the differences and similarities between all its launch events. This helps in understanding how kernel launches vary.
  • The output is now organized by kernel, and written to files based on frame_id and frame_compile_id, similar to before, but the content is now more structured.
• New Helper Functions:
  • _flatten_dict, _unflatten_dict: Utilities to handle nested dictionaries.
  • _to_ranges: Converts a list of indices into a list of continuous ranges, used for summarizing which launches share a particular parameter value.
  • _generate_launch_diff: The core of the new diffing logic. It compares a list of launch events and produces a summary of common and differing parameters. It also has special handling for extracted_args to provide detailed diffs for kernel arguments.

tritonparse/structured_logging.py

• convert: Added special handling for namedtuple to convert it to a dictionary, ensuring field names are preserved during JSON serialization.
• add_launch_metadata: The metadata object (which is a namedtuple) is now converted to a dictionary using _asdict() before being returned.
• LaunchHookImpl.post_launch: The key for the compilation metadata in the trace data has been changed from "metadata" to "compilation_metadata".

[facebook-github-bot]

@FindHao has imported this pull request. If you are a Meta employee, you can view this in D78667008.

[facebook-github-bot]

@FindHao merged this pull request in dff8344.