-
Notifications
You must be signed in to change notification settings - Fork 5
Commit dff8344
Add launch_diff event (#36)
Summary:
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"`.
Pull Request resolved: #36
Reviewed By: davidberard98
Differential Revision: D78667008
Pulled By: FindHao
fbshipit-source-id: 4d526bf8ad441840ffdf58d8d0e86537a0990d5e1 parent c57a3ac commit dff8344Copy full SHA for dff8344
File tree
Expand file treeCollapse file tree
2 files changed
+350
-131
lines changedFilter options
- tritonparse
Expand file treeCollapse file tree
2 files changed
+350
-131
lines changed
0 commit comments