Create info/ Module (Core Query Layer) (#208)

Summary:
This PR creates the `info/` module as the core query layer for kernel information from NDJSON trace files. This is an internal infrastructure PR that will be used by PR4 (reproduce `--kernel`/`--launch-id`) and PR5 (info CLI).

## Changes

- **`tritonparse/info/__init__.py`**: Module initialization, exports core functions
- **`tritonparse/info/kernel_query.py`**: Core query functions:
  - `KernelSummary` dataclass: kernel name, hash, total launches
  - `LaunchInfo` dataclass: launch ID, line index, grid (for PR5)
  - `list_kernels(events)`: List all kernels with their launch counts
  - `find_launch_index_by_kernel(events, kernel_name, launch_id)`: Find 0-based line index for a kernel's N-th launch

- **`tests/test_tritonparse.py`**: Added 6 unit tests in `TestTritonparseCPU`:
  - `test_list_kernels_empty()`: Empty events list
  - `test_list_kernels_single()`: Single kernel with multiple launches
  - `test_list_kernels_multiple()`: Multiple different kernels
  - `test_find_launch_index_valid()`: Valid kernel name and launch_id
  - `test_find_launch_index_kernel_not_found()`: Raises ValueError when kernel not found
  - `test_find_launch_index_out_of_range()`: Raises ValueError with valid range hint

## Testing

Tests use real data from `tests/example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz` when possible, mock data for edge cases (empty list, non-existent kernel, out-of-range launch_id).

## Notes

- This PR does NOT create CLI. It's purely internal infrastructure.
- All indices are 0-based for consistency with Python conventions.
- Kernel name matching is case-sensitive (exact match only).

Pull Request resolved: #208

Reviewed By: wychi

Differential Revision: D88171102

Pulled By: FindHao

fbshipit-source-id: 7c2ac5d74551a1d2b0667dd8b4734972298d7a62

Author: FindHao

tests/test_tritonparse.py

@@ -336,6 +336,152 @@ def test_load_ndjson_gzip_support(self):
 
         print(f"✓ Successfully loaded {len(events)} events from .ndjson.gz file")
 
+    def test_list_kernels_empty(self):
+        """Test listing kernels from empty events list."""
+        from tritonparse.info.kernel_query import list_kernels
+
+        events = []
+        result = list_kernels(events)
+        self.assertEqual(result, [])
+
+    def test_list_kernels_single(self):
+        """Test listing kernels with single kernel and multiple launches."""
+        from pathlib import Path
+
+        from tritonparse.info.kernel_query import list_kernels
+        from tritonparse.tools.prettify_ndjson import load_ndjson
+
+        # Load real test data
+        gz_file = (
+            Path(__file__).parent
+            / "example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz"
+        )
+        events = load_ndjson(gz_file)
+
+        # Filter to only fused_op_kernel launches (4 launches)
+        filtered_events = []
+        for event in events:
+            if event.get("event_type") == "launch":
+                kernel_name = event.get("compilation_metadata", {}).get("name")
+                if kernel_name == "fused_op_kernel":
+                    filtered_events.append(event)
+            else:
+                # Keep non-launch events to test filtering
+                filtered_events.append(event)
+
+        result = list_kernels(filtered_events)
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].name, "fused_op_kernel")
+        self.assertEqual(result[0].total_launches, 4)
+
+    def test_list_kernels_multiple(self):
+        """Test listing kernels with multiple different kernels."""
+        from pathlib import Path
+
+        from tritonparse.info.kernel_query import list_kernels
+        from tritonparse.tools.prettify_ndjson import load_ndjson
+
+        # Load real test data
+        gz_file = (
+            Path(__file__).parent
+            / "example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz"
+        )
+        events = load_ndjson(gz_file)
+
+        result = list_kernels(events)
+        self.assertEqual(len(result), 2)
+
+        # Check that results are sorted by name
+        names = [k.name for k in result]
+        self.assertEqual(names, ["fused_op_kernel", "matmul_kernel"])
+
+        # Check launch counts
+        kernel_dict = {k.name: k for k in result}
+        self.assertEqual(kernel_dict["matmul_kernel"].total_launches, 1553)
+        self.assertEqual(kernel_dict["fused_op_kernel"].total_launches, 4)
+
+    def test_find_launch_index_valid(self):
+        """Test finding valid kernel name and launch_id."""
+        from pathlib import Path
+
+        from tritonparse.info.kernel_query import find_launch_index_by_kernel
+        from tritonparse.tools.prettify_ndjson import load_ndjson
+
+        # Load real test data
+        gz_file = (
+            Path(__file__).parent
+            / "example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz"
+        )
+        events = load_ndjson(gz_file)
+
+        # Test first launch of fused_op_kernel (launch_id=0)
+        index = find_launch_index_by_kernel(events, "fused_op_kernel", 0)
+        self.assertEqual(events[index].get("event_type"), "launch")
+        self.assertEqual(
+            events[index].get("compilation_metadata", {}).get("name"),
+            "fused_op_kernel",
+        )
+
+        # Test second launch of fused_op_kernel (launch_id=1)
+        index = find_launch_index_by_kernel(events, "fused_op_kernel", 1)
+        self.assertEqual(events[index].get("event_type"), "launch")
+        self.assertEqual(
+            events[index].get("compilation_metadata", {}).get("name"),
+            "fused_op_kernel",
+        )
+
+        # Test first launch of matmul_kernel (launch_id=0)
+        index = find_launch_index_by_kernel(events, "matmul_kernel", 0)
+        self.assertEqual(events[index].get("event_type"), "launch")
+        self.assertEqual(
+            events[index].get("compilation_metadata", {}).get("name"),
+            "matmul_kernel",
+        )
+
+    def test_find_launch_index_kernel_not_found(self):
+        """Test that ValueError is raised when kernel not found."""
+        from pathlib import Path
+
+        from tritonparse.info.kernel_query import find_launch_index_by_kernel
+        from tritonparse.tools.prettify_ndjson import load_ndjson
+
+        # Load real test data
+        gz_file = (
+            Path(__file__).parent
+            / "example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz"
+        )
+        events = load_ndjson(gz_file)
+
+        with self.assertRaises(ValueError) as cm:
+            find_launch_index_by_kernel(events, "nonexistent_kernel", 0)
+
+        error_msg = str(cm.exception)
+        self.assertIn("not found", error_msg)
+        self.assertIn("nonexistent_kernel", error_msg)
+
+    def test_find_launch_index_out_of_range(self):
+        """Test that ValueError is raised when launch_id is out of range."""
+        from pathlib import Path
+
+        from tritonparse.info.kernel_query import find_launch_index_by_kernel
+        from tritonparse.tools.prettify_ndjson import load_ndjson
+
+        # Load real test data
+        gz_file = (
+            Path(__file__).parent
+            / "example_output/parsed_output_complex/dedicated_log_triton_trace_findhao__mapped.ndjson.gz"
+        )
+        events = load_ndjson(gz_file)
+
+        # fused_op_kernel has only 4 launches (0-3), test with launch_id=10
+        with self.assertRaises(ValueError) as cm:
+            find_launch_index_by_kernel(events, "fused_op_kernel", 10)
+
+        error_msg = str(cm.exception)
+        self.assertIn("has only 4 launches", error_msg)
+        self.assertIn("--launch-id 10", error_msg)
+        self.assertIn("Valid range: 0 to 3", error_msg)
+
 
 class TestTritonparseCUDA(unittest.TestCase):
     """CUDA tests (require GPU)"""

