From eb0d698f792c6e10cf74f345405c6bf343166504 Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Wed, 26 Aug 2026 12:18:57 -0700 Subject: [PATCH 01/10] create report writing utility file --- .../Reporting/compare_perf_reports_pytorch.py | 22 ++------ ...te_multi_rank_collective_report_pytorch.py | 25 ++------- .../Reporting/generate_perf_report_genesis.py | 33 ++++------- .../Reporting/generate_perf_report_jax.py | 31 +++-------- ...nerate_perf_report_pftrace_hip_activity.py | 43 +++++---------- .../generate_perf_report_pftrace_hip_api.py | 47 +++++----------- ...enerate_perf_report_pftrace_memory_copy.py | 42 +++++--------- .../Reporting/generate_perf_report_pytorch.py | 27 +++------ .../generate_perf_report_pytorch_inference.py | 27 +++------ .../Reporting/generate_perf_report_rocprof.py | 44 +++++---------- TraceLens/Reporting/pftrace_utils.py | 17 ++++++ TraceLens/Reporting/reporting_utils.py | 55 +++++++++++++++++++ tests/test_genesis.py | 2 +- 13 files changed, 171 insertions(+), 244 deletions(-) diff --git a/TraceLens/Reporting/compare_perf_reports_pytorch.py b/TraceLens/Reporting/compare_perf_reports_pytorch.py index 8595dea35..1907e382a 100644 --- a/TraceLens/Reporting/compare_perf_reports_pytorch.py +++ b/TraceLens/Reporting/compare_perf_reports_pytorch.py @@ -13,6 +13,8 @@ import pandas as pd from openpyxl.utils import get_column_letter +from TraceLens.Reporting.reporting_utils import write_report_outputs + # ────────────────────────────────────────────────────────────────────────────── # Configuration # ────────────────────────────────────────────────────────────────────────────── @@ -589,30 +591,18 @@ def generate_compare_perf_reports_pytorch( # ── Write workbook ──────────────────────────────────────────────────────── if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in results.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print( - f"Wrote '{sheet_name}.csv' with {len(df)} rows × {len(df.columns)} columns" - ) + write_report_outputs(results, csvs_dir=output_csvs_dir) if output is not None: with pd.ExcelWriter(output, engine="openpyxl") as xls: for sheet_name, df in results.items(): - if df.empty: - print(f"Sheet '{sheet_name}' is empty (no matching rows)") - df.to_excel( - xls, sheet_name=sheet_name[:31], index=False - ) # Excel 31-char limit + safe = sheet_name[:31] + df.to_excel(xls, sheet_name=safe, index=False) for col in cols_to_hide_xl.get(sheet_name, []): col_idx = df.columns.get_loc(col) + 1 col_letter = get_column_letter(col_idx) - worksheet = xls.sheets[sheet_name[:31]] + worksheet = xls.sheets[safe] worksheet.column_dimensions[col_letter].hidden = True - print( - f"Wrote sheet '{sheet_name}' with {len(df)} rows × {len(df.columns)} columns" - ) return results diff --git a/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py b/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py index a861f683d..f4a867b5f 100644 --- a/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py +++ b/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py @@ -4,7 +4,6 @@ # See LICENSE for license information. ############################################################################### -import importlib.util import os import re import argparse @@ -16,7 +15,7 @@ from TraceLens.Reporting.reporting_utils import ( add_node_span_columns, detect_gpus_per_node, - request_install, + write_report_outputs, ) DEFAULT_RANK_REGEX = r"rank[\[\-_/]?(?P\d+)" @@ -257,25 +256,9 @@ def generate_collective_report( ) # Export DataFrames - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in report_dfs.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - - if output_xlsx_path: - if importlib.util.find_spec("openpyxl") is None: - print("Error importing openpyxl") - request_install("openpyxl") - - print(f"Writing Excel report to {output_xlsx_path}...") - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in report_dfs.items(): - df.to_excel( - writer, sheet_name=sheet_name[:31], index=False - ) # Excel limits sheet names to 31 chars - print(f"Excel report successfully written to {output_xlsx_path}") + write_report_outputs( + report_dfs, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir + ) return report_dfs diff --git a/TraceLens/Reporting/generate_perf_report_genesis.py b/TraceLens/Reporting/generate_perf_report_genesis.py index fc22dbd76..21ac48d1f 100644 --- a/TraceLens/Reporting/generate_perf_report_genesis.py +++ b/TraceLens/Reporting/generate_perf_report_genesis.py @@ -45,30 +45,17 @@ logger = logging.getLogger(__name__) -def _safe_sheet(name: str, used: set) -> str: - base = name[:31] - n = 0 - while base in used: - n += 1 - suffix = f"_{n}" - base = name[: 31 - len(suffix)] + suffix - used.add(base) - return base - - def write_excel(path: Path, sections: Dict[str, Dict[str, pd.DataFrame]]) -> None: - used: set = set() - with pd.ExcelWriter(path, engine="openpyxl") as writer: - for prefix, dfs in sections.items(): - for sheet, df in dfs.items(): - if df is None or df.empty: - continue - # rocprof sheets use short names (no prefix); pftrace keeps prefix for clarity - sheet_label = sheet if prefix == "rocprof" else f"{prefix}_{sheet}" - df.to_excel( - writer, sheet_name=_safe_sheet(sheet_label, used), index=False - ) - logger.info("Wrote %s", path) + from TraceLens.Reporting.reporting_utils import write_report_outputs + + flat: Dict[str, pd.DataFrame] = {} + for prefix, dfs in sections.items(): + for sheet, df in dfs.items(): + if df is None or df.empty: + continue + label = sheet if prefix == "rocprof" else f"{prefix}_{sheet}" + flat[label] = df + write_report_outputs(flat, xlsx_path=str(path)) def _rocprof_sheets_for_excel( diff --git a/TraceLens/Reporting/generate_perf_report_jax.py b/TraceLens/Reporting/generate_perf_report_jax.py index a0e2b81bf..a9f063d9a 100644 --- a/TraceLens/Reporting/generate_perf_report_jax.py +++ b/TraceLens/Reporting/generate_perf_report_jax.py @@ -5,7 +5,6 @@ ############################################################################### import argparse -import importlib.util import os import sys from typing import Optional, Dict @@ -22,8 +21,8 @@ from TraceLens.TreePerf import JaxTreePerfAnalyzer from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, - request_install, resolve_gpu_arch, + write_report_outputs, ) from TraceLens.util import TraceEventUtils @@ -165,26 +164,14 @@ def generate_perf_report_jax( ) # Write all DataFrames to separate sheets in an Excel workbook - if output_csvs_dir: - # Ensure the output directory exists - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - else: - if output_xlsx_path is None: - # split input path at 'xplane.pb' and take the first part and append '.xlsx' - base_path = profile_path.rsplit(".xplane.pb", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - if importlib.util.find_spec("openpyxl") is None: - print("Error importing openpyxl") - request_install("openpyxl") - - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - print(f"DataFrames successfully written to {output_xlsx_path}") + if not output_csvs_dir and output_xlsx_path is None: + base_path = profile_path.rsplit(".xplane.pb", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path if not output_csvs_dir else None, + csvs_dir=output_csvs_dir, + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py index d9f14c9b0..a10ff3255 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py @@ -11,7 +11,6 @@ API↔kernel correlation (see generate_perf_report_pftrace_hip_api for that). """ -import importlib.util import os import argparse import sys @@ -29,7 +28,11 @@ logger = logging.getLogger(__name__) from TraceLens.util import PftraceParser -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) +from TraceLens.Reporting.reporting_utils import write_report_outputs from TraceLens.Reporting.pftrace_hip_activity_analysis import ( PftraceHipActivityAnalyzer, ns_to_ms, @@ -154,33 +157,15 @@ def generate_perf_report_pftrace_hip_activity( dict_name2df["hip_summary"] = analyzer.get_df_hip_summary() logger.info(" - hip_summary (%d rows)", len(dict_name2df["hip_summary"])) - if output_csvs_dir: - logger.info("Writing CSV files to: %s", output_csvs_dir) - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(" - %s.csv (%d rows)", sheet_name, len(df)) - else: - if output_xlsx_path is None: - base = Path(trace_path).resolve() - if base.suffix.lower() == ".pftrace": - base = base.with_suffix("") - elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): - base = base.parent / base.name.replace(".json.gz", "") - else: - base = base.with_suffix("") - output_xlsx_path = str(base) + "_pftrace_activity_report.xlsx" - logger.info("Writing Excel to: %s", output_xlsx_path) - if importlib.util.find_spec("openpyxl") is None: - logger.error("openpyxl required for Excel output") - raise ImportError("openpyxl is required for Excel output") - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dict_name2df.items(): - sn = sheet_name[:31] - df.to_excel(writer, sheet_name=sn, index=False) - logger.info(" - Sheet '%s' (%d rows)", sn, len(df)) - logger.info("Successfully written to %s", output_xlsx_path) + if not output_csvs_dir and output_xlsx_path is None: + output_xlsx_path = derive_pftrace_output_path( + trace_path, "_pftrace_activity_report.xlsx" + ) + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path if not output_csvs_dir else None, + csvs_dir=output_csvs_dir, + ) if output_md_path: logger.info("Writing Markdown to: %s", output_md_path) diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py index 278a1bc88..a243344cf 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py @@ -4,7 +4,6 @@ # See LICENSE for license information. ############################################################################### -import importlib.util import os import re import argparse @@ -24,7 +23,11 @@ logger = logging.getLogger(__name__) from TraceLens.util import PftraceParser -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) +from TraceLens.Reporting.reporting_utils import write_report_outputs from TraceLens.Reporting.pftrace_hip_api_analysis import PftraceHipApiAnalyzer @@ -99,37 +102,15 @@ def generate_perf_report_pftrace_hip_api( ) # Write output - if output_csvs_dir: - logger.info(f"Writing CSV files to: {output_csvs_dir}") - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(f" - {sheet_name}.csv ({len(df)} rows)") - else: - if output_xlsx_path is None: - base = Path(trace_path).resolve() - if base.suffix.lower() == ".pftrace": - base = base.with_suffix("") - elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): - base = base.parent / base.name.replace(".json.gz", "") - else: - base = base.with_suffix("") - output_xlsx_path = str(base) + "_pftrace_hip_api_report.xlsx" - - logger.info(f"Writing Excel file to: {output_xlsx_path}") - if importlib.util.find_spec("openpyxl") is None: - logger.error( - "Error importing openpyxl. Please install: pip install openpyxl" - ) - raise ImportError("openpyxl is required for Excel output") - - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - logger.info(f" - Sheet '{sheet_name}' ({len(df)} rows)") - - logger.info(f"Successfully written to {output_xlsx_path}") + if not output_csvs_dir and output_xlsx_path is None: + output_xlsx_path = derive_pftrace_output_path( + trace_path, "_pftrace_hip_api_report.xlsx" + ) + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path if not output_csvs_dir else None, + csvs_dir=output_csvs_dir, + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py index e612443b0..7a4c433ef 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py @@ -11,7 +11,6 @@ Uses shared pftrace_utils (traceconv) and PftraceParser. """ -import importlib.util import os import argparse import sys @@ -29,7 +28,11 @@ logger = logging.getLogger(__name__) from TraceLens.util import PftraceParser -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) +from TraceLens.Reporting.reporting_utils import write_report_outputs # Event name substrings for direction (ROCm/Perfetto) NAME_HOST_TO_DEVICE = "MEMORY_COPY_HOST_TO_DEVICE" @@ -146,32 +149,15 @@ def generate_perf_report_pftrace_memory_copy( count_df = build_memory_copy_count_df(events) dfs = {"memory_copy_by_copy_bytes": count_df} - if output_csvs_dir: - logger.info("Writing CSV files to: %s", output_csvs_dir) - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dfs.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(" - %s.csv (%d rows)", sheet_name, len(df)) - else: - if output_xlsx_path is None: - base = Path(trace_path).resolve() - if base.suffix.lower() == ".pftrace": - base = base.with_suffix("") - elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): - base = base.parent / base.name.replace(".json.gz", "") - else: - base = base.with_suffix("") - output_xlsx_path = str(base) + "_pftrace_memory_copy_report.xlsx" - logger.info("Writing Excel file to: %s", output_xlsx_path) - if importlib.util.find_spec("openpyxl") is None: - logger.error("openpyxl required for Excel output. pip install openpyxl") - raise ImportError("openpyxl is required for Excel output") - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dfs.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - logger.info(" - Sheet '%s' (%d rows)", sheet_name, len(df)) - logger.info("Successfully written to %s", output_xlsx_path) + if not output_csvs_dir and output_xlsx_path is None: + output_xlsx_path = derive_pftrace_output_path( + trace_path, "_pftrace_memory_copy_report.xlsx" + ) + write_report_outputs( + dfs, + xlsx_path=output_xlsx_path if not output_csvs_dir else None, + csvs_dir=output_csvs_dir, + ) return dfs diff --git a/TraceLens/Reporting/generate_perf_report_pytorch.py b/TraceLens/Reporting/generate_perf_report_pytorch.py index 16bb9b644..e1e6ca4d8 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch.py @@ -19,8 +19,8 @@ from TraceLens.PerfModel.torch_op_mapping import build_sheet_category_to_op_names from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, - request_install, resolve_gpu_arch, + write_report_outputs, ) _WRAPPER_FILE_PATTERNS = frozenset( @@ -1029,25 +1029,12 @@ def _launcher_category(name): print(f"Added {len(additional_dfs)} additional sheets from extension") # Write CSVs and/or Excel (independent options) - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - - if output_xlsx_path is not None or output_csvs_dir is None: - if output_xlsx_path is None: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - if importlib.util.find_spec("openpyxl") is None: - print("Error importing openpyxl") - request_install("openpyxl") - - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - print(f"DataFrames successfully written to {output_xlsx_path}") + if output_xlsx_path is None and output_csvs_dir is None: + base_path = profile_json_path.rsplit(".json", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py index 905193a34..4f6e1a640 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py @@ -25,8 +25,8 @@ from TraceLens.Reporting.generate_perf_report_pytorch import _find_entry_point from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, - request_install, resolve_gpu_arch, + write_report_outputs, ) from TraceLens.util import TraceEventUtils from TraceLens.TraceUtils.annotation_utils import ( @@ -1168,25 +1168,12 @@ def _launcher_category(name): print(f"Added {len(additional_dfs)} additional sheets from extension") # Write CSVs and/or Excel (independent options) - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - - if output_xlsx_path is not None or output_csvs_dir is None: - if output_xlsx_path is None: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - if importlib.util.find_spec("openpyxl") is None: - print("Error importing openpyxl") - request_install("openpyxl") - - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - print(f"DataFrames successfully written to {output_xlsx_path}") + if output_xlsx_path is None and output_csvs_dir is None: + base_path = profile_json_path.rsplit(".json", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_rocprof.py b/TraceLens/Reporting/generate_perf_report_rocprof.py index d1835f341..02ef70cbd 100644 --- a/TraceLens/Reporting/generate_perf_report_rocprof.py +++ b/TraceLens/Reporting/generate_perf_report_rocprof.py @@ -4,7 +4,6 @@ # See LICENSE for license information. ############################################################################### -import importlib.util import os import argparse import sys @@ -22,6 +21,7 @@ from TraceLens.util import RocprofParser from TraceLens.Reporting.rocprof_analysis import RocprofAnalyzer +from TraceLens.Reporting.reporting_utils import write_report_outputs def generate_perf_report_rocprof( @@ -132,37 +132,19 @@ def generate_perf_report_rocprof( ) # 5. Write output - if output_csvs_dir: - logger.info(f"Writing CSV files to: {output_csvs_dir}") - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(f" - {sheet_name}.csv ({len(df)} rows)") - else: - if output_xlsx_path is None: - # Auto-generate output filename - if profile_json_path.endswith("_results.json"): - output_xlsx_path = profile_json_path.replace( - "_results.json", "_perf_report.xlsx" - ) - else: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - - logger.info(f"Writing Excel file to: {output_xlsx_path}") - if importlib.util.find_spec("openpyxl") is None: - logger.error( - "Error importing openpyxl. Please install it with: pip install openpyxl" + if not output_csvs_dir and output_xlsx_path is None: + if profile_json_path.endswith("_results.json"): + output_xlsx_path = profile_json_path.replace( + "_results.json", "_perf_report.xlsx" ) - raise ImportError("openpyxl is required for Excel output") - - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - logger.info(f" - Sheet '{sheet_name}' ({len(df)} rows)") - - logger.info(f"Successfully written to {output_xlsx_path}") + else: + base_path = profile_json_path.rsplit(".json", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path if not output_csvs_dir else None, + csvs_dir=output_csvs_dir, + ) return dict_name2df diff --git a/TraceLens/Reporting/pftrace_utils.py b/TraceLens/Reporting/pftrace_utils.py index b01cd95c7..5906bf121 100644 --- a/TraceLens/Reporting/pftrace_utils.py +++ b/TraceLens/Reporting/pftrace_utils.py @@ -86,3 +86,20 @@ def ensure_trace_json(trace_path: str, traceconv_path: Optional[str] = None) -> raise ValueError( f"Unsupported trace format: {trace_path}. Use .json, .json.gz, or .pftrace." ) + + +def derive_pftrace_output_path(trace_path: str, report_suffix: str) -> str: + """Derive a default output xlsx path from a pftrace input path. + + Strips ``.pftrace``, ``.json.gz``, or other extensions from + *trace_path* and appends *report_suffix* (e.g. + ``"_pftrace_activity_report.xlsx"``). + """ + base = Path(trace_path).resolve() + if base.suffix.lower() == ".pftrace": + base = base.with_suffix("") + elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): + base = base.parent / base.name.replace(".json.gz", "") + else: + base = base.with_suffix("") + return str(base) + report_suffix diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index 10ddc520a..dfd546073 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -8,6 +8,7 @@ import ast import json import logging +import os import re import subprocess import sys @@ -124,6 +125,60 @@ def export_data_df( data_df.to_csv(output_path, index=False) +def _safe_sheet_name(name: str, used: set) -> str: + """Truncate *name* to Excel's 31-char limit with collision avoidance. + + If the truncated name already appears in *used*, a numeric suffix + (``_1``, ``_2``, ...) is appended while staying within the limit. + The final name is added to *used* before returning. + """ + base = name[:31] + n = 0 + while base in used: + n += 1 + suffix = f"_{n}" + base = name[: 31 - len(suffix)] + suffix + used.add(base) + return base + + +def write_report_outputs( + dfs: Dict[str, pd.DataFrame], + *, + xlsx_path: Optional[str] = None, + csvs_dir: Optional[str] = None, +) -> None: + """Write report DataFrames to CSV files and/or an Excel workbook. + + When both *xlsx_path* and *csvs_dir* are provided, both outputs are + written. When neither is provided, nothing happens. Sheet names + are truncated to 31 characters (Excel limit) with collision-safe + suffixes. + """ + if csvs_dir: + os.makedirs(csvs_dir, exist_ok=True) + for sheet_name, df in dfs.items(): + csv_path = os.path.join(csvs_dir, f"{sheet_name}.csv") + df.to_csv(csv_path, index=False) + logger.info("Wrote %s (%d rows)", csv_path, len(df)) + + if xlsx_path: + import importlib.util + + if importlib.util.find_spec("openpyxl") is None: + raise ImportError( + "openpyxl is required for Excel output. Install with: " + "pip install openpyxl" + ) + + used: set = set() + with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer: + for sheet_name, df in dfs.items(): + safe = _safe_sheet_name(sheet_name, used) + df.to_excel(writer, sheet_name=safe, index=False) + logger.info("Wrote %s", xlsx_path) + + def request_install(package_name): """ Prompts the user to install a Python package via pip. If the user agrees, attempts installation. diff --git a/tests/test_genesis.py b/tests/test_genesis.py index d427053c0..e74aa7d79 100644 --- a/tests/test_genesis.py +++ b/tests/test_genesis.py @@ -41,10 +41,10 @@ from TraceLens.Reporting.generate_perf_report_genesis import ( _resolve_steady_state_fallback_s, _rocprof_sheets_for_excel, - _safe_sheet, write_excel, write_genesis_summary_md, ) +from TraceLens.Reporting.reporting_utils import _safe_sheet_name as _safe_sheet ############################################################################### # Shared fixtures — realistic CSV content from actual MI300X genesis traces From 234e536f587ca2b0599552599f8a9e8bf85e76d2 Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Fri, 28 Aug 2026 14:22:07 -0700 Subject: [PATCH 02/10] add columns to hide param --- .../Reporting/compare_perf_reports_pytorch.py | 20 ++++++------------- TraceLens/Reporting/reporting_utils.py | 18 +++++++++++++++++ 2 files changed, 24 insertions(+), 14 deletions(-) diff --git a/TraceLens/Reporting/compare_perf_reports_pytorch.py b/TraceLens/Reporting/compare_perf_reports_pytorch.py index 1907e382a..f7fc00871 100644 --- a/TraceLens/Reporting/compare_perf_reports_pytorch.py +++ b/TraceLens/Reporting/compare_perf_reports_pytorch.py @@ -11,7 +11,6 @@ from typing import Dict, List, Optional, Sequence import pandas as pd -from openpyxl.utils import get_column_letter from TraceLens.Reporting.reporting_utils import write_report_outputs @@ -590,19 +589,12 @@ def generate_compare_perf_reports_pytorch( cols_to_hide_xl[sheet_name] = cols_to_hide # ── Write workbook ──────────────────────────────────────────────────────── - if output_csvs_dir: - write_report_outputs(results, csvs_dir=output_csvs_dir) - - if output is not None: - with pd.ExcelWriter(output, engine="openpyxl") as xls: - for sheet_name, df in results.items(): - safe = sheet_name[:31] - df.to_excel(xls, sheet_name=safe, index=False) - for col in cols_to_hide_xl.get(sheet_name, []): - col_idx = df.columns.get_loc(col) + 1 - col_letter = get_column_letter(col_idx) - worksheet = xls.sheets[safe] - worksheet.column_dimensions[col_letter].hidden = True + write_report_outputs( + results, + xlsx_path=output, + csvs_dir=output_csvs_dir, + hide_columns=cols_to_hide_xl, + ) return results diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index dfd546073..6c3473352 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -147,6 +147,7 @@ def write_report_outputs( *, xlsx_path: Optional[str] = None, csvs_dir: Optional[str] = None, + hide_columns: Optional[Dict[str, List[str]]] = None, ) -> None: """Write report DataFrames to CSV files and/or an Excel workbook. @@ -154,6 +155,14 @@ def write_report_outputs( written. When neither is provided, nothing happens. Sheet names are truncated to 31 characters (Excel limit) with collision-safe suffixes. + + Args: + dfs: Mapping of sheet name -> DataFrame. + xlsx_path: If set, write an ``.xlsx`` workbook here. + csvs_dir: If set, write one CSV per DataFrame into this directory. + hide_columns: Optional mapping of sheet name -> column names to + hide in the Excel output. Hidden columns stay in the file; + names absent from a DataFrame are ignored. """ if csvs_dir: os.makedirs(csvs_dir, exist_ok=True) @@ -171,11 +180,20 @@ def write_report_outputs( "pip install openpyxl" ) + from openpyxl.utils import get_column_letter + + hide_columns = hide_columns or {} used: set = set() with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer: for sheet_name, df in dfs.items(): safe = _safe_sheet_name(sheet_name, used) df.to_excel(writer, sheet_name=safe, index=False) + worksheet = writer.sheets[safe] + for col in hide_columns.get(sheet_name, []): + if col not in df.columns: + continue + col_letter = get_column_letter(df.columns.get_loc(col) + 1) + worksheet.column_dimensions[col_letter].hidden = True logger.info("Wrote %s", xlsx_path) From 841fa1d3cfab40d09d037138099099941393e391 Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Fri, 28 Aug 2026 14:29:30 -0700 Subject: [PATCH 03/10] add skip empty df flag --- TraceLens/Reporting/generate_perf_report_genesis.py | 4 +--- TraceLens/Reporting/reporting_utils.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/TraceLens/Reporting/generate_perf_report_genesis.py b/TraceLens/Reporting/generate_perf_report_genesis.py index 21ac48d1f..f1ed8f952 100644 --- a/TraceLens/Reporting/generate_perf_report_genesis.py +++ b/TraceLens/Reporting/generate_perf_report_genesis.py @@ -51,11 +51,9 @@ def write_excel(path: Path, sections: Dict[str, Dict[str, pd.DataFrame]]) -> Non flat: Dict[str, pd.DataFrame] = {} for prefix, dfs in sections.items(): for sheet, df in dfs.items(): - if df is None or df.empty: - continue label = sheet if prefix == "rocprof" else f"{prefix}_{sheet}" flat[label] = df - write_report_outputs(flat, xlsx_path=str(path)) + write_report_outputs(flat, xlsx_path=str(path), skip_empty=True) def _rocprof_sheets_for_excel( diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index 6c3473352..ff0550477 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -148,13 +148,12 @@ def write_report_outputs( xlsx_path: Optional[str] = None, csvs_dir: Optional[str] = None, hide_columns: Optional[Dict[str, List[str]]] = None, + skip_empty: bool = False, ) -> None: """Write report DataFrames to CSV files and/or an Excel workbook. - When both *xlsx_path* and *csvs_dir* are provided, both outputs are - written. When neither is provided, nothing happens. Sheet names - are truncated to 31 characters (Excel limit) with collision-safe - suffixes. + Sheet names are truncated to 31 characters (Excel limit) with + collision-safe suffixes. Args: dfs: Mapping of sheet name -> DataFrame. @@ -163,7 +162,12 @@ def write_report_outputs( hide_columns: Optional mapping of sheet name -> column names to hide in the Excel output. Hidden columns stay in the file; names absent from a DataFrame are ignored. + skip_empty: If True, ``None`` or empty DataFrames are omitted from + all outputs. """ + if skip_empty: + dfs = {name: df for name, df in dfs.items() if df is not None and not df.empty} + if csvs_dir: os.makedirs(csvs_dir, exist_ok=True) for sheet_name, df in dfs.items(): From f0340254519aa3f1180b8a5733cb420907da63ac Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Fri, 28 Aug 2026 14:33:07 -0700 Subject: [PATCH 04/10] move imports --- TraceLens/Reporting/generate_perf_report_genesis.py | 3 +-- TraceLens/Reporting/reporting_utils.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/TraceLens/Reporting/generate_perf_report_genesis.py b/TraceLens/Reporting/generate_perf_report_genesis.py index f1ed8f952..6882ed55c 100644 --- a/TraceLens/Reporting/generate_perf_report_genesis.py +++ b/TraceLens/Reporting/generate_perf_report_genesis.py @@ -36,6 +36,7 @@ pftrace_to_json, resolve_profile_json, ) +from TraceLens.Reporting.reporting_utils import write_report_outputs logging.basicConfig( stream=sys.stdout, @@ -46,8 +47,6 @@ def write_excel(path: Path, sections: Dict[str, Dict[str, pd.DataFrame]]) -> None: - from TraceLens.Reporting.reporting_utils import write_report_outputs - flat: Dict[str, pd.DataFrame] = {} for prefix, dfs in sections.items(): for sheet, df in dfs.items(): diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index ff0550477..3ccf4f814 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -6,6 +6,7 @@ import argparse import ast +import importlib.util import json import logging import os @@ -176,8 +177,6 @@ def write_report_outputs( logger.info("Wrote %s (%d rows)", csv_path, len(df)) if xlsx_path: - import importlib.util - if importlib.util.find_spec("openpyxl") is None: raise ImportError( "openpyxl is required for Excel output. Install with: " From 3567ebb90b56b232553837786fc6a5fcc5416e64 Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Fri, 28 Aug 2026 14:51:01 -0700 Subject: [PATCH 05/10] add test file --- .../Reporting/generate_perf_report_jax.py | 2 +- ...nerate_perf_report_pftrace_hip_activity.py | 2 +- .../generate_perf_report_pftrace_hip_api.py | 2 +- ...enerate_perf_report_pftrace_memory_copy.py | 2 +- .../Reporting/generate_perf_report_rocprof.py | 2 +- TraceLens/Reporting/reporting_utils.py | 20 +- tests/test_genesis.py | 41 ---- tests/test_reporting_utils.py | 198 +++++++++++++++++- 8 files changed, 211 insertions(+), 58 deletions(-) diff --git a/TraceLens/Reporting/generate_perf_report_jax.py b/TraceLens/Reporting/generate_perf_report_jax.py index a9f063d9a..3170c16cb 100644 --- a/TraceLens/Reporting/generate_perf_report_jax.py +++ b/TraceLens/Reporting/generate_perf_report_jax.py @@ -169,7 +169,7 @@ def generate_perf_report_jax( output_xlsx_path = base_path + "_perf_report.xlsx" write_report_outputs( dict_name2df, - xlsx_path=output_xlsx_path if not output_csvs_dir else None, + xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir, ) diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py index a10ff3255..044de15d2 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py @@ -163,7 +163,7 @@ def generate_perf_report_pftrace_hip_activity( ) write_report_outputs( dict_name2df, - xlsx_path=output_xlsx_path if not output_csvs_dir else None, + xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir, ) diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py index a243344cf..b5c9d564b 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py @@ -108,7 +108,7 @@ def generate_perf_report_pftrace_hip_api( ) write_report_outputs( dict_name2df, - xlsx_path=output_xlsx_path if not output_csvs_dir else None, + xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir, ) diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py index 7a4c433ef..36c96875d 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py @@ -155,7 +155,7 @@ def generate_perf_report_pftrace_memory_copy( ) write_report_outputs( dfs, - xlsx_path=output_xlsx_path if not output_csvs_dir else None, + xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir, ) diff --git a/TraceLens/Reporting/generate_perf_report_rocprof.py b/TraceLens/Reporting/generate_perf_report_rocprof.py index 02ef70cbd..94696bdc1 100644 --- a/TraceLens/Reporting/generate_perf_report_rocprof.py +++ b/TraceLens/Reporting/generate_perf_report_rocprof.py @@ -142,7 +142,7 @@ def generate_perf_report_rocprof( output_xlsx_path = base_path + "_perf_report.xlsx" write_report_outputs( dict_name2df, - xlsx_path=output_xlsx_path if not output_csvs_dir else None, + xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir, ) diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index 3ccf4f814..a82179155 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -6,7 +6,6 @@ import argparse import ast -import importlib.util import json import logging import os @@ -17,6 +16,7 @@ from typing import Dict, List, Optional, Union import pandas as pd +from openpyxl.utils import get_column_letter logger = logging.getLogger(__name__) @@ -177,22 +177,20 @@ def write_report_outputs( logger.info("Wrote %s (%d rows)", csv_path, len(df)) if xlsx_path: - if importlib.util.find_spec("openpyxl") is None: - raise ImportError( - "openpyxl is required for Excel output. Install with: " - "pip install openpyxl" - ) - - from openpyxl.utils import get_column_letter - hide_columns = hide_columns or {} used: set = set() with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer: for sheet_name, df in dfs.items(): safe = _safe_sheet_name(sheet_name, used) df.to_excel(writer, sheet_name=safe, index=False) - worksheet = writer.sheets[safe] - for col in hide_columns.get(sheet_name, []): + cols_to_hide = hide_columns.get(sheet_name, []) + if not cols_to_hide: + continue + # openpyxl may adjust the sheet title (e.g. a case-insensitive + # collision with the default "Sheet"), so grab the worksheet we + # just wrote directly rather than looking it up by name. + worksheet = writer.book.worksheets[-1] + for col in cols_to_hide: if col not in df.columns: continue col_letter = get_column_letter(df.columns.get_loc(col) + 1) diff --git a/tests/test_genesis.py b/tests/test_genesis.py index e74aa7d79..7f9cfe0de 100644 --- a/tests/test_genesis.py +++ b/tests/test_genesis.py @@ -44,7 +44,6 @@ write_excel, write_genesis_summary_md, ) -from TraceLens.Reporting.reporting_utils import _safe_sheet_name as _safe_sheet ############################################################################### # Shared fixtures — realistic CSV content from actual MI300X genesis traces @@ -812,46 +811,6 @@ def test_raises_when_no_data(self): resolve_profile_json(capture_dict, output_dir, include_api=False) -############################################################################### -# generate_perf_report_genesis — _safe_sheet -############################################################################### - - -class TestSafeSheet: - """Validate Excel sheet name deduplication and length limits.""" - - def test_no_collision(self): - used = set() - name = _safe_sheet("gpu_timeline", used) - assert name == "gpu_timeline" - assert "gpu_timeline" in used - - def test_collision_adds_suffix(self): - used = {"gpu_timeline"} - name = _safe_sheet("gpu_timeline", used) - assert name == "gpu_timeline_1" - assert "gpu_timeline_1" in used - - def test_multiple_collisions(self): - used = {"test_sheet", "test_sheet_1", "test_sheet_2"} - name = _safe_sheet("test_sheet", used) - assert name == "test_sheet_3" - - def test_truncates_to_31_chars(self): - used = set() - long_name = "a" * 50 - name = _safe_sheet(long_name, used) - assert len(name) <= 31 - - def test_truncation_with_collision(self): - long_name = "a" * 31 - used = {long_name} - name = _safe_sheet(long_name, used) - assert len(name) <= 31 - assert name != long_name - assert name.endswith("_1") - - ############################################################################### # generate_perf_report_genesis — _rocprof_sheets_for_excel ############################################################################### diff --git a/tests/test_reporting_utils.py b/tests/test_reporting_utils.py index c90a127de..69ca93690 100644 --- a/tests/test_reporting_utils.py +++ b/tests/test_reporting_utils.py @@ -11,12 +11,14 @@ from TraceLens.Reporting.reporting_utils import ( _node_span_for_pg, _parse_pg_ranks, + _safe_sheet_name, add_gpu_arch_cli_args, add_node_span_columns, detect_gpus_per_node, export_data_df, request_install, resolve_gpu_arch, + write_report_outputs, ) from TraceLens.Agent.Analysis.category_analyses import ( analysis_utils as au, @@ -120,7 +122,10 @@ _cleanup_work_dir, generate_perf_report_genesis, ) -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) GPU_ONLY_TRACE = os.path.join( os.path.dirname(__file__), @@ -3165,3 +3170,194 @@ def test_pytorch_report_extension_and_arch(self, tmp_path): include_call_stack=True, ) assert (tmp_path / "ext_out" / "gpu_timeline.csv").exists() + + +############################################################################### +# reporting_utils — _safe_sheet_name +############################################################################### + + +class TestSafeSheetName: + """Validate Excel sheet name deduplication and length limits.""" + + def test_no_collision(self): + used = set() + name = _safe_sheet_name("gpu_timeline", used) + assert name == "gpu_timeline" + assert "gpu_timeline" in used + + def test_collision_adds_suffix(self): + used = {"gpu_timeline"} + name = _safe_sheet_name("gpu_timeline", used) + assert name == "gpu_timeline_1" + assert "gpu_timeline_1" in used + + def test_multiple_collisions(self): + used = {"test_sheet", "test_sheet_1", "test_sheet_2"} + name = _safe_sheet_name("test_sheet", used) + assert name == "test_sheet_3" + + def test_truncates_to_31_chars(self): + used = set() + long_name = "a" * 50 + name = _safe_sheet_name(long_name, used) + assert len(name) <= 31 + + def test_truncation_with_collision(self): + long_name = "a" * 31 + used = {long_name} + name = _safe_sheet_name(long_name, used) + assert len(name) <= 31 + assert name != long_name + assert name.endswith("_1") + + +############################################################################### +# reporting_utils — write_report_outputs +############################################################################### + + +def _read_sheets(xlsx_path): + """Return {sheet_name: DataFrame} for an .xlsx file.""" + return pd.read_excel(xlsx_path, sheet_name=None) + + +class TestWriteReportOutputs: + def _dfs(self): + return { + "alpha": pd.DataFrame({"a": [1, 2], "b": [3, 4]}), + "beta": pd.DataFrame({"x": [5]}), + } + + def test_csvs_only(self, tmp_path): + out = tmp_path / "csvs" + write_report_outputs(self._dfs(), csvs_dir=str(out)) + assert (out / "alpha.csv").exists() + assert (out / "beta.csv").exists() + # No xlsx requested -> none written + assert not list(tmp_path.glob("*.xlsx")) + pd.testing.assert_frame_equal( + pd.read_csv(out / "alpha.csv"), self._dfs()["alpha"] + ) + + def test_xlsx_only(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + write_report_outputs(self._dfs(), xlsx_path=str(xlsx)) + assert xlsx.exists() + assert not list(tmp_path.glob("*.csv")) + sheets = _read_sheets(xlsx) + assert set(sheets) == {"alpha", "beta"} + + def test_both_outputs_written(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + csvs = tmp_path / "csvs" + write_report_outputs(self._dfs(), xlsx_path=str(xlsx), csvs_dir=str(csvs)) + assert xlsx.exists() + assert (csvs / "alpha.csv").exists() + assert (csvs / "beta.csv").exists() + + def test_neither_output_is_noop(self, tmp_path): + write_report_outputs(self._dfs()) + assert not list(tmp_path.iterdir()) + + def test_long_sheet_names_truncated_and_deduped(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + dfs = { + "a" * 40: pd.DataFrame({"c": [1]}), + "a" * 45: pd.DataFrame({"c": [2]}), # truncates to same 31 chars -> deduped + } + write_report_outputs(dfs, xlsx_path=str(xlsx)) + names = list(_read_sheets(xlsx)) + assert all(len(n) <= 31 for n in names) + assert len(names) == 2 # no collision -> both sheets present + + def test_skip_empty_drops_empty_and_none(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + csvs = tmp_path / "csvs" + dfs = { + "keep": pd.DataFrame({"a": [1]}), + "empty": pd.DataFrame(), + "none": None, + } + write_report_outputs( + dfs, xlsx_path=str(xlsx), csvs_dir=str(csvs), skip_empty=True + ) + assert set(_read_sheets(xlsx)) == {"keep"} + assert (csvs / "keep.csv").exists() + assert not (csvs / "empty.csv").exists() + assert not (csvs / "none.csv").exists() + + def test_skip_empty_default_keeps_empty(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + dfs = { + "keep": pd.DataFrame({"a": [1]}), + "empty": pd.DataFrame({"a": []}), + } + write_report_outputs(dfs, xlsx_path=str(xlsx)) + assert set(_read_sheets(xlsx)) == {"keep", "empty"} + + def test_hide_columns_hides_and_keeps_data(self, tmp_path): + from openpyxl import load_workbook + from openpyxl.utils import get_column_letter + + xlsx = tmp_path / "report.xlsx" + df = pd.DataFrame({"keep": [1], "hide_me": [2], "also_keep": [3]}) + write_report_outputs( + {"data": df}, xlsx_path=str(xlsx), hide_columns={"data": ["hide_me"]} + ) + ws = load_workbook(xlsx)["data"] + hidden_col = get_column_letter(df.columns.get_loc("hide_me") + 1) # "B" + assert ws.column_dimensions[hidden_col].hidden is True + # Non-hidden columns are not hidden, and the data is still present + keep_col = get_column_letter(df.columns.get_loc("keep") + 1) + assert ws.column_dimensions[keep_col].hidden is False + assert set(_read_sheets(xlsx)["data"].columns) == { + "keep", + "hide_me", + "also_keep", + } + + def test_hide_columns_ignores_missing_column(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + df = pd.DataFrame({"a": [1]}) + # Should not raise even though "nonexistent" isn't a column + write_report_outputs( + {"data": df}, xlsx_path=str(xlsx), hide_columns={"data": ["nonexistent"]} + ) + assert xlsx.exists() + + def test_sheet_named_like_openpyxl_default(self, tmp_path): + # "sheet" collides case-insensitively with openpyxl's default "Sheet", + # so openpyxl renames it; the worksheet lookup must not rely on the + # requested name. hide_columns exercises that lookup path. + xlsx = tmp_path / "report.xlsx" + df = pd.DataFrame({"a": [1], "b": [2]}) + write_report_outputs( + {"sheet": df}, xlsx_path=str(xlsx), hide_columns={"sheet": ["b"]} + ) + assert xlsx.exists() + + +############################################################################### +# pftrace_utils — derive_pftrace_output_path +############################################################################### + + +class TestDerivePftraceOutputPath: + def test_pftrace_suffix(self): + assert ( + derive_pftrace_output_path("/tmp/trace.pftrace", "_activity_report.xlsx") + == "/tmp/trace_activity_report.xlsx" + ) + + def test_json_gz_suffix(self): + assert ( + derive_pftrace_output_path("/tmp/trace.json.gz", "_hip_api_report.xlsx") + == "/tmp/trace_hip_api_report.xlsx" + ) + + def test_plain_json_suffix(self): + assert ( + derive_pftrace_output_path("/tmp/trace.json", "_memory_copy_report.xlsx") + == "/tmp/trace_memory_copy_report.xlsx" + ) From c376aaf6fe7c73204d10dccd058e61580fb01817 Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Fri, 28 Aug 2026 14:56:43 -0700 Subject: [PATCH 06/10] remove comment --- TraceLens/Reporting/reporting_utils.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index a82179155..6b06fb29a 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -186,9 +186,6 @@ def write_report_outputs( cols_to_hide = hide_columns.get(sheet_name, []) if not cols_to_hide: continue - # openpyxl may adjust the sheet title (e.g. a case-insensitive - # collision with the default "Sheet"), so grab the worksheet we - # just wrote directly rather than looking it up by name. worksheet = writer.book.worksheets[-1] for col in cols_to_hide: if col not in df.columns: From 8ecf92716f151fa30a75052b44ebcff473e1b40f Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Mon, 31 Aug 2026 14:18:26 -0700 Subject: [PATCH 07/10] remove unused imports --- TraceLens/Reporting/generate_perf_report_jax.py | 1 - TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py | 1 - TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py | 1 - 3 files changed, 3 deletions(-) diff --git a/TraceLens/Reporting/generate_perf_report_jax.py b/TraceLens/Reporting/generate_perf_report_jax.py index 3170c16cb..f0bafdfb8 100644 --- a/TraceLens/Reporting/generate_perf_report_jax.py +++ b/TraceLens/Reporting/generate_perf_report_jax.py @@ -5,7 +5,6 @@ ############################################################################### import argparse -import os import sys from typing import Optional, Dict import pandas as pd diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py index b5c9d564b..d8b0cedf7 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py @@ -8,7 +8,6 @@ import re import argparse import sys -from pathlib import Path from typing import Optional, Dict import pandas as pd diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py index 36c96875d..b8dcfffa0 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py @@ -14,7 +14,6 @@ import os import argparse import sys -from pathlib import Path from typing import Optional, Dict, List, Any, Tuple import pandas as pd From 66a256eb726105652cfedabd960ab23a02d4a6c1 Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Mon, 31 Aug 2026 14:18:26 -0700 Subject: [PATCH 08/10] remove unused imports From b7dafc5716b043f972939741c018ebd5286ef41d Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Mon, 31 Aug 2026 17:20:00 -0700 Subject: [PATCH 09/10] merge generate_perf_report_pytorch scripts --- .github/CODEOWNERS | 2 +- .../Reporting/generate_perf_report_pytorch.py | 450 +++++- .../generate_perf_report_pytorch_inference.py | 1402 ----------------- .../generate-perf-report-pytorch-inference.md | 15 +- docs/reference/api-reference.md | 13 +- docs/what-is-tracelens.md | 1 - setup.py | 1 - tests/test_inference_perf_report.py | 4 +- tests/test_pftrace_hip_activity_report.py | 6 +- tests/test_pseudo_ops_extension.py | 2 +- tests/test_reporting_inference_helpers.py | 4 +- tests/test_reporting_utils.py | 13 +- 12 files changed, 451 insertions(+), 1462 deletions(-) delete mode 100644 TraceLens/Reporting/generate_perf_report_pytorch_inference.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 580224a56..20a17b052 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -22,7 +22,7 @@ /TraceLens/Trace2Tree/trace_capture_merge_experimental.py @devalshahamd # Inference reporting, docs, tests, and workflows require Deval review -/TraceLens/Reporting/generate_perf_report_pytorch_inference.py @devalshahamd +/TraceLens/Reporting/generate_perf_report_pytorch.py @devalshahamd /TraceLens/TraceUtils/split_inference_trace_annotation.py @devalshahamd /examples/custom_workflows/inference_analysis/ @devalshahamd /docs/how-to/generate-perf-report-pytorch-inference.md @devalshahamd diff --git a/TraceLens/Reporting/generate_perf_report_pytorch.py b/TraceLens/Reporting/generate_perf_report_pytorch.py index e1e6ca4d8..1fe0843dc 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch.py @@ -6,22 +6,36 @@ import argparse import ast +import collections +import gzip import importlib.util +import json import os import re +import sys import warnings +import zipfile from typing import Dict, Optional import numpy as np import pandas as pd -from TraceLens import NcclAnalyser, TraceDiff, TreePerfAnalyzer +from TraceLens import NcclAnalyser, TraceToTree, TraceDiff, TreePerfAnalyzer from TraceLens.PerfModel.torch_op_mapping import build_sheet_category_to_op_names from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, resolve_gpu_arch, write_report_outputs, ) +from TraceLens.util import TraceEventUtils +from TraceLens.TraceUtils.annotation_utils import ( + CAPTURE_PATTERN, + CaptureAnnotation, + find_events_by_patterns, +) +from TraceLens.Trace2Tree.trace_capture_merge_experimental import ( + merge_capture_trace_into_graph, +) _WRAPPER_FILE_PATTERNS = frozenset( { @@ -180,6 +194,257 @@ def flatten(frames): return empty +def perf_report_sanity_check( + events, + df_gpu_timeline, + df_kernel_launchers, + df_unified_perf, + include_nccl=False, +): + """ + Sanity checks on the performance report DataFrames. + + 1) Total kernel time accounted by df_kernel_launchers and df_unified_perf + should each be >= the computation_time reported in df_gpu_timeline. + 2) Total GPU events in tree events should equal the number of kernels + accounted by df_kernel_launchers and df_unified_perf. + """ + use_time = "busy_time" if include_nccl else "computation_time" + + computation_time_us = ( + df_gpu_timeline.loc[df_gpu_timeline["type"] == use_time, "time ms"].values[0] + * 1e3 + ) + + # --- Check 1: kernel time coverage --- + kl_time_col = ( + "total_direct_kernel_time_sum" + if "total_direct_kernel_time_sum" in df_kernel_launchers.columns + else "total_direct_kernel_time" + ) + up_time_col = ( + "Kernel Time (µs)_sum" + if "Kernel Time (µs)_sum" in df_unified_perf.columns + else "Kernel Time (µs)" + ) + + kl_total_us = df_kernel_launchers[kl_time_col].sum() + up_total_us = df_unified_perf[up_time_col].sum() + + print(f"\n{'='*60}") + print("Perf Report Sanity Check") + print(f"{'='*60}") + print(f" {use_time} (gpu_timeline): {computation_time_us:.2f} µs") + print(f" Kernel time (kernel_launchers): {kl_total_us:.2f} µs") + print(f" Kernel time (unified_perf): {up_total_us:.2f} µs") + + kl_pass = kl_total_us >= computation_time_us + up_pass = up_total_us >= computation_time_us + print(f" kernel_launchers >= computation_time: {'PASS' if kl_pass else 'FAIL'}") + print(f" unified_perf >= computation_time: {'PASS' if up_pass else 'FAIL'}") + + # --- Check 2: per-kernel-name count verification --- + # Build {kernel_name: count} from tree events (ground truth) + tree_kernel_counts = dict( + collections.Counter( + e["name"] + for e in events + if e.get("cat") in {"kernel", "gpu_memcpy", "gpu_memset"} + and ( + not TraceEventUtils.is_communication_string(e.get("name", "")) + or include_nccl + ) + ) + ) + + def _extract_kernel_counts(df, label): + """Extract {kernel_name: count} from a DataFrame's kernel_details column.""" + if "kernel_details_summary" in df.columns: + col = "kernel_details_summary" + elif "kernel_details" in df.columns: + col = "kernel_details" + else: + print(f" WARNING: no kernel_details column in {label}") + return {} + counts = collections.Counter() + for kd in df[col]: + if isinstance(kd, list): + for d in kd: + counts[d["name"]] += d.get("count", 1) + return dict(counts) + + def _print_per_kernel_check(tree_counts, df_counts, label): + """Compare per-kernel counts and print mismatches.""" + df_total = sum(df_counts.values()) + tree_total = sum(tree_counts.values()) + total_pass = tree_total == df_total + print(f"\n --- {label} ---") + print(f" Total GPU events in tree: {tree_total}") + print(f" Kernels accounted: {df_total}") + print(f" Total count match: {'PASS' if total_pass else 'FAIL'}") + + all_names = sorted(set(tree_counts) | set(df_counts)) + mismatches = [] + for name in all_names: + t = tree_counts.get(name, 0) + d = df_counts.get(name, 0) + if t != d: + mismatches.append((name, t, d)) + + if mismatches: + print(f" Per-kernel mismatches ({len(mismatches)}):") + for name, t, d in mismatches: + trunc = name[:80] + "..." if len(name) > 80 else name + print(f" {trunc}") + print(f" tree={t} {label}={d} diff={t - d}") + else: + print(f" Per-kernel detail check: PASS (all match)") + + return df_total, total_pass, mismatches + + # Build {kernel_name: count} from each source + kl_kernel_counts = _extract_kernel_counts(df_kernel_launchers, "kernel_launchers") + up_kernel_counts = _extract_kernel_counts(df_unified_perf, "unified_perf") + + total_gpu_events = sum(tree_kernel_counts.values()) + + kl_kernel_count, kl_count_pass, kl_mismatches = _print_per_kernel_check( + tree_kernel_counts, kl_kernel_counts, "kernel_launchers" + ) + up_kernel_count, up_count_pass, up_mismatches = _print_per_kernel_check( + tree_kernel_counts, up_kernel_counts, "unified_perf" + ) + + print(f"{'='*60}\n") + + return { + "computation_time_us": computation_time_us, + "kl_total_us": kl_total_us, + "up_total_us": up_total_us, + "kl_time_pass": kl_pass, + "up_time_pass": up_pass, + "total_gpu_events": total_gpu_events, + "kl_kernel_count": kl_kernel_count, + "up_kernel_count": up_kernel_count, + "kl_count_pass": kl_count_pass, + "up_count_pass": up_count_pass, + "tree_kernel_counts": dict(tree_kernel_counts), + "kl_kernel_counts": dict(kl_kernel_counts), + "up_kernel_counts": dict(up_kernel_counts), + "kl_mismatches": kl_mismatches, + "up_mismatches": up_mismatches, + } + + +def classify_graph_capture_trace(input_folder: str): + """ + Return {file, batch_size, mode} for a single graph-capture trace file. + Supports .json, .json.gz, and .zip (containing a .json). + """ + execution_details_path = os.path.join(input_folder, "execution_details.json") + if os.path.isfile(execution_details_path): + print( + f"Execution details already exist at {execution_details_path}. Skipping classification." + ) + return + ## vLLM specific dummy run pattern + dummy_run_pattern = re.compile( + r"vllm/v1/worker/gpu_model_runner\.py\(\d+\): _dummy_run" + ) + + def load_trace(path: str) -> dict: + if path.endswith(".zip"): + with zipfile.ZipFile(path, "r") as zf: + json_files = [f for f in zf.namelist() if f.endswith(".json")] + if not json_files: + raise ValueError(f"No .json file found inside {path}") + with zf.open(json_files[0]) as f: + return json.load(f) + if path.endswith(".json.gz"): + with gzip.open(path, "rt") as f: + return json.load(f) + with open(path, "r") as f: + return json.load(f) + + def find_dummy_run_roots(events): + roots = [e for e in events if dummy_run_pattern.match(e.get("name", ""))] + roots.sort(key=lambda x: x.get("ts", 0)) + return roots + + def count_stream_begin_captures(events): + return sum( + 1 + for e in events + if "StreamBeginCapture" in e.get("name", "") + and e.get("cat") == "cuda_runtime" + ) + + def infer_batch_size_from_cpu_ops(events): + first_dims = [] + for e in events: + if e.get("cat") != "cpu_op": + continue + input_dims = e.get("args", {}).get("Input Dims") + if not input_dims: + continue + for dim_list in input_dims: + if isinstance(dim_list, list) and dim_list: + if isinstance(dim_list[0], int): + first_dims.append(dim_list[0]) + if not first_dims: + return None + return collections.Counter(first_dims).most_common(1)[0][0] + + def infer_mode_from_captures(num_captures: int): + return "FULL" if num_captures <= 1 else "PIECEWISE" + + if not os.path.isdir(input_folder): + print(f"Error: {input_folder} is not a directory", file=sys.stderr) + sys.exit(1) + + trace_files = sorted( + os.path.join(input_folder, f) + for f in os.listdir(input_folder) + if f.endswith(".json") or f.endswith(".json.gz") + ) + + if not trace_files: + print(f"No files starting with 'graph_capture_rank_0' found in {input_folder}") + sys.exit(0) + + print(f"Found {len(trace_files)} graph-capture trace file(s) in {input_folder}\n") + + results = [] + for filepath in trace_files: + trace_json = load_trace(filepath) + events = trace_json.get("traceEvents", []) + dummy_roots = find_dummy_run_roots(events) + annotation_roots = find_events_by_patterns(events, [CAPTURE_PATTERN]) + basename = os.path.basename(filepath) + + if annotation_roots and len(annotation_roots) == len(dummy_roots): + cap = CaptureAnnotation(annotation_roots[0]["name"]) + batch_size, mode = cap.batch_size, cap.mode + print( + f"batch_size: {batch_size}, mode: {mode} parsed from annotation, num_captures: {count_stream_begin_captures(events)}" + ) + results.append({"file": basename, "batch_size": batch_size, "mode": mode}) + continue + + num_captures = count_stream_begin_captures(events) + mode = infer_mode_from_captures(num_captures) + batch_size = infer_batch_size_from_cpu_ops(events) + print( + f"batch_size: {batch_size}, mode: {mode} inferred, num_captures: {num_captures}" + ) + + results.append({"file": basename, "batch_size": batch_size, "mode": mode}) + with open(f"{input_folder}/execution_details.json", "w") as f: + json.dump(results, f, indent=2) + print(f"\nResults written to {input_folder}/execution_details.json") + return + + def get_dfs_short_kernels( perf_analyzer, short_kernel_threshold_us=10, histogram_bins=100, topk=None ): @@ -399,6 +664,7 @@ def add_truncated_kernel_details( def generate_perf_report_pytorch( profile_json_path: str, + augmented_tree: TraceToTree = None, output_xlsx_path: Optional[str] = None, output_csvs_dir: Optional[str] = None, # include unlinked kernels in gpu timeline @@ -422,6 +688,7 @@ def generate_perf_report_pytorch( topk_ops: Optional[int] = None, topk_roofline_ops: Optional[int] = None, comparison_json_path: Optional[str] = None, + comparison_augmented_tree: Optional[TraceToTree] = None, extension_file: Optional[str] = None, # for gemm simulator / Origami (Origami requires --enable_origami when arch is set) python_path: Optional[str] = None, @@ -433,6 +700,7 @@ def generate_perf_report_pytorch( enable_origami: bool = False, # activation recompute detection detect_recompute: bool = False, + group_by_parent_module: bool = False, include_call_stack: bool = False, ) -> Dict[str, pd.DataFrame]: gpu_arch_json = resolve_gpu_arch( @@ -440,22 +708,68 @@ def generate_perf_report_pytorch( gpu_arch_platform=gpu_arch_platform, gpu_arch=gpu_arch, ) - add_python_func = True if include_call_stack else False - perf_analyzer = TreePerfAnalyzer.from_file( - profile_filepath=profile_json_path, - arch=gpu_arch_json, - python_path=python_path, - include_unlinked_kernels=include_unlinked_kernels, - enable_pseudo_ops=enable_pseudo_ops, - add_python_func=add_python_func, - detect_recompute=detect_recompute, - enable_origami=enable_origami, - inductor_cache_dir=inductor_cache_dir, + add_python_func = ( + True + if ( + group_by_parent_module + or include_call_stack is True + or augmented_tree is not None + or comparison_augmented_tree is not None + ) + else False ) + if augmented_tree is not None: + perf_analyzer = TreePerfAnalyzer( + tree=augmented_tree, + arch=gpu_arch_json, + python_path=python_path, + include_unlinked_kernels=include_unlinked_kernels, + add_python_func=add_python_func, + enable_pseudo_ops=enable_pseudo_ops, + rebuild_tree=False, + ) + else: + perf_analyzer = TreePerfAnalyzer.from_file( + profile_filepath=profile_json_path, + arch=gpu_arch_json, + python_path=python_path, + include_unlinked_kernels=include_unlinked_kernels, + enable_pseudo_ops=enable_pseudo_ops, + add_python_func=add_python_func, + detect_recompute=detect_recompute, + enable_origami=enable_origami, + inductor_cache_dir=inductor_cache_dir, + ) + + graph_launch_events = [ + event + for event in perf_analyzer.tree.events + if "graphlaunch" in event.get("name", "").lower() + ] + if len(graph_launch_events) > 0: + warnings.warn( + f"There are hipgraph launches (Count: {len(graph_launch_events)}) in this trace, but a graph capture folder not provided, the analysis might be limited", + UserWarning, + ) ## Apply annotation for vLLM eager and replay phase perf_analyzer.tree.apply_annotation( - name_filters=["vllm::unified_attention_with_output"] + name_filters=[ + "vllm::unified_attention_with_output", + "aiter::mha_varlen_fwd", + "pseudo_mla_decode_fwd", + "pseudo_mla_prefill_fwd", + "vllm::gdn_attention_core", + "aiter::fmha_v3_varlen_fwd", + "sglang_profiler::tilelang_kernel_tilelang_sparse_fwd", + "sglang_profiler::attention_paged_attention_ragged", + "aiter::mha_batch_prefill", + "aiter::pa_decode_gluon", + "aiter::v4_attention_with_output", + "pseudo_v4_paged_decode_swa", + "pseudo_v4_paged_decode_csa", + "pseudo_v4_paged_decode_hca", + ] ) if extension_file: @@ -484,6 +798,7 @@ def generate_perf_report_pytorch( df_kernel_launchers_summary_by_category = pd.DataFrame() df_kernel_launchers_unique_args = pd.DataFrame() df_kernel_launchers_unique_args_overlapping_kernels = pd.DataFrame() + df_kernel_launchers = pd.DataFrame() perf_metrics_dfs = {} df_hist = pd.DataFrame() df_short_kernels = pd.DataFrame() @@ -493,20 +808,34 @@ def generate_perf_report_pytorch( df_kernel_launchers = perf_analyzer.get_df_kernel_launchers( include_kernel_details=True, include_first_occurrence_time=include_first_occurrence_time, + include_call_stack=group_by_parent_module, ) - df_kernel_launchers_summary = perf_analyzer.get_df_kernel_launchers_summary( - df_kernel_launchers - ) - df_kernel_launchers_summary_by_category = ( - perf_analyzer.get_df_kernel_launchers_summary_by_category( + if group_by_parent_module: + df_kernel_launchers_summary = ( + perf_analyzer.get_df_kernel_launchers_summary_module( + df_kernel_launchers + ) + ) + df_kernel_launchers_summary_by_category = ( + perf_analyzer.get_df_kernel_launchers_summary_by_category_module( + df_kernel_launchers + ) + ) + else: + df_kernel_launchers_summary = perf_analyzer.get_df_kernel_launchers_summary( df_kernel_launchers ) - ) + df_kernel_launchers_summary_by_category = ( + perf_analyzer.get_df_kernel_launchers_summary_by_category( + df_kernel_launchers + ) + ) df_kernel_launchers_unique_args = ( perf_analyzer.get_df_kernel_launchers_unique_args( df_kernel_launchers, agg_metrics=agg_metrics, include_pct=True, + group_by_parent_module=group_by_parent_module, group_by_num_kernels=group_by_num_kernels, ) ) @@ -522,6 +851,7 @@ def generate_perf_report_pytorch( df_kernel_launchers, agg_metrics=agg_metrics, include_pct=True, + group_by_parent_module=group_by_parent_module, group_by_num_kernels=group_by_num_kernels, include_overlapping_kernels=True, ) @@ -561,7 +891,6 @@ def generate_perf_report_pytorch( "GEMM", "UnaryElementwise", "BinaryElementwise", - "Normalization", ]: # For GEMM: create a single table that covers both fwd and bwd. df_ops_raw = perf_analyzer.build_df_perf_metrics( @@ -802,13 +1131,24 @@ def generate_perf_report_pytorch( # Run TraceDiff when a comparison trace is provided. diff_stats_df is generated _tracediff_diff_stats: Optional[pd.DataFrame] = None if comparison_json_path and not df_unified_perf.empty: - perf_analyzer2 = TreePerfAnalyzer.from_file( - profile_filepath=comparison_json_path, - python_path=perf_analyzer.python_path, - include_unlinked_kernels=perf_analyzer.include_unlinked_kernels, - enable_pseudo_ops=enable_pseudo_ops, - add_python_func=perf_analyzer.add_python_func, - ) + if comparison_augmented_tree is not None: + perf_analyzer2 = TreePerfAnalyzer( + tree=comparison_augmented_tree, + arch=gpu_arch_json, + python_path=python_path, + include_unlinked_kernels=include_unlinked_kernels, + add_python_func=add_python_func, + enable_pseudo_ops=enable_pseudo_ops, + rebuild_tree=False, + ) + else: + perf_analyzer2 = TreePerfAnalyzer.from_file( + profile_filepath=comparison_json_path, + python_path=perf_analyzer.python_path, + include_unlinked_kernels=perf_analyzer.include_unlinked_kernels, + enable_pseudo_ops=enable_pseudo_ops, + add_python_func=perf_analyzer.add_python_func, + ) perf_analyzer2.tree.apply_annotation( name_filters=["vllm::unified_attention_with_output"] ) @@ -906,6 +1246,13 @@ def generate_perf_report_pytorch( # update this dict with the perf_metrics_dfs dict_name2df.update(perf_metrics_dfs) + perf_report_sanity_check( + perf_analyzer.tree.events, + df_gpu_timeline, + df_kernel_launchers, + df_unified_perf, + include_nccl=collective_analysis, + ) # Kernel summary: aggregate per-kernel durations and counts if kernel_summary: @@ -1180,7 +1527,9 @@ def main(): "--include_overlap_info", action="store_true", default=False, - help="Include overlap info in the report. Disabled by default.", + help="Include overlap info in the report. Disabled by default. " + "Adds ops_unique_args_kl_overlap, unified_perf_summary_kl_overlap, and " + "per-category *_kl_overlap / *_fwd_kl_overlap / *_bwd_kl_overlap sheets when data exists.", ) parser.add_argument( "--detect_recompute", @@ -1205,10 +1554,53 @@ def main(): default=False, help="Add call_stack_trimmed and call_stack_full columns to unified_perf_summary.", ) + parser.add_argument( + "--capture_folder", + type=str, + required=False, + help="Path to the graph capture trace folder.", + ) + parser.add_argument( + "--comparison_capture_folder", + type=str, + required=False, + help="Path to the graph capture trace folder for the comparison trace.", + ) + parser.add_argument( + "--group_by_parent_module", + action="store_true", + dest="group_by_parent_module", + default=False, + help="Group kernel launcher summaries by parent module in addition to operation name.", + ) args = parser.parse_args() + if args.comparison_capture_folder and not args.comparison_json_path: + parser.error("--comparison_capture_folder requires --comparison_json_path.") + if args.capture_folder: + metadata_json_path = os.path.join(args.capture_folder, "execution_details.json") + classify_graph_capture_trace(args.capture_folder) + graph_tree = merge_capture_trace_into_graph( + args.capture_folder, + metadata_json_path, + args.profile_json_path, + ) + else: + graph_tree = None + comparison_graph_tree = None + if args.comparison_capture_folder: + comp_metadata = os.path.join( + args.comparison_capture_folder, "execution_details.json" + ) + classify_graph_capture_trace(args.comparison_capture_folder) + comparison_graph_tree = merge_capture_trace_into_graph( + args.comparison_capture_folder, + comp_metadata, + args.comparison_json_path, + ) generate_perf_report_pytorch( profile_json_path=args.profile_json_path, + augmented_tree=graph_tree, output_xlsx_path=args.output_xlsx_path, output_csvs_dir=args.output_csvs_dir, include_unlinked_kernels=args.include_unlinked_kernels, @@ -1225,6 +1617,7 @@ def main(): topk_ops=args.topk_ops, topk_roofline_ops=args.topk_roofline_ops, comparison_json_path=args.comparison_json_path, + comparison_augmented_tree=comparison_graph_tree, extension_file=args.extension_file, python_path=args.python_path, gpu_arch_json_path=args.gpu_arch_json_path, @@ -1233,6 +1626,7 @@ def main(): enable_origami=args.enable_origami, detect_recompute=args.detect_recompute, inductor_cache_dir=args.inductor_cache_dir, + group_by_parent_module=args.group_by_parent_module, include_call_stack=args.include_call_stack, ) diff --git a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py deleted file mode 100644 index 4f6e1a640..000000000 --- a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py +++ /dev/null @@ -1,1402 +0,0 @@ -############################################################################### -# Copyright (c) 2025 - 2026 Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -import argparse -import importlib.util -import json -import os -import sys -import warnings -from typing import Dict, Optional - - -import numpy as np -import pandas as pd -import collections -import gzip -import re -import zipfile - -from TraceLens import NcclAnalyser, TraceToTree, TraceDiff, TreePerfAnalyzer -from TraceLens.PerfModel.torch_op_mapping import build_sheet_category_to_op_names -from TraceLens.Reporting.generate_perf_report_pytorch import _find_entry_point -from TraceLens.Reporting.reporting_utils import ( - add_gpu_arch_cli_args, - resolve_gpu_arch, - write_report_outputs, -) -from TraceLens.util import TraceEventUtils -from TraceLens.TraceUtils.annotation_utils import ( - CAPTURE_PATTERN, - CaptureAnnotation, - find_events_by_patterns, -) -from TraceLens.Trace2Tree.trace_capture_merge_experimental import ( - merge_capture_trace_into_graph, -) - - -def perf_report_sanity_check( - events, - df_gpu_timeline, - df_kernel_launchers, - df_unified_perf, - include_nccl=False, -): - """ - Sanity checks on the performance report DataFrames. - - 1) Total kernel time accounted by df_kernel_launchers and df_unified_perf - should each be >= the computation_time reported in df_gpu_timeline. - 2) Total GPU events in tree events should equal the number of kernels - accounted by df_kernel_launchers and df_unified_perf. - """ - use_time = "busy_time" if include_nccl else "computation_time" - - computation_time_us = ( - df_gpu_timeline.loc[df_gpu_timeline["type"] == use_time, "time ms"].values[0] - * 1e3 - ) - - # --- Check 1: kernel time coverage --- - kl_time_col = ( - "total_direct_kernel_time_sum" - if "total_direct_kernel_time_sum" in df_kernel_launchers.columns - else "total_direct_kernel_time" - ) - up_time_col = ( - "Kernel Time (µs)_sum" - if "Kernel Time (µs)_sum" in df_unified_perf.columns - else "Kernel Time (µs)" - ) - - kl_total_us = df_kernel_launchers[kl_time_col].sum() - up_total_us = df_unified_perf[up_time_col].sum() - - print(f"\n{'='*60}") - print("Perf Report Sanity Check") - print(f"{'='*60}") - print(f" {use_time} (gpu_timeline): {computation_time_us:.2f} µs") - print(f" Kernel time (kernel_launchers): {kl_total_us:.2f} µs") - print(f" Kernel time (unified_perf): {up_total_us:.2f} µs") - - kl_pass = kl_total_us >= computation_time_us - up_pass = up_total_us >= computation_time_us - print(f" kernel_launchers >= computation_time: {'PASS' if kl_pass else 'FAIL'}") - print(f" unified_perf >= computation_time: {'PASS' if up_pass else 'FAIL'}") - - # --- Check 2: per-kernel-name count verification --- - # Build {kernel_name: count} from tree events (ground truth) - tree_kernel_counts = dict( - collections.Counter( - e["name"] - for e in events - if e.get("cat") in {"kernel", "gpu_memcpy", "gpu_memset"} - and ( - not TraceEventUtils.is_communication_string(e.get("name", "")) - or include_nccl - ) - ) - ) - - def _extract_kernel_counts(df, label): - """Extract {kernel_name: count} from a DataFrame's kernel_details column.""" - if "kernel_details_summary" in df.columns: - col = "kernel_details_summary" - elif "kernel_details" in df.columns: - col = "kernel_details" - else: - print(f" WARNING: no kernel_details column in {label}") - return {} - counts = collections.Counter() - for kd in df[col]: - if isinstance(kd, list): - for d in kd: - counts[d["name"]] += d.get("count", 1) - return dict(counts) - - def _print_per_kernel_check(tree_counts, df_counts, label): - """Compare per-kernel counts and print mismatches.""" - df_total = sum(df_counts.values()) - tree_total = sum(tree_counts.values()) - total_pass = tree_total == df_total - print(f"\n --- {label} ---") - print(f" Total GPU events in tree: {tree_total}") - print(f" Kernels accounted: {df_total}") - print(f" Total count match: {'PASS' if total_pass else 'FAIL'}") - - all_names = sorted(set(tree_counts) | set(df_counts)) - mismatches = [] - for name in all_names: - t = tree_counts.get(name, 0) - d = df_counts.get(name, 0) - if t != d: - mismatches.append((name, t, d)) - - if mismatches: - print(f" Per-kernel mismatches ({len(mismatches)}):") - for name, t, d in mismatches: - trunc = name[:80] + "..." if len(name) > 80 else name - print(f" {trunc}") - print(f" tree={t} {label}={d} diff={t - d}") - else: - print(f" Per-kernel detail check: PASS (all match)") - - return df_total, total_pass, mismatches - - # Build {kernel_name: count} from each source - kl_kernel_counts = _extract_kernel_counts(df_kernel_launchers, "kernel_launchers") - up_kernel_counts = _extract_kernel_counts(df_unified_perf, "unified_perf") - - total_gpu_events = sum(tree_kernel_counts.values()) - - kl_kernel_count, kl_count_pass, kl_mismatches = _print_per_kernel_check( - tree_kernel_counts, kl_kernel_counts, "kernel_launchers" - ) - up_kernel_count, up_count_pass, up_mismatches = _print_per_kernel_check( - tree_kernel_counts, up_kernel_counts, "unified_perf" - ) - - print(f"{'='*60}\n") - - return { - "computation_time_us": computation_time_us, - "kl_total_us": kl_total_us, - "up_total_us": up_total_us, - "kl_time_pass": kl_pass, - "up_time_pass": up_pass, - "total_gpu_events": total_gpu_events, - "kl_kernel_count": kl_kernel_count, - "up_kernel_count": up_kernel_count, - "kl_count_pass": kl_count_pass, - "up_count_pass": up_count_pass, - "tree_kernel_counts": dict(tree_kernel_counts), - "kl_kernel_counts": dict(kl_kernel_counts), - "up_kernel_counts": dict(up_kernel_counts), - "kl_mismatches": kl_mismatches, - "up_mismatches": up_mismatches, - } - - -def classify_graph_capture_trace(input_folder: str): - """ - Return {file, batch_size, mode} for a single graph-capture trace file. - Supports .json, .json.gz, and .zip (containing a .json). - """ - execution_details_path = os.path.join(input_folder, "execution_details.json") - if os.path.isfile(execution_details_path): - print( - f"Execution details already exist at {execution_details_path}. Skipping classification." - ) - return - ## vLLM specific dummy run pattern - dummy_run_pattern = re.compile( - r"vllm/v1/worker/gpu_model_runner\.py\(\d+\): _dummy_run" - ) - ## SGLang specific dummy run pattern - ##dummy_run_pattern = re.compile(r"/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py\(\d+\): _capture_graph") - - def load_trace(path: str) -> dict: - if path.endswith(".zip"): - with zipfile.ZipFile(path, "r") as zf: - json_files = [f for f in zf.namelist() if f.endswith(".json")] - if not json_files: - raise ValueError(f"No .json file found inside {path}") - with zf.open(json_files[0]) as f: - return json.load(f) - if path.endswith(".json.gz"): - with gzip.open(path, "rt") as f: - return json.load(f) - with open(path, "r") as f: - return json.load(f) - - def find_dummy_run_roots(events): - roots = [e for e in events if dummy_run_pattern.match(e.get("name", ""))] - roots.sort(key=lambda x: x.get("ts", 0)) - return roots - - def count_stream_begin_captures(events): - return sum( - 1 - for e in events - if "StreamBeginCapture" in e.get("name", "") - and e.get("cat") == "cuda_runtime" - ) - - def infer_batch_size_from_cpu_ops(events): - first_dims = [] - for e in events: - if e.get("cat") != "cpu_op": - continue - input_dims = e.get("args", {}).get("Input Dims") - if not input_dims: - continue - for dim_list in input_dims: - if isinstance(dim_list, list) and dim_list: - if isinstance(dim_list[0], int): - first_dims.append(dim_list[0]) - if not first_dims: - return None - return collections.Counter(first_dims).most_common(1)[0][0] - - def infer_mode_from_captures(num_captures: int): - return "FULL" if num_captures <= 1 else "PIECEWISE" - - if not os.path.isdir(input_folder): - print(f"Error: {input_folder} is not a directory", file=sys.stderr) - sys.exit(1) - - trace_files = sorted( - os.path.join(input_folder, f) - for f in os.listdir(input_folder) - if f.endswith(".json") or f.endswith(".json.gz") - ) - - if not trace_files: - print(f"No files starting with 'graph_capture_rank_0' found in {input_folder}") - sys.exit(0) - - print(f"Found {len(trace_files)} graph-capture trace file(s) in {input_folder}\n") - - results = [] - for filepath in trace_files: - trace_json = load_trace(filepath) - events = trace_json.get("traceEvents", []) - dummy_roots = find_dummy_run_roots(events) - annotation_roots = find_events_by_patterns(events, [CAPTURE_PATTERN]) - basename = os.path.basename(filepath) - - if annotation_roots and len(annotation_roots) == len(dummy_roots): - cap = CaptureAnnotation(annotation_roots[0]["name"]) - batch_size, mode = cap.batch_size, cap.mode - print( - f"batch_size: {batch_size}, mode: {mode} parsed from annotation, num_captures: {count_stream_begin_captures(events)}" - ) - results.append({"file": basename, "batch_size": batch_size, "mode": mode}) - continue - - num_captures = count_stream_begin_captures(events) - mode = infer_mode_from_captures(num_captures) - batch_size = infer_batch_size_from_cpu_ops(events) - print( - f"batch_size: {batch_size}, mode: {mode} inferred, num_captures: {num_captures}" - ) - - results.append({"file": basename, "batch_size": batch_size, "mode": mode}) - with open(f"{input_folder}/execution_details.json", "w") as f: - json.dump(results, f, indent=2) - print(f"\nResults written to {input_folder}/execution_details.json") - return - - -def get_dfs_short_kernels( - perf_analyzer, short_kernel_threshold_us=10, histogram_bins=100, topk=None -): - """ - TODO: move this to the TreePerfAnalyzer class - Analyze short kernel events from the performance data and return two DataFrames: - a histogram of short kernel durations and a summary of top short kernels. - - Args: - perf_analyzer (TreePerfAnalyzer): The performance analyzer object containing kernel data. - short_kernel_threshold_us (int, optional): Threshold in microseconds to classify a kernel as "short". Defaults to 10. - histogram_bins (int, optional): Number of bins for the histogram of short kernel durations. Defaults to 100. - topk (int, optional): Number of top short kernels to include in the summary. If None, include all. Defaults to None. - - Returns: - tuple: A tuple containing: - - pd.DataFrame: Histogram of short kernel durations with columns ['bin_start', 'bin_end', 'count']. - - pd.DataFrame: Summary of top short kernels with detailed statistics and percentage contribution to total time. - """ - df_kernels = perf_analyzer.get_df_kernels() - df_filtered = df_kernels[ - df_kernels["Kernel duration (µs)"] < short_kernel_threshold_us - ] - - # 1. get histogram of these short kernels - if df_filtered.empty: - df_hist = pd.DataFrame(columns=["bin_start", "bin_end", "count"]) - else: - vals = df_filtered["Kernel duration (µs)"].values - counts, bin_edges = np.histogram(vals, bins=histogram_bins) - df_hist = pd.DataFrame( - {"bin_start": bin_edges[:-1], "bin_end": bin_edges[1:], "count": counts} - ) - - # 2. get df short kernels topk by total time - agg_dict = { - "Kernel duration (µs)": ["sum", "count", "mean"], - } - # For GPU-only traces, only group by Kernel name (CPU-related columns don't exist) - # For regular traces, group by all available columns - if perf_analyzer.gpu_only: - groupby_cols = ["Kernel name"] - else: - groupby_cols = [ - "Parent cpu_op", - "Input dims", - "Input strides", - "Concrete Inputs", - "Kernel name", - ] - - # If dataframe is empty, return empty dataframe - if df_filtered.empty: - df_grouped = pd.DataFrame() - else: - df_grouped = df_filtered.groupby( - groupby_cols, - sort=False, - ).agg(agg_dict) - - # Handle empty dataframe case - if df_grouped.empty: - return df_hist, df_grouped - - # Flatten multi-level column names - df_grouped.columns = ["_".join(col).strip() for col in df_grouped.columns] - - # Rename columns for clarity - df_grouped.rename( - columns={ - "Kernel duration (µs)_sum": "Short Kernel duration (µs) sum", - "Kernel duration (µs)_count": "Short Kernel count", - "Kernel duration (µs)_mean": "Short Kernel duration (µs) mean", - }, - inplace=True, - ) - - # Add percentage contribution to total time - df_grouped["Short Kernel duration (µs) percent of total time"] = ( - df_grouped["Short Kernel duration (µs) sum"] - / (perf_analyzer.total_time_ms * 1e3) - * 100 - ) - - # Sort: primary by total short-kernel time (desc), then all other columns for stable order - _sum_col = "Short Kernel duration (µs) sum" - _sort_cols = [_sum_col] + [c for c in df_grouped.columns if c != _sum_col] - _ascending = [False] + [True] * (len(_sort_cols) - 1) - df_grouped.sort_values(by=_sort_cols, ascending=_ascending, inplace=True) - df_grouped.reset_index(inplace=True) - if topk is not None: - df_grouped = df_grouped.head(topk) - return df_hist, df_grouped - - -def apply_extension(perf_analyzer, extension_path): - extension_path = os.path.abspath(extension_path) - extension_name = os.path.splitext(os.path.basename(extension_path))[0] - - from TraceLens.PerfModel.torch_op_mapping import ( - OP_CATEGORY_REGISTRY, - register_op_categories, - register_perf_model_categories, - ) - - spec = importlib.util.spec_from_file_location(extension_name, extension_path) - extension = importlib.util.module_from_spec(spec) - spec.loader.exec_module(extension) - - if hasattr(extension, "tree_postprocess_extension"): - print(f"Applying tree postprocess extension from {extension_path}") - tree_postprocess_extension = getattr(extension, "tree_postprocess_extension") - tree_postprocess_extension(perf_analyzer.tree) - perf_analyzer.tree.label_non_gpu_paths() - - if hasattr(extension, "perf_model_extension"): - print(f"Applying perf model extension from {extension_path}") - perf_model_extension = getattr(extension, "perf_model_extension") - if not isinstance(perf_model_extension, dict): - raise ValueError( - f"Expected perf_model_extension to be a dict, got {type(perf_model_extension)}" - ) - perf_analyzer.op_to_perf_model_class_map.update(perf_model_extension) - register_perf_model_categories( - perf_model_extension, - OP_CATEGORY_REGISTRY, - ) - if hasattr(extension, "op_category_extension"): - print(f"Applying op category extension from {extension_path}") - op_category_extension = getattr(extension, "op_category_extension") - if not isinstance(op_category_extension, dict): - raise ValueError( - f"Expected op_category_extension to be a dict, got {type(op_category_extension)}" - ) - register_op_categories( - op_category_extension, - OP_CATEGORY_REGISTRY, - ) - if hasattr(extension, "dict_cat2names_extension"): - warnings.warn( - "dict_cat2names_extension is deprecated and ignored. Use " - "perf_model_extension for modeled ops or op_category_extension for " - "category-only ops." - ) - - -def trunc_kernel_details(row, kernel_detail_col, trunc_length=64): - """ - Truncates the kernel details in a row to a specified length for readability. - """ - if kernel_detail_col not in row or not row[kernel_detail_col]: - return None # No kernel details available - - truncated_details = [] - for detail in row[kernel_detail_col]: - truncated_name = ( - detail["name"][:trunc_length] + "..." - if len(detail["name"]) > trunc_length - else detail["name"] - ) - truncated_details.append( - { - "name": truncated_name, - "stream": detail.get("stream", None), - "mean_duration_us": round(detail.get("mean_duration_us", 0), 2), - } - ) - - return truncated_details if truncated_details else None - - -def add_truncated_kernel_details( - df: pd.DataFrame, - source_col: str = "kernel_details", - new_col_name: str = None, - trunc_length: int = 64, -) -> pd.DataFrame: - """ - Applies the truncation logic to a DataFrame column and inserts the new - truncated column immediately after the source column for easy comparison. - - Args: - df (pd.DataFrame): The DataFrame to process. - source_col (str): The name of the column containing the full kernel details. - new_col_name (str): The name for the new truncated column. - trunc_length (int): The character length to truncate kernel names to. - - Returns: - pd.DataFrame: A new DataFrame with the added truncated column. - """ - # First, ensure the source column exists. If not, do nothing. - if source_col not in df.columns: - warnings.warn( - f"Source column '{source_col}' not found in DataFrame. Skipping truncation.", - UserWarning, - ) - return df - if new_col_name is None: - new_col_name = f"trunc_{source_col}" - # 1. Create the new column's data. It will be added to the end for now. - df[new_col_name] = df.apply( - lambda row: trunc_kernel_details(row, source_col, trunc_length=trunc_length), - axis=1, - ) - - # 2. Reorder the columns to place the new column next to its source. - cols = df.columns.tolist() - # Pop the new column from the end of the list - new_col = cols.pop(cols.index(new_col_name)) - # Find the position of our source column and insert the new one after it - source_col_idx = cols.index(source_col) - cols.insert(source_col_idx + 1, new_col) - - # Return a new DataFrame with the desired column order - return df[cols] - - -def generate_perf_report_pytorch( - profile_json_path: str, - augmented_tree: TraceToTree = None, - output_xlsx_path: Optional[str] = None, - output_csvs_dir: Optional[str] = None, - # include unlinked kernels in gpu timeline - include_unlinked_kernels: bool = False, - enable_pseudo_ops: bool = False, # pseudo-op generation - # threshold in microseconds for micro idle time - micro_idle_thresh_us: int = None, - # collective analysis - collective_analysis: bool = True, - # overlapping kernel details (optional extra sheets) - include_overlap_info: bool = False, - # kernel summary sheet - kernel_summary: bool = False, - # short kernel study options - short_kernel_study: bool = False, - short_kernel_threshold_us: int = 10, - short_kernel_histogram_bins: int = 100, - topk_short_kernels: Optional[int] = None, # include all below thresh by default - topk_ops: Optional[int] = None, - topk_roofline_ops: Optional[int] = None, - comparison_json_path: Optional[str] = None, - comparison_augmented_tree: Optional[TraceToTree] = None, - extension_file: Optional[str] = None, - # for gemm simulator / Origami (Origami requires --enable_origami when arch is set) - python_path: Optional[str] = None, - gpu_arch_json_path: Optional[str] = None, - gpu_arch_platform: Optional[str] = None, - gpu_arch: Optional[dict] = None, - enable_origami: bool = False, - group_by_parent_module: bool = False, - group_by_num_kernels: bool = False, - include_call_stack: bool = False, -) -> Dict[str, pd.DataFrame]: - gpu_arch_json = resolve_gpu_arch( - gpu_arch_json_path=gpu_arch_json_path, - gpu_arch_platform=gpu_arch_platform, - gpu_arch=gpu_arch, - ) - add_python_func = ( - True - if ( - group_by_parent_module - or include_call_stack is True - or augmented_tree is not None - or comparison_augmented_tree is not None - ) - else False - ) - if augmented_tree is not None: - perf_analyzer = TreePerfAnalyzer( - tree=augmented_tree, - arch=gpu_arch_json, - python_path=python_path, - include_unlinked_kernels=include_unlinked_kernels, - add_python_func=add_python_func, - enable_pseudo_ops=enable_pseudo_ops, - rebuild_tree=False, - ) - else: - perf_analyzer = TreePerfAnalyzer.from_file( - profile_filepath=profile_json_path, - arch=gpu_arch_json, - python_path=python_path, - include_unlinked_kernels=include_unlinked_kernels, - add_python_func=add_python_func, - enable_pseudo_ops=enable_pseudo_ops, - ) - - graph_launch_events = [ - event - for event in perf_analyzer.tree.events - if "graphlaunch" in event.get("name", "").lower() - ] - if len(graph_launch_events) > 0: - warnings.warn( - f"There are hipgraph launches (Count: {len(graph_launch_events)}) in this trace, but a graph capture folder not provided, the analysis might be limited", - UserWarning, - ) - - ## Apply annotation for vLLM eager and replay phase - perf_analyzer.tree.apply_annotation( - name_filters=[ - "vllm::unified_attention_with_output", - "aiter::mha_varlen_fwd", - "pseudo_mla_decode_fwd", - "pseudo_mla_prefill_fwd", - "vllm::gdn_attention_core", - "aiter::fmha_v3_varlen_fwd", - "sglang_profiler::tilelang_kernel_tilelang_sparse_fwd", - "sglang_profiler::attention_paged_attention_ragged", - "aiter::mha_batch_prefill", - "aiter::pa_decode_gluon", - "aiter::v4_attention_with_output", - "pseudo_v4_paged_decode_swa", - "pseudo_v4_paged_decode_csa", - "pseudo_v4_paged_decode_hca", - ] - ) - - if extension_file: - apply_extension(perf_analyzer, extension_file) - - # Detect GPU-only trace early and inform user - if perf_analyzer.gpu_only: - print( - "Detected GPU-only trace. Skipping CPU-dependent analysis and generating only GPU timeline and kernel summary." - ) - agg_metrics = ["mean", "median", "std", "min", "max"] - - # Generate base DataFrames - df_gpu_timeline = perf_analyzer.get_df_gpu_timeline( - micro_idle_thresh_us=micro_idle_thresh_us - ) - - # TODO: move this to the TreePerfAnalyzer class - total_time_row = df_gpu_timeline[df_gpu_timeline["type"] == "total_time"] - total_time_ms = total_time_row["time ms"].values[0] - perf_analyzer.total_time_ms = total_time_ms - - # Initialize empty DataFrames for GPU-only traces to avoid NameError - df_kernel_launchers_summary = pd.DataFrame() - df_kernel_launchers_summary_by_category = pd.DataFrame() - df_kernel_launchers_unique_args = pd.DataFrame() - df_kernel_launchers_unique_args_overlapping_kernels = pd.DataFrame() - df_kernel_launchers = pd.DataFrame() - perf_metrics_dfs = {} - df_hist = pd.DataFrame() - df_short_kernels = pd.DataFrame() - - # Only process CPU-dependent analysis for non-GPU-only traces - if not perf_analyzer.gpu_only: - df_kernel_launchers = perf_analyzer.get_df_kernel_launchers( - include_kernel_details=True, - include_call_stack=group_by_parent_module, - ) - df_kernel_launchers_summary = ( - perf_analyzer.get_df_kernel_launchers_summary_module(df_kernel_launchers) - ) - df_kernel_launchers_summary_by_category = ( - perf_analyzer.get_df_kernel_launchers_summary_by_category_module( - df_kernel_launchers - ) - ) - df_kernel_launchers_unique_args = ( - perf_analyzer.get_df_kernel_launchers_unique_args( - df_kernel_launchers, - agg_metrics=agg_metrics, - include_pct=True, - group_by_parent_module=group_by_parent_module, - group_by_num_kernels=group_by_num_kernels, - ) - ) - df_kernel_launchers_unique_args = add_truncated_kernel_details( - df_kernel_launchers_unique_args, - source_col="kernel_details_summary", - new_col_name="trunc_kernel_details", - ) - df_kernel_launchers_unique_args_overlapping_kernels = pd.DataFrame() - if include_overlap_info: - df_kernel_launchers_unique_args_overlapping_kernels = ( - perf_analyzer.get_df_kernel_launchers_unique_args( - df_kernel_launchers, - agg_metrics=agg_metrics, - include_pct=True, - group_by_parent_module=group_by_parent_module, - group_by_num_kernels=group_by_num_kernels, - include_overlapping_kernels=True, - ) - ) - df_kernel_launchers_unique_args_overlapping_kernels = ( - add_truncated_kernel_details( - df_kernel_launchers_unique_args_overlapping_kernels, - source_col="kernel_details_summary", - new_col_name="trunc_kernel_details", - ) - ) - df_kernel_launchers_unique_args_overlapping_kernels = ( - add_truncated_kernel_details( - df_kernel_launchers_unique_args_overlapping_kernels, - source_col="overlapping_kernels_details_summary", - new_col_name="trunc_overlapping_kernels_details", - ) - ) - # Dictionary to hold the op-specific DataFrames - perf_metrics_dfs = {} - sheet_category_to_op_names = build_sheet_category_to_op_names( - perf_analyzer.op_to_perf_model_class_map - ) - for sheet_category, op_names in sheet_category_to_op_names.items(): - # Filter events belonging to the current legacy sheet category - op_events = [ - event - for event in perf_analyzer.tree.events - if event["name"] in op_names - ] - if len(op_events) == 0: - continue - # Skip categories with no events - if sheet_category in ["GEMM", "UnaryElementwise", "BinaryElementwise"]: - # For GEMM: create a single table that covers both fwd and bwd. - df_ops_raw = perf_analyzer.build_df_perf_metrics( - op_events, bwd=False, include_kernel_details=True, include_args=True - ) - df_ops = perf_analyzer.summarize_df_perf_metrics( - df_ops_raw, - agg_metrics, - group_by_num_kernels=group_by_num_kernels, - ) - df_ops = add_truncated_kernel_details( - df_ops, - source_col="kernel_details__summarize_kernel_stats", - new_col_name="trunc_kernel_details", - ) - if not df_ops.empty: - perf_metrics_dfs[sheet_category] = df_ops - if include_overlap_info: - df_ops_overlapping_kernels = ( - perf_analyzer.summarize_df_perf_metrics( - df_ops_raw, - agg_metrics, - group_by_num_kernels=group_by_num_kernels, - include_overlapping_kernels=True, - ) - ) - df_ops_overlapping_kernels = add_truncated_kernel_details( - df_ops_overlapping_kernels, - source_col="kernel_details__summarize_kernel_stats", - new_col_name="trunc_kernel_details", - ) - df_ops_overlapping_kernels = add_truncated_kernel_details( - df_ops_overlapping_kernels, - source_col="overlapping_kernels_details__summarize_kernel_stats", - new_col_name="trunc_overlapping_kernels_details", - ) - if not df_ops_overlapping_kernels.empty: - perf_metrics_dfs[f"{sheet_category}_kl_overlap"] = ( - df_ops_overlapping_kernels - ) - else: - # For FLASH_ATTN and CONV: create separate tables for forward and backward passes. - df_ops_fwd_raw = perf_analyzer.build_df_perf_metrics( - op_events, bwd=False, include_kernel_details=True, include_args=True - ) - df_ops_fwd = perf_analyzer.summarize_df_perf_metrics( - df_ops_fwd_raw, - agg_metrics, - group_by_num_kernels=group_by_num_kernels, - ) - df_ops_fwd = add_truncated_kernel_details( - df_ops_fwd, - source_col="kernel_details__summarize_kernel_stats", - new_col_name="trunc_kernel_details", - ) - # For now, flash_attention_varlen_backward and aten::convolution_backward are processed with bwd=True, - # so we need a workaround to extract them from the fwd df and append them to the bwd df. - filtered_df_bwd_ops = None - df_ops_bwd_raw = None - if not df_ops_fwd.empty: - # Filter out backward operations that were incorrectly included in forward - bwd_op_names = [ - "flash_attn::_flash_attn_varlen_backward", - "aten::convolution_backward", - ] - filtered_df_bwd_ops = df_ops_fwd[ - df_ops_fwd["name"].isin(bwd_op_names) - ] - df_ops_fwd = df_ops_fwd[~df_ops_fwd["name"].isin(bwd_op_names)] - df_ops_fwd = df_ops_fwd[ - df_ops_fwd["name"] != "flash_attn::_flash_attn_varlen_backward" - ] - - op_events = [] - if len(op_events) > 0: - df_ops_bwd_raw = perf_analyzer.build_df_perf_metrics( - op_events, - bwd=True, - include_kernel_details=True, - include_args=True, - ) - df_ops_bwd = perf_analyzer.summarize_df_perf_metrics( - df_ops_bwd_raw, - agg_metrics, - group_by_num_kernels=group_by_num_kernels, - ) - df_ops_bwd = add_truncated_kernel_details( - df_ops_bwd, - source_col="kernel_details__summarize_kernel_stats", - new_col_name="trunc_kernel_details", - ) - if filtered_df_bwd_ops is not None: - df_ops_bwd = pd.concat([df_ops_bwd, filtered_df_bwd_ops]) - if not df_ops_bwd.empty: - perf_metrics_dfs[f"{sheet_category}_bwd"] = df_ops_bwd - if not df_ops_fwd.empty: - perf_metrics_dfs[f"{sheet_category}_fwd"] = df_ops_fwd - - if include_overlap_info: - df_ops_fwd_overlapping_kernels = ( - perf_analyzer.summarize_df_perf_metrics( - df_ops_fwd_raw, - agg_metrics, - group_by_num_kernels=group_by_num_kernels, - include_overlapping_kernels=True, - ) - ) - df_ops_fwd_overlapping_kernels = add_truncated_kernel_details( - df_ops_fwd_overlapping_kernels, - source_col="kernel_details__summarize_kernel_stats", - new_col_name="trunc_kernel_details", - ) - df_ops_fwd_overlapping_kernels = add_truncated_kernel_details( - df_ops_fwd_overlapping_kernels, - source_col="overlapping_kernels_details__summarize_kernel_stats", - new_col_name="trunc_overlapping_kernels_details", - ) - filtered_df_bwd_ops_overlapping_kernels = None - if not df_ops_fwd_overlapping_kernels.empty: - bwd_op_names = [ - "flash_attn::_flash_attn_varlen_backward", - "aten::convolution_backward", - ] - filtered_df_bwd_ops_overlapping_kernels = ( - df_ops_fwd_overlapping_kernels[ - df_ops_fwd_overlapping_kernels["name"].isin( - bwd_op_names - ) - ] - ) - df_ops_fwd_overlapping_kernels = df_ops_fwd_overlapping_kernels[ - ~df_ops_fwd_overlapping_kernels["name"].isin(bwd_op_names) - ] - df_ops_fwd_overlapping_kernels = df_ops_fwd_overlapping_kernels[ - df_ops_fwd_overlapping_kernels["name"] - != "flash_attn::_flash_attn_varlen_backward" - ] - - df_ops_bwd_overlapping_kernels = pd.DataFrame() - if df_ops_bwd_raw is not None: - df_ops_bwd_overlapping_kernels = ( - perf_analyzer.summarize_df_perf_metrics( - df_ops_bwd_raw, - agg_metrics, - group_by_num_kernels=group_by_num_kernels, - include_overlapping_kernels=True, - ) - ) - df_ops_bwd_overlapping_kernels = add_truncated_kernel_details( - df_ops_bwd_overlapping_kernels, - source_col="kernel_details__summarize_kernel_stats", - new_col_name="trunc_kernel_details", - ) - df_ops_bwd_overlapping_kernels = add_truncated_kernel_details( - df_ops_bwd_overlapping_kernels, - source_col="overlapping_kernels_details__summarize_kernel_stats", - new_col_name="trunc_overlapping_kernels_details", - ) - if filtered_df_bwd_ops_overlapping_kernels is not None: - df_ops_bwd_overlapping_kernels = pd.concat( - [ - df_ops_bwd_overlapping_kernels, - filtered_df_bwd_ops_overlapping_kernels, - ] - ) - if not df_ops_bwd_overlapping_kernels.empty: - perf_metrics_dfs[f"{sheet_category}_bwd_kl_overlap"] = ( - df_ops_bwd_overlapping_kernels - ) - if not df_ops_fwd_overlapping_kernels.empty: - perf_metrics_dfs[f"{sheet_category}_fwd_kl_overlap"] = ( - df_ops_fwd_overlapping_kernels - ) - - # Short kernel study (works for both GPU-only and regular traces) - if short_kernel_study: - df_hist, df_short_kernels = get_dfs_short_kernels( - perf_analyzer, - short_kernel_threshold_us=short_kernel_threshold_us, - histogram_bins=short_kernel_histogram_bins, - topk=topk_short_kernels, - ) - - # Build dict_name2df - only include sheets that have data - dict_name2df = {"gpu_timeline": df_gpu_timeline} - df_unified_perf: pd.DataFrame = pd.DataFrame() - - # Add CPU-dependent sheets only if not GPU-only - if not perf_analyzer.gpu_only: - if not df_kernel_launchers_summary_by_category.empty: - dict_name2df["ops_summary_by_category"] = ( - df_kernel_launchers_summary_by_category - ) - if not df_kernel_launchers_summary.empty: - dict_name2df["ops_summary"] = df_kernel_launchers_summary - if not df_kernel_launchers_unique_args.empty: - dict_name2df["ops_unique_args"] = df_kernel_launchers_unique_args - if ( - include_overlap_info - and not df_kernel_launchers_unique_args_overlapping_kernels.empty - ): - dict_name2df["ops_unique_args_kl_overlap"] = ( - df_kernel_launchers_unique_args_overlapping_kernels - ) - - # Add unified perf metrics table (ops with perf models + leaf ops with GPU kernels) - df_unified_perf = perf_analyzer.build_df_unified_perf_table( - include_nccl=collective_analysis - ) - - # Run TraceDiff when a comparison trace is provided. diff_stats_df is generated - _tracediff_diff_stats: Optional[pd.DataFrame] = None - if comparison_json_path and not df_unified_perf.empty: - if comparison_augmented_tree is not None: - perf_analyzer2 = TreePerfAnalyzer( - tree=comparison_augmented_tree, - arch=gpu_arch_json, - python_path=python_path, - include_unlinked_kernels=include_unlinked_kernels, - add_python_func=add_python_func, - enable_pseudo_ops=enable_pseudo_ops, - rebuild_tree=False, - ) - else: - perf_analyzer2 = TreePerfAnalyzer.from_file( - profile_filepath=comparison_json_path, - python_path=perf_analyzer.python_path, - include_unlinked_kernels=perf_analyzer.include_unlinked_kernels, - enable_pseudo_ops=enable_pseudo_ops, - add_python_func=perf_analyzer.add_python_func, - ) - perf_analyzer2.tree.apply_annotation( - name_filters=["vllm::unified_attention_with_output"] - ) - td = TraceDiff(perf_analyzer.tree, perf_analyzer2.tree) - td.generate_tracediff_report() - _tracediff_diff_stats = td.diff_stats_df - - if not df_unified_perf.empty: - df_unified_perf_summary = perf_analyzer.summarize_df_unified_perf_table( - df_unified_perf, - agg_metrics=agg_metrics, - include_pct=True, - group_by_num_kernels=group_by_num_kernels, - include_call_stack=include_call_stack, - tree=perf_analyzer.tree, - ) - if not df_unified_perf_summary.empty: - df_unified_perf_summary = add_truncated_kernel_details( - df_unified_perf_summary, - source_col="kernel_details_summary", - new_col_name="trunc_kernel_details", - ) - if "call_stack_full" in df_unified_perf_summary.columns: - cs_col = df_unified_perf_summary.columns.get_loc("call_stack_full") - ep_results = df_unified_perf_summary.apply( - lambda row: _find_entry_point( - row["call_stack_full"], row["name"] - ), - axis=1, - ) - df_unified_perf_summary.insert( - cs_col, - "entry_point", - ep_results.apply(lambda x: x["entry_point"]), - ) - if os.environ.get("TRACELENS_DEBUG"): - df_unified_perf_summary.insert( - cs_col + 1, - "num_wrappers", - ep_results.apply(lambda x: x["num_wrappers"]), - ) - df_unified_perf_summary.insert( - cs_col + 2, - "traversal", - ep_results.apply(lambda x: x["traversal"]), - ) - df_unified_perf_summary.insert( - cs_col + 3, - "wrappers", - ep_results.apply(lambda x: x["wrappers"]), - ) - dict_name2df["unified_perf_summary"] = df_unified_perf_summary - - if _tracediff_diff_stats is not None and not _tracediff_diff_stats.empty: - from TraceLens.Reporting.tracediff_comparison_extension import ( - enrich_perf_report_dict_inplace, - ) - - dict_name2df = enrich_perf_report_dict_inplace( - dict_name2df, - _tracediff_diff_stats, - df_unified_perf=df_unified_perf, - ) - dict_name2df["diff_stats"] = _tracediff_diff_stats - - if include_overlap_info: - df_unified_perf_summary_overlapping_kernels = ( - perf_analyzer.summarize_df_unified_perf_table( - df_unified_perf, - agg_metrics=agg_metrics, - include_pct=True, - group_by_num_kernels=group_by_num_kernels, - include_overlapping_kernels=True, - ) - ) - if not df_unified_perf_summary_overlapping_kernels.empty: - df_unified_perf_summary_overlapping_kernels = ( - add_truncated_kernel_details( - df_unified_perf_summary_overlapping_kernels, - source_col="kernel_details_summary", - new_col_name="trunc_kernel_details", - ) - ) - df_unified_perf_summary_overlapping_kernels = ( - add_truncated_kernel_details( - df_unified_perf_summary_overlapping_kernels, - source_col="overlapping_kernels_details_summary", - new_col_name="trunc_overlapping_kernels_details", - ) - ) - if not df_unified_perf_summary_overlapping_kernels.empty: - dict_name2df["unified_perf_summary_kl_overlap"] = ( - df_unified_perf_summary_overlapping_kernels - ) - - # update this dict with the perf_metrics_dfs - dict_name2df.update(perf_metrics_dfs) - perf_report_sanity_check( - perf_analyzer.tree.events, - df_gpu_timeline, - df_kernel_launchers, - df_unified_perf, - include_nccl=collective_analysis, - ) - - # Kernel summary: aggregate per-kernel durations and counts - if kernel_summary: - try: - df_kernels = perf_analyzer.get_df_kernels(launcher_detail=True) - except Exception: - df_kernels = pd.DataFrame() - if not df_kernels.empty and "Kernel duration (µs)" in df_kernels.columns: - # Fallback: If Parent cpu_op is missing, fill it from Launcher (for display purposes) - if ( - "Parent cpu_op" in df_kernels.columns - and "Launcher" in df_kernels.columns - ): - mask_missing_parent = df_kernels["Parent cpu_op"].isna() - if mask_missing_parent.any(): - df_kernels.loc[mask_missing_parent, "Parent cpu_op"] = ( - df_kernels.loc[mask_missing_parent, "Launcher"] - ) - - # Fallback categorization for graph/runtime launched kernels with no cpu_op - # Note: Basic 'Parent op category' is added by get_kernel_details() in tree_perf.py - # This adds categorization for kernels that don't have a parent cpu_op - if "Parent op category" not in df_kernels.columns: - df_kernels["Parent op category"] = np.nan - - if "Launcher" in df_kernels.columns: - mask_missing_cat = df_kernels["Parent op category"].isna() - if mask_missing_cat.any(): - - def _launcher_category(name): - s = str(name).lower() - if "cudagraph" in s or "graphlaunch" in s: - return "graph" - return "runtime" if s and s != "nan" else np.nan - - df_kernels.loc[mask_missing_cat, "Parent op category"] = ( - df_kernels.loc[mask_missing_cat, "Launcher"].apply( - _launcher_category - ) - ) - - # Group by category/cpu_op along with kernel identifiers when available - group_cols = [] - for col in [ - "Parent op category", - "Parent cpu_op", - "Kernel name", - "Kernel stream", - ]: - if col in df_kernels.columns: - group_cols.append(col) - if not group_cols: - group_cols = ( - ["Kernel name"] if "Kernel name" in df_kernels.columns else [] - ) - - agg_dict = {"Kernel duration (µs)": ["sum", "count", "mean", "min", "max"]} - df_kernel_summary = df_kernels.groupby(group_cols, dropna=False).agg( - agg_dict - ) - df_kernel_summary.columns = [ - "_".join(col).strip() for col in df_kernel_summary.columns.values - ] - df_kernel_summary.reset_index(inplace=True) - - # Percent columns: - # 1) Percent of kernels time: sums to ~100% across rows - total_kernels_us = df_kernels["Kernel duration (µs)"].sum() - if total_kernels_us > 0: - df_kernel_summary["Percent of kernels time (%)"] = ( - df_kernel_summary["Kernel duration (µs)_sum"] / total_kernels_us - ) * 100 - else: - df_kernel_summary["Percent of kernels time (%)"] = np.nan - # 2) Percent of total time (GPU timeline baseline; includes idle/non-kernel) - total_us = ( - perf_analyzer.total_time_ms * 1e3 - if hasattr(perf_analyzer, "total_time_ms") - else None - ) - if total_us: - df_kernel_summary["Percent of total time (%)"] = ( - df_kernel_summary["Kernel duration (µs)_sum"] / total_us - ) * 100 - else: - df_kernel_summary["Percent of total time (%)"] = np.nan - - df_kernel_summary.sort_values( - by="Kernel duration (µs)_sum", ascending=False, inplace=True - ) - df_kernel_summary.reset_index(drop=True, inplace=True) - dict_name2df["kernel_summary"] = df_kernel_summary - - if short_kernel_study: - dict_name2df["short_kernel_histogram"] = df_hist - dict_name2df["short_kernels_summary"] = df_short_kernels - - # Skip collective analysis for GPU-only traces (no CPU ops means no collectives) - if collective_analysis and not perf_analyzer.gpu_only: - nccl_analyser = NcclAnalyser([profile_json_path], None) - df_nccl_summary = nccl_analyser.build_df_summary_long() - if not df_nccl_summary.empty: - dict_name2df["coll_analysis"] = df_nccl_summary - - # Get additional DataFrames from extension if available - if extension_file: - extension_path = os.path.abspath(extension_file) - extension_name = os.path.splitext(os.path.basename(extension_path))[0] - spec = importlib.util.spec_from_file_location(extension_name, extension_path) - extension = importlib.util.module_from_spec(spec) - spec.loader.exec_module(extension) - - if hasattr(extension, "get_additional_dataframes_extension"): - print(f"Getting additional DataFrames from extension: {extension_path}") - get_additional_dfs = getattr( - extension, "get_additional_dataframes_extension" - ) - additional_dfs = get_additional_dfs(perf_analyzer.tree) - if additional_dfs: - dict_name2df.update(additional_dfs) - print(f"Added {len(additional_dfs)} additional sheets from extension") - - # Write CSVs and/or Excel (independent options) - if output_xlsx_path is None and output_csvs_dir is None: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - write_report_outputs( - dict_name2df, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir - ) - - return dict_name2df - - -def main(): - - parser = argparse.ArgumentParser( - description="Process a JSON trace profile and generate performance report tables." - ) - parser.add_argument( - "--profile_json_path", - type=str, - required=True, - help="Path to the profile.json or .json.gz file", - ) - parser.add_argument( - "--output_xlsx_path", - type=str, - default=None, - help="Path to the output Excel file", - ) - parser.add_argument( - "--output_csvs_dir", - type=str, - default=None, - help="Directory to save output CSV files", - ) - - # Optional arguments - parser.add_argument( - "--include_unlinked_kernels", - action="store_true", - help="Include unlinked kernels in the GPU timeline analysis.", - ) - parser.add_argument( - "--micro_idle_thresh_us", - type=int, - default=None, - help="Threshold in microseconds to classify idle interval as micro idle in GPU timeline analysis. " - "Default is None and all idle times are included in one category.", - ) - parser.add_argument( - "--disable_coll_analysis", - action="store_false", - dest="collective_analysis", - default=False, - help="Disable collective analysis section in the report. Enabled by default.", - ) - parser.add_argument( - "--enable_kernel_summary", - action="store_true", - dest="kernel_summary", - default=False, - help="Enable kernel summary sheet in the report. Disabled by default.", - ) - - parser.add_argument( - "--group_by_parent_module", - action="store_true", - dest="group_by_parent_module", - default=False, - help="Group kernel launcher summaries by parent module in addition to operation name.", - ) - parser.add_argument( - "--short_kernel_study", - action="store_true", - help="Include short kernel study in the report.", - ) - parser.add_argument( - "--short_kernel_threshold_us", - type=int, - default=10, - help='Threshold in microseconds to classify a kernel as "short". Defaults to 10 us.', - ) - parser.add_argument( - "--short_kernel_histogram_bins", - type=int, - default=100, - help="Number of bins for the short-kernel histogram.", - ) - parser.add_argument( - "--topk_short_kernels", - type=int, - default=None, - help="Rows to keep in the short-kernel table.", - ) - parser.add_argument( - "--enable_pseudo_ops", - action="store_true", - default=False, - help="Enable automatic pseudo-op augmentation to tree to isolate specific kernels (e.g., FusedMoE).", - ) - parser.add_argument( - "--topk_ops", - type=int, - default=None, - help="Rows to keep in the unique-args launcher table.", - ) - parser.add_argument( - "--topk_roofline_ops", - type=int, - default=None, - help="Rows to keep in the roofline table.", - ) - - parser.add_argument( - "--comparison_json_path", - type=str, - default=None, - help=( - "Path to a second trace to compare against the primary trace. " - "Runs TraceDiff and adds speedup, delta, and LCA columns to " - "unified_perf_summary, plus a diff_stats sheet." - ), - ) - - parser.add_argument( - "--extension_file", - type=str, - default=None, - help="Path to the extension file containing custom extensions for TraceTree and PerfModel.", - ) - - parser.add_argument( - "--python_path", - type=str, - default=None, - help="Path to the python executable for gemm simulator", - ) - add_gpu_arch_cli_args(parser) - parser.add_argument( - "--enable-origami", - action="store_true", - default=False, - help="Use Origami for simulated GEMM/SDPA times when a GPU arch JSON is provided", - ) - - parser.add_argument( - "--capture_folder", - type=str, - required=False, - help="Path to the capture trace folder", - ) - parser.add_argument( - "--comparison_capture_folder", - type=str, - required=False, - help="Path to the capture trace folder for the comparison trace", - ) - parser.add_argument( - "--group_by_num_kernels", - action="store_true", - default=False, - help="Group by number of kernels in summary tables.", - ) - parser.add_argument( - "--include_call_stack", - action="store_true", - default=False, - help="Include callstack in the report.", - ) - parser.add_argument( - "--include_overlap_info", - action="store_true", - default=False, - help="Include overlap info in the report. Disabled by default. " - "Adds ops_unique_args_kl_overlap, unified_perf_summary_kl_overlap, and " - "per-category *_kl_overlap / *_fwd_kl_overlap / *_bwd_kl_overlap sheets when data exists.", - ) - - args = parser.parse_args() - if args.comparison_capture_folder and not args.comparison_json_path: - parser.error("--comparison_capture_folder requires --comparison_json_path.") - if args.capture_folder: - metadata_json_path = os.path.join(args.capture_folder, "execution_details.json") - classify_graph_capture_trace(args.capture_folder) - graph_tree = merge_capture_trace_into_graph( - args.capture_folder, - metadata_json_path, - args.profile_json_path, - ) - else: - graph_tree = None - comparison_graph_tree = None - if args.comparison_capture_folder: - comp_metadata = os.path.join( - args.comparison_capture_folder, "execution_details.json" - ) - classify_graph_capture_trace(args.comparison_capture_folder) - comparison_graph_tree = merge_capture_trace_into_graph( - args.comparison_capture_folder, - comp_metadata, - args.comparison_json_path, - ) - generate_perf_report_pytorch( - profile_json_path=args.profile_json_path, - augmented_tree=graph_tree, - output_xlsx_path=args.output_xlsx_path, - output_csvs_dir=args.output_csvs_dir, - include_unlinked_kernels=args.include_unlinked_kernels, - enable_pseudo_ops=args.enable_pseudo_ops, - micro_idle_thresh_us=args.micro_idle_thresh_us, - collective_analysis=args.collective_analysis, - include_overlap_info=args.include_overlap_info, - kernel_summary=args.kernel_summary, - short_kernel_study=args.short_kernel_study, - short_kernel_threshold_us=args.short_kernel_threshold_us, - short_kernel_histogram_bins=args.short_kernel_histogram_bins, - topk_short_kernels=args.topk_short_kernels, - topk_ops=args.topk_ops, - topk_roofline_ops=args.topk_roofline_ops, - comparison_json_path=args.comparison_json_path, - comparison_augmented_tree=comparison_graph_tree, - extension_file=args.extension_file, - python_path=args.python_path, - gpu_arch_json_path=args.gpu_arch_json_path, - gpu_arch_platform=args.gpu_arch_platform, - enable_origami=args.enable_origami, - group_by_parent_module=args.group_by_parent_module, - group_by_num_kernels=args.group_by_num_kernels, - include_call_stack=args.include_call_stack, - ) - - -if __name__ == "__main__": - main() diff --git a/docs/how-to/generate-perf-report-pytorch-inference.md b/docs/how-to/generate-perf-report-pytorch-inference.md index 519eb4092..48564106d 100644 --- a/docs/how-to/generate-perf-report-pytorch-inference.md +++ b/docs/how-to/generate-perf-report-pytorch-inference.md @@ -12,11 +12,12 @@ See LICENSE for license information. ``` -`TraceLens_generate_perf_report_pytorch_inference` is the inference-oriented -variant of the PyTorch report. It targets inference traces from frameworks such -as vLLM, SGLang, ATOM, and xDiT that run in CUDA/HIP graph mode, and can merge the -graph-capture traces back into the graph-replay trace to recover the call-stack -and input-shape metadata that graph execution drops. +`TraceLens_generate_perf_report_pytorch` is a unified PyTorch report generator +for both training and inference traces. For inference traces from frameworks such +as vLLM, SGLang, ATOM, and xDiT that run in CUDA/HIP graph mode, use +`--capture_folder` to merge graph-capture traces back into the graph-replay +trace and recover the call-stack and input-shape metadata that graph execution +drops, and `--group_by_parent_module` for module-level grouping. This topic covers the end-to-end inference workflow: collecting traces, splitting them into steady-state windows, and generating the report. For training or @@ -429,7 +430,7 @@ Report generation is supported for both eager-mode and graph-mode (capture + replay) traces. Pass the graph-replay trace to generate the default Excel report: ```bash -TraceLens_generate_perf_report_pytorch_inference \ +TraceLens_generate_perf_report_pytorch \ --profile_json_path tests/traces/inference/graph_full/graph_execution.json.gz ``` @@ -446,7 +447,7 @@ traces back into the replay trace and restore that metadata for richer operator and roofline analysis: ```bash -TraceLens_generate_perf_report_pytorch_inference \ +TraceLens_generate_perf_report_pytorch \ --profile_json_path tests/traces/inference/graph_full/graph_execution.json.gz \ --capture_folder tests/traces/inference/graph_full/capture_traces ``` diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 349022b2c..7ef5c403c 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -50,20 +50,15 @@ Generate a multi-sheet Excel report from a PyTorch (`torch.profiler`) trace. **Output:** an `.xlsx` workbook with the GPU-timeline, operator-category, operator, unique-argument, and roofline sheets. -### TraceLens_generate_perf_report_pytorch_inference - -Inference-oriented variant of the PyTorch report. +For inference traces (vLLM/SGLang/ATOM/xDiT), additional flags are available: | Argument | Default | Description | |----------|---------|-------------| -| `--profile_json_path` | required | Path to the trace. | | `--group_by_parent_module` | off | Group kernel-launcher summaries by parent `nn.Module`. | -| `--capture_folder` | None | Path to the capture-trace folder. | -| `--include_overlap_info` | off | Add `*_kl_overlap` sheets when data exists. | +| `--capture_folder` | None | Path to the graph-capture-trace folder (merges capture metadata into the replay tree). | +| `--comparison_capture_folder` | None | Capture folder for the comparison trace. | -Shares most options with `TraceLens_generate_perf_report_pytorch` (output -paths, short-kernel study, roofline/Origami, comparison, call stack). Run with -`--help` for the full list. +Run with `--help` for the full list. ### TraceLens_generate_perf_report_jax diff --git a/docs/what-is-tracelens.md b/docs/what-is-tracelens.md index ecd4850bf..895263b62 100644 --- a/docs/what-is-tracelens.md +++ b/docs/what-is-tracelens.md @@ -82,7 +82,6 @@ TraceLens supports the following trace formats: | Format | Source tool | Report CLI | |--------|-------------|------------| | PyTorch | `torch.profiler` | `TraceLens_generate_perf_report_pytorch` | -| PyTorch (inference) | `torch.profiler` | `TraceLens_generate_perf_report_pytorch_inference` | | JAX | XPlane protobuf | `TraceLens_generate_perf_report_jax` | | rocprofv3 JSON | AMD ROCm ROCprofiler-SDK | `TraceLens_generate_perf_report_rocprof` | | rocprofv3 pftrace (Perfetto-style) | `rocprofv3 --output-format pftrace` | `TraceLens_generate_perf_report_pftrace_hip_activity`, `..._pftrace_hip_api`, `..._pftrace_memory_copy` | diff --git a/setup.py b/setup.py index 3a924c4d7..d9fd1f6ba 100755 --- a/setup.py +++ b/setup.py @@ -84,7 +84,6 @@ def _wheel_version(): "console_scripts": [ "TraceLens_generate_perf_report_jax = TraceLens.Reporting.generate_perf_report_jax:main", "TraceLens_generate_perf_report_pytorch = TraceLens.Reporting.generate_perf_report_pytorch:main", - "TraceLens_generate_perf_report_pytorch_inference = TraceLens.Reporting.generate_perf_report_pytorch_inference:main", "TraceLens_generate_perf_report_rocprof = TraceLens.Reporting.generate_perf_report_rocprof:main", "TraceLens_compare_perf_reports_pytorch = TraceLens.Reporting.compare_perf_reports_pytorch:main", "TraceLens_generate_multi_rank_collective_report_pytorch = TraceLens.Reporting.generate_multi_rank_collective_report_pytorch:main", diff --git a/tests/test_inference_perf_report.py b/tests/test_inference_perf_report.py index 8fe3ae9f0..a71f070a4 100644 --- a/tests/test_inference_perf_report.py +++ b/tests/test_inference_perf_report.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -# Regression tests for generate_perf_report_pytorch_inference. +# Regression tests for generate_perf_report_pytorch (inference mode). # Each test case is a subdirectory under tests/traces/inference/ containing: # - A .json.gz trace file # - A perf_csvs/ folder with reference CSV files (one per output sheet) @@ -13,7 +13,7 @@ import os, numpy as np, pandas as pd, pytest, ast, re, gzip, json, glob from pandas.api.types import is_float_dtype -from TraceLens.Reporting.generate_perf_report_pytorch_inference import ( +from TraceLens.Reporting.generate_perf_report_pytorch import ( classify_graph_capture_trace, generate_perf_report_pytorch, generate_perf_report_pytorch as gen_inf, diff --git a/tests/test_pftrace_hip_activity_report.py b/tests/test_pftrace_hip_activity_report.py index ef1967020..69afcc75e 100644 --- a/tests/test_pftrace_hip_activity_report.py +++ b/tests/test_pftrace_hip_activity_report.py @@ -381,11 +381,13 @@ def test_pftrace_analyzer_and_report(self, tmp_path): def test_inference_report_main(tmp_path): trace = _write_trace(tmp_path, [("aten::mm", "gemm_kernel", 80)]) out_dir = tmp_path / "inf_csvs" - import TraceLens.Reporting.generate_perf_report_pytorch_inference as mod + import importlib + + mod = importlib.import_module("TraceLens.Reporting.generate_perf_report_pytorch") old_argv = sys.argv sys.argv = [ - "generate_perf_report_pytorch_inference", + "generate_perf_report_pytorch", "--profile_json_path", trace, "--output_csvs_dir", diff --git a/tests/test_pseudo_ops_extension.py b/tests/test_pseudo_ops_extension.py index bc3decab8..492a8f010 100644 --- a/tests/test_pseudo_ops_extension.py +++ b/tests/test_pseudo_ops_extension.py @@ -38,7 +38,7 @@ ) from tests.fixtures.traces import NORM_TRACE, TRACES_ROOT from tests.fixtures.treeperf import _make_gpu_event, _mk_ac2g -from TraceLens.Reporting.generate_perf_report_pytorch_inference import ( +from TraceLens.Reporting.generate_perf_report_pytorch import ( generate_perf_report_pytorch as generate_inference_report, ) from tests.fixtures.reporting import _mk_ac2g, _mk_event diff --git a/tests/test_reporting_inference_helpers.py b/tests/test_reporting_inference_helpers.py index 2ba0c48c6..9715eb78c 100644 --- a/tests/test_reporting_inference_helpers.py +++ b/tests/test_reporting_inference_helpers.py @@ -4,10 +4,10 @@ # See LICENSE for license information. ############################################################################### -"""Unit tests for helper functions in generate_perf_report_pytorch_inference.""" +"""Unit tests for helper functions in generate_perf_report_pytorch.""" import os, pandas as pd, pytest -from TraceLens.Reporting.generate_perf_report_pytorch_inference import ( +from TraceLens.Reporting.generate_perf_report_pytorch import ( add_truncated_kernel_details, get_dfs_short_kernels, perf_report_sanity_check, diff --git a/tests/test_reporting_utils.py b/tests/test_reporting_utils.py index 69ca93690..be17bcbf1 100644 --- a/tests/test_reporting_utils.py +++ b/tests/test_reporting_utils.py @@ -26,7 +26,6 @@ ) from TraceLens.Reporting import ( generate_multi_rank_collective_report_pytorch as coll_mod, - generate_perf_report_pytorch_inference as inf_mod, reporting_utils as ru, tracediff_comparison_extension as tde, ) @@ -51,7 +50,7 @@ TRACES_ROOT, _discover_inference_cases, ) -from TraceLens.Reporting.generate_perf_report_pytorch_inference import ( +from TraceLens.Reporting.generate_perf_report_pytorch import ( add_truncated_kernel_details as add_truncated_inference, add_truncated_kernel_details as add_truncated_kernel_details_inference, apply_extension as apply_extension_inference, @@ -1068,7 +1067,7 @@ def test_inference_report_main_cli(tmp_path): old_argv = sys.argv sys.argv = [ - "generate_perf_report_pytorch_inference", + "generate_perf_report_pytorch", "--profile_json_path", trace, "--output_csvs_dir", @@ -1080,7 +1079,9 @@ def test_inference_report_main_cli(tmp_path): "--group_by_parent_module", ] try: - inf_mod.main() + importlib.import_module( + "TraceLens.Reporting.generate_perf_report_pytorch" + ).main() finally: sys.argv = old_argv assert xlsx.exists() @@ -1644,12 +1645,12 @@ def test_generate_perf_report_pytorch_main(self, tmp_path): def test_generate_perf_report_inference_main(self, tmp_path): mod = importlib.import_module( - "TraceLens.Reporting.generate_perf_report_pytorch_inference" + "TraceLens.Reporting.generate_perf_report_pytorch" ) trace = _write_trace(tmp_path, [("aten::mm", "gemm_kernel", 100)], "inf.json") old_argv = sys.argv sys.argv = [ - "generate_perf_report_pytorch_inference", + "generate_perf_report_pytorch", "--profile_json_path", trace, "--output_csvs_dir", From dffbaecf3ff55c8c4605156de511246b56ad66ba Mon Sep 17 00:00:00 2001 From: Kyle Hoffmeyer Date: Mon, 31 Aug 2026 17:52:40 -0700 Subject: [PATCH 10/10] update imports --- tests/test_pftrace_hip_activity_report.py | 9 ++++----- tests/test_reporting_utils.py | 14 ++++---------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/tests/test_pftrace_hip_activity_report.py b/tests/test_pftrace_hip_activity_report.py index 69afcc75e..25d10f793 100644 --- a/tests/test_pftrace_hip_activity_report.py +++ b/tests/test_pftrace_hip_activity_report.py @@ -11,6 +11,9 @@ _write_markdown_report, generate_perf_report_pftrace_hip_activity, ) +from TraceLens.Reporting.generate_perf_report_pytorch import ( + main as generate_perf_report_pytorch_main, +) from TraceLens.Reporting.pftrace_hip_activity_analysis import ( Event, PftraceHipActivityAnalyzer, @@ -381,10 +384,6 @@ def test_pftrace_analyzer_and_report(self, tmp_path): def test_inference_report_main(tmp_path): trace = _write_trace(tmp_path, [("aten::mm", "gemm_kernel", 80)]) out_dir = tmp_path / "inf_csvs" - import importlib - - mod = importlib.import_module("TraceLens.Reporting.generate_perf_report_pytorch") - old_argv = sys.argv sys.argv = [ "generate_perf_report_pytorch", @@ -396,7 +395,7 @@ def test_inference_report_main(tmp_path): "--enable_kernel_summary", ] try: - mod.main() + generate_perf_report_pytorch_main() finally: sys.argv = old_argv assert (out_dir / "gpu_timeline.csv").exists() diff --git a/tests/test_reporting_utils.py b/tests/test_reporting_utils.py index be17bcbf1..c25416f61 100644 --- a/tests/test_reporting_utils.py +++ b/tests/test_reporting_utils.py @@ -59,6 +59,7 @@ generate_perf_report_pytorch as gen_inf, generate_perf_report_pytorch as generate_inference_report, get_dfs_short_kernels as get_dfs_short_kernels_inference, + main as generate_perf_report_pytorch_main, perf_report_sanity_check, ) from TraceLens.Reporting.compare_perf_reports_pytorch import ( @@ -1079,9 +1080,7 @@ def test_inference_report_main_cli(tmp_path): "--group_by_parent_module", ] try: - importlib.import_module( - "TraceLens.Reporting.generate_perf_report_pytorch" - ).main() + generate_perf_report_pytorch_main() finally: sys.argv = old_argv assert xlsx.exists() @@ -1644,9 +1643,6 @@ def test_generate_perf_report_pytorch_main(self, tmp_path): assert (tmp_path / "csv" / "gpu_timeline.csv").exists() def test_generate_perf_report_inference_main(self, tmp_path): - mod = importlib.import_module( - "TraceLens.Reporting.generate_perf_report_pytorch" - ) trace = _write_trace(tmp_path, [("aten::mm", "gemm_kernel", 100)], "inf.json") old_argv = sys.argv sys.argv = [ @@ -1659,7 +1655,7 @@ def test_generate_perf_report_inference_main(self, tmp_path): str(tmp_path / "out.xlsx"), ] try: - mod.main() + generate_perf_report_pytorch_main() finally: sys.argv = old_argv assert (tmp_path / "csv" / "gpu_timeline.csv").exists() @@ -1995,8 +1991,6 @@ def test_pytorch_report_main(tmp_path): ) out_dir = tmp_path / "py_csvs" xlsx = tmp_path / "py.xlsx" - mod = importlib.import_module("TraceLens.Reporting.generate_perf_report_pytorch") - old_argv = sys.argv sys.argv = [ "generate_perf_report_pytorch", @@ -2012,7 +2006,7 @@ def test_pytorch_report_main(tmp_path): "--group_by_num_kernels", ] try: - mod.main() + generate_perf_report_pytorch_main() finally: sys.argv = old_argv assert xlsx.exists()