PerfModel: unify kernel library classification

README.md

@@ -31,6 +31,8 @@ TraceLens is a Python library for **automated performance analysis of training a
 
 **TraceLens Agent**: Receive a prioritized human-readable optimization report, derived through an agentic workflow, covering compute kernels, system bottlenecks, and kernel fusion opportunities with root-cause reasoning and concrete resolutions.
 
+**Trace corpus indexing**: Build a searchable SQLite catalog of TraceLens reports so you can find traces by op, category, or kernel name without reopening every raw file. Scanner and importer sit behind a storage interface; SQLite is the first backend.
+
 ---
 
 ## Quick Start
@@ -79,6 +81,16 @@ TraceLens_compare_perf_reports_pytorch \
     -o comparison.xlsx
 ```
 
+Index traces and existing TraceLens CSV reports (see [Index a corpus of traces](docs/how-to/trace-index.md)):
+
+```bash
+TraceLens_trace_index --db trace_index.sqlite append \
+    --trace-path /path/to/rank0_trace.json.gz \
+    --report-dir path/to/perf_report_csvs
+TraceLens_trace_index --db trace_index.sqlite build --traces-file traces.txt
+TraceLens_trace_index --db trace_index.sqlite search Cijk
+```
+
 For multi-rank runs, generate a collective-communication report across ranks (see [Generate a collective-communication report](docs/how-to/collective-report.md)):
 
 ```bash
@@ -139,6 +151,7 @@ Each format's linked doc covers its full CLI reference. For PyTorch report compa
 | Multi-Rank Collective Report | [docs/how-to/collective-report.md](docs/how-to/collective-report.md)                                                             |
 | Performance Report Columns   | [docs/reference/perf-report-columns.md](docs/reference/perf-report-columns.md)                                                   |
 | TraceLens Agent              | [docs/how-to/agent.md](docs/how-to/agent.md)                                                                                     |
+| TraceIndex                   | [docs/how-to/trace-index.md](docs/how-to/trace-index.md)                                                                         |
 
 ---
 

TraceLens/Agent/Analysis/category_analyses/analysis_utils.py

@@ -27,6 +27,7 @@
 import numpy as np
 import pandas as pd
 
+from TraceLens.PerfModel.kernel_library import classify_kernel_library
 from TraceLens.PerfModel.utils import torch_dtype_map
 
 TARGET_HIGH = 100.0
@@ -44,26 +45,6 @@
 # Efficiency Boundary for grouping
 _EFF_BUCKET_BOUNDARIES = (30, 60)
 
-_OP_NAME_LIBRARY_RULES = [
-    ("aiter::", "AITER"),
-    ("rocm_aiter", "AITER"),
-    ("fbgemm", "FBGEMM"),
-    ("miopen", "MIOpen"),
-    ("triton", "Triton"),
-]
-_KERNEL_NAME_LIBRARY_RULES = [
-    ("aiter", "AITER"),
-    ("ck_tile::", "CK"),
-    ("ck_tile6kentry", "CK"),
-    ("FmhaFwd", "CK"),
-    ("FmhaBwd", "CK"),
-    ("Cijk_", "Tensile"),
-    ("wvSplitK", "rocBLAS"),
-    ("splitKreduce", "rocBLAS"),
-    ("rocprim::", "rocPRIM"),
-    ("triton_", "Triton"),
-    ("void at::native::", "PyTorch Native"),
-]
 _COMPARATIVE_SPEEDUP_COL = "speedup (trace2/trace1)"
 _COMPARATIVE_DELTA_COL = "delta_us (trace2 - trace1)"
 _COMPARATIVE_T2_TIME_COL = "lca_total_kernel_time_trace2_us"