tritonparse/info/__init__.py

@@ -0,0 +1,24 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Info module for querying kernel information from NDJSON trace files.
+
+This module provides core query functions for kernel information:
+- Listing all kernels with their launch counts
+- Finding launch events by kernel name and launch ID
+- Querying launch information for specific kernels
+"""
+
+from tritonparse.info.kernel_query import (
+    find_launch_index_by_kernel,
+    KernelSummary,
+    LaunchInfo,
+    list_kernels,
+)
+
+__all__ = [
+    "KernelSummary",
+    "LaunchInfo",
+    "list_kernels",
+    "find_launch_index_by_kernel",
+]

tritonparse/info/kernel_query.py

@@ -0,0 +1,107 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Core query functions for kernel information from NDJSON trace files.
+
+This module provides functions to query kernel launch information from parsed
+event lists. It supports both raw log files and parsed ndjson files (with launch_diff events).
+"""
+
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+
+@dataclass
+class KernelSummary:
+    """Summary information about a kernel."""
+
+    name: str
+    hash: str
+    total_launches: int
+
+
+@dataclass
+class LaunchInfo:
+    """Information about a specific kernel launch."""
+
+    launch_id: int  # 0-based
+    line_index: int  # 0-based (index in events list)
+    grid: List[int]
+
+
+def list_kernels(events: List[Dict[str, Any]]) -> List[KernelSummary]:
+    """
+    List all kernels with their launch counts.
+
+    Args:
+        events: List of parsed event dictionaries from NDJSON file
+
+    Returns:
+        List of KernelSummary objects, sorted by kernel name
+    """
+    # Count launches per kernel
+    kernel_counts: Dict[str, Dict[str, Any]] = defaultdict(
+        lambda: {"hash": "", "count": 0}
+    )
+
+    for event in events:
+        if event.get("event_type") != "launch":
+            continue
+
+        comp_meta = event.get("compilation_metadata", {})
+        kernel_name = comp_meta.get("name")
+        kernel_hash = comp_meta.get("hash", "")
+
+        if kernel_name:
+            kernel_counts[kernel_name]["hash"] = kernel_hash
+            kernel_counts[kernel_name]["count"] += 1
+
+    # Convert to KernelSummary list
+    summaries = [
+        KernelSummary(name=name, hash=info["hash"], total_launches=info["count"])
+        for name, info in kernel_counts.items()
+    ]
+
+    # Sort by kernel name for consistent output
+    summaries.sort(key=lambda x: x.name)
+
+    return summaries
+
+
+def find_launch_index_by_kernel(
+    events: List[Dict[str, Any]], kernel_name: str, launch_id: int
+) -> int:
+    """
+    Find the 0-based line index for a kernel's N-th launch.
+
+    Args:
+        events: List of parsed event dictionaries
+        kernel_name: Exact kernel name to match (case-sensitive)
+        launch_id: 0-based launch index for the kernel
+
+    Returns:
+        0-based line index (index in events list)
+
+    Raises:
+        ValueError: If kernel not found or launch_id out of range
+    """
+    count = 0
+    for i, event in enumerate(events):
+        if event.get("event_type") != "launch":
+            continue
+
+        comp_meta = event.get("compilation_metadata", {})
+        name = comp_meta.get("name")
+        if name == kernel_name:
+            if count == launch_id:
+                return i
+            count += 1
+
+    if count == 0:
+        raise ValueError(f"Kernel '{kernel_name}' not found")
+    else:
+        raise ValueError(
+            f"Kernel '{kernel_name}' has only {count} launches, "
+            f"but --launch-id {launch_id} was requested. Valid range: 0 to {count - 1}"
+        )