Add info CLI Subcommand

[FindHao]

Summary

This PR adds the info subcommand to query kernel information from NDJSON trace files. It automatically detects raw logs and parses them, enabling users to list kernels and their launches without manual parsing.

Changes

• info/cli.py (NEW):

  • _add_info_args(): Add CLI arguments for info subcommand
  • info_command(): Main command implementation
  • Auto-detects raw logs (no launch_diff events) and parses them automatically
  • Supports listing all kernels or launches for a specific kernel
  • Provides fuzzy matching suggestions when kernel not found

• info/parse_helper.py (NEW):

  • parse_and_compress_raw_log(): Parse and compress raw log files
  • Reuses parse_single_file() and gzip_single_file() from parse module
  • Handles file naming for various input formats (.ndjson, .ndjson.gz, .bin.ndjson)

• info/kernel_query.py:

  • list_launches_for_kernel(): List all launches for a specific kernel with grid info
  • find_similar_kernels(): Fuzzy matching using difflib.get_close_matches()
  • list_kernels_fast(): Optimized listing using launch_diff events when available
    • Merges kernels with same name (sums up launches)
    • Falls back to list_kernels() for raw logs

• info/__init__.py:

  • Export new functions: list_launches_for_kernel, find_similar_kernels, list_kernels_fast

• cli.py:

  • Add info subcommand integration
  • Import _add_info_args and info_command from info.cli
  • Update help examples

• tests/test_tritonparse.py:

  • test_info_kernel_query_functions(): Combined unit test for 3 query functions
  • test_info_list_kernels(): Integration test for listing all kernels
  • test_info_kernel_launches(): Integration test for listing kernel launches
  • test_info_kernel_not_found(): Integration test for error handling with suggestions

Usage

# List all kernels with launch counts
tritonparseoss info trace.ndjson

# Output example:
# Kernels in trace.ndjson:
#   fused_op_kernel        4 launches (id: 0-3)
#   matmul_kernel       1553 launches (id: 0-1552)

# List launches for a specific kernel
tritonparseoss info trace.ndjson --kernel matmul_kernel

# Output example:
# Launches for 'matmul_kernel':
#   id=  0  line   123  grid=[128, 1, 1]
#   id=  1  line   456  grid=[256, 1, 1]
#   ...

# Works with raw logs (auto-parses)
tritonparseoss info raw_trace.ndjson
# Automatically parses and compresses, then queries

Features

• Auto-parsing: Automatically detects and parses raw logs (no launch_diff events)
• Compression: Parsed files are automatically compressed using parse module's gzip functionality
• Performance: Uses launch_diff events for fast kernel listing when available
• Error handling: Suggests similar kernel names when kernel not found
• Temporary files: Parsed files saved to temp directory (not cleaned up, user can reuse)

Testing

• 1 combined unit test covering 3 query functions
• 3 integration tests covering CLI functionality
• All tests use real data from test fixtures
• Tests verify both fast path (launch_diff) and slow path (full traversal)

Notes

• Temporary directories are not cleaned up (as per design requirement)
• Parsed files are compressed (.ndjson.gz) for consistency with parse module
• Kernel name matching is case-sensitive (exact match only)
• Fuzzy matching uses 0.6 cutoff threshold for similarity

Test Plan

% python -m unittest tests.test_tritonparse.TestTritonparseCPU.test_info_kernel_query_functions tests.test_tritonparse.TestTritonparseCPU.test_info_list_kernels tests.test_tritonparse.TestTritonparseCPU.test_info_kernel_launches tests.test_tritonparse.TestTritonparseCPU.test_info_kernel_not_found  -v
test_info_kernel_query_functions (tests.test_tritonparse.TestTritonparseCPU.test_info_kernel_query_functions)
Test info module kernel query functions. ... ok
test_info_list_kernels (tests.test_tritonparse.TestTritonparseCPU.test_info_list_kernels)
Integration test: info command lists all kernels. ... ok
test_info_kernel_launches (tests.test_tritonparse.TestTritonparseCPU.test_info_kernel_launches)
Integration test: info command lists launches for specific kernel. ... ok
test_info_kernel_not_found (tests.test_tritonparse.TestTritonparseCPU.test_info_kernel_not_found)
Integration test: info command handles kernel not found. ... ok

----------------------------------------------------------------------
Ran 4 tests in 0.708s

OK

[meta-codesync]

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

[meta-codesync]

@FindHao merged this pull request in ab34d43.