@@ -1093,15 +1074,5 @@ def run_category_analysis(
     print(f"Metrics written to: {output_path}")
 
 
-# Category-specific helper functions
-def classify_kernel_library(op_name: str, kernel_details: str = "") -> Optional[str]:
-    """Identify the GPU library backing an operation from its name or kernel strings."""
-    op_lower = op_name.lower()
-    for marker, lib in _OP_NAME_LIBRARY_RULES:
-        if marker in op_lower:
-            return lib
-    kd = str(kernel_details) if kernel_details and not pd.isna(kernel_details) else ""
-    for marker, lib in _KERNEL_NAME_LIBRARY_RULES:
-        if marker in kd:
-            return lib
-    return None
+# Category-specific helper functions — classify_kernel_library lives in
+# TraceLens.PerfModel.kernel_library and is imported above.

TraceLens/PerfModel/kernel_library.py

@@ -0,0 +1,79 @@
+###############################################################################
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Shared kernel-library classification for perf reports, agents, and TraceIndex."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from TraceLens.PerfModel import kernel_name_parser as knp
+
+_OP_NAME_LIBRARY_RULES = [
+    ("aiter::", "AITER"),
+    ("rocm_aiter", "AITER"),
+    ("fbgemm", "FBGEMM"),
+    ("miopen", "MIOpen"),
+    ("triton", "Triton"),
+]
+
+_KERNEL_NAME_LIBRARY_RULES = [
+    ("aiter", "AITER"),
+    ("ck_tile::", "CK"),
+    ("ck_tile6kentry", "CK"),
+    ("FmhaFwd", "CK"),
+    ("FmhaBwd", "CK"),
+    ("Cijk_", "Tensile"),
+    ("wvSplitK", "rocBLAS"),
+    ("splitKreduce", "rocBLAS"),
+    ("rocprim::", "rocPRIM"),
+    ("triton_", "Triton"),
+    ("void at::native::", "PyTorch Native"),
+    ("nccl", "RCCL/NCCL"),
+    ("rccl", "RCCL/NCCL"),
+    ("composable", "CK"),
+]
+
+# Ordered GEMM detectors; extend when kernel_name_parser gains more libraries (#805).
+_GEMM_LIBRARY_CHECKS = (
+    (knp.is_rocm_gemm, "Tensile"),
+    (knp.is_cuda_gemm, "nvjet"),
+)
+
+
+def _coerce_kernel_details(kernel_details: Any) -> str:
+    if kernel_details in (None, ""):
+        return ""
+    if isinstance(kernel_details, float) and kernel_details != kernel_details:
+        return ""
+    return str(kernel_details)
+
+
+def _library_from_gemm_kernel_name(kernel_name: str) -> Optional[str]:
+    for detector, library in _GEMM_LIBRARY_CHECKS:
+        if detector(kernel_name):
+            return library
+    return None
+
+
+def classify_kernel_library(op_name: str, kernel_details: Any = "") -> Optional[str]:
+    """Identify the GPU library backing an operation from its name or kernel strings."""
+    op_lower = (op_name or "").lower()
+    for marker, library in _OP_NAME_LIBRARY_RULES:
+        if marker in op_lower:
+            return library
+
+    kd = _coerce_kernel_details(kernel_details)
+    if kd:
+        gemm_library = _library_from_gemm_kernel_name(kd)
+        if gemm_library is not None:
+            return gemm_library
+
+        kd_lower = kd.lower()
+        for marker, library in _KERNEL_NAME_LIBRARY_RULES:
+            if marker.lower() in kd_lower:
+                return library
+    return None

TraceLens/TraceIndex/__init__.py

@@ -0,0 +1,31 @@
+###############################################################################
+# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+"""Trace corpus indexing helpers."""
+
+from .core import (
+    append_trace,
+    build_traces,
+    execute_read_query,
+    generate_report_and_import,
+    search_index,
+)
+from .models import SearchHit, TraceRecord, TraceReport
+from .sqlite_store import SQLiteTraceIndexStore
+from .store import TraceIndexStore
+
+__all__ = [
+    "TraceIndexStore",
+    "SQLiteTraceIndexStore",
+    "SearchHit",
+    "TraceRecord",
+    "TraceReport",
+    "append_trace",
+    "build_traces",
+    "execute_read_query",
+    "generate_report_and_import",
+    "search_index",
+]