Add info CLI Subcommand (#210)

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

```bash
# 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

Pull Request resolved: #210

Test Plan:
```bash
% 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
```

Reviewed By: wychi

Differential Revision: D88171129

Pulled By: FindHao

fbshipit-source-id: ff464052f8d80819912b65551a6d87ba8a184fa2

Author: FindHao

tests/test_tritonparse.py

@@ -614,6 +614,114 @@ def test_reproduce_launch_id_out_of_range(self):
         finally:
             self.cleanup_temp_reproduce_dir(temp_dir)
 
+    def test_info_kernel_query_functions(self):
+        """Test info module kernel query functions."""
+        from tritonparse.info.kernel_query import (
+            find_similar_kernels,
+            list_kernels,
+            list_kernels_fast,
+            list_launches_for_kernel,
+        )
+        from tritonparse.tools.prettify_ndjson import load_ndjson
+
+        gz_file = self._get_test_ndjson_file()
+        events = load_ndjson(gz_file)
+
+        # Test list_launches_for_kernel
+        launches = list_launches_for_kernel(events, "fused_op_kernel")
+        self.assertGreater(len(launches), 0)
+        self.assertEqual(launches[0].launch_id, 0)
+        self.assertIsInstance(launches[0].grid, list)
+
+        # Test list_launches_for_kernel with non-existent kernel
+        with self.assertRaises(ValueError) as cm:
+            list_launches_for_kernel(events, "nonexistent_kernel")
+        self.assertIn("not found", str(cm.exception))
+
+        # Test find_similar_kernels
+        similar = find_similar_kernels(events, "fused_op", n=3)
+        self.assertGreater(len(similar), 0)
+        self.assertIn("fused_op_kernel", similar)
+
+        similar = find_similar_kernels(events, "fused_op_kernel", n=3)
+        self.assertIn("fused_op_kernel", similar)
+
+        similar = find_similar_kernels(events, "xyz_abc_123", n=3)
+        self.assertEqual(len(similar), 0)
+
+        # Test list_kernels_fast (should use launch_diff and match list_kernels)
+        kernels_fast = list_kernels_fast(events)
+        self.assertGreater(len(kernels_fast), 0)
+
+        kernels_slow = list_kernels(events)
+        fast_dict = {k.name: k.total_launches for k in kernels_fast}
+        slow_dict = {k.name: k.total_launches for k in kernels_slow}
+        self.assertEqual(fast_dict, slow_dict)
+
+    def test_info_list_kernels(self):
+        """Integration test: info command lists all kernels."""
+        import sys
+        from io import StringIO
+
+        from tritonparse.info.cli import info_command
+
+        gz_file = self._get_test_ndjson_file()
+
+        # Capture stdout
+        old_stdout = sys.stdout
+        sys.stdout = captured_output = StringIO()
+
+        try:
+            info_command(str(gz_file), kernel_name=None)
+            output = captured_output.getvalue()
+            self.assertIn("Kernels in", output)
+            self.assertIn("launches", output)
+        finally:
+            sys.stdout = old_stdout
+
+    def test_info_kernel_launches(self):
+        """Integration test: info command lists launches for specific kernel."""
+        import sys
+        from io import StringIO
+
+        from tritonparse.info.cli import info_command
+
+        gz_file = self._get_test_ndjson_file()
+
+        # Capture stdout
+        old_stdout = sys.stdout
+        sys.stdout = captured_output = StringIO()
+
+        try:
+            info_command(str(gz_file), kernel_name="fused_op_kernel")
+            output = captured_output.getvalue()
+            self.assertIn("Launches for 'fused_op_kernel'", output)
+            self.assertIn("id=", output)
+            self.assertIn("line", output)
+        finally:
+            sys.stdout = old_stdout
+
+    def test_info_kernel_not_found(self):
+        """Integration test: info command handles kernel not found."""
+        import sys
+        from io import StringIO
+
+        from tritonparse.info.cli import info_command
+
+        gz_file = self._get_test_ndjson_file()
+
+        # Capture stdout
+        old_stdout = sys.stdout
+        sys.stdout = captured_output = StringIO()
+
+        try:
+            with self.assertRaises(ValueError):
+                info_command(str(gz_file), kernel_name="nonexistent_kernel")
+            output = captured_output.getvalue()
+            self.assertIn("not found", output)
+        finally:
+            sys.stdout = old_stdout
+
 
 class TestTritonparseCUDA(unittest.TestCase):
     """CUDA tests (require GPU)"""

tritonparse/cli.py

@@ -4,6 +4,7 @@
 from importlib.metadata import PackageNotFoundError, version
 
 from .common import is_fbcode
+from .info.cli import _add_info_args, info_command
 from .reproducer.cli import _add_reproducer_args
 from .reproducer.orchestrator import reproduce
 from .utils import _add_parse_args, unified_parse
@@ -31,6 +32,8 @@ def main():
             "Examples:\n"
             f"  {prog_name} parse /path/to/logs --out parsed_output\n"
             f"  {prog_name} reproduce /path/to/trace.ndjson --line 1 --out-dir repro_output\n"
+            f"  {prog_name} info /path/to/trace.ndjson\n"
+            f"  {prog_name} info /path/to/trace.ndjson --kernel matmul_kernel\n"
         ),
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
@@ -60,6 +63,14 @@ def main():
     _add_reproducer_args(repro_parser)
     repro_parser.set_defaults(func="reproduce")
 
+    # info subcommand
+    info_parser = subparsers.add_parser(
+        "info",
+        help="Query kernel information from trace file",
+    )
+    _add_info_args(info_parser)
+    info_parser.set_defaults(func="info")
+
     args = parser.parse_args()
 
     if args.func == "parse":
@@ -89,6 +100,8 @@ def main():
             kernel_import=args.kernel_import,
             replacer=replacer,
         )
+    elif args.func == "info":
+        info_command(input_path=args.input, kernel_name=args.kernel)
     else:
         raise RuntimeError(f"Unknown command: {args.func}")
 

tritonparse/info/__init__.py

@@ -11,14 +11,20 @@
 
 from tritonparse.info.kernel_query import (
     find_launch_index_by_kernel,
+    find_similar_kernels,
     KernelSummary,
     LaunchInfo,
     list_kernels,
+    list_kernels_fast,
+    list_launches_for_kernel,
 )
 
 __all__ = [
     "KernelSummary",
     "LaunchInfo",
     "list_kernels",
+    "list_kernels_fast",
+    "list_launches_for_kernel",
     "find_launch_index_by_kernel",
+    "find_similar_kernels",
 ]

tritonparse/info/cli.py

@@ -0,0 +1,121 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+CLI implementation for the info subcommand.
+
+This module provides command-line interface for querying kernel information
+from NDJSON trace files.
+"""
+
+import argparse
+import tempfile
+from typing import Optional
+
+from tritonparse.info.kernel_query import (
+    find_similar_kernels,
+    list_kernels_fast,
+    list_launches_for_kernel,
+)
+from tritonparse.info.parse_helper import parse_and_compress_raw_log
+from tritonparse.tools.prettify_ndjson import load_ndjson
+
+
+def _add_info_args(parser: argparse.ArgumentParser) -> None:
+    """Add arguments for the info subcommand."""
+    parser.add_argument(
+        "input",
+        help="Path to ndjson/ndjson.gz/.bin.ndjson file",
+    )
+    parser.add_argument(
+        "--kernel",
+        type=str,
+        default=None,
+        help="Kernel name to list launches for",
+    )
+
+
+def info_command(input_path: str, kernel_name: Optional[str] = None) -> None:
+    """
+    Main function for the info command.
+
+    Args:
+        input_path: Path to ndjson file
+        kernel_name: Optional kernel name to list launches for
+    """
+    # 1. Load and detect type
+    events = load_ndjson(input_path)
+    has_launch_diff = any(e.get("event_type") == "launch_diff" for e in events)
+
+    # 2. If no launch_diff, auto-parse
+    if not has_launch_diff:
+        print(
+            f"Input file '{input_path}' appears to be raw log (no launch_diff events)."
+        )
+        print("Parsing automatically to generate launch_diff events...")
+
+        temp_dir = tempfile.mkdtemp(prefix="tritonparse_info_")
+
+        try:
+            # Parse and compress (reuses parse module's functions)
+            parsed_file = parse_and_compress_raw_log(
+                input_path,
+                output_dir=temp_dir,
+                split_inductor_compilations=False,
+                verbose=False,
+            )
+
+            # Load compressed file (load_ndjson supports .ndjson.gz)
+            events = load_ndjson(parsed_file)
+
+            print(f"✓ Parsed and compressed file: {parsed_file}")
+            print(f"  (Temporary directory: {temp_dir})")
+        except Exception as e:
+            raise RuntimeError(f"Failed to parse input file '{input_path}': {e}") from e
+    else:
+        print(f"Using parsed trace file: {input_path}")
+
+    # 3. Process query
+    if kernel_name:
+        # List launches for specific kernel
+        try:
+            launches = list_launches_for_kernel(events, kernel_name)
+            print(f"\nLaunches for '{kernel_name}':")
+            print("-" * 60)
+            for launch in launches:
+                grid_str = str(launch.grid) if launch.grid else "N/A"
+                print(
+                    f"  id={launch.launch_id:3d}  line {launch.line_index:5d}  grid={grid_str}"
+                )
+        except ValueError as e:
+            error_msg = str(e)
+            print(f"\nError: {error_msg}")
+            # Try to suggest similar kernels
+            try:
+                similar = find_similar_kernels(events, kernel_name, n=3)
+                if similar:
+                    print("\nDid you mean one of these?")
+                    all_kernels = list_kernels_fast(
+                        events
+                    )  # Use fast path for consistency
+                    kernel_dict = {k.name: k for k in all_kernels}
+                    for name in similar:
+                        count = kernel_dict[name].total_launches
+                        print(f"  - {name} ({count} launches)")
+                    print("\nUse 'tritonparseoss info <file>' to list all kernels.")
+            except Exception:
+                pass  # Ignore errors in suggestion
+            raise
+    else:
+        # List all kernels
+        kernels = list_kernels_fast(events)
+        print(f"\nKernels in {input_path}:")
+        print("-" * 60)
+        for kernel in kernels:
+            if kernel.total_launches > 0:
+                max_id = kernel.total_launches - 1
+                print(
+                    f"  {kernel.name:30s} {kernel.total_launches:3d} launches "
+                    f"(id: 0-{max_id})"
+                )
+            else:
+                print(f"  {kernel.name:30s} {kernel.total_launches:3d} launches")